Refactor stackify

crates/codegen/src/cfg_edit.rs

@@ -335,25 +335,53 @@ impl<'f> CfgEditor<'f> {
         (from, new_block)
     }
 
-    pub fn split_edge(&mut self, from: BlockId, to: BlockId) -> BlockId {
+    fn insert_edge_block(&mut self, to: BlockId) -> BlockId {
+        let mid = self.func.dfg.make_block();
+        if self.func.layout.entry_block() == Some(to) {
+            // Splitting an edge whose destination is the entry block (e.g. a multiway self-loop
+            // on the entry): inserting `mid` immediately before `to` would make `mid` the new
+            // entry, because `Layout::insert_block_before` reassigns `entry_block` when `before`
+            // has no predecessor in the layout. Keep `to` as the entry by placing `mid` after it;
+            // layout order past the entry is only a fallthrough hint and is recomputed by later
+            // analyses.
+            self.func.layout.insert_block_after(mid, to);
+        } else {
+            self.func.layout.insert_block_before(mid, to);
+        }
+        let mut cursor = InstInserter::at_location(CursorLocation::BlockTop(mid));
+        cursor.append_inst_data(self.func, Jump::new(self.func.dfg.inst_set().jump(), to));
+
+        mid
+    }
+
+    /// Split one outgoing branch edge slot through its own new jump block.
+    ///
+    /// Phi nodes are keyed by predecessor block rather than edge slot. If another parallel edge
+    /// from `from` to the same target remains, its incoming value is copied for the new block;
+    /// otherwise the predecessor label is moved from `from` to the new block.
+    pub fn split_out_edge(&mut self, from: BlockId, edge_idx: usize) -> BlockId {
         assert!(self.func.layout.is_block_inserted(from));
-        assert!(self.func.layout.is_block_inserted(to));
 
         let term = self.branch_terminator(from);
         let branch_info = self.func.dfg.branch_info(term).unwrap();
-        assert!(
-            branch_info.dests().into_iter().any(|dest| dest == to),
-            "edge {from:?} -> {to:?} does not exist"
-        );
+        let dests = branch_info.dests();
+        let to = *dests
+            .get(edge_idx)
+            .unwrap_or_else(|| panic!("outgoing edge index out of bounds: {edge_idx}"));
+        assert!(self.func.layout.is_block_inserted(to));
 
-        let mid = self.func.dfg.make_block();
-        let mut cursor = InstInserter::at_location(CursorLocation::BlockTop(to));
-        cursor.insert_block_before(self.func, mid);
-        cursor.set_location(CursorLocation::BlockTop(mid));
-        cursor.append_inst_data(self.func, Jump::new(self.func.dfg.inst_set().jump(), to));
+        let has_parallel_edge = dests
+            .iter()
+            .enumerate()
+            .any(|(other_idx, &dest)| other_idx != edge_idx && dest == to);
+        let mid = self.insert_edge_block(to);
+        self.func.dfg.rewrite_branch_edge_dest(term, edge_idx, mid);
 
-        self.func.dfg.rewrite_branch_edges_to_block(term, to, mid);
-        replace_phi_incoming_block(self.func, to, from, mid);
+        if has_parallel_edge {
+            copy_phi_incoming_block(self.func, to, from, mid);
+        } else {
+            replace_phi_incoming_block(self.func, to, from, mid);
+        }
 
         self.recompute_cfg();
         mid
@@ -1001,6 +1029,33 @@ fn replace_phi_incoming_block(
     }
 }
 
+fn copy_phi_incoming_block(
+    func: &mut Function,
+    block: BlockId,
+    old_pred: BlockId,
+    new_pred: BlockId,
+) {
+    let phi_inputs = iter_phis_in_block(func, block)
+        .map(|phi_inst| {
+            let phi = func.dfg.cast_phi(phi_inst).unwrap();
+            let mut incoming = phi
+                .args()
+                .iter()
+                .filter(|(_, pred)| *pred == old_pred)
+                .map(|(value, _)| *value);
+            let value = incoming.next().unwrap_or_else(|| {
+                panic!("phi {phi_inst:?} in {block:?} missing incoming from {old_pred:?}")
+            });
+            assert!(
+                incoming.next().is_none(),
+                "phi {phi_inst:?} in {block:?} has duplicate incoming from {old_pred:?}"
+            );
+            (phi_inst, value)
+        })
+        .collect::<Vec<_>>();
+    append_phi_inputs_for_new_pred(func, block, new_pred, &phi_inputs);
+}
+
 pub(crate) fn simplify_trivial_phis_in_block(func: &mut Function, block: BlockId) -> bool {
     let mut changed = false;
     let mut next_inst = func.layout.first_inst_of(block);
@@ -1332,6 +1387,53 @@ block2:
         });
     }
 
