[spv-out] implement array value indexing

gfx-rs/naga#723 back from dead

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

naga/src/back/spv/block.rs

@@ -346,14 +346,32 @@ 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 (id, variable) = self.writer.promote_access_expression_to_variable(
+                            &self.ir_module.types,
+                            result_type_id,
+                            base_id,
+                            &self.fun_info[base].ty,
+                            index_id,
+                            ty_element,
+                            block,
+                        )?;
+                        self.function.internal_variables.push(variable);
+                        id
+                    }
+                    // TODO: 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

@@ -143,6 +143,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,62 @@ 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_resolution: &TypeResolution,
+        index_id: Word,
+        element_ty: Handle<crate::Type>,
+        block: &mut Block,
+    ) -> Result<(Word, LocalVariable), Error> {
+        let container_type_id = self.get_expression_type_id(container_resolution);
+        let pointer_type_id = self.id_gen.next();
+        Instruction::type_pointer(
+            pointer_type_id,
+            spirv::StorageClass::Function,
+            container_type_id,
+        )
+        .to_words(&mut self.logical_layout.declarations);
+
+        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/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];
+}