CFG-based loop body detection, remove Theta union, fix DCE used_ids

Three fixes for remaining CI difftest failures:

1. Loop body detection: Replace contiguous index range (header..merge)
   with BFS from loop header following branch targets, stopping at the
   merge block. SPIR-V does not guarantee contiguous block layout, so
   the old range missed non-contiguous loop body blocks that should be
   protected from RVSDG transformation. Fixes control_flow_complex.

2. Remove Theta-value union: The union of id{N} with Theta(true, id{N},
   FConst(0.0)) puts the synthetic zero init value in the same e-class
   as the actual computation, allowing the extractor to pick FConst(0.0)
   instead of the real value. The Theta term still exists in the egraph
   for LoopInvariant propagation without the union. Fixes math_ops.

3. DCE used_ids: Add operand references from types_global_values to
   used_ids set. OpConstantComposite and other globals may reference IDs
   that were aliased by the optimizer. Fixes matrix_ops undefined id.

Author: LegNeato

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

@@ -220,16 +220,43 @@ pub fn optimize_module_direct(module: &Module) -> Result<Module, EgglogOptError>
                     ) = (inst.operands.first(), inst.operands.get(1))
                     {
                         let label_map = &func_block_labels[func_idx];
-                        if let Some(&merge_idx) = label_map.get(merge_label) {
-                            let continue_idx = label_map.get(continue_label).copied();
-                            // Loop body spans from header to just before merge block
-                            let body_indices: Vec<usize> = (block_idx..merge_idx).collect();
-                            loop_constructs.push(LoopInfo {
-                                body_block_indices: body_indices,
-                                continue_block_idx: continue_idx,
-                                func_idx,
-                            });
+                        let continue_idx = label_map.get(continue_label).copied();
+                        // Collect loop body via CFG traversal from header, stopping
+                        // at the merge block. This handles non-contiguous block layouts
+                        // that a simple (header..merge) index range would miss.
+                        let mut body_indices: Vec<usize> = Vec::new();
+                        let mut visited: HashSet<usize> = HashSet::new();
+                        let mut worklist: Vec<usize> = vec![block_idx];
+                        let merge_idx = label_map.get(merge_label).copied();
+                        while let Some(idx) = worklist.pop() {
+                            if !visited.insert(idx) {
+                                continue;
+                            }
+                            // Don't include the merge block itself
+                            if Some(idx) == merge_idx {
+                                continue;
+                            }
+                            body_indices.push(idx);
+                            // Follow all branch targets from this block
+                            if let Some(blk) = func.blocks.get(idx) {
+                                for bi in &blk.instructions {
+                                    for op in &bi.operands {
+                                        if let Some(target_label) = op.id_ref_any() {
+                                            if let Some(&target_idx) =
+                                                label_map.get(&target_label)
+                                            {
+                                                worklist.push(target_idx);
+                                            }
+                                        }
+                                    }
+                                }
+                            }
                         }
+                        loop_constructs.push(LoopInfo {
+                            body_block_indices: body_indices,
+                            continue_block_idx: continue_idx,
+                            func_idx,
+                        });
                     }
                 }
             }
@@ -467,12 +494,12 @@ pub fn optimize_module_direct(module: &Module) -> Result<Module, EgglogOptError>
                             .map_err(|e| EgglogOptError::ExecutionError(e.to_string()))?;
                         theta_bound_ids.insert(id);
 
-                        // Union original ID with Theta - after saturation, the egraph will
-                        // have propagated LoopInvariant through the expression if applicable
-                        let union_cmd = format!("(union id{} theta_{})", id, id);
-                        egraph
-                            .parse_and_run_program(None, &union_cmd)
-                            .map_err(|e| EgglogOptError::ExecutionError(e.to_string()))?;
+                        // NOTE: We intentionally do NOT union id{N} with theta_{N}.
+                        // Doing so puts FConst(0.0)/Const(0)/BoolConst(0) init values
+                        // into the same e-class as the actual computation, which can
+                        // cause the extractor to pick the zero constant instead of
+                        // the real value. The Theta term exists in the egraph so
+                        // LoopInvariant rules can still detect and mark invariants.
                     }
                 }
             }
@@ -1870,6 +1897,14 @@ fn remove_dead_instructions(module: &mut Module, true_roots: &HashSet<Word>) ->
                     used_ids.insert(id);
                 }
             }
+            // Also mark operand references from types_global_values as used.
+            // OpConstantComposite and other globals may reference function-body
+            // IDs that were aliased/rewritten by the optimizer.
+            for op in &inst.operands {
+                if let Some(ref_id) = op.id_ref_any() {
+                    used_ids.insert(ref_id);
+                }
+            }
         }
 
         let mut removed_any = false;