Rework br_table to use BlockCall (#5731)

Rework br_table to use BlockCall, allowing us to avoid adding new nodes during ssa construction to hold block arguments. Additionally, many places where we previously matched on InstructionData to extract branch destinations can be replaced with a use of branch_destination or branch_destination_mut.

Cargo.lock

@@ -93,12 +93,6 @@ dependencies = [
  "derive_arbitrary",
 ]
 
-[[package]]
-name = "arrayvec"
-version = "0.7.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8da52d66c7071e2e3fa2a1e5c6d088fec47b593032b254f5e980de8ea54454d6"
-
 [[package]]
 name = "async-trait"
 version = "0.1.53"
@@ -549,7 +543,6 @@ name = "cranelift-codegen"
 version = "0.94.0"
 dependencies = [
  "anyhow",
- "arrayvec",
  "bincode",
  "bumpalo",
  "capstone",

cranelift/codegen/Cargo.toml

@@ -13,7 +13,6 @@ build = "build.rs"
 edition.workspace = true
 
 [dependencies]
-arrayvec = "0.7"
 anyhow = { workspace = true, optional = true }
 bumpalo = "3"
 capstone = { workspace = true, optional = true }

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

@@ -287,6 +287,7 @@ fn verify_format(inst_name: &str, operands_in: &[Operand], format: &InstructionF
     // - its number and names of input immediate operands,
     // - whether it has a value list or not.
     let mut num_values = 0;
+    let mut num_blocks = 0;
     let mut num_immediates = 0;
 
     for operand in operands_in.iter() {
@@ -301,7 +302,9 @@ fn verify_format(inst_name: &str, operands_in: &[Operand], format: &InstructionF
         if operand.is_value() {
             num_values += 1;
         }
-        if operand.is_immediate_or_entityref() {
+        if operand.kind.is_block() {
+            num_blocks += 1;
+        } else if operand.is_immediate_or_entityref() {
             if let Some(format_field) = format.imm_fields.get(num_immediates) {
                 assert_eq!(
                     format_field.kind.rust_field_name,
@@ -325,6 +328,13 @@ fn verify_format(inst_name: &str, operands_in: &[Operand], format: &InstructionF
         inst_name, format.name
     );
 
+    assert_eq!(
+        num_blocks, format.num_block_operands,
+        "inst {} doesn't have as many block input operands as its format {} declares; you may need \
+        to use a different format.",
+        inst_name, format.name,
+    );
+
     assert_eq!(
         num_immediates,
         format.imm_fields.len(),

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

@@ -79,7 +79,8 @@ fn define_control_flow(
 
         Use ``x`` as an unsigned index into the jump table ``JT``. If a jump
         table entry is found, branch to the corresponding block. If no entry was
-        found or the index is out-of-bounds, branch to the given default block.
+        found or the index is out-of-bounds, branch to the default block of the
+        table.
 
         Note that this branch instruction can't pass arguments to the targeted
         blocks. Split critical edges as needed to work around this.

cranelift/codegen/src/inst_predicates.rs

@@ -191,16 +191,23 @@ pub(crate) fn visit_block_succs<F: FnMut(Inst, Block, bool)>(
             }
 
             ir::InstructionData::BranchTable { table, .. } => {
+                let pool = &f.dfg.value_lists;
                 let table = &f.stencil.dfg.jump_tables[*table];
 
                 // The default block is reached via a direct conditional branch,
-                // so it is not part of the table. We visit the default block first
+                // so it is not part of the table. We visit the default block
-                // explicitly, as some callers of visit_block_succs depend on that
+                // first explicitly, to mirror the traversal order of
-                // ordering.
+                // `JumpTableData::all_branches`, and transitively the order of
-                visit(inst, table.default_block(), false);
+                // `InstructionData::branch_destination`.
+                //
+                // Additionally, this case is why we are unable to replace this
+                // whole function with a loop over `branch_destination`: we need
+                // to report which branch targets come from the table vs the
+                // default.
+                visit(inst, table.default_block().block(pool), false);
 
-                for &dest in table.as_slice() {
+                for dest in table.as_slice() {
-                    visit(inst, dest, true);
+                    visit(inst, dest.block(pool), true);
                 }
             }
 

cranelift/codegen/src/ir/dfg.rs

@@ -718,7 +718,7 @@ impl DataFlowGraph {
             .iter()
             .chain(
                 self.insts[inst]
-                    .branch_destination()
+                    .branch_destination(&self.jump_tables)
                     .into_iter()
                     .flat_map(|branch| branch.args_slice(&self.value_lists).iter()),
             )
@@ -735,10 +735,10 @@ impl DataFlowGraph {
             self.inst_args_mut(inst)[i] = body(self, arg);
         }
 
-        for block_ix in 0..self.insts[inst].branch_destination().len() {
+        for block_ix in 0..self.insts[inst].branch_destination(&self.jump_tables).len() {
             // We aren't changing the size of the args list, so we won't need to write the branch
             // back to the instruction.
-            let mut block = self.insts[inst].branch_destination()[block_ix];
+            let mut block = self.insts[inst].branch_destination(&self.jump_tables)[block_ix];
             for i in 0..block.args_slice(&self.value_lists).len() {
                 let arg = block.args_slice(&self.value_lists)[i];
                 block.args_slice_mut(&mut self.value_lists)[i] = body(self, arg);
@@ -757,8 +757,8 @@ impl DataFlowGraph {
             *arg = values.next().unwrap();
         }
 
-        for block_ix in 0..self.insts[inst].branch_destination().len() {
+        for block_ix in 0..self.insts[inst].branch_destination(&self.jump_tables).len() {
-            let mut block = self.insts[inst].branch_destination()[block_ix];
+            let mut block = self.insts[inst].branch_destination(&self.jump_tables)[block_ix];
             for arg in block.args_slice_mut(&mut self.value_lists) {
                 *arg = values.next().unwrap();
             }

cranelift/codegen/src/ir/function.rs

@@ -6,9 +6,9 @@
 use crate::entity::{PrimaryMap, SecondaryMap};
 use crate::ir::{
     self, Block, DataFlowGraph, DynamicStackSlot, DynamicStackSlotData, DynamicStackSlots,
-    DynamicType, ExtFuncData, FuncRef, GlobalValue, GlobalValueData, Inst, InstructionData,
+    DynamicType, ExtFuncData, FuncRef, GlobalValue, GlobalValueData, Inst, JumpTable,
-    JumpTable, JumpTableData, Layout, Opcode, SigRef, Signature, SourceLocs, StackSlot,
+    JumpTableData, Layout, Opcode, SigRef, Signature, SourceLocs, StackSlot, StackSlotData,
-    StackSlotData, StackSlots, Table, TableData, Type,
+    StackSlots, Table, TableData, Type,
 };
 use crate::isa::CallConv;
 use crate::value_label::ValueLabelsRanges;
@@ -273,37 +273,10 @@ impl FunctionStencil {
     /// Rewrite the branch destination to `new_dest` if the destination matches `old_dest`.
     /// Does nothing if called with a non-jump or non-branch instruction.
     pub fn rewrite_branch_destination(&mut self, inst: Inst, old_dest: Block, new_dest: Block) {
-        match &mut self.dfg.insts[inst] {
+        for dest in self.dfg.insts[inst].branch_destination_mut(&mut self.dfg.jump_tables) {
-            InstructionData::Jump {
+            if dest.block(&self.dfg.value_lists) == old_dest {
-                destination: dest, ..
+                dest.set_block(new_dest, &mut self.dfg.value_lists)
-            } => {
-                if dest.block(&self.dfg.value_lists) == old_dest {
-                    dest.set_block(new_dest, &mut self.dfg.value_lists)
-                }
             }
-
-            InstructionData::Brif {
-                blocks: [block_then, block_else],
-                ..
-            } => {
-                if block_then.block(&self.dfg.value_lists) == old_dest {
-                    block_then.set_block(new_dest, &mut self.dfg.value_lists);
-                }
-
-                if block_else.block(&self.dfg.value_lists) == old_dest {
-                    block_else.set_block(new_dest, &mut self.dfg.value_lists);
-                }
-            }
-
-            InstructionData::BranchTable { table, .. } => {
-                for entry in self.dfg.jump_tables[*table].all_branches_mut() {
-                    if *entry == old_dest {
-                        *entry = new_dest;
-                    }
-                }
-            }
-
-            inst => debug_assert!(!inst.opcode().is_branch()),
         }
     }
 

cranelift/codegen/src/ir/instructions.rs

@@ -304,13 +304,13 @@ impl InstructionData {
     /// Get the destinations of this instruction, if it's a branch.
     ///
     /// `br_table` returns the empty slice.
-    pub fn branch_destination(&self) -> &[BlockCall] {
+    pub fn branch_destination<'a>(&'a self, jump_tables: &'a ir::JumpTables) -> &[BlockCall] {
         match self {
             Self::Jump {
                 ref destination, ..
             } => std::slice::from_ref(destination),
-            Self::Brif { blocks, .. } => blocks,
+            Self::Brif { blocks, .. } => blocks.as_slice(),
-            Self::BranchTable { .. } => &[],
+            Self::BranchTable { table, .. } => jump_tables.get(*table).unwrap().all_branches(),
             _ => {
                 debug_assert!(!self.opcode().is_branch());
                 &[]
@@ -321,14 +321,19 @@ impl InstructionData {
     /// Get a mutable slice of the destinations of this instruction, if it's a branch.
     ///
     /// `br_table` returns the empty slice.
-    pub fn branch_destination_mut(&mut self) -> &mut [BlockCall] {
+    pub fn branch_destination_mut<'a>(
+        &'a mut self,
+        jump_tables: &'a mut ir::JumpTables,
+    ) -> &mut [BlockCall] {
         match self {
             Self::Jump {
                 ref mut destination,
                 ..
             } => std::slice::from_mut(destination),
-            Self::Brif { blocks, .. } => blocks,
+            Self::Brif { blocks, .. } => blocks.as_mut_slice(),
-            Self::BranchTable { .. } => &mut [],
+            Self::BranchTable { table, .. } => {
+                jump_tables.get_mut(*table).unwrap().all_branches_mut()
+            }
             _ => {
                 debug_assert!(!self.opcode().is_branch());
                 &mut []

cranelift/codegen/src/ir/jumptable.rs

@@ -3,7 +3,8 @@
 //! Jump tables are declared in the preamble and assigned an `ir::entities::JumpTable` reference.
 //! The actual table of destinations is stored in a `JumpTableData` struct defined in this module.
 
-use crate::ir::entities::Block;
+use crate::ir::instructions::ValueListPool;
+use crate::ir::BlockCall;
 use alloc::vec::Vec;
 use core::fmt::{self, Display, Formatter};
 use core::slice::{Iter, IterMut};
@@ -23,57 +24,57 @@ use serde::{Deserialize, Serialize};
 #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
 pub struct JumpTableData {
     // Table entries.
-    table: Vec<Block>,
+    table: Vec<BlockCall>,
 }
 
 impl JumpTableData {
     /// Create a new jump table with the provided blocks
-    pub fn new(def: Block, table: &[Block]) -> Self {
+    pub fn new(def: BlockCall, table: &[BlockCall]) -> Self {
         Self {
             table: std::iter::once(def).chain(table.iter().copied()).collect(),
         }
     }
 
     /// Fetch the default block for this jump table.
-    pub fn default_block(&self) -> Block {
+    pub fn default_block(&self) -> BlockCall {
         *self.table.first().unwrap()
     }
 
     /// Mutable access to the default block of this jump table.
-    pub fn default_block_mut(&mut self) -> &mut Block {
+    pub fn default_block_mut(&mut self) -> &mut BlockCall {
         self.table.first_mut().unwrap()
     }
 
     /// The jump table and default block as a single slice. The default block will always be first.
-    pub fn all_branches(&self) -> &[Block] {
+    pub fn all_branches(&self) -> &[BlockCall] {
         self.table.as_slice()
     }
 
     /// The jump table and default block as a single mutable slice. The default block will always
     /// be first.
-    pub fn all_branches_mut(&mut self) -> &mut [Block] {
+    pub fn all_branches_mut(&mut self) -> &mut [BlockCall] {
         self.table.as_mut_slice()
     }
 
     /// Access the jump table as a slice. This excludes the default block.
-    pub fn as_slice(&self) -> &[Block] {
+    pub fn as_slice(&self) -> &[BlockCall] {
         &self.table.as_slice()[1..]
     }
 
     /// Access the jump table as a mutable slice. This excludes the default block.
-    pub fn as_mut_slice(&mut self) -> &mut [Block] {
+    pub fn as_mut_slice(&mut self) -> &mut [BlockCall] {
         &mut self.table.as_mut_slice()[1..]
     }
 
     /// Returns an iterator to the jump table, excluding the default block.
     #[deprecated(since = "7.0.0", note = "please use `.as_slice()` instead")]
-    pub fn iter(&self) -> Iter<Block> {
+    pub fn iter(&self) -> Iter<BlockCall> {
         self.as_slice().iter()
     }
 
     /// Returns an iterator that allows modifying each value, excluding the default block.
     #[deprecated(since = "7.0.0", note = "please use `.as_mut_slice()` instead")]
-    pub fn iter_mut(&mut self) -> IterMut<Block> {
+    pub fn iter_mut(&mut self) -> IterMut<BlockCall> {
         self.as_mut_slice().iter_mut()
     }
 
@@ -81,15 +82,26 @@ impl JumpTableData {
     pub fn clear(&mut self) {
         self.table.drain(1..);
     }
+
+    /// Return a value that can display the contents of this jump table.
+    pub fn display<'a>(&'a self, pool: &'a ValueListPool) -> DisplayJumpTable<'a> {
+        DisplayJumpTable { jt: self, pool }
+    }
 }
 
-impl Display for JumpTableData {
+/// A wrapper for the context required to display a [JumpTableData].
-    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
+pub struct DisplayJumpTable<'a> {
-        write!(fmt, "{}, [", self.default_block())?;
+    jt: &'a JumpTableData,
-        if let Some((first, rest)) = self.as_slice().split_first() {
+    pool: &'a ValueListPool,
-            write!(fmt, "{}", first)?;
+}
+
+impl<'a> Display for DisplayJumpTable<'a> {
+    fn fmt(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
+        write!(fmt, "{}, [", self.jt.default_block().display(self.pool))?;
+        if let Some((first, rest)) = self.jt.as_slice().split_first() {
+            write!(fmt, "{}", first.display(self.pool))?;
             for block in rest {
-                write!(fmt, ", {}", block)?;
+                write!(fmt, ", {}", block.display(self.pool))?;
             }
         }
         write!(fmt, "]")
@@ -100,12 +112,14 @@ impl Display for JumpTableData {
 mod tests {
     use super::JumpTableData;
     use crate::entity::EntityRef;
-    use crate::ir::Block;
+    use crate::ir::instructions::ValueListPool;
-    use alloc::string::ToString;
+    use crate::ir::{Block, BlockCall, Value};
+    use std::string::ToString;
 
     #[test]
     fn empty() {
-        let def = Block::new(0);
+        let mut pool = ValueListPool::default();
+        let def = BlockCall::new(Block::new(0), &[], &mut pool);
 
         let jt = JumpTableData::new(def, &[]);
 
@@ -114,7 +128,7 @@ mod tests {
         assert_eq!(jt.as_slice().get(0), None);
         assert_eq!(jt.as_slice().get(10), None);
 
-        assert_eq!(jt.to_string(), "block0, []");
+        assert_eq!(jt.display(&pool).to_string(), "block0, []");
 
         assert_eq!(jt.all_branches(), [def]);
         assert_eq!(jt.as_slice(), []);
@@ -122,16 +136,33 @@ mod tests {
 
     #[test]
     fn insert() {
-        let def = Block::new(0);
+        let mut pool = ValueListPool::default();
+
+        let v0 = Value::new(0);
+        let v1 = Value::new(1);
+
+        let e0 = Block::new(0);
         let e1 = Block::new(1);
         let e2 = Block::new(2);
 
-        let jt = JumpTableData::new(def, &[e1, e2, e1]);
+        let def = BlockCall::new(e0, &[], &mut pool);
+        let b1 = BlockCall::new(e1, &[v0], &mut pool);
+        let b2 = BlockCall::new(e2, &[], &mut pool);
+        let b3 = BlockCall::new(e1, &[v1], &mut pool);
+
+        let jt = JumpTableData::new(def, &[b1, b2, b3]);
 
         assert_eq!(jt.default_block(), def);
-        assert_eq!(jt.to_string(), "block0, [block1, block2, block1]");
+        assert_eq!(
+            jt.display(&pool).to_string(),
+            "block0, [block1(v0), block2, block1(v1)]"
+        );
 
-        assert_eq!(jt.all_branches(), [def, e1, e2, e1]);
+        assert_eq!(jt.all_branches(), [def, b1, b2, b3]);
-        assert_eq!(jt.as_slice(), [e1, e2, e1]);
+        assert_eq!(jt.as_slice(), [b1, b2, b3]);
+
+        assert_eq!(jt.as_slice()[0].args_slice(&pool), [v0]);
+        assert_eq!(jt.as_slice()[1].args_slice(&pool), []);
+        assert_eq!(jt.as_slice()[2].args_slice(&pool), [v1]);
     }
 }

cranelift/codegen/src/isa/aarch64/mod.rs

@@ -377,9 +377,14 @@ mod test {
         let mut pos = FuncCursor::new(&mut func);
 
         pos.insert_block(bb0);
-        let jt = pos
+        let jt_data = JumpTableData::new(
-            .func
+            pos.func.dfg.block_call(bb3, &[]),
-            .create_jump_table(JumpTableData::new(bb3, &[bb1, bb2]));
+            &[
+                pos.func.dfg.block_call(bb1, &[]),
+                pos.func.dfg.block_call(bb2, &[]),
+            ],
+        );
+        let jt = pos.func.create_jump_table(jt_data);
         pos.ins().br_table(arg0, jt);
 
         pos.insert_block(bb1);

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

@@ -418,9 +418,14 @@ mod test {
         let mut pos = FuncCursor::new(&mut func);
 
         pos.insert_block(bb0);
-        let jt = pos
+        let jt_data = JumpTableData::new(
-            .func
+            pos.func.dfg.block_call(bb3, &[]),
-            .create_jump_table(JumpTableData::new(bb3, &[bb1, bb2]));
+            &[
+                pos.func.dfg.block_call(bb1, &[]),
+                pos.func.dfg.block_call(bb2, &[]),
+            ],
+        );
+        let jt = pos.func.create_jump_table(jt_data);
         pos.ins().br_table(arg0, jt);
 
         pos.insert_block(bb1);

cranelift/codegen/src/machinst/lower.rs

@@ -942,29 +942,13 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
             // Avoid immutable borrow by explicitly indexing.
             let (inst, succ) = self.vcode.block_order().succ_indices(block)[succ_idx];
 
-            // Get branch args and convert to Regs.
+            // The use of `succ_idx` to index `branch_destination` is valid on the assumption that
-            let branch_args = match &self.f.dfg.insts[inst] {
+            // the traversal order defined in `visit_block_succs` mirrors the order returned by
-                InstructionData::Jump {
+            // `branch_destination`. If that assumption is violated, the branch targets returned
-                    destination: block, ..
+            // here will not match the clif.
-                } => block.args_slice(&self.f.dfg.value_lists),
+            let branches = self.f.dfg.insts[inst].branch_destination(&self.f.dfg.jump_tables);
-                InstructionData::Brif {
+            let branch_args = branches[succ_idx].args_slice(&self.f.dfg.value_lists);
-                    blocks: [then_block, else_block],
+
-                    ..
-                } => {
-                    // NOTE: `succ_idx == 0` implying that we're traversing the `then_block` is
-                    // enforced by the traversal order defined in `visit_block_succs`. Eventually
-                    // we should traverse the `branch_destination` slice there, which would
-                    // simplify computing the branch args significantly.
-                    if succ_idx == 0 {
-                        then_block.args_slice(&self.f.dfg.value_lists)
-                    } else {
-                        assert!(succ_idx == 1);
-                        else_block.args_slice(&self.f.dfg.value_lists)
-                    }
-                }
-                InstructionData::BranchTable { .. } => &[],
-                _ => unreachable!(),
-            };
             let mut branch_arg_vregs: SmallVec<[Reg; 16]> = smallvec![];
             for &arg in branch_args {
                 let arg = self.f.dfg.resolve_aliases(arg);

cranelift/codegen/src/remove_constant_phis.rs

@@ -7,7 +7,6 @@ use crate::ir;
 use crate::ir::Function;
 use crate::ir::{Block, BlockCall, Inst, Value};
 use crate::timing;
-use arrayvec::ArrayVec;
 use bumpalo::Bump;
 use cranelift_entity::SecondaryMap;
 use smallvec::SmallVec;
@@ -171,9 +170,10 @@ struct BlockSummary<'a> {
     /// We don't bother to include transfers that pass zero parameters
     /// since that makes more work for the solver for no purpose.
     ///
-    /// Note that, because blocks used with `br_table`s cannot have block
+    /// We optimize for the case where a branch instruction has up to two
-    /// arguments, there are at most two outgoing edges from these blocks.
+    /// outgoing edges, as unconditional jumps and conditional branches are
-    dests: ArrayVec<OutEdge<'a>, 2>,
+    /// more prominent than br_table.
+    dests: SmallVec<[OutEdge<'a>; 2]>,
 }
 
 impl<'a> BlockSummary<'a> {
@@ -239,7 +239,11 @@ pub fn do_remove_constant_phis(func: &mut Function, domtree: &mut DominatorTree)
         let mut summary = BlockSummary::new(&bump, formals);
 
         for inst in func.layout.block_insts(b) {
-            for (ix, dest) in func.dfg.insts[inst].branch_destination().iter().enumerate() {
+            for (ix, dest) in func.dfg.insts[inst]
+                .branch_destination(&func.dfg.jump_tables)
+                .iter()
+                .enumerate()
+            {
                 if let Some(edge) = OutEdge::new(&bump, &func.dfg, inst, ix, *dest) {
                     summary.dests.push(edge);
                 }
@@ -382,8 +386,8 @@ pub fn do_remove_constant_phis(func: &mut Function, domtree: &mut DominatorTree)
             }
 
             let dfg = &mut func.dfg;
-            let block =
+            let dests = dfg.insts[edge.inst].branch_destination_mut(&mut dfg.jump_tables);
-                &mut dfg.insts[edge.inst].branch_destination_mut()[edge.branch_index as usize];
+            let block = &mut dests[edge.branch_index as usize];
 
             old_actuals.extend(block.args_slice(&dfg.value_lists));
 

cranelift/codegen/src/verifier/mod.rs

@@ -849,8 +849,9 @@ impl<'a> Verifier<'a> {
                 format!("invalid jump table reference {}", j),
             ))
         } else {
-            for &block in self.func.stencil.dfg.jump_tables[j].all_branches() {
+            let pool = &self.func.stencil.dfg.value_lists;
-                self.verify_block(inst, block, errors)?;
+            for block in self.func.stencil.dfg.jump_tables[j].all_branches() {
+                self.verify_block(inst, block.block(pool), errors)?;
             }
             Ok(())
         }
@@ -1285,53 +1286,19 @@ impl<'a> Verifier<'a> {
         errors: &mut VerifierErrors,
     ) -> VerifierStepResult<()> {
         match &self.func.dfg.insts[inst] {
-            ir::InstructionData::Jump {
+            ir::InstructionData::Jump { destination, .. } => {
-                destination: block, ..
+                self.typecheck_block_call(inst, destination, errors)?;
-            } => {
-                let iter = self
-                    .func
-                    .dfg
-                    .block_params(block.block(&self.func.dfg.value_lists))
-                    .iter()
-                    .map(|&v| self.func.dfg.value_type(v));
-                let args = block.args_slice(&self.func.dfg.value_lists);
-                self.typecheck_variable_args_iterator(inst, iter, args, errors)?;
             }
             ir::InstructionData::Brif {
                 blocks: [block_then, block_else],
                 ..
             } => {
-                let iter = self
+                self.typecheck_block_call(inst, block_then, errors)?;
-                    .func
+                self.typecheck_block_call(inst, block_else, errors)?;
-                    .dfg
-                    .block_params(block_then.block(&self.func.dfg.value_lists))
-                    .iter()
-                    .map(|&v| self.func.dfg.value_type(v));
-                let args_then = block_then.args_slice(&self.func.dfg.value_lists);
-                self.typecheck_variable_args_iterator(inst, iter, args_then, errors)?;
-
-                let iter = self
-                    .func
-                    .dfg
-                    .block_params(block_else.block(&self.func.dfg.value_lists))
-                    .iter()
-                    .map(|&v| self.func.dfg.value_type(v));
-                let args_else = block_else.args_slice(&self.func.dfg.value_lists);
-                self.typecheck_variable_args_iterator(inst, iter, args_else, errors)?;
             }
             ir::InstructionData::BranchTable { table, .. } => {
                 for block in self.func.stencil.dfg.jump_tables[*table].all_branches() {
-                    let arg_count = self.func.dfg.num_block_params(*block);
+                    self.typecheck_block_call(inst, block, errors)?;
-                    if arg_count != 0 {
-                        return errors.nonfatal((
-                            inst,
-                            self.context(inst),
-                            format!(
-                                "takes no arguments, but had target {} with {} arguments",
-                                block, arg_count,
-                            ),
-                        ));
-                    }
                 }
             }
             inst => debug_assert!(!inst.opcode().is_branch()),
@@ -1358,6 +1325,23 @@ impl<'a> Verifier<'a> {
         Ok(())
     }
 
+    fn typecheck_block_call(
+        &self,
+        inst: Inst,
+        block: &ir::BlockCall,
+        errors: &mut VerifierErrors,
+    ) -> VerifierStepResult<()> {
+        let pool = &self.func.dfg.value_lists;
+        let iter = self
+            .func
+            .dfg
+            .block_params(block.block(pool))
+            .iter()
+            .map(|&v| self.func.dfg.value_type(v));
+        let args = block.args_slice(pool);
+        self.typecheck_variable_args_iterator(inst, iter, args, errors)
+    }
+
     fn typecheck_variable_args_iterator<I: Iterator<Item = Type>>(
         &self,
         inst: Inst,

cranelift/codegen/src/write.rs

@@ -421,7 +421,9 @@ pub fn write_operands(w: &mut dyn Write, dfg: &DataFlowGraph, inst: Inst) -> fmt
             write!(w, " {}, {}", arg, block_then.display(pool))?;
             write!(w, ", {}", block_else.display(pool))
         }
-        BranchTable { arg, table, .. } => write!(w, " {}, {}", arg, jump_tables[table]),
+        BranchTable { arg, table, .. } => {
+            write!(w, " {}, {}", arg, jump_tables[table].display(pool))
+        }
         Call {
             func_ref, ref args, ..
         } => write!(w, " {}({})", func_ref, DisplayValues(args.as_slice(pool))),

cranelift/filetests/filetests/isa/x64/branches.clif

@@ -727,3 +727,185 @@ block2:
 ;   popq %rbp
 ;   retq
 
+function %br_table_i32(i32) -> i32 {
+block0(v0: i32):
+  br_table v0, block4, [block1, block2, block2, block3]
+
+block1:
+  v1 = iconst.i32 1
+  jump block5(v1)
+
+block2:
+  v2 = iconst.i32 2
+  jump block5(v2)
+
+block3:
+  v3 = iconst.i32 3
+  jump block5(v3)
+
+block4:
+  v4 = iconst.i32 4
+  jump block5(v4)
+
+block5(v5: i32):
+  v6 = iadd.i32 v0, v5
+  return v6
+}
+
+; VCode:
+;   pushq   %rbp
+;   movq    %rsp, %rbp
+; block0:
+;   cmpl    $4, %edi
+;   br_table %rdi, %rcx, %rdx
+; block1:
+;   movl    $4, %edx
+;   jmp     label2
+; block2:
+;   jmp     label10
+; block3:
+;   movl    $1, %edx
+;   jmp     label4
+; block4:
+;   jmp     label10
+; block5:
+;   jmp     label7
+; block6:
+;   jmp     label7
+; block7:
+;   movl    $2, %edx
+;   jmp     label10
+; block8:
+;   movl    $3, %edx
+;   jmp     label9
+; block9:
+;   jmp     label10
+; block10:
+;   movq    %rdi, %rax
+;   addl    %eax, %edx, %eax
+;   movq    %rbp, %rsp
+;   popq    %rbp
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   pushq %rbp
+;   movq %rsp, %rbp
+; block1: ; offset 0x4
+;   cmpl $4, %edi
+;   jae 0x39
+;   movl %edi, %edx
+;   movl $0, %ecx
+;   cmovaeq %rcx, %rdx
+;   leaq 0xa(%rip), %rcx
+;   movslq (%rcx, %rdx, 4), %rdx
+;   addq %rdx, %rcx
+;   jmpq *%rcx
+;   sbbb (%rax), %al
+;   addb %al, (%rax)
+;   andb $0, %al
+;   addb %al, (%rax)
+;   andb $0, %al
+;   addb %al, (%rax)
+;   addb %al, %cs:(%rax)
+; block2: ; offset 0x39
+;   movl $4, %edx
+; block3: ; offset 0x3e
+;   jmp 0x5c
+; block4: ; offset 0x43
+;   movl $1, %edx
+; block5: ; offset 0x48
+;   jmp 0x5c
+; block6: ; offset 0x4d
+;   movl $2, %edx
+;   jmp 0x5c
+; block7: ; offset 0x57
+;   movl $3, %edx
+; block8: ; offset 0x5c
+;   movq %rdi, %rax
+;   addl %edx, %eax
+;   movq %rbp, %rsp
+;   popq %rbp
+;   retq
+
+function %br_table_i32_inline(i32) -> i32 {
+block0(v0: i32):
+  v1 = iconst.i32 1
+  v2 = iconst.i32 2
+  v3 = iconst.i32 3
+  v4 = iconst.i32 4
+  br_table v0, block1(v4), [block1(v1), block1(v2), block1(v2), block1(v3)]
+
+block1(v5: i32):
+  return v5
+}
+
+; VCode:
+;   pushq   %rbp
+;   movq    %rsp, %rbp
+; block0:
+;   movl    $1, %edx
+;   movl    $2, %r8d
+;   movl    $3, %r9d
+;   movl    $4, %eax
+;   cmpl    $4, %edi
+;   br_table %rdi, %r11, %r10
+; block1:
+;   jmp     label6
+; block2:
+;   movq    %rdx, %rax
+;   jmp     label6
+; block3:
+;   movq    %r8, %rax
+;   jmp     label6
+; block4:
+;   movq    %r8, %rax
+;   jmp     label6
+; block5:
+;   movq    %r9, %rax
+;   jmp     label6
+; block6:
+;   movq    %rbp, %rsp
+;   popq    %rbp
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   pushq %rbp
+;   movq %rsp, %rbp
+; block1: ; offset 0x4
+;   movl $1, %edx
+;   movl $2, %r8d
+;   movl $3, %r9d
+;   movl $4, %eax
+;   cmpl $4, %edi
+;   jae 0x72
+;   movl %edi, %r10d
+;   movl $0, %r11d
+;   cmovaeq %r11, %r10
+;   leaq 0xb(%rip), %r11
+;   movslq (%r11, %r10, 4), %r10
+;   addq %r10, %r11
+;   jmpq *%r11
+;   adcl $0x1d000000, %eax
+;   addb %al, (%rax)
+;   addb %ah, 0x2d000000(%rip)
+;   addb %al, (%rax)
+; block2: ; offset 0x52
+;   jmp 0x72
+; block3: ; offset 0x57
+;   movq %rdx, %rax
+;   jmp 0x72
+; block4: ; offset 0x5f
+;   movq %r8, %rax
+;   jmp 0x72
+; block5: ; offset 0x67
+;   movq %r8, %rax
+;   jmp 0x72
+; block6: ; offset 0x6f
+;   movq %r9, %rax
+; block7: ; offset 0x72
+;   movq %rbp, %rsp
+;   popq %rbp
+;   retq
+

cranelift/filetests/filetests/runtests/br_table.clif

@@ -56,3 +56,47 @@ block2:
 }
 ; run: %br_table_cold_block(0) == 0
 ; run: %br_table_cold_block(1) == 0
+
+function %br_table_i32_inline(i32) -> i32 {
+block0(v0: i32):
+  v1 = iconst.i32 1
+  v2 = iconst.i32 2
+  v3 = iconst.i32 3
+  v4 = iconst.i32 4
+  br_table v0, block1(v4), [block1(v1), block1(v2), block1(v2), block1(v3)]
+
+block1(v5: i32):
+  return v5
+}
+
+; run: %br_table_i32_inline(0) == 1
+; run: %br_table_i32_inline(1) == 2
+; run: %br_table_i32_inline(2) == 2
+; run: %br_table_i32_inline(3) == 3
+; run: %br_table_i32_inline(4) == 4
+; run: %br_table_i32_inline(297) == 4
+; run: %br_table_i32_inline(65535) == 4
+
+function %br_table_i32_inline_varied(i32) -> i32 {
+block0(v0: i32):
+  v1 = iconst.i32 1
+  v2 = iconst.i32 2
+  v3 = iconst.i32 3
+  v4 = iconst.i32 4
+  br_table v0, block1(v4), [block1(v1), block2(v2, v4), block2(v4, v3), block1(v3)]
+
+block2(v6: i32, v7: i32):
+  v8 = iadd v6, v7
+  jump block1(v8)
+
+block1(v5: i32):
+  return v5
+}
+
+; run: %br_table_i32_inline_varied(0) == 1
+; run: %br_table_i32_inline_varied(1) == 6
+; run: %br_table_i32_inline_varied(2) == 7
+; run: %br_table_i32_inline_varied(3) == 3
+; run: %br_table_i32_inline_varied(4) == 4
+; run: %br_table_i32_inline_varied(297) == 4
+; run: %br_table_i32_inline_varied(65535) == 4

cranelift/filetests/filetests/verifier/type_check.clif

@@ -63,12 +63,10 @@ function %fn_call_incorrect_arg_type(i64) {
         return
 }
 
-; TODO: Should we instead just verify that jump tables contain no blocks that take arguments? This
-; error doesn't occur if no instruction uses the jump table.
 function %jump_table_args() {
     block0:
         v0 = iconst.i32 0
-        br_table v0, block2, [block1] ; error: takes no arguments, but had target block1 with 1 arguments
+        br_table v0, block2, [block1] ; error: mismatched argument count
 
     block1(v5: i32):
         return

cranelift/frontend/src/frontend.rs

@@ -130,6 +130,8 @@ impl<'short, 'long> InstBuilderBase<'short> for FuncInstBuilder<'short, 'long> {
             }
 
             ir::InstructionData::BranchTable { table, .. } => {
+                let pool = &self.builder.func.dfg.value_lists;
+
                 // Unlike all other jumps/branches, jump tables are
                 // capable of having the same successor appear
                 // multiple times, so we must deduplicate.
@@ -144,7 +146,8 @@ impl<'short, 'long> InstBuilderBase<'short> for FuncInstBuilder<'short, 'long> {
                     .expect("you are referencing an undeclared jump table")
                     .all_branches()
                 {
-                    if !unique.insert(*dest_block) {
+                    let block = dest_block.block(pool);
+                    if !unique.insert(block) {
                         continue;
                     }
 
@@ -153,7 +156,7 @@ impl<'short, 'long> InstBuilderBase<'short> for FuncInstBuilder<'short, 'long> {
                     self.builder
                         .func_ctx
                         .ssa
-                        .declare_block_predecessor(*dest_block, inst);
+                        .declare_block_predecessor(block, inst);
                 }
             }
 
@@ -691,7 +694,7 @@ impl<'a> FunctionBuilder<'a> {
     /// other jump instructions.
     pub fn change_jump_destination(&mut self, inst: Inst, old_block: Block, new_block: Block) {
         let dfg = &mut self.func.dfg;
-        for block in dfg.insts[inst].branch_destination_mut() {
+        for block in dfg.insts[inst].branch_destination_mut(&mut dfg.jump_tables) {
             if block.block(&dfg.value_lists) == old_block {
                 self.func_ctx.ssa.remove_block_predecessor(old_block, inst);
                 block.set_block(new_block, &mut dfg.value_lists);

cranelift/frontend/src/ssa.rs

@@ -15,7 +15,7 @@ use cranelift_codegen::cursor::{Cursor, FuncCursor};
 use cranelift_codegen::entity::{EntityList, EntitySet, ListPool, SecondaryMap};
 use cranelift_codegen::ir::immediates::{Ieee32, Ieee64};
 use cranelift_codegen::ir::types::{F32, F64, I128, I64};
-use cranelift_codegen::ir::{Block, Function, Inst, InstBuilder, InstructionData, Type, Value};
+use cranelift_codegen::ir::{Block, Function, Inst, InstBuilder, Type, Value};
 use cranelift_codegen::packed_option::PackedOption;
 
 /// Structure containing the data relevant the construction of SSA for a given function.
@@ -488,7 +488,6 @@ impl SSABuilder {
         &mut self,
         func: &mut Function,
         sentinel: Value,
-        var: Variable,
         dest_block: Block,
     ) -> Value {
         // Determine how many predecessors are yielding unique, non-temporary Values. If a variable
@@ -545,71 +544,23 @@ impl SSABuilder {
             // There is disagreement in the predecessors on which value to use so we have
             // to keep the block argument.
             let mut preds = self.ssa_blocks[dest_block].predecessors;
-            let var_defs = &mut self.variables[var];
+            let dfg = &mut func.stencil.dfg;
             for (idx, &val) in results.as_slice().iter().enumerate() {
                 let pred = preds.get_mut(idx, &mut self.inst_pool).unwrap();
                 let branch = *pred;
-                if let Some((new_block, new_branch)) =
-                    Self::append_jump_argument(func, branch, dest_block, val)
-                {
-                    *pred = new_branch;
-                    let old_block = func.layout.inst_block(branch).unwrap();
-                    self.ssa_blocks[new_block] = SSABlockData {
-                        predecessors: EntityList::from_slice(&[branch], &mut self.inst_pool),
-                        sealed: Sealed::Yes,
-                        single_predecessor: PackedOption::from(old_block),
-                    };
-                    var_defs[new_block] = PackedOption::from(val);
-                    self.side_effects.split_blocks_created.push(new_block);
-                }
-            }
-            sentinel
-        }
-    }
 
-    /// Appends a jump argument to a jump instruction, returns block created in case of
+                let dests = dfg.insts[branch].branch_destination_mut(&mut dfg.jump_tables);
-    /// critical edge splitting.
+                assert!(
-    fn append_jump_argument(
+                    !dests.is_empty(),
-        func: &mut Function,
+                    "you have declared a non-branch instruction as a predecessor to a block!"
-        branch: Inst,
+                );
-        dest_block: Block,
+                for block in dests {
-        val: Value,
-    ) -> Option<(Block, Inst)> {
-        let dfg = &mut func.stencil.dfg;
-        match &mut dfg.insts[branch] {
-            // For a single destination appending a jump argument to the instruction
-            // is sufficient.
-            InstructionData::Jump { destination, .. } => {
-                destination.append_argument(val, &mut dfg.value_lists);
-                None
-            }
-            InstructionData::Brif { blocks, .. } => {
-                for block in blocks {
                     if block.block(&dfg.value_lists) == dest_block {
                         block.append_argument(val, &mut dfg.value_lists);
                     }
                 }
-                None
-            }
-            InstructionData::BranchTable { table: jt, .. } => {
-                // In the case of a jump table, the situation is tricky because br_table doesn't
-                // support arguments. We have to split the critical edge.
-                let middle_block = dfg.blocks.add();
-                func.stencil.layout.append_block(middle_block);
-
-                for block in dfg.jump_tables[*jt].all_branches_mut() {
-                    if *block == dest_block {
-                        *block = middle_block;
-                    }
-                }
-
-                let mut cur = FuncCursor::new(func).at_bottom(middle_block);
-                let middle_jump_inst = cur.ins().jump(dest_block, &[val]);
-                Some((middle_block, middle_jump_inst))
-            }
-            _ => {
-                panic!("you have declared a non-branch instruction as a predecessor to a block");
             }
+            sentinel
         }
     }
 
@@ -644,7 +595,7 @@ impl SSABuilder {
                     self.use_var_nonlocal(func, var, ty, block);
                 }
                 Call::FinishPredecessorsLookup(sentinel, dest_block) => {
-                    let val = self.finish_predecessors_lookup(func, sentinel, var, dest_block);
+                    let val = self.finish_predecessors_lookup(func, sentinel, dest_block);
                     self.results.push(val);
                 }
             }
@@ -1010,7 +961,13 @@ mod tests {
         ssa.def_var(x_var, x1, block0);
         ssa.use_var(&mut func, x_var, I32, block0).0;
         let br_table = {
-            let jump_table = JumpTableData::new(block2, &[block2, block1]);
+            let jump_table = JumpTableData::new(
+                func.dfg.block_call(block2, &[]),
+                &[
+                    func.dfg.block_call(block2, &[]),
+                    func.dfg.block_call(block1, &[]),
+                ],
+            );
             let jt = func.create_jump_table(jump_table);
             let mut cur = FuncCursor::new(&mut func).at_bottom(block0);
             cur.ins().br_table(x1, jt)

cranelift/frontend/src/switch.rs

@@ -220,7 +220,13 @@ impl Switch {
             "Jump tables bigger than 2^32-1 are not yet supported"
         );
 
-        let jt_data = JumpTableData::new(otherwise, blocks);
+        let jt_data = JumpTableData::new(
+            bx.func.dfg.block_call(otherwise, &[]),
+            &blocks
+                .iter()
+                .map(|block| bx.func.dfg.block_call(*block, &[]))
+                .collect::<Vec<_>>(),
+        );
         let jump_table = bx.create_jump_table(jt_data);
 
         let discr = if first_index == 0 {

cranelift/fuzzgen/src/function_generator.rs

@@ -18,15 +18,6 @@ use std::collections::HashMap;
 use std::ops::RangeInclusive;
 use target_lexicon::{Architecture, Triple};
 
-/// Generates a Vec with `len` elements comprised of `options`
-fn arbitrary_vec<T: Clone>(
-    u: &mut Unstructured,
-    len: usize,
-    options: &[T],
-) -> arbitrary::Result<Vec<T>> {
-    (0..len).map(|_| u.choose(options).cloned()).collect()
-}
-
 type BlockSignature = Vec<Type>;
 
 fn insert_opcode(
@@ -1610,7 +1601,15 @@ where
             }
             BlockTerminator::BrTable(default, targets) => {
                 // Create jump tables on demand
-                let jt = builder.create_jump_table(JumpTableData::new(default, &targets));
+                let mut jt = Vec::with_capacity(targets.len());
+                for block in targets {
+                    let args = self.generate_values_for_block(builder, block)?;
+                    jt.push(builder.func.dfg.block_call(block, &args))
+                }
+
+                let args = self.generate_values_for_block(builder, default)?;
+                let jt_data = JumpTableData::new(builder.func.dfg.block_call(default, &args), &jt);
+                let jt = builder.create_jump_table(jt_data);
 
                 // br_table only supports I32
                 let val = builder.use_var(self.get_variable_of_type(I32)?);
@@ -1799,21 +1798,21 @@ where
                     // If we have more than one block we can allow terminators that target blocks.
                     // TODO: We could add some kind of BrReturn here, to explore edges where we
                     // exit in the middle of the function
-                    valid_terminators
-                        .extend_from_slice(&[BlockTerminatorKind::Jump, BlockTerminatorKind::Br]);
-                }
-
-                // BrTable and the Switch interface only allow targeting blocks without params
-                // we also need to ensure that the next block has no params, since that one is
-                // guaranteed to be picked in either case.
-                if has_paramless_targets && next_block_is_paramless {
                     valid_terminators.extend_from_slice(&[
+                        BlockTerminatorKind::Jump,
+                        BlockTerminatorKind::Br,
                         BlockTerminatorKind::BrTable,
-                        BlockTerminatorKind::Switch,
                     ]);
                 }
 
-                let terminator = self.u.choose(&valid_terminators[..])?;
+                // As the Switch interface only allows targeting blocks without params we need
+                // to ensure that the next block has no params, since that one is guaranteed to be
+                // picked in either case.
+                if has_paramless_targets && next_block_is_paramless {
+                    valid_terminators.push(BlockTerminatorKind::Switch);
+                }
+
+                let terminator = self.u.choose(&valid_terminators)?;
 
                 // Choose block targets for the terminators that we picked above
                 Ok(match terminator {
@@ -1829,10 +1828,8 @@ where
                         let default = next_block;
 
                         let target_count = self.param(&self.config.jump_table_entries)?;
-                        let targets = arbitrary_vec(
+                        let targets = Result::from_iter(
-                            self.u,
+                            (0..target_count).map(|_| self.generate_target_block(block)),
-                            target_count,
-                            self.resources.forward_blocks_without_params(block),
                         )?;
 
                         BlockTerminator::BrTable(default, targets)

cranelift/interpreter/src/step.rs

@@ -349,8 +349,11 @@ where
     // Interpret a Cranelift instruction.
     Ok(match inst.opcode() {
         Opcode::Jump => {
-            let block = inst.branch_destination()[0];
+            if let InstructionData::Jump { destination, .. } = inst {
-            continue_at(block)?
+                continue_at(destination)?
+            } else {
+                unreachable!()
+            }
         }
         Opcode::Brif => {
             if let InstructionData::Brif {
@@ -383,7 +386,7 @@ where
                     .copied()
                     .unwrap_or(jt_data.default_block());
 
-                ControlFlow::ContinueAt(jump_target, SmallVec::new())
+                continue_at(jump_target)?
             } else {
                 unreachable!()
             }

cranelift/reader/src/parser.rs

@@ -1783,9 +1783,13 @@ impl<'a> Parser<'a> {
 
     // Parse a jump table literal.
     //
-    // jump-table-lit ::= "[" block {"," block } "]"
+    // jump-table-lit ::= "[" block(args) {"," block(args) } "]"
     //                  | "[]"
-    fn parse_jump_table(&mut self, def: Block) -> ParseResult<JumpTableData> {
+    fn parse_jump_table(
+        &mut self,
+        ctx: &mut Context,
+        def: ir::BlockCall,
+    ) -> ParseResult<ir::JumpTable> {
         self.match_token(Token::LBracket, "expected '[' before jump table contents")?;
 
         let mut data = Vec::new();
@@ -1793,7 +1797,8 @@ impl<'a> Parser<'a> {
         match self.token() {
             Some(Token::Block(dest)) => {
                 self.consume();
-                data.push(dest);
+                let args = self.parse_opt_value_list()?;
+                data.push(ctx.function.dfg.block_call(dest, &args));
 
                 loop {
                     match self.token() {
@@ -1801,7 +1806,8 @@ impl<'a> Parser<'a> {
                             self.consume();
                             if let Some(Token::Block(dest)) = self.token() {
                                 self.consume();
-                                data.push(dest);
+                                let args = self.parse_opt_value_list()?;
+                                data.push(ctx.function.dfg.block_call(dest, &args));
                             } else {
                                 return err!(self.loc, "expected jump_table entry");
                             }
@@ -1817,7 +1823,11 @@ impl<'a> Parser<'a> {
 
         self.consume();
 
-        Ok(JumpTableData::new(def, &data))
+        Ok(ctx
+            .function
+            .dfg
+            .jump_tables
+            .push(JumpTableData::new(def, &data)))
     }
 
     // Parse a constant decl.
@@ -2588,9 +2598,10 @@ impl<'a> Parser<'a> {
                 let arg = self.match_value("expected SSA value operand")?;
                 self.match_token(Token::Comma, "expected ',' between operands")?;
                 let block_num = self.match_block("expected branch destination block")?;
+                let args = self.parse_opt_value_list()?;
+                let destination = ctx.function.dfg.block_call(block_num, &args);
                 self.match_token(Token::Comma, "expected ',' between operands")?;
-                let table_data = self.parse_jump_table(block_num)?;
+                let table = self.parse_jump_table(ctx, destination)?;
-                let table = ctx.function.dfg.jump_tables.push(table_data);
                 InstructionData::BranchTable { opcode, arg, table }
             }
             InstructionFormat::TernaryImm8 => {

cranelift/src/bugpoint.rs

@@ -422,7 +422,7 @@ impl Mutator for ReplaceBlockParamWithConst {
         // Remove parameters in branching instructions that point to this block
         for pred in cfg.pred_iter(self.block) {
             let dfg = &mut func.dfg;
-            for branch in dfg.insts[pred.inst].branch_destination_mut().into_iter() {
+            for branch in dfg.insts[pred.inst].branch_destination_mut(&mut dfg.jump_tables) {
                 if branch.block(&dfg.value_lists) == self.block {
                     branch.remove(param_index, &mut dfg.value_lists);
                 }
@@ -711,7 +711,7 @@ impl Mutator for MergeBlocks {
 
         // If the branch instruction that lead us to this block wasn't an unconditional jump, then
         // we have a conditional jump sequence that we should not break.
-        let branch_dests = func.dfg.insts[pred.inst].branch_destination();
+        let branch_dests = func.dfg.insts[pred.inst].branch_destination(&func.dfg.jump_tables);
         if branch_dests.len() != 1 {
             return Some((
                 func,

cranelift/wasm/src/code_translator.rs

@@ -526,7 +526,7 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
                         frame.set_branched_to_exit();
                         frame.br_destination()
                     };
-                    data.push(block);
+                    data.push(builder.func.dfg.block_call(block, &[]));
                 }
                 let block = {
                     let i = state.control_stack.len() - 1 - (default as usize);
@@ -534,6 +534,7 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
                     frame.set_branched_to_exit();
                     frame.br_destination()
                 };
+                let block = builder.func.dfg.block_call(block, &[]);
                 let jt = builder.create_jump_table(JumpTableData::new(block, &data));
                 builder.ins().br_table(val, jt);
             } else {
@@ -552,7 +553,7 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
                             *entry.insert(block)
                         }
                     };
-                    data.push(branch_block);
+                    data.push(builder.func.dfg.block_call(branch_block, &[]));
                 }
                 let default_branch_block = match dest_block_map.entry(default as usize) {
                     hash_map::Entry::Occupied(entry) => *entry.get(),
@@ -562,6 +563,7 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
                         *entry.insert(block)
                     }
                 };
+                let default_branch_block = builder.func.dfg.block_call(default_branch_block, &[]);
                 let jt = builder.create_jump_table(JumpTableData::new(default_branch_block, &data));
                 builder.ins().br_table(val, jt);
                 for (depth, dest_block) in dest_block_sequence {