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.