Cranelift: fix regalloc2 integration bug wrt blockparam branch args. (#4042)

Previously, the block successor accumulation and the blockparam branch
arg setup were decoupled. The lowering backend implicitly specified
the order of successor edges via its `MachTerminator` enum on the last
instruction in the block, while the `Lower` toplevel
machine-independent driver set up blockparam branch args in the edge
order seen in CLIF.

In some cases, these orders did not match -- for example, when the
conditional branch depended on an FP condition that was implemented by
swapping taken/not-taken edges and inverting the condition code.

This PR refactors the successor handling to be centralized in `Lower`
rather than flow through the terminator `MachInst`, and adds a
successor block and its blockparam args at the same time, ensuring the
orders match.

cranelift/codegen/src/machinst/lower.rs

@@ -745,6 +745,9 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
     fn lower_clif_branches<B: LowerBackend<MInst = I>>(
         &mut self,
         backend: &B,
+        // Lowered block index:
+        bindex: BlockIndex,
+        // Original CLIF block:
         block: Block,
         branches: &SmallVec<[Inst; 2]>,
         targets: &SmallVec<[MachLabel; 2]>,
@@ -762,12 +765,15 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
         let loc = self.srcloc(branches[0]);
         self.finish_ir_inst(loc);
         // Add block param outputs for current block.
-        self.lower_branch_blockparam_args(block);
+        self.lower_branch_blockparam_args(bindex);
         Ok(())
     }
 
-    fn lower_branch_blockparam_args(&mut self, block: Block) {
-        visit_block_succs(self.f, block, |inst, _succ| {
+    fn lower_branch_blockparam_args(&mut self, block: BlockIndex) {
+        for succ_idx in 0..self.vcode.block_order().succ_indices(block).len() {
+            // Avoid immutable borrow by explicitly indexing.
+            let (inst, succ) = self.vcode.block_order().succ_indices(block)[succ_idx];
+            // Get branch args and convert to Regs.
             let branch_args = self.f.dfg.inst_variable_args(inst);
             let mut branch_arg_vregs: SmallVec<[Reg; 16]> = smallvec![];
             for &arg in branch_args {
@@ -778,8 +784,8 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
                     branch_arg_vregs.push(vreg.into());
                 }
             }
-            self.vcode.add_branch_args_for_succ(&branch_arg_vregs[..]);
-        });
+            self.vcode.add_succ(succ, &branch_arg_vregs[..]);
+        }
         self.finish_ir_inst(SourceLoc::default());
     }
 
@@ -849,7 +855,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
             if let Some(bb) = lb.orig_block() {
                 self.collect_branches_and_targets(bindex, bb, &mut branches, &mut targets);
                 if branches.len() > 0 {
-                    self.lower_clif_branches(backend, bb, &branches, &targets)?;
+                    self.lower_clif_branches(backend, bindex, bb, &branches, &targets)?;
                     self.finish_ir_inst(self.srcloc(branches[0]));
                 }
             } else {
@@ -875,7 +881,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
                             .add_block_param(vreg, self.vcode.get_vreg_type(vreg));
                     }
                 }
-                self.vcode.add_branch_args_for_succ(&branch_arg_vregs[..]);
+                self.vcode.add_succ(succ, &branch_arg_vregs[..]);
 
                 self.emit(I::gen_jump(MachLabel::from_block(succ)));
                 self.finish_ir_inst(SourceLoc::default());

cranelift/codegen/src/machinst/mod.rs

@@ -94,7 +94,7 @@ pub trait MachInst: Clone + Debug {
 
     /// Is this a terminator (branch or ret)? If so, return its type
     /// (ret/uncond/cond) and target if applicable.
-    fn is_term<'a>(&'a self) -> MachTerminator<'a>;
+    fn is_term(&self) -> MachTerminator;
 
     /// Returns true if the instruction is an epilogue placeholder.
     fn is_epilogue_placeholder(&self) -> bool;
@@ -221,18 +221,22 @@ pub trait MachInstLabelUse: Clone + Copy + Debug + Eq {
 
 /// Describes a block terminator (not call) in the vcode, when its branches
 /// have not yet been finalized (so a branch may have two targets).
+///
+/// Actual targets are not included: the single-source-of-truth for
+/// those is the VCode itself, which holds, for each block, successors
+/// and outgoing branch args per successor.
 #[derive(Clone, Debug, PartialEq, Eq)]
-pub enum MachTerminator<'a> {
+pub enum MachTerminator {
     /// Not a terminator.
     None,
     /// A return instruction.
     Ret,
     /// An unconditional branch to another block.
-    Uncond(MachLabel),
+    Uncond,
     /// A conditional branch to one of two other blocks.
-    Cond(MachLabel, MachLabel),
+    Cond,
     /// An indirect branch with known possible targets.
-    Indirect(&'a [MachLabel]),
+    Indirect,
 }
 
 /// A trait describing the ability to encode a MachInst into binary machine code.

cranelift/codegen/src/machinst/vcode.rs

@@ -371,7 +371,7 @@ impl<I: VCodeInst> VCodeBuilder<I> {
         self.vcode.block_params.push(param.into());
     }
 
-    pub fn add_branch_args_for_succ(&mut self, args: &[Reg]) {
+    fn add_branch_args_for_succ(&mut self, args: &[Reg]) {
         let start = self.vcode.branch_block_args.len();
         self.vcode
             .branch_block_args
@@ -385,34 +385,16 @@ impl<I: VCodeInst> VCodeBuilder<I> {
     /// Push an instruction for the current BB and current IR inst
     /// within the BB.
     pub fn push(&mut self, insn: I) {
-        match insn.is_term() {
-            MachTerminator::None | MachTerminator::Ret => {}
-            MachTerminator::Uncond(target) => {
-                self.vcode
-                    .block_succs_preds
-                    .push(BlockIndex::new(target.get() as usize));
-            }
-            MachTerminator::Cond(true_branch, false_branch) => {
-                self.vcode
-                    .block_succs_preds
-                    .push(BlockIndex::new(true_branch.get() as usize));
-                self.vcode
-                    .block_succs_preds
-                    .push(BlockIndex::new(false_branch.get() as usize));
-            }
-            MachTerminator::Indirect(targets) => {
-                for target in targets {
-                    self.vcode
-                        .block_succs_preds
-                        .push(BlockIndex::new(target.get() as usize));
-                }
-            }
-        }
-
         self.vcode.insts.push(insn);
         self.vcode.srclocs.push(self.cur_srcloc);
     }
 
+    /// Add a successor block with branch args.
+    pub fn add_succ(&mut self, block: BlockIndex, args: &[Reg]) {
+        self.vcode.block_succs_preds.push(block);
+        self.add_branch_args_for_succ(args);
+    }
+
     /// Set the current source location.
     pub fn set_srcloc(&mut self, srcloc: SourceLoc) {
         self.cur_srcloc = srcloc;
@@ -1191,9 +1173,7 @@ impl<I: VCodeInst> RegallocFunction for VCode<I> {
 
     fn is_branch(&self, insn: InsnIndex) -> bool {
         match self.insts[insn.index()].is_term() {
-            MachTerminator::Cond(..)
-            | MachTerminator::Uncond(..)
-            | MachTerminator::Indirect(..) => true,
+            MachTerminator::Cond | MachTerminator::Uncond | MachTerminator::Indirect => true,
             _ => false,
         }
     }