[wgsl-in, spv-out] Allow dynamic indexing of arrays by value.

Bring gfx-rs/naga#723 back from the dead.

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>
Co-authored-by: Dzmitry Malyshau <kvarkus@gmail.com>
Co-authored-by: Jim Blandy <jimb@red-bean.com>

naga/src/back/spv/block.rs

@@ -343,14 +343,40 @@ impl<'w> BlockContext<'w> {
 
                         load_id
                     }
+                    crate::TypeInner::Array {
+                        base: ty_element, ..
+                    } => {
+                        let index_id = self.cached[index];
+                        let base_id = self.cached[base];
+                        let base_ty = match self.fun_info[base].ty {
+                            TypeResolution::Handle(handle) => handle,
+                            TypeResolution::Value(_) => {
+                                return Err(Error::Validation(
+                                    "Array types should always be in the arena",
+                                ))
+                            }
+                        };
+                        let (id, variable) = self.writer.promote_access_expression_to_variable(
+                            &self.ir_module.types,
+                            result_type_id,
+                            base_id,
+                            base_ty,
+                            index_id,
+                            ty_element,
+                            block,
+                        )?;
+                        self.function.internal_variables.push(variable);
+                        id
+                    }
+                    // wgpu#4337: Support `crate::TypeInner::Matrix`
                     ref other => {
                         log::error!(
                             "Unable to access base {:?} of type {:?}",
                             self.ir_function.expressions[base],
                             other
                         );
                         return Err(Error::Validation(
-                            "only vectors may be dynamically indexed by value",
+                            "only vectors and arrays may be dynamically indexed by value",
                         ));
                     }
                 }

naga/src/back/spv/mod.rs

@@ -144,6 +144,7 @@ struct Function {
     signature: Option<Instruction>,
     parameters: Vec<FunctionArgument>,
     variables: crate::FastHashMap<Handle<crate::LocalVariable>, LocalVariable>,
+    internal_variables: Vec<LocalVariable>,
     blocks: Vec<TerminatedBlock>,
     entry_point_context: Option<EntryPointContext>,
 }

naga/src/back/spv/writer.rs

