Add support for any, fixed-reg and stack-only branch arguments being defined in its branch instruction

Author: d-sonuga

doc/FASTALLOC.md

@@ -128,6 +128,10 @@ arguments will be in their dedicated spillslots.
 4. At the beginning of a block, all branch parameters and livein 
 virtual registers will be in their dedicated spillslots.
 
+There is an exception to invariant 3: if a branch instruction defines
+the VReg used as a branch arg, then there may be no opportunity for
+the VReg to be placed in its spillslot.
+
 # Instruction Allocation
 
 To allocate a single instruction, the first step is to reset the
@@ -283,6 +287,15 @@ It's after these edits have been inserted that the parallel move
 resolver is then used to generate and insert edits to move from
 those spillslots to the spillslots of the block parameters.
 
+There is an exception to the invariant - it's possible that the
+branch argument is defined in the same branch instruction.
+If the branch argument VReg has a fixed-reg constraint, the move
+will have to be done in the successor.
+If it has an stack or anywhere constraint, it is allocated directly
+into the block param's spillslot, so there is no need to insert moves.
+The other constraints, reuse and any-reg, are not supported in this
+case.
+
 # Across Blocks
 
 When a block completes processing, some VRegs will still be live.
@@ -297,6 +310,20 @@ to be in from the first instruction.
 All block parameters are freed, just like defs, and liveins' current
 allocations in `vreg_allocs` are set to their spillslots.
 
+Any block parameter that receives a branch argument from a predecessor
+where the argument VReg was defined in the branch instruction will
+also need moves inserted at the block beginning because the predecessor
+couldn't have inserted the required moves.
+All predecessors branch arguments to the block are checked to see if any
+are defined in the same branch instruction. For all branch arguments that
+are defined in the branch instruction and have fixed-reg constraints, a 
+move will be inserted from the fixed-reg to the block param's spillslot
+at the beginning of the block. In the case of stack and anywhere constraints,
+nothing is done, because in that case, the VRegs used as the branch arguments
+will be defined directly into the block param's spillslot. Reuse and any-reg
+constraints are not supported and aren't handled.
+
+
 # Edits Order
 
 `regalloc2`'s outward interface guarantees that edits are in

src/fastalloc/mod.rs

