changes

cranelift/codegen/src/machinst/blockorder.rs

@@ -81,15 +81,6 @@ pub struct BlockLoweringOrder {
     /// Ranges in `lowered_succ_indices` giving the successor lists for each lowered
     /// block. Indexed by lowering-order index (`BlockIndex`).
     lowered_succ_ranges: Vec<(Option<Inst>, std::ops::Range<usize>)>,
-    /// Cold blocks. These blocks are not reordered in the
-    /// `lowered_order` above; the lowered order must respect RPO
-    /// (uses after defs) in order for lowering to be
-    /// correct. Instead, this set is used to provide `is_cold()`,
-    /// which is used by VCode emission to sink the blocks at the last
-    /// moment (when we actually emit bytes into the MachBuffer).
-    cold_blocks: FxHashSet<BlockIndex>,
-    /// Lowered blocks that are indirect branch targets.
-    indirect_branch_targets: FxHashSet<BlockIndex>,
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
@@ -98,6 +89,15 @@ pub enum LoweredBlock {
     Orig {
         /// Original CLIF block.
         block: Block,
+        /// Cold block. These blocks are not reordered in the
+        /// `lowered_order` above; the lowered order must respect RPO
+        /// (uses after defs) in order for lowering to be
+        /// correct. Instead, this bool is used to provide `is_cold()`,
+        /// which is used by VCode emission to sink the blocks at the last
+        /// moment (when we actually emit bytes into the MachBuffer).
+        cold: bool,
+        /// Target of indirect branch
+        indirect_branch_target: bool,
     },
 
     /// Critical edge between two CLIF blocks.
@@ -112,6 +112,11 @@ pub enum LoweredBlock {
         /// indexing order as `inst_predicates::visit_block_succs`. This is used to distinguish
         /// multiple edges between the same CLIF blocks.
         succ_idx: u32,
+
+        /// See [LoweredBlock::orig]
+        cold: bool,
+        /// See [LoweredBlock::orig]
+        indirect_branch_target: bool,
     },
 }
 
@@ -119,7 +124,7 @@ impl LoweredBlock {
     /// Unwrap an `Orig` block.
     pub fn orig_block(&self) -> Option<Block> {
         match self {
-            &LoweredBlock::Orig { block } => Some(block),
+            &LoweredBlock::Orig { block, .. } => Some(block),
             &LoweredBlock::CriticalEdge { .. } => None,
         }
     }
@@ -141,6 +146,20 @@ impl LoweredBlock {
             &LoweredBlock::Orig { .. } => None,
         }
     }
+
+    pub fn is_cold(&self) -> bool {
+        match self {
+            &LoweredBlock::CriticalEdge { cold, .. } => cold,
+            &LoweredBlock::Orig { cold, .. } => cold
+        }
+    }
+
+    pub fn is_indirect_branch_target(&self) -> bool {
+        match self {
+            &LoweredBlock::CriticalEdge { indirect_branch_target, .. } => indirect_branch_target,
+            &LoweredBlock::Orig { indirect_branch_target, .. } => indirect_branch_target
+        }
+    }
 }
 
 impl BlockLoweringOrder {
@@ -165,7 +184,7 @@ impl BlockLoweringOrder {
             visit_block_succs(f, block, |_, succ, from_table| {
                 block_out_count[block] += 1;
                 block_in_count[succ] += 1;
-                block_succs.push(LoweredBlock::Orig { block: succ });
+                block_succs.push(LoweredBlock::Orig { block: succ, cold: f.layout.is_cold(succ), indirect_branch_target: false });
 
                 if from_table {
                     indirect_branch_target_clif_blocks.insert(succ);
@@ -192,7 +211,7 @@ impl BlockLoweringOrder {
         let mut lowered_order = Vec::new();
 
         for &block in domtree.cfg_postorder().iter().rev() {
-            let lb = LoweredBlock::Orig { block };
+            let lb = LoweredBlock::Orig { block, cold: f.layout.is_cold(block), indirect_branch_target: indirect_branch_target_clif_blocks.contains(&block) };
             let bindex = BlockIndex::new(lowered_order.len());
             lb_to_bindex.insert(lb.clone(), bindex);
             lowered_order.push(lb);
@@ -208,6 +227,10 @@ impl BlockLoweringOrder {
                             pred: block,
                             succ,
                             succ_idx: succ_ix as u32,
+                            // Edges inherit indirect branch and cold block metadata from their
+                            // successor.
+                            cold: f.layout.is_cold(succ),
+                            indirect_branch_target: indirect_branch_target_clif_blocks.contains(&succ)
                         };
                         let bindex = BlockIndex::new(lowered_order.len());
                         lb_to_bindex.insert(*lb, bindex);
@@ -221,8 +244,6 @@ impl BlockLoweringOrder {
         // during the creation of `lowering_order`, as we need `lb_to_bindex` to be fully populated
         // first.
         let mut lowered_succ_indices = Vec::new();
-        let mut cold_blocks = FxHashSet::default();
-        let mut indirect_branch_targets = FxHashSet::default();
         let lowered_succ_ranges =
             Vec::from_iter(lowered_order.iter().enumerate().map(|(ix, lb)| {
                 let bindex = BlockIndex::new(ix);
@@ -230,18 +251,24 @@ impl BlockLoweringOrder {
                 let opt_inst = match lb {
                     // Block successors are pulled directly over, as they'll have been mutated when
                     // determining the block order already.
-                    &LoweredBlock::Orig { block } => {
+                    &LoweredBlock::Orig { block, .. } => {
                         let range = block_succ_range[block].clone();
                         lowered_succ_indices
-                            .extend(block_succs[range].iter().map(|lb| lb_to_bindex[lb]));
-
-                        if f.layout.is_cold(block) {
-                            cold_blocks.insert(bindex);
-                        }
-
-                        if indirect_branch_target_clif_blocks.contains(&block) {
-                            indirect_branch_targets.insert(bindex);
-                        }
+                            .extend(block_succs[range].iter().map(|lb| {
+                                match lb {
+                                    LoweredBlock::Orig { block, cold, .. } => {
+                                        match lb_to_bindex.get(&LoweredBlock::Orig { block: *block, cold: *cold, indirect_branch_target: indirect_branch_target_clif_blocks.contains(block) } ) {
+                                            Some(idx) => *idx,
+                                            None => {
+                                                panic!("Trying to hashlookup {:?} in HashMap: {:?}", lb, lb_to_bindex);
+                                            }
+                                        }
+                                    },
+                                    LoweredBlock::CriticalEdge { .. } => {
+                                        lb_to_bindex[&lb]
+                                    }
+                                }
+                            }));
 
                         let last = f.layout.last_inst(block).unwrap();
                         let opcode = f.dfg.insts[last].opcode();
@@ -253,21 +280,10 @@ impl BlockLoweringOrder {
 
                     // Critical edges won't have successor information in block_succ_range, but
                     // they only have a single known successor to record anyway.
-                    &LoweredBlock::CriticalEdge { succ, .. } => {
-                        let succ_index = lb_to_bindex[&LoweredBlock::Orig { block: succ }];
+                    &LoweredBlock::CriticalEdge { succ, cold, indirect_branch_target, .. } => {
+                        let succ_index = lb_to_bindex[&LoweredBlock::Orig { block: succ, cold, indirect_branch_target }];
                         lowered_succ_indices.push(succ_index);
 
-                        // Edges inherit indirect branch and cold block metadata from their
-                        // successor.
-
-                        if f.layout.is_cold(succ) {
-                            cold_blocks.insert(bindex);
-                        }
-
-                        if indirect_branch_target_clif_blocks.contains(&succ) {
-                            indirect_branch_targets.insert(bindex);
-                        }
-
                         None
                     }
                 };
@@ -278,9 +294,7 @@ impl BlockLoweringOrder {
         let result = BlockLoweringOrder {
             lowered_order,
             lowered_succ_indices,
-            lowered_succ_ranges,
-            cold_blocks,
-            indirect_branch_targets,
+            lowered_succ_ranges
         };
 
         trace!("BlockLoweringOrder: {:#?}", result);
@@ -300,13 +314,13 @@ impl BlockLoweringOrder {
 
     /// Determine whether the given lowered-block index is cold.
     pub fn is_cold(&self, block: BlockIndex) -> bool {
-        self.cold_blocks.contains(&block)
+        self.lowered_order[block.index()].is_cold()
     }
 
     /// Determine whether the given lowered block index is an indirect branch
     /// target.
     pub fn is_indirect_branch_target(&self, block: BlockIndex) -> bool {
-        self.indirect_branch_targets.contains(&block)
+        self.lowered_order[block.index()].is_indirect_branch_target()
     }
 }