fix: bug in visualization construction

Author: aPere3

zhc_ir/src/visualization/placement/annotation.rs

@@ -199,6 +199,17 @@ impl PlacementSolution {
     }
 }
 
+/// Convert raw layer positions to ranks (0..n-1) based on sorted order.
+fn positions_to_ranks(positions: SmallVec<Place>) -> SmallVec<Place> {
+    let mut indexed = positions.iter().enumerate().map(|(i, p)| (i, *p)).covec();
+    indexed.sort_by_key(|(_, pos)| *pos);
+    let mut ranks = vec![Place(0.0); positions.len()];
+    for (rank, (orig_idx, _)) in indexed.iter().enumerate() {
+        ranks[*orig_idx] = Place(rank as f64);
+    }
+    ranks.into_iter().collect()
+}
+
 fn resolve(input: OpMap<PlacementVariable>) -> OpMap<PlacementSolution> {
     input.map(|v| match v {
         PlacementVariable::NonGroup { op, .. } => PlacementSolution::NonGroup { op: op.get_val() },
@@ -208,12 +219,18 @@ fn resolve(input: OpMap<PlacementVariable>) -> OpMap<PlacementSolution> {
             inputs,
             outputs,
             ..
-        } => PlacementSolution::Group {
-            op: op.get_val(),
-            inputs: inputs.into_iter().map(|a| a.get_val()).collect(),
-            outputs: outputs.into_iter().map(|a| a.get_val()).collect(),
-            maps: (resolve(maps.0), maps.1.map(|_| ())),
-        },
+        } => {
+            // Convert raw layer positions to ranks, since other ops on the same
+            // layer can create gaps in the position sequence.
+            let input_positions = inputs.into_iter().map(|a| a.get_val()).cosvec();
+            let output_positions = outputs.into_iter().map(|a| a.get_val()).cosvec();
+            PlacementSolution::Group {
+                op: op.get_val(),
+                inputs: positions_to_ranks(input_positions),
+                outputs: positions_to_ranks(output_positions),
+                maps: (resolve(maps.0), maps.1.map(|_| ())),
+            }
+        }
     })
 }
 

zhc_ir/src/visualization/placement/solver.rs