@@ -827,6 +827,17 @@ impl<'a, F: Function> Env<'a, F> {
                 .iter()
                 .enumerate()
             {
+                if self
+                    .func
+                    .inst_operands(inst)
+                    .iter()
+                    .find(|op| op.vreg() == *vreg && op.kind() == OperandKind::Def)
+                    .is_some()
+                {
+                    // vreg is defined in this instruction, so it's dead already.
+                    // Can't move it.
+                    continue;
+                }
                 let succ_params = self.func.block_params(*succ);
                 let succ_param_vreg = succ_params[pos];
                 if self.vreg_spillslots[succ_param_vreg.vreg()].is_invalid() {
@@ -1061,6 +1072,41 @@ impl<'a, F: Function> Env<'a, F> {
                     Operand::new(op.vreg(), reused_op.constraint(), op.kind(), op.pos());
                 trace!("allocating reuse op {op} as {new_reuse_op}");
                 self.process_operand_allocation(inst, new_reuse_op, op_idx)?;
+            } else if self.func.is_branch(inst) {
+                // If the defined vreg is used as a branch arg and it has an
+                // any or stack constraint, define it into the block param spillslot
+                let mut param_spillslot = None;
+                'outer: for (succ_idx, succ) in
+                    self.func.block_succs(block).iter().cloned().enumerate()
+                {
+                    for (param_idx, branch_arg_vreg) in self
+                        .func
+                        .branch_blockparams(block, inst, succ_idx)
+                        .iter()
+                        .cloned()
+                        .enumerate()
+                    {
+                        if op.vreg() == branch_arg_vreg {
+                            if matches!(
+                                op.constraint(),
+                                OperandConstraint::Any | OperandConstraint::Stack
+                            ) {
+                                let block_param = self.func.block_params(succ)[param_idx];
+                                param_spillslot = Some(self.get_spillslot(block_param));
+                            }
+                            break 'outer;
+                        }
+                    }
+                }
+                if let Some(param_spillslot) = param_spillslot {
+                    let spillslot = self.vreg_spillslots[op.vreg().vreg()];
+                    self.vreg_spillslots[op.vreg().vreg()] = param_spillslot;
+                    let op = Operand::new(op.vreg(), OperandConstraint::Stack, op.kind(), op.pos());
+                    self.process_operand_allocation(inst, op, op_idx)?;
+                    self.vreg_spillslots[op.vreg().vreg()] = spillslot;
+                } else {
+                    self.process_operand_allocation(inst, op, op_idx)?;
+                }
             } else {
                 self.process_operand_allocation(inst, op, op_idx)?;
             }
@@ -1159,13 +1205,10 @@ impl<'a, F: Function> Env<'a, F> {
                 // be none at this point.
                 continue;
             }
-            if self.vreg_spillslots[vreg.vreg()].is_invalid() {
-                self.vreg_spillslots[vreg.vreg()] = self.stack.allocstack(vreg.class());
-            }
             // The allocation where the vreg is expected to be before
             // the first instruction.
             let prev_alloc = self.vreg_allocs[vreg.vreg()];
-            let slot = Allocation::stack(self.vreg_spillslots[vreg.vreg()]);
+            let slot = Allocation::stack(self.get_spillslot(vreg));
             self.vreg_to_live_inst_range[vreg.vreg()].2 = slot;
             self.vreg_to_live_inst_range[vreg.vreg()].0 = ProgPoint::before(first_inst);
             trace!("{} is a block param. Freeing it", vreg);
@@ -1260,6 +1303,67 @@ impl<'a, F: Function> Env<'a, F> {
                 InstPosition::Before,
             );
         }
+        // Reset this, in case a fixed reg used by a branch arg defined on the branch
+        // is used as a scratch reg in the previous loop.
+        self.edits.scratch_regs = self.edits.dedicated_scratch_regs.clone();
+
+        let get_succ_idx_of_pred = |pred, func: &F| {
+            for (idx, pred_succ) in func.block_succs(pred).iter().enumerate() {
+                if *pred_succ == block {
+                    return idx;
+                }
+            }
+            unreachable!(
+                "{:?} was not found in the successor list of its predecessor {:?}",
+                block, pred
+            );
+        };
+        trace!(
+            "Checking for predecessor branch args defined in the branch with fixed-reg constraint"
+        );
+        for (param_idx, block_param) in self.func.block_params(block).iter().cloned().enumerate() {
+            // Block param is never used. Don't bother.
+            if self.vreg_spillslots[block_param.vreg()].is_invalid() {
+                continue;
+            }
+            for pred in self.func.block_preds(block).iter().cloned() {
+                let pred_last_inst = self.func.block_insns(pred).last();
+                let curr_block_succ_idx = get_succ_idx_of_pred(pred, self.func);
+                let branch_arg_for_param =
+                    self.func
+                        .branch_blockparams(pred, pred_last_inst, curr_block_succ_idx)[param_idx];
+                // If the branch arg is defined in the branch instruction, the move will have to be done
+                // here, instead of at the end of the predecessor block.
+                let move_from = self.func.inst_operands(pred_last_inst)
+                    .iter()
+                    .find_map(|op| if op.kind() == OperandKind::Def && op.vreg() == branch_arg_for_param {
+                        match op.constraint() {
+                            OperandConstraint::FixedReg(reg) => {
+                                trace!("Found one for branch arg {branch_arg_for_param} and param {block_param} in reg {reg}");
+                                Some(Allocation::reg(reg))
+                            },
+                            // In these cases, the vreg is defined directly into the block param
+                            // spillslot.
+                            OperandConstraint::Stack | OperandConstraint::Any => None,
+                            constraint => panic!("fastalloc does not support using any-reg or reuse constraints ({}) defined on a branch instruction as a branch arg on the same instruction", constraint),
+                        }
+                    } else {
+                        None
+                    });
+                if let Some(from) = move_from {
+                    let to = Allocation::stack(self.vreg_spillslots[block_param.vreg()]);
+                    trace!("Inserting edit to move from {from} to {to}");
+                    self.add_move(
+                        first_inst,
+                        from,
+                        to,
+                        block_param.class(),
+                        InstPosition::Before,
+                    )?;
+                }
+                break;
+            }
+        }
         if trace_enabled!() {
             self.log_post_reload_at_begin_state(block);
         }