Helper function for identifying Cx2 matrices in a uniform buffer

Author: andyleiserson

naga/src/back/hlsl/writer.rs

@@ -2891,37 +2891,13 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                     //
                     // Note that this only works for `Load`s and we handle
                     // `Store`s differently in `Statement::Store`.
-                    let cx2_columns;
                     if let Some(MatrixType {
                         columns,
                         rows: crate::VectorSize::Bi,
                         width: 4,
                     }) = get_inner_matrix_of_struct_array_member(module, base, func_ctx, true)
+                        .or_else(|| get_global_uniform_matrix(module, base, func_ctx))
                     {
-                        cx2_columns = Some(columns);
-                    } else {
-                        let base_tr = func_ctx
-                            .resolve_type(base, &module.types)
-                            .pointer_base_type();
-                        let base_ty = base_tr.as_ref().map(|tr| tr.inner_with(&module.types));
-                        match (&func_ctx.expressions[base], base_ty) {
-                            (
-                                &Expression::GlobalVariable(handle),
-                                Some(&TypeInner::Matrix {
-                                    columns,
-                                    rows: crate::VectorSize::Bi,
-                                    ..
-                                }),
-                            ) if module.global_variables[handle].space
-                                == crate::AddressSpace::Uniform =>
-                            {
-                                cx2_columns = Some(columns);
-                            }
-                            _ => cx2_columns = None,
-                        }
-                    }
-
-                    if let Some(columns) = cx2_columns {
                         write!(self.out, "__get_col_of_mat{}x2(", columns as u8)?;
                         self.write_expr(module, base, func_ctx)?;
                         write!(self.out, ", ")?;
@@ -3033,38 +3009,16 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 {
                     // do nothing, the chain is written on `Load`/`Store`
                 } else {
-                    // We write the matrix column access in a special way since
-                    // the type of `base` is our special __matCx2 struct.
-                    let is_cx2;
+                    // See if we need to write the matrix column access in a
+                    // special way since the type of `base` is our special
+                    // __matCx2 struct.
                     if let Some(MatrixType {
                         rows: crate::VectorSize::Bi,
                         width: 4,
                         ..
                     }) = get_inner_matrix_of_struct_array_member(module, base, func_ctx, true)
+                        .or_else(|| get_global_uniform_matrix(module, base, func_ctx))
                     {
-                        is_cx2 = true;
-                    } else {
-                        let base_tr = func_ctx
-                            .resolve_type(base, &module.types)
-                            .pointer_base_type();
-                        let base_ty = base_tr.as_ref().map(|tr| tr.inner_with(&module.types));
-                        match (&func_ctx.expressions[base], base_ty) {
-                            (
-                                &Expression::GlobalVariable(handle),
-                                Some(&TypeInner::Matrix {
-                                    rows: crate::VectorSize::Bi,
-                                    ..
-                                }),
-                            ) if module.global_variables[handle].space
-                                == crate::AddressSpace::Uniform =>
-                            {
-                                is_cx2 = true;
-                            }
-                            _ => is_cx2 = false,
-                        }
-                    }
-
-                    if is_cx2 {
                         self.write_expr(module, base, func_ctx)?;
                         write!(self.out, "._{index}")?;
                         return Ok(());
@@ -4497,6 +4451,36 @@ pub(super) fn get_inner_matrix_of_struct_array_member(
     None
 }
 
+/// Simpler version of get_inner_matrix_of_global_uniform that only looks at the
+/// immediate expression, rather than traversing an access chain.
+fn get_global_uniform_matrix(
+    module: &Module,
+    base: Handle<crate::Expression>,
+    func_ctx: &back::FunctionCtx<'_>,
+) -> Option<MatrixType> {
+    let base_tr = func_ctx
+        .resolve_type(base, &module.types)
+        .pointer_base_type();
+    let base_ty = base_tr.as_ref().map(|tr| tr.inner_with(&module.types));
+    match (&func_ctx.expressions[base], base_ty) {
+        (
+            &crate::Expression::GlobalVariable(handle),
+            Some(&TypeInner::Matrix {
+                columns,
+                rows,
+                scalar,
+            }),
+        ) if module.global_variables[handle].space == crate::AddressSpace::Uniform => {
+            Some(MatrixType {
+                columns,
+                rows,
+                width: scalar.width,
+            })
+        }
+        _ => None,
+    }
+}
+
 /// Returns the matrix data if the access chain starting at `base`:
 /// - starts with an expression with resolved type of [`TypeInner::Matrix`]
 /// - contains zero or more expressions with resolved type of [`TypeInner::Array`] of [`TypeInner::Matrix`]