linker/duplicates: handle all decorations instead of special-casing zombies.

crates/rustc_codegen_spirv/src/linker/duplicates.rs

@@ -1,4 +1,3 @@
-use crate::decorations::{CustomDecoration, ZombieDecoration};
 use rspirv::binary::Assemble;
 use rspirv::dr::{Instruction, Module, Operand};
 use rspirv::spirv::{Op, Word};
@@ -73,15 +72,20 @@ fn make_annotation_key(inst: &Instruction) -> Vec<u32> {
 fn gather_annotations(annotations: &[Instruction]) -> FxHashMap<Word, Vec<u32>> {
     let mut map = FxHashMap::default();
     for inst in annotations {
-        if inst.class.opcode == Op::Decorate || inst.class.opcode == Op::MemberDecorate {
-            match map.entry(inst.operands[0].id_ref_any().unwrap()) {
+        match inst.class.opcode {
+            Op::Decorate
+            | Op::DecorateId
+            | Op::DecorateString
+            | Op::MemberDecorate
+            | Op::MemberDecorateString => match map.entry(inst.operands[0].id_ref_any().unwrap()) {
                 hash_map::Entry::Vacant(entry) => {
                     entry.insert(vec![make_annotation_key(inst)]);
                 }
                 hash_map::Entry::Occupied(mut entry) => {
                     entry.get_mut().push(make_annotation_key(inst));
                 }
-            }
+            },
+            _ => {}
         }
     }
     map.into_iter()
@@ -110,7 +114,6 @@ fn gather_names(debug_names: &[Instruction]) -> FxHashMap<Word, String> {
 fn make_dedupe_key(
     inst: &Instruction,
     unresolved_forward_pointers: &FxHashSet<Word>,
-    zombies: &FxHashSet<Word>,
     annotations: &FxHashMap<Word, Vec<u32>>,
     names: &FxHashMap<Word, String>,
 ) -> Vec<u32> {
@@ -136,22 +139,16 @@ fn make_dedupe_key(
         }
     }
     if let Some(id) = inst.result_id {
-        data.push(if zombies.contains(&id) {
-            if inst.result_type.is_some() {
-                id
-            } else {
-                1
-            }
-        } else {
-            0
-        });
         if let Some(annos) = annotations.get(&id) {
             data.extend_from_slice(annos);
         }
         if inst.class.opcode == Op::Variable {
             // Names only matter for OpVariable.
             if let Some(name) = names.get(&id) {
                 // Jump through some hoops to shove a String into a Vec<u32>.
+                //
+                // FIXME(eddyb) this should `.assemble_into(&mut data)` the
+                // `Operand::LiteralString(...)` from the original `Op::Name`.
                 for chunk in name.as_bytes().chunks(4) {
                     let slice = match *chunk {
                         [a] => [a, 0, 0, 0],
@@ -184,11 +181,6 @@ fn rewrite_inst_with_rules(inst: &mut Instruction, rules: &FxHashMap<u32, u32>)
 pub fn remove_duplicate_types(module: &mut Module) {
     // Keep in mind, this algorithm requires forward type references to not exist - i.e. it's a valid spir-v module.
 
-    // Include zombies in the key to not merge zombies with non-zombies
-    let zombies: FxHashSet<Word> = ZombieDecoration::decode_all(module)
-        .map(|(z, _)| z)
-        .collect();
-
     // When a duplicate type is encountered, then this is a map from the deleted ID, to the new, deduplicated ID.
     let mut rewrite_rules = FxHashMap::default();
     // Instructions are encoded into "keys": their opcode, followed by arguments, then annotations.
@@ -224,13 +216,7 @@ pub fn remove_duplicate_types(module: &mut Module) {
         // all_inst_iter_mut pass below. However, the code is a lil bit cleaner this way I guess.
         rewrite_inst_with_rules(inst, &rewrite_rules);
 
-        let key = make_dedupe_key(
-            inst,
-            &unresolved_forward_pointers,
-            &zombies,
-            &annotations,
-            &names,
-        );
+        let key = make_dedupe_key(inst, &unresolved_forward_pointers, &annotations, &names);
 
         match key_to_result_id.entry(key) {
             hash_map::Entry::Vacant(entry) => {

crates/rustc_codegen_spirv/src/linker/import_export_link.rs

@@ -1,4 +1,5 @@
 use super::Result;
+use crate::decorations::{CustomDecoration, ZombieDecoration};
 use rspirv::dr::{Instruction, Module};
 use rspirv::spirv::{Capability, Decoration, LinkageType, Op, Word};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
@@ -26,6 +27,27 @@ fn find_import_export_pairs_and_killed_params(
     let mut rewrite_rules = FxHashMap::default();
     let mut killed_parameters = FxHashSet::default();
 
+    // HACK(eddyb) collect all zombies, and anything that transitively mentions
+    // them (or is "infected"), to ignore them when doing type checks, as they
+    // will either be removed later, or become an error of their own.
+    let mut zombie_infected: FxHashSet<Word> = ZombieDecoration::decode_all(module)
+        .map(|(z, _)| z)
+        .collect();
+    for inst in module.global_inst_iter() {
+        if let Some(result_id) = inst.result_id {
+            if !zombie_infected.contains(&result_id) {
+                let mut id_operands = inst.operands.iter().filter_map(|o| o.id_ref_any());
+                // NOTE(eddyb) this takes advantage of the fact that the module
+                // is ordered def-before-use (with the minor exception of forward
+                // references for recursive data, which are not fully supported),
+                // to be able to propagate "zombie infection" in one pass.
+                if id_operands.any(|id| zombie_infected.contains(&id)) {
+                    zombie_infected.insert(result_id);
+                }
+            }
+        }
+    }
+
     // First, collect all the exports.
     for annotation in &module.annotations {
         let (id, name) = match get_linkage_inst(annotation) {
@@ -53,7 +75,14 @@ fn find_import_export_pairs_and_killed_params(
         };
         let import_type = *type_map.get(&import_id).expect("Unexpected op");
         // Make sure the import/export pair has the same type.
-        check_tys_equal(sess, module, name, import_type, export_type)?;
+        check_tys_equal(
+            sess,
+            module,
+            name,
+            import_type,
+            export_type,
+            &zombie_infected,
+        )?;
         rewrite_rules.insert(import_id, export_id);
         if let Some(params) = fn_parameters.get(&import_id) {
             for &param in params {
@@ -111,8 +140,17 @@ fn check_tys_equal(
     name: &str,
     import_type: Word,
     export_type: Word,
+    zombie_infected: &FxHashSet<Word>,
 ) -> Result<()> {
-    if import_type == export_type {
+    let allowed = import_type == export_type || {
+        // HACK(eddyb) zombies can cause types to differ in definition, due to
+        // requiring multiple different instances (usually different `Span`s),
+        // so we ignore them, as `find_import_export_pairs_and_killed_params`'s
+        // own comment explains (zombies will cause errors or be removed, *later*).
+        zombie_infected.contains(&import_type) && zombie_infected.contains(&export_type)
+    };
+
+    if allowed {
         Ok(())
     } else {
         // We have an error. It's okay to do something really slow now to report the error.