+    #[test]
+    fn split_out_edge_preserves_parallel_phi_inputs() {
+        let module = parse_test_module(
+            r#"
+target = "evm-ethereum-osaka"
+
+func private %f() -> i32 {
+block0:
+    br_table 0.i8 block1 (0.i8 block1) (1.i8 block2);
+
+block1:
+    v0.i32 = phi (7.i32 block0);
+    return v0;
+
+block2:
+    return 9.i32;
+}
+"#,
+        );
+        let func_ref = module.funcs()[0];
+        module.func_store.modify(func_ref, |func| {
+            let blocks: Vec<_> = func.layout.iter_block().collect();
+            let [b0, b1, b2] = blocks.as_slice() else {
+                panic!("expected three blocks");
+            };
+
+            let phi_inst = func.layout.first_inst_of(*b1).unwrap();
+            let incoming = func.dfg.cast_phi(phi_inst).unwrap().args()[0].0;
+            let mut editor = CfgEditor::new(func, CleanupMode::Strict);
+            let first_mid = editor.split_out_edge(*b0, 0);
+            let second_mid = editor.split_out_edge(*b0, 1);
+
+            let term = editor.func().layout.last_inst_of(*b0).unwrap();
+            let dests = editor.func().dfg.branch_info(term).unwrap().dests();
+            assert_eq!(dests.as_slice(), &[first_mid, second_mid, *b2]);
+
+            let phi = editor.func().dfg.cast_phi(phi_inst).unwrap();
+            assert_eq!(phi.args().len(), 2);
+            assert!(phi.args().iter().all(
+                |&(value, pred)| value == incoming && (pred == first_mid || pred == second_mid)
+            ));
+            assert_eq!(editor.cfg().preds_as_slice(*b1).len(), 2);
+            assert!(editor.cfg().preds_of(*b1).any(|&pred| pred == first_mid));
+            assert!(editor.cfg().preds_of(*b1).any(|&pred| pred == second_mid));
+        });
+    }
+
     #[test]
     fn truncate_block_from_inst_and_append_inst_with_result_rebuild_tail() {
         let mb = test_module_builder();

crates/codegen/src/critical_edge.rs

@@ -1,71 +1,47 @@
-use sonatina_ir::{BlockId, ControlFlowGraph, Function, InstId};
+use sonatina_ir::{BlockId, ControlFlowGraph, Function};
 
 use crate::cfg_edit::{CfgEditor, CleanupMode};
 
-#[derive(Debug)]
-pub struct CriticalEdgeSplitter {
-    critical_edges: Vec<CriticalEdge>,
-}
-
-impl Default for CriticalEdgeSplitter {
-    fn default() -> Self {
-        Self::new()
-    }
-}
+#[derive(Debug, Default)]
+pub struct CriticalEdgeSplitter;
 
 impl CriticalEdgeSplitter {
     pub fn new() -> Self {
-        Self {
-            critical_edges: Vec::default(),
-        }
+        Self
     }
 
     pub fn run(&mut self, func: &mut Function, cfg: &mut ControlFlowGraph) {
-        self.clear();
-
+        let mut critical_edges = Vec::<(BlockId, usize)>::new();
         for block in func.layout.iter_block() {
-            if let Some(last_inst) = func.layout.last_inst_of(block) {
-                self.add_critical_edges(last_inst, func, cfg);
+            if cfg.succ_num_of(block) < 2 {
+                continue;
             }
-        }
 
-        let edges = std::mem::take(&mut self.critical_edges);
-        let mut editor = CfgEditor::new(func, CleanupMode::Strict);
-        for edge in edges {
-            let from = editor.func().layout.inst_block(edge.inst);
-            editor.split_edge(from, edge.to);
+            // Preserve the historical destination ordering for ordinary distinct edges so bridge
+            // creation does not perturb later fallthrough placement. Parallel slots for one
+            // target are still kept distinct and ordered by their branch slot.
+            for &to in cfg.succs_of(block) {
+                if cfg.pred_num_of(to) < 2 {
+                    continue;
+                }
+                critical_edges.extend(
+                    cfg.succ_edges_of(block)
+                        .map(|&edge| cfg.edge_data(edge))
+                        .filter(|edge| edge.to == to)
+                        .map(|edge| (block, edge.branch_slot)),
+                );
+            }
         }
 
-        cfg.compute(editor.func());
-    }
-
-    pub fn clear(&mut self) {
-        self.critical_edges.clear();
-    }
-
-    fn add_critical_edges(&mut self, inst_id: InstId, func: &Function, cfg: &ControlFlowGraph) {
-        let block = func.layout.inst_block(inst_id);
-        if cfg.succ_num_of(block) < 2 {
+        if critical_edges.is_empty() {
             return;
         }
 
-        for &succ in cfg.succs_of(block) {
-            if cfg.pred_num_of(succ) > 1 {
-                self.critical_edges.push(CriticalEdge::new(inst_id, succ));
-            }
+        let mut editor = CfgEditor::new(func, CleanupMode::Strict);
+        for (from, branch_slot) in critical_edges {
+            editor.split_out_edge(from, branch_slot);
         }
-    }
-}
-
-#[derive(Debug)]
-struct CriticalEdge {
-    inst: InstId,
-    to: BlockId,
-}
-
-impl CriticalEdge {
-    fn new(inst: InstId, to: BlockId) -> Self {
-        Self { inst, to }
+        cfg.clone_from(editor.cfg());
     }
 }
 
@@ -80,6 +56,7 @@ mod tests {
         },
         isa::Isa,
     };
+    use sonatina_parser::parse_module;
 
     use super::*;
 
@@ -414,4 +391,78 @@ mod tests {
             .view(func_ref, |func| cfg_split.compute(func));
         assert_eq!(cfg, cfg_split);
     }
+
+    #[test]
+    fn critical_edge_br_table_splits_duplicate_target_slots_individually() {
+        let parsed = parse_module(
+            r#"
+target = "evm-ethereum-osaka"
+
+func private %f(v0.i1, v1.i8) -> i8 {
+block0:
+    br v0 block1 block2;
+
+block1:
+    br_table v1 block3 (0.i8 block3) (1.i8 block4);
+
+block2:
+    jump block3;
+
+block3:
+    v2.i8 = phi (7.i8 block1) (9.i8 block2);
+    return v2;
+
+block4:
+    return 11.i8;
+}
+"#,
+        )
+        .expect("module parses");
+        let func_ref = parsed.module.funcs()[0];
+        parsed.module.func_store.modify(func_ref, |func| {
+            let blocks: Vec<_> = func.layout.iter_block().collect();
+            let [_, duplicate_pred, other_pred, target, distinct_target] = blocks.as_slice() else {
+                panic!("expected five blocks");
+            };
+            let phi_inst = func.layout.first_inst_of(*target).unwrap();
+            let duplicate_incoming = func
+                .dfg
+                .cast_phi(phi_inst)
+                .unwrap()
+                .args()
+                .iter()
+                .find(|(_, pred)| pred == duplicate_pred)
+                .unwrap()
+                .0;
+
+            let mut cfg = ControlFlowGraph::default();
+            cfg.compute(func);
+            CriticalEdgeSplitter::new().run(func, &mut cfg);
+
+            let term = func.layout.last_inst_of(*duplicate_pred).unwrap();
+            let dests = func.dfg.branch_info(term).unwrap().dests();
+            assert_eq!(dests.len(), 3);
+            assert_ne!(dests[0], dests[1]);
+            assert_eq!(dests[2], *distinct_target);
+            assert!(
+                dests[..2]
+                    .iter()
+                    .all(|&mid| { cfg.succs_of(mid).eq(std::iter::once(target)) })
+            );
+
+            let phi = func.dfg.cast_phi(phi_inst).unwrap();
+            assert_eq!(phi.args().len(), 3);
+            assert!(!phi.args().iter().any(|(_, pred)| pred == duplicate_pred));
+            assert!(phi.args().iter().any(|(_, pred)| pred == other_pred));
+            assert_eq!(
+                phi.args()
+                    .iter()
+                    .filter(|(value, pred)| {
+                        *value == duplicate_incoming && *pred != *other_pred
+                    })
+                    .count(),
+                2
+            );
+        });
+    }
 }

crates/codegen/src/isa/evm/backend.rs

@@ -265,7 +265,7 @@ impl EvmBackend {
                 if let Some(alloc) = prepared.function_plan(func).map(|plan| &plan.alloc) {
                     module.func_store.view(func, |function| {
                         for v in function.dfg.value_ids() {
-                            let Some(slot) = alloc.scratch_slot_of_value[v] else {
+                            let Some(slot) = alloc.scratch_slot(v) else {
                                 continue;
                             };
                             scratch_spills.push((v, slot));