Fix Issue 4857 Case Two

naga/src/back/spv/block.rs

@@ -423,8 +423,19 @@ impl<'w> BlockContext<'w> {
                         };
 
                         let binding_type_id = self.get_type_id(LookupType::Handle(binding_type));
-
                         let load_id = self.gen_id();
+
+                        // Map `binding_type_id` back to the original pointer if it is an opaque
+                        // type.
+                        match self.ir_module.types[binding_type].inner {
+                            crate::TypeInner::Image { .. }
+                            | crate::TypeInner::Sampler { .. }
+                            | crate::TypeInner::AccelerationStructure => {
+                                self.function_arg_ids.insert(load_id, result_id);
+                            }
+                            _ => {}
+                        }
+
                         block.body.push(Instruction::load(
                             binding_type_id,
                             load_id,
@@ -514,8 +525,19 @@ impl<'w> BlockContext<'w> {
                         };
 
                         let binding_type_id = self.get_type_id(LookupType::Handle(binding_type));
-
                         let load_id = self.gen_id();
+
+                        // Map `binding_type_id` back to the original pointer if it is an opaque
+                        // type.
+                        match self.ir_module.types[binding_type].inner {
+                            crate::TypeInner::Image { .. }
+                            | crate::TypeInner::Sampler { .. }
+                            | crate::TypeInner::AccelerationStructure => {
+                                self.function_arg_ids.insert(load_id, result_id);
+                            }
+                            _ => {}
+                        }
+
                         block.body.push(Instruction::load(
                             binding_type_id,
                             load_id,
@@ -2668,7 +2690,13 @@ impl<'w> BlockContext<'w> {
                     let id = self.gen_id();
                     self.temp_list.clear();
                     for &argument in arguments {
-                        self.temp_list.push(self.cached[argument]);
+                        // Check if we should use the `argument_id` directly or the pointer to it.
+                        let argument_id = self.cached[argument];
+                        let argument_id = self
+                            .function_arg_ids
+                            .get(&argument_id)
+                            .map_or(argument_id, |&pointer_id| pointer_id);
+                        self.temp_list.push(argument_id);
                     }
 
                     let type_id = match result {

naga/src/back/spv/mod.rs

@@ -493,6 +493,7 @@ impl CachedExpressions {
         self.ids.resize(length, 0);
     }
 }
+
 impl ops::Index<Handle<crate::Expression>> for CachedExpressions {
     type Output = Word;
     fn index(&self, h: Handle<crate::Expression>) -> &Word {
@@ -689,6 +690,10 @@ struct BlockContext<'w> {
     /// SPIR-V ids for expressions we've evaluated.
     cached: CachedExpressions,
 
+    /// The pointers of the cached expressions' SPIR-V ids from [`Block::Context::cached`].
+    /// Only used when loaded opaque types need to be passed to a function call.
+    function_arg_ids: crate::FastIndexMap<Word, Word>,
+
     /// The `Writer`'s temporary vector, for convenience.
     temp_list: Vec<Word>,
 
@@ -762,6 +767,11 @@ pub struct Writer {
     // retain the table here between functions to save heap allocations.
     saved_cached: CachedExpressions,
 
+    // Maps the expression ids from `saved_cached` to the pointer id they were loaded from.
+    // Only used when opaque types need to be passed to a function call.
+    // This goes alongside `saved_cached`, so it too is only meaningful within a BlockContext.
+    function_arg_ids: crate::FastIndexMap<Word, Word>,
+
     gl450_ext_inst_id: Word,
 
     // Just a temporary list of SPIR-V ids

naga/src/back/spv/recyclable.rs

@@ -59,6 +59,13 @@ impl<K, S: Clone> Recyclable for indexmap::IndexSet<K, S> {
     }
 }
 
+impl<K, V, S: Clone> Recyclable for indexmap::IndexMap<K, V, S> {
+    fn recycle(mut self) -> Self {
+        self.clear();
+        self
+    }
+}
+
 impl<K: Ord, V> Recyclable for std::collections::BTreeMap<K, V> {
     fn recycle(mut self) -> Self {
         self.clear();

naga/src/back/spv/writer.rs

@@ -79,6 +79,7 @@ impl Writer {
             global_variables: HandleVec::new(),
             binding_map: options.binding_map.clone(),
             saved_cached: CachedExpressions::default(),
+            function_arg_ids: crate::FastIndexMap::default(),
             gl450_ext_inst_id,
             temp_list: Vec::new(),
         })
@@ -130,6 +131,7 @@ impl Writer {
             cached_constants: take(&mut self.cached_constants).recycle(),
             global_variables: take(&mut self.global_variables).recycle(),
             saved_cached: take(&mut self.saved_cached).recycle(),
+            function_arg_ids: take(&mut self.function_arg_ids).recycle(),
             temp_list: take(&mut self.temp_list).recycle(),
         };
 
@@ -585,6 +587,7 @@ impl Writer {
             function: &mut function,
             // Re-use the cached expression table from prior functions.
             cached: std::mem::take(&mut self.saved_cached),
+            function_arg_ids: std::mem::take(&mut self.function_arg_ids),
 
             // Steal the Writer's temp list for a bit.
             temp_list: std::mem::take(&mut self.temp_list),

naga/src/front/spv/mod.rs

@@ -4406,6 +4406,17 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                             crate::Expression::FunctionArgument(i) => {
                                 fun_parameter_sampling[i as usize] |= flags;
                             }
+                            crate::Expression::Access { base, .. } => match expressions[base] {
+                                crate::Expression::GlobalVariable(handle) => {
+                                    if let Some(sampling) = self.handle_sampling.get_mut(&handle) {
+                                        *sampling |= flags
+                                    }
+                                }
+                                crate::Expression::FunctionArgument(i) => {
+                                    fun_parameter_sampling[i as usize] |= flags;
+                                }
+                                ref other => return Err(Error::InvalidGlobalVar(other.clone())),
+                            },
                             ref other => return Err(Error::InvalidGlobalVar(other.clone())),
                         }
                     }