@@ -32,6 +32,9 @@ impl Function {
                 for local_var in self.variables.values() {
                     local_var.instruction.to_words(sink);
                 }
+                for internal_var in self.internal_variables.iter() {
+                    internal_var.instruction.to_words(sink);
+                }
             }
             for instruction in block.body.iter() {
                 instruction.to_words(sink);
@@ -135,6 +138,56 @@ impl Writer {
         self.capabilities_used.insert(spirv::Capability::Shader);
     }
 
+    #[allow(clippy::too_many_arguments)]
+    pub(super) fn promote_access_expression_to_variable(
+        &mut self,
+        ir_types: &UniqueArena<crate::Type>,
+        result_type_id: Word,
+        container_id: Word,
+        container_ty: Handle<crate::Type>,
+        index_id: Word,
+        element_ty: Handle<crate::Type>,
+        block: &mut Block,
+    ) -> Result<(Word, LocalVariable), Error> {
+        let pointer_type_id =
+            self.get_pointer_id(ir_types, container_ty, spirv::StorageClass::Function)?;
+
+        let variable = {
+            let id = self.id_gen.next();
+            LocalVariable {
+                id,
+                instruction: Instruction::variable(
+                    pointer_type_id,
+                    id,
+                    spirv::StorageClass::Function,
+                    None,
+                ),
+            }
+        };
+        block
+            .body
+            .push(Instruction::store(variable.id, container_id, None));
+
+        let element_pointer_id = self.id_gen.next();
+        let element_pointer_type_id =
+            self.get_pointer_id(ir_types, element_ty, spirv::StorageClass::Function)?;
+        block.body.push(Instruction::access_chain(
+            element_pointer_type_id,
+            element_pointer_id,
+            variable.id,
+            &[index_id],
+        ));
+        let id = self.id_gen.next();
+        block.body.push(Instruction::load(
+            result_type_id,
+            id,
+            element_pointer_id,
+            None,
+        ));
+
+        Ok((id, variable))
+    }
+
     /// Indicate that the code requires any one of the listed capabilities.
     ///
     /// If nothing in `capabilities` appears in the available capabilities

naga/src/lib.rs

@@ -1402,21 +1402,20 @@ pub enum Expression {
     /// ## Dynamic indexing restrictions
     ///
     /// To accommodate restrictions in some of the shader languages that Naga
-    /// targets, it is not permitted to subscript a matrix or array with a
-    /// dynamically computed index unless that matrix or array appears behind a
-    /// pointer. In other words, if the inner type of `base` is [`Array`] or
-    /// [`Matrix`], then `index` must be a constant. But if the type of `base`
-    /// is a [`Pointer`] to an array or matrix or a [`ValuePointer`] with a
-    /// `size`, then the index may be any expression of integer type.
+    /// targets, it is not permitted to subscript a matrix with a dynamically
+    /// computed index unless that matrix appears behind a pointer. In other
+    /// words, if the inner type of `base` is [`Matrix`], then `index` must be a
+    /// constant. But if the type of `base` is a [`Pointer`] to an matrix, then
+    /// the index may be any expression of integer type.
     ///
     /// You can use the [`Expression::is_dynamic_index`] method to determine
-    /// whether a given index expression requires matrix or array base operands
-    /// to be behind a pointer.
+    /// whether a given index expression requires matrix base operands to be
+    /// behind a pointer.
     ///
     /// (It would be simpler to always require the use of `AccessIndex` when
-    /// subscripting arrays and matrices that are not behind pointers, but to
-    /// accommodate existing front ends, Naga also permits `Access`, with a
-    /// restricted `index`.)
+    /// subscripting matrices that are not behind pointers, but to accommodate
+    /// existing front ends, Naga also permits `Access`, with a restricted
+    /// `index`.)
     ///
     /// [`Vector`]: TypeInner::Vector
     /// [`Matrix`]: TypeInner::Matrix

naga/src/proc/mod.rs

@@ -521,12 +521,12 @@ impl crate::Expression {
         }
     }
 
-    /// Return true if this expression is a dynamic array index, for [`Access`].
+    /// Return true if this expression is a dynamic array/vector/matrix index,
+    /// for [`Access`].
     ///
     /// This method returns true if this expression is a dynamically computed
-    /// index, and as such can only be used to index matrices and arrays when
-    /// they appear behind a pointer. See the documentation for [`Access`] for
-    /// details.
+    /// index, and as such can only be used to index matrices when they appear
+    /// behind a pointer. See the documentation for [`Access`] for details.
     ///
     /// Note, this does not check the _type_ of the given expression. It's up to
     /// the caller to establish that the `Access` expression is well-typed

naga/src/proc/typifier.rs

@@ -92,6 +92,13 @@ pub enum TypeResolution {
     /// available in the associated arena. However, the `TypeInner` itself may
     /// contain `Handle<Type>` values referring to types from the arena.
     ///
+    /// The inner type must only be one of the following variants:
+    /// - TypeInner::Pointer
+    /// - TypeInner::ValuePointer
+    /// - TypeInner::Matrix (generated by matrix multiplication)
+    /// - TypeInner::Vector
+    /// - TypeInner::Scalar
+    ///
     /// [`TypeInner`]: crate::TypeInner
     Value(crate::TypeInner),
 }

naga/src/valid/expression.rs

@@ -240,9 +240,10 @@ impl super::Validator {
                 let base_type = &resolver[base];
                 // See the documentation for `Expression::Access`.
                 let dynamic_indexing_restricted = match *base_type {
-                    Ti::Vector { .. } => false,
-                    Ti::Matrix { .. } | Ti::Array { .. } => true,
-                    Ti::Pointer { .. }
+                    Ti::Matrix { .. } => true,
+                    Ti::Vector { .. }
+                    | Ti::Array { .. }
+                    | Ti::Pointer { .. }
                     | Ti::ValuePointer { size: Some(_), .. }
                     | Ti::BindingArray { .. } => false,
                     ref other => {

naga/tests/in/access.wgsl

@@ -167,3 +167,14 @@ fn assign_through_ptr() {
 	var arr = array<vec4<f32>, 2>(vec4(6.0), vec4(7.0));
     assign_array_through_ptr_fn(&arr);
 }
+
+@vertex
+fn foo(@builtin(vertex_index) vi: u32) -> @builtin(position) vec4<f32> {
+	let arr = array<i32, 5>(1, 2, 3, 4, 5);
+	let value = arr[vi];
+	return vec4<f32>(vec4<i32>(value));
+}
+
+fn array_by_value(a: array<i32, 5>, i: i32) -> i32 {
+    return a[i];
+}