Move jump tables to the DataFlowGraph (#5745)

Move the storage for jump tables off of FunctionStencil and onto DataFlowGraph. This change is in service of #5731, making it easier to access the jump table data in the context of helpers like inst_values.

Author: elliottt

cranelift/codegen/src/dominator_tree.rs

@@ -367,7 +367,7 @@ impl DominatorTree {
                     destination: dest,
                     ..
                 } => {
-                    for succ in func.jump_tables[*jt].iter() {
+                    for succ in func.stencil.dfg.jump_tables[*jt].iter() {
                         self.push_if_unseen(*succ);
                     }
                     self.push_if_unseen(*dest);

cranelift/codegen/src/flowgraph.rs

@@ -137,7 +137,7 @@ impl ControlFlowGraph {
                 } => {
                     self.add_edge(block, inst, *dest);
 
-                    for dest in func.jump_tables[*jt].iter() {
+                    for dest in func.stencil.dfg.jump_tables[*jt].iter() {
                         self.add_edge(block, inst, *dest);
                     }
                 }

cranelift/codegen/src/inst_predicates.rs

@@ -199,7 +199,7 @@ pub(crate) fn visit_block_succs<F: FnMut(Inst, Block, bool)>(
                 // so it is not part of the table.
                 visit(inst, *dest, false);
 
-                for &dest in f.jump_tables[*table].as_slice() {
+                for &dest in f.stencil.dfg.jump_tables[*table].as_slice() {
                     visit(inst, dest, true);
                 }
             }

cranelift/codegen/src/ir/dfg.rs

@@ -5,12 +5,11 @@ use crate::ir;
 use crate::ir::builder::ReplaceBuilder;
 use crate::ir::dynamic_type::{DynamicTypeData, DynamicTypes};
 use crate::ir::instructions::{CallInfo, InstructionData};
-use crate::ir::{types, ConstantData, ConstantPool, Immediate};
 use crate::ir::{
-    Block, BlockCall, DynamicType, FuncRef, Inst, SigRef, Signature, Type, Value,
+    types, Block, BlockCall, ConstantData, ConstantPool, DynamicType, ExtFuncData, FuncRef,
+    Immediate, Inst, JumpTables, RelSourceLoc, SigRef, Signature, Type, Value,
     ValueLabelAssignments, ValueList, ValueListPool,
 };
-use crate::ir::{ExtFuncData, RelSourceLoc};
 use crate::packed_option::ReservedValue;
 use crate::write::write_operands;
 use core::fmt;
@@ -144,6 +143,9 @@ pub struct DataFlowGraph {
 
     /// Stores large immediates that otherwise will not fit on InstructionData
     pub immediates: PrimaryMap<Immediate, ConstantData>,
+
+    /// Jump tables used in this function.
+    pub jump_tables: JumpTables,
 }
 
 impl DataFlowGraph {
@@ -162,6 +164,7 @@ impl DataFlowGraph {
             values_labels: None,
             constants: ConstantPool::new(),
             immediates: PrimaryMap::new(),
+            jump_tables: JumpTables::new(),
         }
     }
 
@@ -179,6 +182,7 @@ impl DataFlowGraph {
         self.values_labels = None;
         self.constants.clear();
         self.immediates.clear();
+        self.jump_tables.clear();
     }
 
     /// Get the total number of instructions created in this function, whether they are currently

cranelift/codegen/src/ir/function.rs

@@ -7,7 +7,7 @@ use crate::entity::{PrimaryMap, SecondaryMap};
 use crate::ir::{
     self, Block, DataFlowGraph, DynamicStackSlot, DynamicStackSlotData, DynamicStackSlots,
     DynamicType, ExtFuncData, FuncRef, GlobalValue, GlobalValueData, Inst, InstructionData,
-    JumpTable, JumpTableData, JumpTables, Layout, Opcode, SigRef, Signature, SourceLocs, StackSlot,
+    JumpTable, JumpTableData, Layout, Opcode, SigRef, Signature, SourceLocs, StackSlot,
     StackSlotData, StackSlots, Table, TableData, Type,
 };
 use crate::isa::CallConv;
@@ -170,9 +170,6 @@ pub struct FunctionStencil {
     /// Tables referenced.
     pub tables: PrimaryMap<ir::Table, ir::TableData>,
 
-    /// Jump tables used in this function.
-    pub jump_tables: JumpTables,
-
     /// Data flow graph containing the primary definition of all instructions, blocks and values.
     pub dfg: DataFlowGraph,
 
@@ -200,7 +197,6 @@ impl FunctionStencil {
         self.dynamic_stack_slots.clear();
         self.global_values.clear();
         self.tables.clear();
-        self.jump_tables.clear();
         self.dfg.clear();
         self.layout.clear();
         self.srclocs.clear();
@@ -209,7 +205,7 @@ impl FunctionStencil {
 
     /// Creates a jump table in the function, to be used by `br_table` instructions.
     pub fn create_jump_table(&mut self, data: JumpTableData) -> JumpTable {
-        self.jump_tables.push(data)
+        self.dfg.jump_tables.push(data)
     }
 
     /// Creates a sized stack slot in the function, to be used by `stack_load`, `stack_store`
@@ -304,7 +300,7 @@ impl FunctionStencil {
                 destination: default_dest,
                 ..
             } => {
-                self.jump_tables[*table].iter_mut().for_each(|entry| {
+                self.dfg.jump_tables[*table].iter_mut().for_each(|entry| {
                     if *entry == old_dest {
                         *entry = new_dest;
                     }
@@ -433,7 +429,6 @@ impl Function {
                 dynamic_stack_slots: DynamicStackSlots::new(),
                 global_values: PrimaryMap::new(),
                 tables: PrimaryMap::new(),
-                jump_tables: PrimaryMap::new(),
                 dfg: DataFlowGraph::new(),
                 layout: Layout::new(),
                 srclocs: SecondaryMap::new(),

cranelift/codegen/src/unreachable_code.rs

@@ -21,7 +21,7 @@ pub fn eliminate_unreachable_code(
 ) {
     let _tt = timing::unreachable_code();
     let mut pos = FuncCursor::new(func);
-    let mut used_tables = EntitySet::with_capacity(pos.func.jump_tables.len());
+    let mut used_tables = EntitySet::with_capacity(pos.func.stencil.dfg.jump_tables.len());
     while let Some(block) = pos.next_block() {
         if domtree.is_reachable(block) {
             let inst = pos.func.layout.last_inst(block).unwrap();
@@ -50,7 +50,7 @@ pub fn eliminate_unreachable_code(
         pos.func.layout.remove_block(block);
     }
 
-    for (table, jt_data) in func.jump_tables.iter_mut() {
+    for (table, jt_data) in func.stencil.dfg.jump_tables.iter_mut() {
         if !used_tables.contains(table) {
             jt_data.clear();
         }

cranelift/codegen/src/verifier/mod.rs

@@ -447,7 +447,7 @@ impl<'a> Verifier<'a> {
     }
 
     fn verify_jump_tables(&self, errors: &mut VerifierErrors) -> VerifierStepResult<()> {
-        for (jt, jt_data) in &self.func.jump_tables {
+        for (jt, jt_data) in &self.func.stencil.dfg.jump_tables {
             for &block in jt_data.iter() {
                 self.verify_block(jt, block, errors)?;
             }
@@ -854,7 +854,7 @@ impl<'a> Verifier<'a> {
         j: JumpTable,
         errors: &mut VerifierErrors,
     ) -> VerifierStepResult<()> {
-        if !self.func.jump_tables.is_valid(j) {
+        if !self.func.stencil.dfg.jump_tables.is_valid(j) {
             errors.nonfatal((
                 inst,
                 self.context(inst),
@@ -1344,7 +1344,7 @@ impl<'a> Verifier<'a> {
                         ),
                     ));
                 }
-                for block in self.func.jump_tables[*table].iter() {
+                for block in self.func.stencil.dfg.jump_tables[*table].iter() {
                     let arg_count = self.func.dfg.num_block_params(*block);
                     if arg_count != 0 {
                         return errors.nonfatal((

cranelift/codegen/src/write.rs

@@ -82,7 +82,7 @@ pub trait FuncWriter {
             }
         }
 
-        for (jt, jt_data) in &func.jump_tables {
+        for (jt, jt_data) in &func.stencil.dfg.jump_tables {
             any = true;
             self.write_entity_definition(w, func, jt.into(), jt_data)?;
         }

cranelift/frontend/src/frontend.rs

@@ -139,6 +139,8 @@ impl<'short, 'long> InstBuilderBase<'short> for FuncInstBuilder<'short, 'long> {
                 for dest_block in self
                     .builder
                     .func
+                    .stencil
+                    .dfg
                     .jump_tables
                     .get(*table)
                     .expect("you are referencing an undeclared jump table")

cranelift/frontend/src/ssa.rs

@@ -601,7 +601,7 @@ impl SSABuilder {
                 let middle_block = dfg.blocks.add();
                 func.stencil.layout.append_block(middle_block);
 
-                let table = &mut func.stencil.jump_tables[*jt];
+                let table = &mut dfg.jump_tables[*jt];
                 for block in table.iter_mut() {
                     if *block == dest_block {
                         *block = middle_block;