Legalize br_table to a BB-like format.

cranelift/codegen/src/legalizer/mod.rs

@@ -224,10 +224,29 @@ fn expand_br_table_jt(
         _ => panic!("Expected br_table: {}", func.dfg.display_inst(inst, None)),
     };
 
+    // Rewrite:
+    //
+    //     br_table $idx, default_ebb, $jt
+    //
+    // To:
+    //
+    //     $oob = ifcmp_imm $idx, len($jt)
+    //     brif uge $oob, default_ebb
+    //     jump fallthrough_ebb
+    //
+    //   fallthrough_ebb:
+    //     $base = jump_table_base.i64 $jt
+    //     $rel_addr = jump_table_entry.i64 $idx, $base, 4, $jt
+    //     $addr = iadd $base, $rel_addr
+    //     indirect_jump_table_br $addr, $jt
+
     let table_size = func.jump_tables[table].len();
     let addr_ty = isa.pointer_type();
     let entry_ty = I32;
 
+    let ebb = func.layout.pp_ebb(inst);
+    let jump_table_ebb = func.dfg.make_ebb();
+
     let mut pos = FuncCursor::new(func).at_inst(inst);
     pos.use_srcloc(inst);
 
@@ -237,6 +256,8 @@ fn expand_br_table_jt(
         .icmp_imm(IntCC::UnsignedGreaterThanOrEqual, arg, table_size as i64);
 
     pos.ins().brnz(oob, default_ebb, &[]);
+    pos.ins().jump(jump_table_ebb, &[]);
+    pos.insert_ebb(jump_table_ebb);
 
     let base_addr = pos.ins().jump_table_base(addr_ty, table);
     let entry = pos
@@ -246,9 +267,9 @@ fn expand_br_table_jt(
     let addr = pos.ins().iadd(base_addr, entry);
     pos.ins().indirect_jump_table_br(addr, table);
 
-    let ebb = pos.current_ebb().unwrap();
     pos.remove_inst();
     cfg.recompute_ebb(pos.func, ebb);
+    cfg.recompute_ebb(pos.func, jump_table_ebb);
 }
 
 /// Expand br_table to series of conditionals.
@@ -270,8 +291,15 @@ fn expand_br_table_conds(
         _ => panic!("Expected br_table: {}", func.dfg.display_inst(inst, None)),
     };
 
+    let ebb = func.layout.pp_ebb(inst);
+
     // This is a poor man's jump table using just a sequence of conditional branches.
     let table_size = func.jump_tables[table].len();
+    let mut cond_failed_ebb = std::vec::Vec::with_capacity(table_size - 1);
+    for _ in 0..table_size - 1 {
+        cond_failed_ebb.push(func.dfg.make_ebb());
+    }
+
     let mut pos = FuncCursor::new(func).at_inst(inst);
     pos.use_srcloc(inst);
 
@@ -279,14 +307,21 @@ fn expand_br_table_conds(
         let dest = pos.func.jump_tables[table].as_slice()[i];
         let t = pos.ins().icmp_imm(IntCC::Equal, arg, i as i64);
         pos.ins().brnz(t, dest, &[]);
+        // Jump to the next case.
+        if i < table_size - 1 {
+            pos.ins().jump(cond_failed_ebb[i], &[]);
+            pos.insert_ebb(cond_failed_ebb[i]);
+        }
     }
 
     // `br_table` jumps to the default destination if nothing matches
     pos.ins().jump(default_ebb, &[]);
 
-    let ebb = pos.current_ebb().unwrap();
     pos.remove_inst();
     cfg.recompute_ebb(pos.func, ebb);
+    for failed_ebb in cond_failed_ebb.into_iter() {
+        cfg.recompute_ebb(pos.func, failed_ebb);
+    }
 }
 
 /// Expand the select instruction.

cranelift/filetests/filetests/isa/x86/legalize-br-table.clif

@@ -1,6 +1,8 @@
 test compile
 
 target x86_64
+; regex: V=v\d+
+; regex: EBB=ebb\d+
 
 function u0:0(i64) system_v {
     ss0 = explicit_slot 1
@@ -10,6 +12,14 @@ ebb0(v0: i64):
     v1 = stack_addr.i64 ss0
     v2 = load.i8 v1
     br_table v2, ebb2, jt0
+; check:     $(oob=$V) = ifcmp_imm $(idx=$V), 1
+; nextln:    brif uge $oob, ebb2
+; nextln:    fallthrough $(inb=$EBB)
+; check:   $inb:
+; nextln:    $(base=$V) = jump_table_base.i64 jt0
+; nextln:    $(rel_addr=$V) = jump_table_entry.i64 $idx, $base, 4, jt0
+; nextln:    $(addr=$V) = iadd $base, $rel_addr
+; nextln:    indirect_jump_table_br $addr, jt0
 
 ebb2:
     jump ebb1

cranelift/filetests/filetests/legalizer/br_table_cond.clif

@@ -4,6 +4,8 @@ set jump_tables_enabled=false
 target x86_64
 
 ; Test that when jump_tables_enables is false, all jump tables are eliminated.
+; regex: V=v\d+
+; regex: EBB=ebb\d+
 
 function u0:0(i64 vmctx) baldrdash {
     gv0 = vmctx
@@ -20,10 +22,32 @@ ebb5:
     trapnz v2, interrupt
     v3 = iconst.i32 0
     br_table v3, ebb3, jt0
+; check:  ebb5:
+; check:    $(val0=$V) = iconst.i32 0
+; nextln:   $(cmp0=$V) = icmp_imm eq $val0, 0
+; nextln:   brnz $cmp0, ebb2
+; nextln:   jump $(fail0=$EBB)
+; check:  $fail0:
+; nextln:   $(cmp1=$V) = icmp_imm.i32 eq $val0, 1
+; nextln:   brnz $cmp1, ebb2
+; nextln:   jump $(fail1=$EBB)
+; check:  $fail1:
+; nextln:   $(cmp2=$V) = icmp_imm.i32 eq $val0, 2
+; nextln:   brnz $cmp2, ebb7
+; nextln:   jump ebb3
 
 ebb7:
     v4 = iconst.i32 0
     br_table v4, ebb3, jt1
+; check:  ebb7:
+; check:    $(val1=$V) = iconst.i32 0
+; nextln:   $(cmp3=$V) = icmp_imm eq $val1, 0
+; nextln:   brnz $cmp3, ebb8
+; nextln:   jump $(fail3=$EBB)
+; check:  $fail3:
+; nextln:   $(cmp4=$V) = icmp_imm.i32 eq $val1, 1
+; nextln:   brnz $cmp4, ebb8
+; nextln:   jump ebb3
 
 ebb8:
     jump ebb5