riscv64: Clear the top 32bits in the br_table index (#5831)

We were unintentionally relying on these to be zeroed when jumping.

cranelift/codegen/src/isa/riscv64/inst.isle

@@ -213,6 +213,7 @@
     (BrTable
       (index Reg)
       (tmp1 WritableReg)
+      (tmp2 WritableReg)
       (targets VecBranchTarget))
 
     ;; atomic compare and set operation

cranelift/codegen/src/isa/riscv64/inst/emit.rs

@@ -984,11 +984,13 @@ impl MachInstEmit for Inst {
             &Inst::BrTable {
                 index,
                 tmp1,
+                tmp2,
                 ref targets,
             } => {
                 let index = allocs.next(index);
                 let tmp1 = allocs.next_writable(tmp1);
-                let tmp2 = writable_spilltmp_reg();
+                let tmp2 = allocs.next_writable(tmp2);
+                let ext_index = writable_spilltmp_reg();
 
                 // The default target is passed in as the 0th element of `targets`
                 // separate it here for clarity.
@@ -1001,15 +1003,18 @@ impl MachInstEmit for Inst {
                 // that offset. Each jump table entry is a regular auipc+jalr which we emit sequentially.
                 //
                 // Build the following sequence:
+                //
+                // extend_index:
+                //     zext.w  ext_index, index
                 // bounds_check:
                 //     li      tmp, n_labels
-                //     bltu    index, tmp, compute_target
+                //     bltu    ext_index, tmp, compute_target
                 // jump_to_default_block:
                 //     auipc   pc, 0
                 //     jalr    zero, pc, default_block
                 // compute_target:
                 //     auipc   pc, 0
-                //     slli    tmp, index, 3
+                //     slli    tmp, ext_index, 3
                 //     add     pc, pc, tmp
                 //     jalr    zero, pc, 0x10
                 // jump_table:
@@ -1017,6 +1022,19 @@ impl MachInstEmit for Inst {
                 //     auipc   pc, 0
                 //     jalr    zero, pc, block_target
 
+                // Extend the index to 64 bits.
+                //
+                // This prevents us branching on the top 32 bits of the index, which
+                // are undefined.
+                Inst::Extend {
+                    rd: ext_index,
+                    rn: index,
+                    signed: false,
+                    from_bits: 32,
+                    to_bits: 64,
+                }
+                .emit(&[], sink, emit_info, state);
+
                 // Bounds check.
                 //
                 // Check if the index passed in is larger than the number of jumptable
@@ -1030,7 +1048,7 @@ impl MachInstEmit for Inst {
                     not_taken: BranchTarget::zero(),
                     kind: IntegerCompare {
                         kind: IntCC::UnsignedLessThan,
-                        rs1: index,
+                        rs1: ext_index.to_reg(),
                         rs2: tmp2.to_reg(),
                     },
                 }
@@ -1059,7 +1077,7 @@ impl MachInstEmit for Inst {
                 Inst::AluRRImm12 {
                     alu_op: AluOPRRI::Slli,
                     rd: tmp2,
-                    rs: index,
+                    rs: ext_index.to_reg(),
                     imm12: Imm12::from_bits(3),
                 }
                 .emit(&[], sink, emit_info, state);

cranelift/codegen/src/isa/riscv64/inst/mod.rs

@@ -338,9 +338,12 @@ fn riscv64_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut Operan
     match inst {
         &Inst::Nop0 => {}
         &Inst::Nop4 => {}
-        &Inst::BrTable { index, tmp1, .. } => {
+        &Inst::BrTable {
+            index, tmp1, tmp2, ..
+        } => {
             collector.reg_use(index);
             collector.reg_early_def(tmp1);
+            collector.reg_early_def(tmp2);
         }
         &Inst::Auipc { rd, .. } => collector.reg_def(rd),
         &Inst::Lui { rd, .. } => collector.reg_def(rd),
@@ -1171,15 +1174,17 @@ impl Inst {
             &Inst::BrTable {
                 index,
                 tmp1,
+                tmp2,
                 ref targets,
             } => {
                 let targets: Vec<_> = targets.iter().map(|x| x.as_label().unwrap()).collect();
                 format!(
-                    "{} {},{}##tmp1={}",
+                    "{} {},{}##tmp1={},tmp2={}",
                     "br_table",
                     format_reg(index, allocs),
                     format_labels(&targets[..]),
                     format_reg(tmp1.to_reg(), allocs),
+                    format_reg(tmp2.to_reg(), allocs),
                 )
             }
             &Inst::Auipc { rd, imm } => {

cranelift/codegen/src/isa/riscv64/lower/isle.rs

@@ -378,6 +378,7 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Riscv64Backend> {
 
     fn lower_br_table(&mut self, index: Reg, targets: &VecMachLabel) -> Unit {
         let tmp1 = self.temp_writable_reg(I64);
+        let tmp2 = self.temp_writable_reg(I64);
         let targets: Vec<BranchTarget> = targets
             .into_iter()
             .copied()
@@ -386,6 +387,7 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Riscv64Backend> {
         self.emit(&MInst::BrTable {
             index,
             tmp1,
+            tmp2,
             targets,
         });
     }