@@ -94,12 +94,25 @@ fn place_once_top_down<'ir, 'ann>(ir: AnnIRView<'ir, 'ann, LayoutDialect, Placem
 
         // Last, for the group operations of the layer, we reconcile the rets positions with the
         // constraints coming from inside of the group. That is, we reorder the rets
-        // position following the order of the group outputs
+        // position following the order of the group outputs.
+        // Note: we use ranks (relative order among outputs) rather than raw layer positions,
+        // since other ops on the same layer can create gaps in the position sequence.
         for op in layer_ops.iter_mut() {
             if let PlacementVariable::Group { outputs, rets, .. } = op.get_annotation() {
                 let rets_original_places = rets.iter().map(|r| r.get_val()).cosvec();
-                for (ret, val) in (rets.iter(), outputs.iter()).mzip() {
-                    ret.set_val(rets_original_places[val.get_val().0 as usize]);
+                // Compute rank of each output based on sorted position order
+                let mut indexed = outputs
+                    .iter()
+                    .enumerate()
+                    .map(|(i, o)| (i, o.get_val()))
+                    .cosvec();
+                indexed.sort_by_key(|(_, pos)| *pos);
+                let mut ranks = vec![0usize; outputs.len()];
+                for (rank, (orig_idx, _)) in indexed.iter().enumerate() {
+                    ranks[*orig_idx] = rank;
+                }
+                for (ret, rank) in (rets.iter(), ranks.iter()).mzip() {
+                    ret.set_val(rets_original_places[*rank]);
                 }
             }
         }
@@ -198,11 +211,24 @@ fn place_once_bottom_up<'ir, 'ann>(ir: AnnIRView<'ir, 'ann, LayoutDialect, Place
         // Last, for the group operations of the layer, we reconcile the args positions with the
         // constraints coming from inside of the group. That is, we reorder the args
         // position following the order of the group inputs.
+        // Note: we use ranks (relative order among inputs) rather than raw layer positions,
+        // since other ops on the same layer can create gaps in the position sequence.
         for op in layer_ops.iter_mut() {
             if let PlacementVariable::Group { inputs, args, .. } = op.get_annotation() {
                 let args_original_places = args.iter().map(|r| r.get_val()).cosvec();
-                for (arg, val) in (args.iter(), inputs.iter()).mzip() {
-                    arg.set_val(args_original_places[val.get_val().0 as usize]);
+                // Compute rank of each input based on sorted position order
+                let mut indexed = inputs
+                    .iter()
+                    .enumerate()
+                    .map(|(i, o)| (i, o.get_val()))
+                    .cosvec();
+                indexed.sort_by_key(|(_, pos)| *pos);
+                let mut ranks = vec![0usize; inputs.len()];
+                for (rank, (orig_idx, _)) in indexed.iter().enumerate() {
+                    ranks[*orig_idx] = rank;
+                }
+                for (arg, rank) in (args.iter(), ranks.iter()).mzip() {
+                    arg.set_val(args_original_places[*rank]);
                 }
             }
         }

zhc_ir/src/visualization/test.rs

@@ -1076,3 +1076,80 @@ fn test_cross_group_reorder() {
     op_annotations.insert(r, root.clone());
     draw_ir_html(&ir, op_annotations, "test33.html");
 }
+
+/// Bug reproducer: GroupInput shares layer 1 with an IntInput inside the group.
+/// This triggers the indexing bug in place_once_bottom_up (line 203) where
+/// the layer position is used as an index into args_original_places.
+///
+/// Forces crossing: two external inputs A, B go into the group. An internal source C
+/// also exists on layer 1. The median heuristic will reorder based on downstream users,
+/// potentially placing C between the GroupInputs.
+#[test]
+fn test_mixed_first_layer_in_group() {
+    // Structure:
+    //   A, B (root) -> [group: gi_A, gi_B, internal_C (all layer 1) -> use them -> outputs] -> ret
+    // With careful wiring, reordering can place internal_C at position 1, pushing gi_B to position
+    // 2. args_original_places has 2 elements, so index 2 is out of bounds.
+    let mut ir: IR<TestLang> = IR::empty();
+    let (op_a, a) = ir.add_op(TestInstructionSet::IntInput { pos: 0 }, svec![]);
+    let (op_b, b) = ir.add_op(TestInstructionSet::IntInput { pos: 1 }, svec![]);
+    // Inside group_a:
+    let (op_c, c) = ir.add_op(TestInstructionSet::IntInput { pos: 2 }, svec![]); // internal source
+    // Wire so that C is used first (lower position), then B, then A - forces reorder
+    let (op_add1, add1) = ir.add_op(TestInstructionSet::Add, svec![c[0], b[0]]); // uses C and B
+    let (op_add2, add2) = ir.add_op(TestInstructionSet::Add, svec![add1[0], a[0]]); // uses result and A
+    // Back at root:
+    let (op_ret, _) = ir.add_op(TestInstructionSet::Return, svec![add2[0]]);
+
+    let root = Hierarchy::new();
+    let mut group_a = root.clone();
+    group_a.push("group_a");
+
+    let mut op_annotations = ir.empty_opmap();
+    op_annotations.insert(op_a, root.clone());
+    op_annotations.insert(op_b, root.clone());
+    op_annotations.insert(op_c, group_a.clone()); // internal source, layer 1
+    op_annotations.insert(op_add1, group_a.clone());
+    op_annotations.insert(op_add2, group_a.clone());
+    op_annotations.insert(op_ret, root.clone());
+    draw_ir_html(&ir, op_annotations, "test34.html");
+}
+
+/// Bug reproducer: GroupOutput shares last layer with a Return inside the group.
+/// This triggers the indexing bug in place_once_top_down (line 100) where
+/// the layer position is used as an index into rets_original_places.
+///
+/// Strategy: Single GroupOutput + Return on last layer. If Return ends up at
+/// position 0 and GroupOutput at position 1, indexing with 1 into a 1-element
+/// rets_original_places array causes out of bounds.
+/// We use two inputs where the one feeding the internal Return has lower position.
+#[test]
+fn test_mixed_last_layer_in_group() {
+    let mut ir: IR<TestLang> = IR::empty();
+    // inp0 at position 0, inp1 at position 1
+    let (op_inp0, inp0) = ir.add_op(TestInstructionSet::IntInput { pos: 0 }, svec![]);
+    let (op_inp1, inp1) = ir.add_op(TestInstructionSet::IntInput { pos: 1 }, svec![]);
+    // Inside group_a:
+    // inc1 from inp0 (position 0) -> feeds internal Return (will want position 0 on last layer)
+    let (op_inc1, inc1) = ir.add_op(TestInstructionSet::Inc, svec![inp0[0]]);
+    // inc2 from inp1 (position 1) -> exits group (will want position 1 on last layer)
+    let (op_inc2, inc2) = ir.add_op(TestInstructionSet::Inc, svec![inp1[0]]);
+    // Internal Return fed by inc1 - should get position 0 on last layer
+    let (op_internal_ret, _) = ir.add_op(TestInstructionSet::Return, svec![inc1[0]]);
+    // inc2 exits the group - GroupOutput should get position 1 on last layer
+    // But rets_original_places has only 1 element!
+    let (op_ret, _) = ir.add_op(TestInstructionSet::Return, svec![inc2[0]]);
+
+    let root = Hierarchy::new();
+    let mut group_a = root.clone();
+    group_a.push("group_a");
+
+    let mut op_annotations = ir.empty_opmap();
+    op_annotations.insert(op_inp0, root.clone());
+    op_annotations.insert(op_inp1, root.clone());
+    op_annotations.insert(op_inc1, group_a.clone());
+    op_annotations.insert(op_inc2, group_a.clone());
+    op_annotations.insert(op_internal_ret, group_a.clone());
+    op_annotations.insert(op_ret, root.clone());
+    draw_ir_html(&ir, op_annotations, "test35.html");
+}