refactor: streamline branch resolution by introducing resolve_branch_info method and enhance phi handling

Author: Arnaud Riess

crates/herkos/src/ir/builder/core.rs

@@ -326,11 +326,6 @@ impl ControlFrame {
         }
     }
 
-    /// Returns `true` if this is a Loop frame (backward-branch target).
-    pub(super) fn is_loop(&self) -> bool {
-        matches!(self, ControlFrame::Loop { .. })
-    }
-
     /// Loop phi vars (one per Wasm local); empty slice for non-Loop frames.
     pub(super) fn loop_phi_vars(&self) -> &[UseVar] {
         match self {
@@ -505,7 +500,9 @@ impl IrBuilder {
         if let Some(block) = self.blocks.iter_mut().find(|b| b.id == self.current_block) {
             block.instructions.push(instr);
         } else {
-            // Current block doesn't exist yet, create it
+            // Current block doesn't exist yet, create it as a fallback.
+            // This handles cases where instructions are emitted before explicit block creation,
+            // which is valid in the IR builder's lazy block creation model.
             self.blocks.push(IrBlock {
                 id: self.current_block,
                 instructions: vec![instr],
@@ -732,6 +729,32 @@ impl IrBuilder {
         Ok(target)
     }
 
+    /// Resolve branch target and metadata for relative depth N.
+    ///
+    /// Returns `(target_block, is_loop, frame_idx)` needed for recording phi predecessors
+    /// and determining terminator type.
+    pub(super) fn resolve_branch_info(&self, depth: u32) -> Result<(BlockId, bool, usize)> {
+        let frame_idx = self
+            .control_stack
+            .len()
+            .checked_sub(depth as usize + 1)
+            .ok_or_else(|| {
+                anyhow::anyhow!(
+                    "Branch depth {} exceeds control stack depth {}",
+                    depth,
+                    self.control_stack.len()
+                )
+            })?;
+
+        let frame = &self.control_stack[frame_idx];
+        let (target, is_loop) = match frame {
+            ControlFrame::Loop { start_block, .. } => (*start_block, true),
+            _ => (frame.end_block(), false),
+        };
+
+        Ok((target, is_loop, frame_idx))
+    }
+
     /// Start a new block (create and switch to it).
     pub(super) fn start_block(&mut self, block_id: BlockId) {
         self.current_block = block_id;
@@ -799,10 +822,10 @@ impl IrBuilder {
         &mut self,
         join_block: BlockId,
         predecessors: &[(BlockId, Vec<UseVar>)],
-    ) {
+    ) -> Result<()> {
         let num_locals = self.local_vars.len();
         if predecessors.is_empty() || num_locals == 0 {
-            return;
+            return Ok(());
         }
 
         // Collect phis to insert: allocate dest vars before touching self.blocks.
@@ -838,12 +861,18 @@ impl IrBuilder {
         self.local_vars = new_locals;
 
         if !phi_instrs.is_empty() {
-            if let Some(block) = self.blocks.iter_mut().find(|b| b.id == join_block) {
-                let old = std::mem::take(&mut block.instructions);
-                block.instructions = phi_instrs;
-                block.instructions.extend(old);
-            }
+            let block = self
+                .blocks
+                .iter_mut()
+                .find(|b| b.id == join_block)
+                .ok_or_else(|| {
+                    anyhow::anyhow!("join block {:?} not found in blocks", join_block)
+                })?;
+            let old = std::mem::take(&mut block.instructions);
+            block.instructions = phi_instrs;
+            block.instructions.extend(old);
         }
+        Ok(())
     }
 
     /// Emit phi instructions for a loop frame into its header block.
@@ -906,11 +935,20 @@ impl IrBuilder {
             });
         }
 
-        if let Some(block) = self.blocks.iter_mut().find(|b| b.id == start_block) {
-            let old = std::mem::take(&mut block.instructions);
-            block.instructions = phi_instrs;
-            block.instructions.extend(old);
-        }
+        let block = self
+            .blocks
+            .iter_mut()
+            .find(|b| b.id == start_block)
+            .unwrap_or_else(|| {
+                panic!(
+                    "Loop header block {:?} not found when emitting loop phis. \
+                     This indicates a bug in the IR builder's block or control frame management.",
+                    start_block
+                );
+            });
+        let old = std::mem::take(&mut block.instructions);
+        block.instructions = phi_instrs;
+        block.instructions.extend(old);
     }
 }
 

crates/herkos/src/ir/builder/translate.rs

@@ -471,7 +471,7 @@ impl IrBuilder {
                             if branch_incoming.is_empty() {
                                 self.dead_code = true;
                             } else {
-                                self.insert_phis_at_join(end_block, &branch_incoming);
+                                self.insert_phis_at_join(end_block, &branch_incoming)?;
                             }
                         }
 
@@ -495,7 +495,7 @@ impl IrBuilder {
                                 preds.push((then_block, then_locals));
                             }
                             self.push_predecessor_if_live(&mut preds);
-                            preds.extend(branch_incoming.iter().cloned());
+                            preds.extend(branch_incoming.clone());
 
                             // Assign result if needed (else-branch fall-through)
                             if let Some(result_var) = rv {
@@ -516,7 +516,7 @@ impl IrBuilder {
                             if preds.is_empty() {
                                 self.dead_code = true;
                             } else {
-                                self.insert_phis_at_join(end_block, &preds);
+                                self.insert_phis_at_join(end_block, &preds)?;
                             }
                         }
 
@@ -553,7 +553,7 @@ impl IrBuilder {
                             if branch_incoming.is_empty() {
                                 self.dead_code = true;
                             } else {
-                                self.insert_phis_at_join(end_block, &branch_incoming);
+                                self.insert_phis_at_join(end_block, &branch_incoming)?;
                             }
                         }
 
@@ -588,7 +588,7 @@ impl IrBuilder {
                             if branch_incoming.is_empty() {
                                 self.dead_code = true;
                             } else {
-                                self.insert_phis_at_join(end_block, &branch_incoming);
+                                self.insert_phis_at_join(end_block, &branch_incoming)?;
                             }
                         }
                     }
@@ -980,22 +980,7 @@ impl IrBuilder {
                 // Before terminating the block we snapshot `local_vars` and record it
                 // as a phi predecessor for the target frame.  This snapshot is used by
                 // `insert_phis_at_join` (or `emit_loop_phis`) to build the phi sources.
-                let depth = *relative_depth as usize;
-                let frame_idx =
-                    self.control_stack
-                        .len()
-                        .checked_sub(depth + 1)
-                        .ok_or_else(|| {
-                            anyhow::anyhow!("br: depth {} exceeds control stack", relative_depth)
-                        })?;
-
-                let (target, is_loop) = {
-                    let frame = &self.control_stack[frame_idx];
-                    match frame {
-                        super::core::ControlFrame::Loop { start_block, .. } => (*start_block, true),
-                        _ => (frame.end_block(), false),
-                    }
-                };
+                let (target, is_loop, frame_idx) = self.resolve_branch_info(*relative_depth)?;
 
                 // Record snapshot *before* terminating so local_vars is still valid.
                 if is_loop {
@@ -1035,22 +1020,7 @@ impl IrBuilder {
                     .ok_or_else(|| anyhow::anyhow!("Stack underflow for br_if"))?
                     .var_id();
 
-                let depth = *relative_depth as usize;
-                let frame_idx =
-                    self.control_stack
-                        .len()
-                        .checked_sub(depth + 1)
-                        .ok_or_else(|| {
-                            anyhow::anyhow!("br_if: depth {} exceeds control stack", relative_depth)
-                        })?;
-
-                let (target, is_loop) = {
-                    let frame = &self.control_stack[frame_idx];
-                    match frame {
-                        super::core::ControlFrame::Loop { start_block, .. } => (*start_block, true),
-                        _ => (frame.end_block(), false),
-                    }
-                };
+                let (target, is_loop, frame_idx) = self.resolve_branch_info(*relative_depth)?;
 
                 // Record the taken branch as a phi predecessor (snapshot before termination).
                 if is_loop {
@@ -1091,18 +1061,17 @@ impl IrBuilder {
                 let target_depths: Vec<u32> = targets.targets().collect::<Result<Vec<_>, _>>()?;
                 let default_depth = targets.default();
 
-                let stack_len = self.control_stack.len();
+                // Collect all branch targets and record phi predecessors, deduplicating by frame_idx
+                // (multiple table entries may point to the same target frame).
                 let mut recorded: std::collections::HashSet<usize> =
                     std::collections::HashSet::new();
                 for depth in target_depths
                     .iter()
                     .copied()
                     .chain(std::iter::once(default_depth))
                 {
-                    let depth = depth as usize;
-                    let frame_idx = stack_len.saturating_sub(depth + 1);
+                    let (_, is_loop, frame_idx) = self.resolve_branch_info(depth)?;
                     if recorded.insert(frame_idx) {
-                        let is_loop = self.control_stack[frame_idx].is_loop();
                         if is_loop {
                             self.record_loop_back_branch(frame_idx);
                         } else {