cranelift: Rework block instructions to use BlockCall (#5464)

Add a new type BlockCall that represents the pair of a block name with arguments to be passed to it. (The mnemonic here is that it looks a bit like a function call.) Rework the implementation of jump, brz, and brnz to use BlockCall instead of storing the block arguments as varargs in the instruction's ValueList.

To ensure that we're processing block arguments from BlockCall values in instructions, three new functions have been introduced on DataFlowGraph that both sets of arguments:

inst_values - returns an iterator that traverses values in the instruction and block arguments
map_inst_values - applies a function to each value in the instruction and block arguments
overwrite_inst_values - overwrite all values in an instruction and block arguments with values from the iterator

Co-authored-by: Jamey Sharp <jamey@minilop.net>

cranelift/codegen/src/ir/instructions.rs

@@ -31,6 +31,95 @@ pub type ValueList = entity::EntityList<Value>;
 /// Memory pool for holding value lists. See `ValueList`.
 pub type ValueListPool = entity::ListPool<Value>;
 
+/// A pair of a Block and its arguments, stored in a single EntityList internally.
+///
+/// NOTE: We don't expose either value_to_block or block_to_value outside of this module because
+/// this operation is not generally safe. However, as the two share the same underlying layout,
+/// they can be stored in the same value pool.
+///
+/// BlockCall makes use of this shared layout by storing all of its contents (a block and its
+/// argument) in a single EntityList. This is a bit better than introducing a new entity type for
+/// the pair of a block name and the arguments entity list, as we don't pay any indirection penalty
+/// to get to the argument values -- they're stored in-line with the block in the same list.
+///
+/// The BlockCall::new function guarantees this layout by requiring a block argument that's written
+/// in as the first element of the EntityList. Any subsequent entries are always assumed to be real
+/// Values.
+#[derive(Debug, Clone, Copy, PartialEq, Hash)]
+#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
+pub struct BlockCall {
+    /// The underlying storage for the BlockCall. The first element of the values EntityList is
+    /// guaranteed to always be a Block encoded as a Value via BlockCall::block_to_value.
+    /// Consequently, the values entity list is never empty.
+    values: entity::EntityList<Value>,
+}
+
+impl BlockCall {
+    // NOTE: the only uses of this function should be internal to BlockCall. See the block comment
+    // on BlockCall for more context.
+    fn value_to_block(val: Value) -> Block {
+        Block::from_u32(val.as_u32())
+    }
+
+    // NOTE: the only uses of this function should be internal to BlockCall. See the block comment
+    // on BlockCall for more context.
+    fn block_to_value(block: Block) -> Value {
+        Value::from_u32(block.as_u32())
+    }
+
+    /// Construct a BlockCall with the given block and arguments.
+    pub fn new(block: Block, args: &[Value], pool: &mut ValueListPool) -> Self {
+        let mut values = ValueList::default();
+        values.push(Self::block_to_value(block), pool);
+        values.extend(args.iter().copied(), pool);
+        Self { values }
+    }
+
+    /// Return the block for this BlockCall.
+    pub fn block(&self, pool: &ValueListPool) -> Block {
+        let val = self.values.first(pool).unwrap();
+        Self::value_to_block(val)
+    }
+
+    /// Replace the block for this BlockCall.
+    pub fn set_block(&mut self, block: Block, pool: &mut ValueListPool) {
+        *self.values.get_mut(0, pool).unwrap() = Self::block_to_value(block);
+    }
+
+    /// Append an argument to the block args.
+    pub fn append_argument(&mut self, arg: Value, pool: &mut ValueListPool) {
+        self.values.push(arg, pool);
+    }
+
+    /// Return a slice for the arguments of this block.
+    pub fn args_slice<'a>(&self, pool: &'a ValueListPool) -> &'a [Value] {
+        &self.values.as_slice(pool)[1..]
+    }
+
+    /// Return a slice for the arguments of this block.
+    pub fn args_slice_mut<'a>(&'a mut self, pool: &'a mut ValueListPool) -> &'a mut [Value] {
+        &mut self.values.as_mut_slice(pool)[1..]
+    }
+
+    /// Remove the argument at ix from the argument list.
+    pub fn remove(&mut self, ix: usize, pool: &mut ValueListPool) {
+        self.values.remove(1 + ix, pool)
+    }
+
+    /// Clear out the arguments list.
+    pub fn clear(&mut self, pool: &mut ValueListPool) {
+        self.values.truncate(1, pool)
+    }
+
+    /// Appends multiple elements to the arguments.
+    pub fn extend<I>(&mut self, elements: I, pool: &mut ValueListPool)
+    where
+        I: IntoIterator<Item = Value>,
+    {
+        self.values.extend(elements, pool)
+    }
+}
+
 // Include code generated by `cranelift-codegen/meta/src/gen_inst.rs`. This file contains:
 //
 // - The `pub enum InstructionFormat` enum with all the instruction formats.
@@ -178,18 +267,10 @@ impl InstructionData {
     ///
     /// Any instruction that can transfer control to another block reveals its possible destinations
     /// here.
-    pub fn analyze_branch<'a>(&'a self, pool: &'a ValueListPool) -> BranchInfo<'a> {
+    pub fn analyze_branch(&self) -> BranchInfo {
         match *self {
-            Self::Jump {
-                destination,
-                ref args,
-                ..
-            } => BranchInfo::SingleDest(destination, args.as_slice(pool)),
-            Self::Branch {
-                destination,
-                ref args,
-                ..
-            } => BranchInfo::SingleDest(destination, &args.as_slice(pool)[1..]),
+            Self::Jump { destination, .. } => BranchInfo::SingleDest(destination),
+            Self::Branch { destination, .. } => BranchInfo::SingleDest(destination),
             Self::BranchTable {
                 table, destination, ..
             } => BranchInfo::Table(table, Some(destination)),
@@ -204,7 +285,7 @@ impl InstructionData {
     /// branch or jump.
     ///
     /// Multi-destination branches like `br_table` return `None`.
-    pub fn branch_destination(&self) -> Option<Block> {
+    pub fn branch_destination(&self) -> Option<BlockCall> {
         match *self {
             Self::Jump { destination, .. } | Self::Branch { destination, .. } => Some(destination),
             Self::BranchTable { .. } => None,
@@ -219,7 +300,7 @@ impl InstructionData {
     /// single destination branch or jump.
     ///
     /// Multi-destination branches like `br_table` return `None`.
-    pub fn branch_destination_mut(&mut self) -> Option<&mut Block> {
+    pub fn branch_destination_mut(&mut self) -> Option<&mut BlockCall> {
         match *self {
             Self::Jump {
                 ref mut destination,
@@ -366,14 +447,14 @@ impl InstructionData {
 }
 
 /// Information about branch and jump instructions.
-pub enum BranchInfo<'a> {
+pub enum BranchInfo {
     /// This is not a branch or jump instruction.
     /// This instruction will not transfer control to another block in the function, but it may still
     /// affect control flow by returning or trapping.
     NotABranch,
 
     /// This is a branch or jump to a single destination block, possibly taking value arguments.
-    SingleDest(Block, &'a [Value]),
+    SingleDest(BlockCall),
 
     /// This is a jump table branch which can have many destination blocks and maybe one default block.
     Table(JumpTable, Option<Block>),