Move default blocks into jump tables

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

@@ -15,10 +15,6 @@ pub(crate) struct EntityRefs {
     /// This is primarily used in control flow instructions.
     pub(crate) block_call: OperandKind,
 
-    /// A reference to a basic block in the same function.
-    /// This is primarily used in control flow instructions.
-    pub(crate) label: OperandKind,
-
     /// A reference to a basic block in the same function, with its arguments provided.
     /// This is primarily used in control flow instructions.
     pub(crate) block_then: OperandKind,
@@ -63,12 +59,6 @@ impl EntityRefs {
                 "a basic block in the same function, with its arguments provided.",
             ),
 
-            label: new(
-                "destination",
-                "ir::Block",
-                "a basic block in the same function.",
-            ),
-
             block_then: new(
                 "block_then",
                 "ir::BlockCall",

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

@@ -117,7 +117,6 @@ impl Formats {
 
             branch_table: Builder::new("BranchTable")
                 .value()
-                .imm(&entities.label)
                 .imm(&entities.jump_table)
                 .build(),
 

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

@@ -19,7 +19,6 @@ fn define_control_flow(
 ) {
     let block_call = &Operand::new("block_call", &entities.block_call)
         .with_doc("Destination basic block, with its arguments provided");
-    let label = &Operand::new("label", &entities.label).with_doc("Destination basic block");
 
     ig.push(
         Inst::new(
@@ -93,7 +92,7 @@ fn define_control_flow(
         "#,
                 &formats.branch_table,
             )
-            .operands_in(vec![x, label, JT])
+            .operands_in(vec![x, JT])
             .branches(),
         );
     }

cranelift/codegen/src/inst_predicates.rs

@@ -190,16 +190,14 @@ pub(crate) fn visit_block_succs<F: FnMut(Inst, Block, bool)>(
                 visit(inst, block_else.block(&f.dfg.value_lists), false);
             }
 
-            ir::InstructionData::BranchTable {
-                table,
-                destination: dest,
-                ..
-            } => {
+            ir::InstructionData::BranchTable { table, .. } => {
+                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.
-                visit(inst, *dest, false);
+                visit(inst, table.default_block(), false);
 
-                for &dest in f.stencil.dfg.jump_tables[*table].as_slice() {
+                for &dest in table.as_slice() {
                     visit(inst, dest, true);
                 }
             }

cranelift/codegen/src/ir/function.rs

@@ -295,19 +295,11 @@ impl FunctionStencil {
                 }
             }
 
-            InstructionData::BranchTable {
-                table,
-                destination: default_dest,
-                ..
-            } => {
-                self.dfg.jump_tables[*table].iter_mut().for_each(|entry| {
+            InstructionData::BranchTable { table, .. } => {
+                for entry in self.dfg.jump_tables[*table].all_branches_mut() {
                     if *entry == old_dest {
                         *entry = new_dest;
                     }
-                });
-
-                if *default_dest == old_dest {
-                    *default_dest = new_dest;
                 }
             }
 

cranelift/codegen/src/ir/jumptable.rs

@@ -14,6 +14,10 @@ use serde::{Deserialize, Serialize};
 /// Contents of a jump table.
 ///
 /// All jump tables use 0-based indexing and are densely populated.
+///
+/// The default block for the jump table is stored as the last element of the underlying vector,
+/// and is not included in the length of the jump table. It can be accessed through the
+/// `default_block` function. The default block is iterated via the `iter` and `iter_mut` methods.
 #[derive(Clone, PartialEq, Hash)]
 #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
 pub struct JumpTableData {
@@ -22,72 +26,71 @@ pub struct JumpTableData {
 }
 
 impl JumpTableData {
-    /// Create a new empty jump table.
-    pub fn new() -> Self {
-        Self { table: Vec::new() }
-    }
-
-    /// Create a new empty jump table with the specified capacity.
-    pub fn with_capacity(capacity: usize) -> Self {
-        Self {
-            table: Vec::with_capacity(capacity),
-        }
-    }
     /// Create a new jump table with the provided blocks
-    pub fn with_blocks(table: Vec<Block>) -> Self {
+    pub fn new(def: Block, mut table: Vec<Block>) -> Self {
+        table.push(def);
         Self { table }
     }
 
-    /// Get the number of table entries.
-    pub fn len(&self) -> usize {
-        self.table.len()
+    /// Fetch the default block for this jump table.
+    pub fn default_block(&self) -> Block {
+        *self.table.last().unwrap()
     }
 
-    /// Append a table entry.
-    pub fn push_entry(&mut self, dest: Block) {
-        self.table.push(dest)
+    /// Mutable access to the default block of this jump table.
+    pub fn default_block_mut(&mut self) -> &mut Block {
+        self.table.last_mut().unwrap()
     }
 
-    /// Checks if any of the entries branch to `block`.
-    pub fn branches_to(&self, block: Block) -> bool {
-        self.table.iter().any(|target_block| *target_block == block)
+    /// The jump table and default block as a single slice. The default block will always be last.
+    pub fn all_branches(&self) -> &[Block] {
+        self.table.as_slice()
     }
 
-    /// Access the whole table as a slice.
+    /// The jump table and default block as a single mutable slice. The default block will always
+    /// be last.
+    pub fn all_branches_mut(&mut self) -> &mut [Block] {
+        self.table.as_mut_slice()
+    }
+
+    /// Access the jump table as a slice. This excludes the default block.
     pub fn as_slice(&self) -> &[Block] {
-        self.table.as_slice()
+        let last = self.table.len() - 1;
+        &self.table.as_slice()[0..last]
     }
 
-    /// Access the whole table as a mutable slice.
+    /// Access the jump table as a mutable slice. This excludes the default block.
     pub fn as_mut_slice(&mut self) -> &mut [Block] {
-        self.table.as_mut_slice()
+        let last = self.table.len() - 1;
+        &mut self.table.as_mut_slice()[0..last]
     }
 
-    /// Returns an iterator over the table.
+    /// 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> {
-        self.table.iter()
+        self.as_slice().iter()
     }
 
-    /// Returns an iterator that allows modifying each value.
+    /// 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> {
-        self.table.iter_mut()
+        self.as_mut_slice().iter_mut()
     }
 
-    /// Clears all entries in this jump table.
+    /// Clears all entries in this jump table, except for the default block.
     pub fn clear(&mut self) {
-        self.table.clear();
+        self.table.drain(0..self.table.len() - 1);
     }
 }
 
 impl Display for JumpTableData {
     fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
-        write!(fmt, "[")?;
-        match self.table.first() {
-            None => (),
-            Some(first) => write!(fmt, "{}", first)?,
-        }
-        for block in self.table.iter().skip(1) {
-            write!(fmt, ", {}", block)?;
+        write!(fmt, "{}, [", self.default_block())?;
+        if let Some((first, rest)) = self.as_slice().split_first() {
+            write!(fmt, "{}", first)?;
+            for block in rest {
+                write!(fmt, ", {}", block)?;
+            }
         }
         write!(fmt, "]")
     }
@@ -102,31 +105,33 @@ mod tests {
 
     #[test]
     fn empty() {
-        let jt = JumpTableData::new();
+        let def = Block::new(0);
+
+        let jt = JumpTableData::new(def, vec![]);
+
+        assert_eq!(jt.all_branches().get(0), Some(&def));
 
         assert_eq!(jt.as_slice().get(0), None);
         assert_eq!(jt.as_slice().get(10), None);
 
-        assert_eq!(jt.to_string(), "[]");
+        assert_eq!(jt.to_string(), "block0, []");
 
-        let v = jt.as_slice();
-        assert_eq!(v, []);
+        assert_eq!(jt.all_branches(), [def]);
+        assert_eq!(jt.as_slice(), []);
     }
 
     #[test]
     fn insert() {
+        let def = Block::new(0);
         let e1 = Block::new(1);
         let e2 = Block::new(2);
 
-        let mut jt = JumpTableData::new();
-
-        jt.push_entry(e1);
-        jt.push_entry(e2);
-        jt.push_entry(e1);
+        let jt = JumpTableData::new(def, vec![e1, e2, e1]);
 
-        assert_eq!(jt.to_string(), "[block1, block2, block1]");
+        assert_eq!(jt.default_block(), def);
+        assert_eq!(jt.to_string(), "block0, [block1, block2, block1]");
 
-        let v = jt.as_slice();
-        assert_eq!(v, [e1, e2, e1]);
+        assert_eq!(jt.all_branches(), [e1, e2, e1, def]);
+        assert_eq!(jt.as_slice(), [e1, e2, e1]);
     }
 }

cranelift/codegen/src/isa/aarch64/lower.isle

@@ -2515,7 +2515,7 @@
 
 ;; `targets` contains the default target with the list of branch targets
 ;; concatenated.
-(rule (lower_branch (br_table idx _ _) targets)
+(rule (lower_branch (br_table idx _) targets)
       (let ((jt_size u32 (targets_jt_size targets))
             (_ InstOutput (side_effect
                   (emit_island (targets_jt_space targets))))

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

@@ -361,11 +361,10 @@ mod test {
         let mut pos = FuncCursor::new(&mut func);
 
         pos.insert_block(bb0);
-        let mut jt_data = JumpTableData::new();
-        jt_data.push_entry(bb1);
-        jt_data.push_entry(bb2);
-        let jt = pos.func.create_jump_table(jt_data);
-        pos.ins().br_table(arg0, bb3, jt);
+        let jt = pos
+            .func
+            .create_jump_table(JumpTableData::new(bb3, vec![bb1, bb2]));
+        pos.ins().br_table(arg0, jt);
 
         pos.insert_block(bb1);
         let v1 = pos.ins().iconst(I32, 1);

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

@@ -1942,7 +1942,7 @@
 (extern constructor lower_br_table lower_br_table)
 
 (rule
-  (lower_branch (br_table index _ _) targets)
+  (lower_branch (br_table index _) targets)
   (lower_br_table index targets))
 
 (decl load_ra () Reg)

cranelift/codegen/src/isa/s390x/lower.isle

@@ -3725,7 +3725,7 @@
 
 ;; Jump table.  `targets` contains the default target followed by the
 ;; list of branch targets per index value.
-(rule (lower_branch (br_table val_idx _ _) targets)
+(rule (lower_branch (br_table val_idx _) targets)
       (let ((idx Reg (put_in_reg_zext64 val_idx))
             ;; Bounds-check the index and branch to default.
             ;; This is an internal branch that is not a terminator insn.

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

@@ -2942,7 +2942,7 @@
 
 ;; Rules for `br_table` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
-(rule (lower_branch (br_table idx @ (value_type ty) _ _) (jump_table_targets default_target jt_targets))
+(rule (lower_branch (br_table idx @ (value_type ty) _) (jump_table_targets default_target jt_targets))
       (emit_side_effect (jmp_table_seq ty idx default_target jt_targets)))
 
 ;; Rules for `select_spectre_guard` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

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

@@ -408,11 +408,10 @@ mod test {
         let mut pos = FuncCursor::new(&mut func);
 
         pos.insert_block(bb0);
-        let mut jt_data = JumpTableData::new();
-        jt_data.push_entry(bb1);
-        jt_data.push_entry(bb2);
-        let jt = pos.func.create_jump_table(jt_data);
-        pos.ins().br_table(arg0, bb3, jt);
+        let jt = pos
+            .func
+            .create_jump_table(JumpTableData::new(bb3, vec![bb1, bb2]));
+        pos.ins().br_table(arg0, jt);
 
         pos.insert_block(bb1);
         let v1 = pos.ins().iconst(I32, 1);

cranelift/codegen/src/verifier/mod.rs

@@ -579,10 +579,7 @@ impl<'a> Verifier<'a> {
                 self.verify_block(inst, block_then.block(&self.func.dfg.value_lists), errors)?;
                 self.verify_block(inst, block_else.block(&self.func.dfg.value_lists), errors)?;
             }
-            BranchTable {
-                table, destination, ..
-            } => {
-                self.verify_block(inst, destination, errors)?;
+            BranchTable { table, .. } => {
                 self.verify_jump_table(inst, table, errors)?;
             }
             Call {
@@ -852,7 +849,7 @@ impl<'a> Verifier<'a> {
                 format!("invalid jump table reference {}", j),
             ))
         } else {
-            for &block in self.func.stencil.dfg.jump_tables[j].as_slice() {
+            for &block in self.func.stencil.dfg.jump_tables[j].all_branches() {
                 self.verify_block(inst, block, errors)?;
             }
             Ok(())
@@ -1322,23 +1319,8 @@ impl<'a> Verifier<'a> {
                 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,
-                destination: block,
-                ..
-            } => {
-                let arg_count = self.func.dfg.num_block_params(*block);
-                if arg_count != 0 {
-                    return errors.nonfatal((
-                        inst,
-                        self.context(inst),
-                        format!(
-                            "takes no arguments, but had target {} with {} arguments",
-                            block, arg_count,
-                        ),
-                    ));
-                }
-                for block in self.func.stencil.dfg.jump_tables[*table].iter() {
+            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);
                     if arg_count != 0 {
                         return errors.nonfatal((

cranelift/codegen/src/write.rs

@@ -421,12 +421,7 @@ 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,
-            destination,
-            table,
-            ..
-        } => write!(w, " {}, {}, {}", arg, destination, jump_tables[table]),
+        BranchTable { arg, table, .. } => write!(w, " {}, {}", arg, jump_tables[table]),
         Call {
             func_ref, ref args, ..
         } => write!(w, " {}({})", func_ref, DisplayValues(args.as_slice(pool))),