Inliner: rewrite IDs on decorations fix

CHANGELOG.md

@@ -38,6 +38,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
     you may be able to `cargo update -p naga` to update to a fixed `naga` version  
     (`0.11.1` for `wgpu 0.15`, `0.12.1` for `wgpu 0.16`, and any later versions)
 
+### Fixed 🩹
+- [PR#1059](https://github.com/EmbarkStudios/rust-gpu/pull/1059) fixed the `inline` pass not copying `OpDecorate`s in the callee (which led to their loss).
+
 ## [0.7.0]
 
 ### Added ⭐

crates/rustc_codegen_spirv/src/linker/inline.rs

@@ -71,6 +71,7 @@ pub fn inline(sess: &Session, module: &mut Module) -> super::Result<()> {
     let mut inliner = Inliner {
         header: module.header.as_mut().unwrap(),
         types_global_values: &mut module.types_global_values,
+        annotations: &mut module.annotations,
         void,
         functions: &functions,
         legal_globals: &legal_globals,
@@ -337,6 +338,7 @@ fn should_inline(
 struct Inliner<'m, 'map> {
     header: &'m mut ModuleHeader,
     types_global_values: &'m mut Vec<Instruction>,
+    annotations: &'m mut Vec<Instruction>,
     void: Word,
     functions: &'map FunctionMap,
     legal_globals: &'map FxHashMap<Word, LegalGlobal>,
@@ -350,6 +352,24 @@ impl Inliner<'_, '_> {
         result
     }
 
+    /// Applies all rewrite rules to the decorations in the header.
+    fn apply_rewrite_for_decorations(&mut self, rewrite_rules: &FxHashMap<Word, Word>) {
+        // NOTE(siebencorgie): We don't care *what* decoration we rewrite atm. AFAIK there is no case where rewriting
+        // the decoration on inline wouldn't be valid.
+        for annotation_idx in 0..self.annotations.len() {
+            if self.annotations[annotation_idx].class.opcode == Op::Decorate {
+                if let Some(id) = self.annotations[annotation_idx].operands[0].id_ref_any_mut() {
+                    if let Some(&rewrite) = rewrite_rules.get(id) {
+                        // Copy decoration instruction and push it.
+                        let mut instcpy = self.annotations[annotation_idx].clone();
+                        *instcpy.operands[0].id_ref_any_mut().unwrap() = rewrite;
+                        self.annotations.push(instcpy);
+                    }
+                }
+            }
+        }
+    }
+
     fn ptr_ty(&mut self, pointee: Word) -> Word {
         // TODO: This is horribly slow, fix this
         let existing = self.types_global_values.iter().find(|inst| {
@@ -448,6 +468,7 @@ impl Inliner<'_, '_> {
         // fn is inlined multiple times.
         self.add_clone_id_rules(&mut rewrite_rules, &inlined_callee_blocks);
         apply_rewrite_rules(&rewrite_rules, &mut inlined_callee_blocks);
+        self.apply_rewrite_for_decorations(&rewrite_rules);
 
         // Split the block containing the `OpFunctionCall` into pre-call vs post-call.
         let pre_call_block_idx = block_idx;