Resolve composite element types during emission, protect loop continue blocks

- Build composite-to-element type map from OpTypeVector/Matrix/Array
  definitions and use element type for CompositeConstruct components
  instead of inheriting the composite result type.
- Skip RVSDG transform for selections whose branch targets overlap
  with loop continue blocks.

Author: LegNeato

rust/spirv-tools-opt/src/direct/emit.rs

@@ -129,6 +129,8 @@ pub struct EmitCtx<'a> {
     pub glsl_ext_id: Option<Word>,
     pub type_widths: &'a HashMap<Word, u32>,
     pub id_to_type: &'a HashMap<Word, Word>,
+    /// Maps composite type IDs (OpTypeVector, OpTypeMatrix, OpTypeArray) to their element type ID.
+    pub composite_element_types: &'a HashMap<Word, Word>,
 }
 
 // ---------------------------------------------------------------------------
@@ -1641,9 +1643,14 @@ fn emit_composite_construct(
     }
     let mut synth = Vec::new();
     let mut operands = Vec::new();
+    // For CompositeConstruct, the element type is derived from the composite type.
+    let element_type = ctx
+        .composite_element_types
+        .get(&result_type)
+        .copied()
+        .unwrap_or(result_type);
     for arg in &args[..expected] {
-        // Each component's type differs from result_type (composite type).
-        let component_type = resolve_term_type(arg, ctx).unwrap_or(result_type);
+        let component_type = resolve_term_type(arg, ctx).unwrap_or(element_type);
         let (arg_id, mut s) = emit_term(arg, component_type, ctx)?;
         synth.append(&mut s);
         operands.push(rspirv::dr::Operand::IdRef(arg_id));
@@ -1667,8 +1674,13 @@ fn flatten_expr_list(
     match term {
         Term::App { op, .. } if op == "ENil" => Some((Vec::new(), Vec::new())),
         Term::App { op, args } if op == "ECons" && args.len() >= 2 => {
-            // Each element's type differs from result_type (composite type).
-            let head_type = resolve_term_type(&args[0], ctx).unwrap_or(result_type);
+            // Element type from the composite type for correct component emission.
+            let element_type = ctx
+                .composite_element_types
+                .get(&result_type)
+                .copied()
+                .unwrap_or(result_type);
+            let head_type = resolve_term_type(&args[0], ctx).unwrap_or(element_type);
             let (head_id, head_synth) = emit_term(&args[0], head_type, ctx)?;
             let (mut rest_ids, mut rest_synth) = flatten_expr_list(&args[1], result_type, ctx)?;
             let mut ids = vec![head_id];
@@ -1863,6 +1875,7 @@ mod tests {
             glsl_ext_id: None,
             type_widths: &type_widths,
             id_to_type: &HashMap::new(),
+            composite_element_types: &HashMap::new(),
         };
         let term = parse_sexpr("(Sym \"id5\")").unwrap();
         let (result_id, synth) = emit_term(&term, 10, &mut ctx).unwrap();
@@ -1888,6 +1901,7 @@ mod tests {
             glsl_ext_id: None,
             type_widths: &type_widths,
             id_to_type: &HashMap::new(),
+            composite_element_types: &HashMap::new(),
         };
         let term = parse_sexpr("(Const 42)").unwrap();
         let (result_id, synth) = emit_term(&term, 10, &mut ctx).unwrap();
@@ -1915,6 +1929,7 @@ mod tests {
             glsl_ext_id: None,
             type_widths: &type_widths,
             id_to_type: &HashMap::new(),
+            composite_element_types: &HashMap::new(),
         };
         let term = parse_sexpr("(Add (Const 3) (Const 5))").unwrap();
         let (result_id, synth) = emit_term(&term, 10, &mut ctx).unwrap();
@@ -1946,6 +1961,7 @@ mod tests {
             glsl_ext_id: None,
             type_widths: &type_widths,
             id_to_type: &HashMap::new(),
+            composite_element_types: &HashMap::new(),
         };
         let term = parse_sexpr("(IntToExpr (Sym \"id5\"))").unwrap();
         let (result_id, synth) = emit_term(&term, 10, &mut ctx).unwrap();
@@ -1974,6 +1990,11 @@ mod tests {
             glsl_ext_id: None,
             type_widths,
             id_to_type: EMPTY_ID_TO_TYPE.get_or_init(HashMap::new),
+            composite_element_types: {
+                static EMPTY_COMP: std::sync::OnceLock<HashMap<Word, Word>> =
+                    std::sync::OnceLock::new();
+                EMPTY_COMP.get_or_init(HashMap::new)
+            },
         }
     }
 

rust/spirv-tools-opt/src/direct/mod.rs

@@ -559,12 +559,35 @@ pub fn optimize_module_direct(module: &Module) -> Result<Module, EgglogOptError>
         }
     }
 
+    // Build a set of (func_idx, block_idx) pairs that are loop continue blocks.
+    // Selections whose branch targets overlap these must also be skipped.
+    let mut continue_block_set: HashSet<(usize, usize)> = HashSet::new();
+    for loop_info in &loop_constructs {
+        if let Some(continue_idx) = loop_info.continue_block_idx {
+            continue_block_set.insert((loop_info.func_idx, continue_idx));
+        }
+    }
+
     // For each selection construct, convert to RVSDG EffGamma
     for (sel_idx, sel) in selection_constructs.iter().enumerate() {
         // Skip selection constructs inside loop bodies to avoid breaking continue block reachability
         if loop_block_set.contains(&(sel.func_idx, sel.header_block_idx)) {
             continue;
         }
+        // Skip selections whose branch targets or merge block overlap with loop continue blocks
+        {
+            let label_map = &func_block_labels[sel.func_idx];
+            let touches_continue = [&sel.then_label, &sel.else_label, &sel.merge_label]
+                .iter()
+                .any(|label| {
+                    label_map
+                        .get(label)
+                        .map_or(false, |&idx| continue_block_set.contains(&(sel.func_idx, idx)))
+                });
+            if touches_continue {
+                continue;
+            }
+        }
         let func = &module.functions[sel.func_idx];
 
         // Find the header block to get the condition
@@ -854,6 +877,21 @@ pub fn optimize_module_direct(module: &Module) -> Result<Module, EgglogOptError>
     let float32_type = find_spirv_type(module, Op::TypeFloat, Some(32));
     let float64_type = find_spirv_type(module, Op::TypeFloat, Some(64));
 
+    // Build composite → element type mapping for CompositeConstruct emission
+    let mut composite_element_types: HashMap<Word, Word> = HashMap::new();
+    for inst in &module.types_global_values {
+        match inst.class.opcode {
+            Op::TypeVector | Op::TypeMatrix | Op::TypeArray | Op::TypeRuntimeArray => {
+                if let (Some(composite_id), Some(rspirv::dr::Operand::IdRef(element_id))) =
+                    (inst.result_id, inst.operands.first())
+                {
+                    composite_element_types.insert(composite_id, *element_id);
+                }
+            }
+            _ => {}
+        }
+    }
+
     // Only extract from IDs that are both:
     // 1. True roots (operands of side effects) - these are the outputs we need
     // 2. Live (reachable via liveness propagation in the e-graph)
@@ -939,6 +977,7 @@ pub fn optimize_module_direct(module: &Module) -> Result<Module, EgglogOptError>
                         glsl_ext_id: ctx.glsl_ext_id(),
                         type_widths: &type_widths,
                         id_to_type: &ctx.id_to_type,
+                        composite_element_types: &composite_element_types,
                     };
                     if let Some((final_id, new_insts)) =
                         emit::emit_term(term_tree, corrected_type, &mut emit_ctx)