Make DataFlowGraph::blocks public (#5740)

Similar to when we exposed the DataFlowGraph::insts field through a restrictive newtype, expose DataFlowGraph::blocks through an interface that allows a restrictive set of operations. This field being public now allows us to avoid a rematch in ssa construction, and simplifies the implementation of adding a block argument to a block referenced by a br_table instruction.

Author: elliottt

cranelift/codegen/src/ir/dfg.rs

@@ -45,6 +45,45 @@ impl IndexMut<Inst> for Insts {
     }
 }
 
+/// Storage for basic blocks within the DFG.
+#[derive(Clone, PartialEq, Hash)]
+#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
+pub struct Blocks(PrimaryMap<Block, BlockData>);
+
+impl Blocks {
+    /// Create a new basic block.
+    pub fn add(&mut self) -> Block {
+        self.0.push(BlockData::new())
+    }
+
+    /// Get the total number of basic blocks created in this function, whether they are
+    /// currently inserted in the layout or not.
+    ///
+    /// This is intended for use with `SecondaryMap::with_capacity`.
+    pub fn len(&self) -> usize {
+        self.0.len()
+    }
+
+    /// Returns `true` if the given block reference is valid.
+    pub fn is_valid(&self, block: Block) -> bool {
+        self.0.is_valid(block)
+    }
+}
+
+impl Index<Block> for Blocks {
+    type Output = BlockData;
+
+    fn index(&self, block: Block) -> &BlockData {
+        &self.0[block]
+    }
+}
+
+impl IndexMut<Block> for Blocks {
+    fn index_mut(&mut self, block: Block) -> &mut BlockData {
+        &mut self.0[block]
+    }
+}
+
 /// A data flow graph defines all instructions and basic blocks in a function as well as
 /// the data flow dependencies between them. The DFG also tracks values which can be either
 /// instruction results or block parameters.
@@ -70,7 +109,7 @@ pub struct DataFlowGraph {
     ///
     /// This map is not in program order. That is handled by `Layout`, and so is the sequence of
     /// instructions contained in each block.
-    blocks: PrimaryMap<Block, BlockData>,
+    pub blocks: Blocks,
 
     /// Dynamic types created.
     pub dynamic_types: DynamicTypes,
@@ -113,7 +152,7 @@ impl DataFlowGraph {
         Self {
             insts: Insts(PrimaryMap::new()),
             results: SecondaryMap::new(),
-            blocks: PrimaryMap::new(),
+            blocks: Blocks(PrimaryMap::new()),
             dynamic_types: DynamicTypes::new(),
             value_lists: ValueListPool::new(),
             values: PrimaryMap::new(),
@@ -130,7 +169,7 @@ impl DataFlowGraph {
     pub fn clear(&mut self) {
         self.insts.0.clear();
         self.results.clear();
-        self.blocks.clear();
+        self.blocks.0.clear();
         self.dynamic_types.clear();
         self.value_lists.clear();
         self.values.clear();
@@ -1084,17 +1123,17 @@ impl DataFlowGraph {
 impl DataFlowGraph {
     /// Create a new basic block.
     pub fn make_block(&mut self) -> Block {
-        self.blocks.push(BlockData::new())
+        self.blocks.add()
     }
 
     /// Get the number of parameters on `block`.
     pub fn num_block_params(&self, block: Block) -> usize {
-        self.blocks[block].params.len(&self.value_lists)
+        self.blocks[block].params(&self.value_lists).len()
     }
 
     /// Get the parameters on `block`.
     pub fn block_params(&self, block: Block) -> &[Value] {
-        self.blocks[block].params.as_slice(&self.value_lists)
+        self.blocks[block].params(&self.value_lists)
     }
 
     /// Get the types of the parameters on `block`.
@@ -1250,7 +1289,7 @@ impl DataFlowGraph {
 /// match the function arguments.
 #[derive(Clone, PartialEq, Hash)]
 #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
-struct BlockData {
+pub struct BlockData {
     /// List of parameters to this block.
     params: ValueList,
 }
@@ -1261,6 +1300,11 @@ impl BlockData {
             params: ValueList::new(),
         }
     }
+
+    /// Get the parameters on `block`.
+    pub fn params<'a>(&self, pool: &'a ValueListPool) -> &'a [Value] {
+        self.params.as_slice(pool)
+    }
 }
 
 /// Object that can display an instruction.

cranelift/codegen/src/ir/mod.rs

@@ -33,7 +33,7 @@ pub use crate::ir::builder::{
     InsertBuilder, InstBuilder, InstBuilderBase, InstInserterBase, ReplaceBuilder,
 };
 pub use crate::ir::constant::{ConstantData, ConstantPool};
-pub use crate::ir::dfg::{DataFlowGraph, ValueDef};
+pub use crate::ir::dfg::{BlockData, DataFlowGraph, ValueDef};
 pub use crate::ir::dynamic_type::{dynamic_to_fixed, DynamicTypeData, DynamicTypes};
 pub use crate::ir::entities::{
     Block, Constant, DynamicStackSlot, DynamicType, FuncRef, GlobalValue, Immediate, Inst,

cranelift/frontend/src/ssa.rs

@@ -15,9 +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, JumpTableData, Type, Value,
-};
+use cranelift_codegen::ir::{Block, Function, Inst, InstBuilder, InstructionData, Type, Value};
 use cranelift_codegen::packed_option::PackedOption;
 
 /// Structure containing the data relevant the construction of SSA for a given function.
@@ -577,58 +575,41 @@ impl SSABuilder {
         dest_block: Block,
         val: Value,
     ) -> Option<(Block, Inst)> {
-        match &func.dfg.insts[branch] {
+        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 { .. } => {
-                let dfg = &mut func.dfg;
-                for dest in dfg.insts[branch].branch_destination_mut() {
-                    dest.append_argument(val, &mut dfg.value_lists);
-                }
+            InstructionData::Jump { destination, .. } => {
+                destination.append_argument(val, &mut dfg.value_lists);
                 None
             }
-            InstructionData::Brif { .. } => {
-                let dfg = &mut func.dfg;
-                for block in dfg.insts[branch].branch_destination_mut() {
+            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: mut jt, .. } => {
+            InstructionData::BranchTable {
+                table: jt,
+                destination,
+                ..
+            } => {
                 // 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 = func.dfg.make_block();
-                func.layout.append_block(middle_block);
-
-                let table = &func.jump_tables[jt];
-                let mut copied = JumpTableData::with_capacity(table.len());
-                let mut changed = false;
-                for &destination in table.iter() {
-                    if destination == dest_block {
-                        copied.push_entry(middle_block);
-                        changed = true;
-                    } else {
-                        copied.push_entry(destination);
-                    }
-                }
+                let middle_block = dfg.blocks.add();
+                func.stencil.layout.append_block(middle_block);
 
-                if changed {
-                    jt = func.create_jump_table(copied);
+                let table = &mut func.stencil.jump_tables[*jt];
+                for block in table.iter_mut() {
+                    if *block == dest_block {
+                        *block = middle_block;
+                    }
                 }
 
-                // Redo the match from above, but this time capture mutable references
-                match &mut func.dfg.insts[branch] {
-                    InstructionData::BranchTable {
-                        destination, table, ..
-                    } => {
-                        if *destination == dest_block {
-                            *destination = middle_block;
-                        }
-                        *table = jt;
-                    }
-                    _ => unreachable!(),
+                if *destination == dest_block {
+                    *destination = middle_block;
                 }
 
                 let mut cur = FuncCursor::new(func).at_bottom(middle_block);