cranelift: Restrict br_table to i32 indices (#4510)

* cranelift: Restrict `br_table` to `i32` indices

In #4498 it was proposed that we should only accept `i32` indices
to `br_table`. The rationale for this is that larger types lead the
users to a false sense of flexibility (since we don't support jump
tables larger than u32's), and narrower types are not well tested
paths that would be safer if we removed them.

* cranelift: Reduce directly from i128 to i32 in Switch

cranelift/codegen/meta/src/shared/instructions.rs

@@ -157,7 +157,12 @@ fn define_control_flow(
     }
 
     {
-        let x = &Operand::new("x", iB).with_doc("index into jump table");
+        let _i32 = &TypeVar::new(
+            "i32",
+            "A 32 bit scalar integer type",
+            TypeSetBuilder::new().ints(32..32).build(),
+        );
+        let x = &Operand::new("x", _i32).with_doc("i32 index into jump table");
         let JT = &Operand::new("JT", &entities.jump_table);
 
         ig.push(

cranelift/codegen/src/isa/x64/lower.rs

@@ -3244,19 +3244,13 @@ impl LowerBackend for X64Backend {
                     let jt_size = jt_size as u32;
 
                     let ty = ctx.input_ty(branches[0], 0);
-                    let ext_spec = match ty {
-                        types::I128 => panic!("BrTable unimplemented for I128"),
-                        types::I64 => ExtSpec::ZeroExtendTo64,
-                        _ => ExtSpec::ZeroExtendTo32,
-                    };
-
                     let idx = extend_input_to_reg(
                         ctx,
                         InsnInput {
                             insn: branches[0],
                             input: 0,
                         },
-                        ext_spec,
+                        ExtSpec::ZeroExtendTo32,
                     );
 
                     // Emit the compound instruction that does:

cranelift/filetests/filetests/isa/aarch64/jumptable.clif

@@ -2,30 +2,30 @@ test compile precise-output
 set unwind_info=false
 target aarch64
 
-function %f(i64) -> i64 {
+function %f(i32) -> i32 {
   jt0 = jump_table [block1, block2, block3]
 
-block0(v0: i64):
+block0(v0: i32):
   br_table v0, block4, jt0
 
 block1:
-  v1 = iconst.i64 1
+  v1 = iconst.i32 1
   jump block5(v1)
 
 block2:
-  v2 = iconst.i64 2
+  v2 = iconst.i32 2
   jump block5(v2)
 
 block3:
-  v3 = iconst.i64 3
+  v3 = iconst.i32 3
   jump block5(v3)
 
 block4:
-  v4 = iconst.i64 4
+  v4 = iconst.i32 4
   jump block5(v4)
 
-block5(v5: i64):
+block5(v5: i32):
-  v6 = iadd.i64 v0, v5
+  v6 = iadd.i32 v0, v5
   return v6
 }
 
@@ -54,6 +54,5 @@ block5(v5: i64):
 ; block8:
 ;   b label9
 ; block9:
-;   add x0, x0, x5
+;   add w0, w0, w5
 ;   ret
-

cranelift/filetests/filetests/isa/s390x/jumptable.clif

@@ -1,59 +1,60 @@
 test compile precise-output
 target s390x
 
-function %f(i64) -> i64 {
+function %f(i32) -> i32 {
   jt0 = jump_table [block1, block2, block3]
 
-block0(v0: i64):
+block0(v0: i32):
   br_table v0, block4, jt0
 
 block1:
-  v1 = iconst.i64 1
+  v1 = iconst.i32 1
   jump block5(v1)
 
 block2:
-  v2 = iconst.i64 2
+  v2 = iconst.i32 2
   jump block5(v2)
 
 block3:
-  v3 = iconst.i64 3
+  v3 = iconst.i32 3
   jump block5(v3)
 
 block4:
-  v4 = iconst.i64 4
+  v4 = iconst.i32 4
   jump block5(v4)
 
-block5(v5: i64):
+block5(v5: i32):
-  v6 = iadd.i64 v0, v5
+  v6 = iadd.i32 v0, v5
   return v6
 }
 
 ; block0:
-;   clgfi %r2, 3
+;   llgfr %r3, %r2
+;   clgfi %r3, 3
 ;   jghe label1
-;   sllg %r5, %r2, 2
+;   sllg %r3, %r3, 2
-;   larl %r1, 14 ; agf %r1, 0(%r1, %r5) ; br %r1 ; jt_entries label3 label5 label7
+;   larl %r1, 14 ; agf %r1, 0(%r1, %r3) ; br %r1 ; jt_entries label3 label5 label7
 ; block1:
-;   lghi %r4, 4
+;   lhi %r5, 4
 ;   jg label2
 ; block2:
 ;   jg label9
 ; block3:
-;   lghi %r4, 1
+;   lhi %r5, 1
 ;   jg label4
 ; block4:
 ;   jg label9
 ; block5:
-;   lghi %r4, 2
+;   lhi %r5, 2
 ;   jg label6
 ; block6:
 ;   jg label9
 ; block7:
-;   lghi %r4, 3
+;   lhi %r5, 3
 ;   jg label8
 ; block8:
 ;   jg label9
 ; block9:
-;   agr %r2, %r4
+;   ar %r2, %r5
 ;   br %r14
 

cranelift/filetests/filetests/isa/x64/atomic-cas-bug.clif

@@ -78,11 +78,11 @@ function u0:31(i64, i32, i32, i8, i8) -> i32, i32 system_v {
 @0001                               v8 = stack_load.i8 ss5
 @0001                               stack_store v8, ss3+1
 @0001                               v9 = stack_load.i8 ss3
-@0001                               v10 = uextend.i64 v9
+@0001                               v10 = uextend.i32 v9
 @0005                               jump block37
 
                                 block37:
-@0005                               br_table.i64 v10, block36, jt0
+@0005                               br_table v10, block36, jt0
 
                                 block2:
 @0001                               v11 = stack_load.i8 ss3+1

cranelift/filetests/filetests/licm/rewrite-jump-table.clif

@@ -5,7 +5,7 @@ function %rewrite_jump_table() {
     jt0 = jump_table [block1, block2]
 
     block0:
-        v0 = iconst.i64 1
+        v0 = iconst.i32 1
         br_table v0, block1, jt0
 
     block1:

cranelift/filetests/filetests/runtests/br_table.clif

@@ -4,79 +4,6 @@ target aarch64
 target x86_64
 target s390x
 
-
-function %br_table_i8(i8) -> i8 {
-  jt0 = jump_table [block1, block2, block2, block3]
-
-block0(v0: i8):
-  br_table v0, block4, jt0
-
-block1:
-  v1 = iconst.i8 1
-  jump block5(v1)
-
-block2:
-  v2 = iconst.i8 2
-  jump block5(v2)
-
-block3:
-  v3 = iconst.i8 3
-  jump block5(v3)
-
-block4:
-  v4 = iconst.i8 4
-  jump block5(v4)
-
-block5(v5: i8):
-  v6 = iadd.i8 v0, v5
-  return v6
-}
-; run: %br_table_i8(0) == 1
-; run: %br_table_i8(1) == 3
-; run: %br_table_i8(2) == 4
-; run: %br_table_i8(3) == 6
-; run: %br_table_i8(4) == 8
-; run: %br_table_i8(5) == 9
-; run: %br_table_i8(6) == 10
-; run: %br_table_i8(-1) == 3
-
-
-function %br_table_i16(i16) -> i16 {
-  jt0 = jump_table [block1, block2, block2, block3]
-
-block0(v0: i16):
-  br_table v0, block4, jt0
-
-block1:
-  v1 = iconst.i16 1
-  jump block5(v1)
-
-block2:
-  v2 = iconst.i16 2
-  jump block5(v2)
-
-block3:
-  v3 = iconst.i16 3
-  jump block5(v3)
-
-block4:
-  v4 = iconst.i16 4
-  jump block5(v4)
-
-block5(v5: i16):
-  v6 = iadd.i16 v0, v5
-  return v6
-}
-; run: %br_table_i16(0) == 1
-; run: %br_table_i16(1) == 3
-; run: %br_table_i16(2) == 4
-; run: %br_table_i16(3) == 6
-; run: %br_table_i16(4) == 8
-; run: %br_table_i16(5) == 9
-; run: %br_table_i16(6) == 10
-; run: %br_table_i16(-1) == 3
-
-
 function %br_table_i32(i32) -> i32 {
   jt0 = jump_table [block1, block2, block2, block3]
 
@@ -111,39 +38,3 @@ block5(v5: i32):
 ; run: %br_table_i32(5) == 9
 ; run: %br_table_i32(6) == 10
 ; run: %br_table_i32(-1) == 3
-
-
-function %br_table_i64(i64) -> i64 {
-  jt0 = jump_table [block1, block2, block2, block3]
-
-block0(v0: i64):
-  br_table v0, block4, jt0
-
-block1:
-  v1 = iconst.i64 1
-  jump block5(v1)
-
-block2:
-  v2 = iconst.i64 2
-  jump block5(v2)
-
-block3:
-  v3 = iconst.i64 3
-  jump block5(v3)
-
-block4:
-  v4 = iconst.i64 4
-  jump block5(v4)
-
-block5(v5: i64):
-  v6 = iadd.i64 v0, v5
-  return v6
-}
-; run: %br_table_i64(0) == 1
-; run: %br_table_i64(1) == 3
-; run: %br_table_i64(2) == 4
-; run: %br_table_i64(3) == 6
-; run: %br_table_i64(4) == 8
-; run: %br_table_i64(5) == 9
-; run: %br_table_i64(6) == 10
-; run: %br_table_i64(-1) == 3

cranelift/filetests/filetests/verifier/jump_table.clif

@@ -1,19 +1,19 @@
 test verifier
 
-function %br_invalid_default(i64) {
+function %br_invalid_default(i32) {
     jt0 = jump_table [block1, block1]
 
-block0(v0: i64):
+block0(v0: i32):
-    br_table.i64 v0, block2, jt0 ; error: invalid block reference block2
+    br_table v0, block2, jt0 ; error: invalid block reference block2
 block1:
     return
 }
 
-function %br(i64) {
+function %br(i32) {
     jt0 = jump_table [block1, block2] ; error: invalid block reference block2
 
-block0(v0: i64):
+block0(v0: i32):
-    br_table.i64 v0, block1, jt0
+    br_table v0, block1, jt0
 block1:
     return
 }

cranelift/frontend/src/switch.rs

@@ -264,21 +264,23 @@ impl Switch {
                 }
             };
 
-            let discr = if bx.func.dfg.value_type(discr).bits() > 64 {
+            let discr = match bx.func.dfg.value_type(discr).bits() {
-                // Check for overflow of cast to u32. This is the max supported jump table entries.
+                bits if bits > 32 => {
-                let new_block = bx.create_block();
+                    // Check for overflow of cast to u32. This is the max supported jump table entries.
-                let bigger_than_u32 =
+                    let new_block = bx.create_block();
-                    bx.ins()
+                    let bigger_than_u32 =
-                        .icmp_imm(IntCC::UnsignedGreaterThan, discr, u32::MAX as i64);
+                        bx.ins()
-                bx.ins().brnz(bigger_than_u32, otherwise, &[]);
+                            .icmp_imm(IntCC::UnsignedGreaterThan, discr, u32::MAX as i64);
-                bx.ins().jump(new_block, &[]);
+                    bx.ins().brnz(bigger_than_u32, otherwise, &[]);
-                bx.seal_block(new_block);
+                    bx.ins().jump(new_block, &[]);
-                bx.switch_to_block(new_block);
+                    bx.seal_block(new_block);
+                    bx.switch_to_block(new_block);
 
-                // Cast to u64, as br_table is not implemented for i128
+                    // Cast to i32, as br_table is not implemented for i64/i128
-                bx.ins().isplit(discr).0
+                    bx.ins().ireduce(types::I32, discr)
-            } else {
+                }
-                discr
+                bits if bits < 32 => bx.ins().uextend(types::I32, discr),
+                _ => discr,
             };
 
             bx.ins().br_table(discr, otherwise, jump_table);
@@ -419,7 +421,8 @@ block0:
     jump block3
 
 block3:
-    br_table.i8 v0, block0, jt0"
+    v1 = uextend.i32 v0
+    br_table v1, block0, jt0"
         );
     }
 
@@ -470,11 +473,13 @@ block8:
     jump block12
 
 block12:
-    br_table.i8 v0, block0, jt0
+    v5 = uextend.i32 v0
+    br_table v5, block0, jt0
 
 block10:
-    v5 = iadd_imm.i8 v0, -10
+    v6 = iadd_imm.i8 v0, -10
-    br_table v5, block0, jt1"
+    v7 = uextend.i32 v6
+    br_table v7, block0, jt1"
         );
     }
 
@@ -528,7 +533,8 @@ block0:
     jump block4
 
 block4:
-    br_table.i8 v0, block0, jt0"
+    v2 = uextend.i32 v0
+    br_table v2, block0, jt0"
         );
     }
 
@@ -620,7 +626,13 @@ block0:
     jump block4
 
 block4:
-    br_table.i64 v0, block3, jt0"
+    v1 = icmp_imm.i64 ugt v0, 0xffff_ffff
+    brnz v1, block3
+    jump block5
+
+block5:
+    v2 = ireduce.i32 v0
+    br_table v2, block3, jt0"
         );
     }
 
@@ -660,7 +672,7 @@ block4:
     jump block5
 
 block5:
-    v2, v3 = isplit.i128 v0
+    v2 = ireduce.i32 v0
     br_table v2, block3, jt0"
         );
     }

cranelift/fuzzgen/src/function_generator.rs

@@ -419,8 +419,7 @@ where
 
     /// Generates a br_table into a random block
     fn generate_br_table(&mut self, builder: &mut FunctionBuilder) -> Result<()> {
-        let _type = *self.u.choose(&[I8, I16, I32, I64][..])?;
+        let var = self.get_variable_of_type(I32)?; // br_table only supports I32
-        let var = self.get_variable_of_type(_type)?;
         let val = builder.use_var(var);
 
         let valid_blocks = self.generate_valid_jumptable_target_blocks();