cranelift: Remove redundant dominator-tree structure (#7948)

* Don't keep extra copies of root block ID

We already know the entry block everywhere we use this.

* Replace DomTreeWithChildren with DominatorTreePreorder

These two types had nearly identical interfaces except that the latter
also supports constant-time dominance checks.

* Use dom-tree preorder for O(1) dominance checks

It turns out this was the only thing the egraph pass used the original
dominator tree for, so we can stop passing that around.

We could also require that the caller of EgraphPass::new construct the
DominatorTreePreorder and avoid passing the original tree into the
egraph pass at all, but I've chosen to stop refactoring here.

We should review all uses of DominatorTree::dominates to see if it makes
sense to pass a DominatorTreePreorder around more places.

* review comments: Rename/document away from "domtree with children"

cranelift/codegen/src/egraph.rs

@@ -3,8 +3,7 @@
 use crate::alias_analysis::{AliasAnalysis, LastStores};
 use crate::ctxhash::{CtxEq, CtxHash, CtxHashMap};
 use crate::cursor::{Cursor, CursorPosition, FuncCursor};
-use crate::dominator_tree::DominatorTree;
-use crate::egraph::domtree::DomTreeWithChildren;
+use crate::dominator_tree::{DominatorTree, DominatorTreePreorder};
 use crate::egraph::elaborate::Elaborator;
 use crate::fx::FxHashSet;
 use crate::inst_predicates::{is_mergeable_for_egraph, is_pure_for_egraph};
@@ -22,7 +21,6 @@ use smallvec::SmallVec;
 use std::hash::Hasher;
 
 mod cost;
-mod domtree;
 mod elaborate;
 
 /// Pass over a Function that does the whole aegraph thing.
@@ -46,14 +44,12 @@ mod elaborate;
 pub struct EgraphPass<'a> {
     /// The function we're operating on.
     func: &'a mut Function,
-    /// Dominator tree, used for elaboration pass.
-    domtree: &'a DominatorTree,
+    /// Dominator tree for the CFG, used to visit blocks in pre-order
+    /// so we see value definitions before their uses, and also used for
+    /// O(1) dominance checks.
+    domtree: DominatorTreePreorder,
     /// Alias analysis, used during optimization.
     alias_analysis: &'a mut AliasAnalysis<'a>,
-    /// "Domtree with children": like `domtree`, but with an explicit
-    /// list of children, complementing the parent pointers stored
-    /// in `domtree`.
-    domtree_children: DomTreeWithChildren,
     /// Loop analysis results, used for built-in LICM during
     /// elaboration.
     loop_analysis: &'a LoopAnalysis,
@@ -401,16 +397,16 @@ impl<'a> EgraphPass<'a> {
     /// Create a new EgraphPass.
     pub fn new(
         func: &'a mut Function,
-        domtree: &'a DominatorTree,
+        raw_domtree: &'a DominatorTree,
         loop_analysis: &'a LoopAnalysis,
         alias_analysis: &'a mut AliasAnalysis<'a>,
     ) -> Self {
         let num_values = func.dfg.num_values();
-        let domtree_children = DomTreeWithChildren::new(func, domtree);
+        let mut domtree = DominatorTreePreorder::new();
+        domtree.compute(raw_domtree, &func.layout);
         Self {
             func,
             domtree,
-            domtree_children,
             loop_analysis,
             alias_analysis,
             stats: Stats::default(),
@@ -497,7 +493,7 @@ impl<'a> EgraphPass<'a> {
 
         // In domtree preorder, visit blocks. (TODO: factor out an
         // iterator from this and elaborator.)
-        let root = self.domtree_children.root();
+        let root = cursor.layout().entry_block().unwrap();
         enum StackEntry {
             Visit(Block),
             Pop,
@@ -509,8 +505,7 @@ impl<'a> EgraphPass<'a> {
                     // We popped this block; push children
                     // immediately, then process this block.
                     block_stack.push(StackEntry::Pop);
-                    block_stack
-                        .extend(self.domtree_children.children(block).map(StackEntry::Visit));
+                    block_stack.extend(self.domtree.children(block).map(StackEntry::Visit));
                     effectful_gvn_map.increment_depth();
 
                     trace!("Processing block {}", block);
@@ -615,8 +610,7 @@ impl<'a> EgraphPass<'a> {
     fn elaborate(&mut self) {
         let mut elaborator = Elaborator::new(
             self.func,
-            self.domtree,
-            &self.domtree_children,
+            &self.domtree,
             self.loop_analysis,
             &mut self.remat_values,
             &mut self.stats,

cranelift/codegen/src/egraph/domtree.rs

@@ -1,74 +0,0 @@
-//! Extended domtree with various traversal support.
-
-use crate::dominator_tree::DominatorTree;
-use crate::ir::{Block, Function};
-use cranelift_entity::{packed_option::PackedOption, SecondaryMap};
-
-#[derive(Clone, Debug)]
-/// Like a [`DominatorTree`], but with an explicit list of children,
-/// rather than parent pointers.
-pub(crate) struct DomTreeWithChildren {
-    nodes: SecondaryMap<Block, DomTreeNode>,
-    root: Block,
-}
-
-#[derive(Clone, Copy, Debug, Default)]
-struct DomTreeNode {
-    /// Points to the first child of the node, and implicitly to an entire
-    /// linked list of the children.
-    children: PackedOption<Block>,
-    /// Points to the next sibling, if any.
-    next: PackedOption<Block>,
-}
-
-impl DomTreeWithChildren {
-    pub(crate) fn new(func: &Function, domtree: &DominatorTree) -> DomTreeWithChildren {
-        let mut nodes: SecondaryMap<Block, DomTreeNode> =
-            SecondaryMap::with_capacity(func.dfg.num_blocks());
-
-        for block in func.layout.blocks() {
-            let Some(idom_inst) = domtree.idom(block) else {
-                continue;
-            };
-            let idom = func
-                .layout
-                .inst_block(idom_inst)
-                .expect("Dominating instruction should be part of a block");
-
-            // Insert at the front of nodes[idom].children
-            nodes[block].next = nodes[idom].children;
-            nodes[idom].children = block.into();
-        }
-
-        let root = func.layout.entry_block().unwrap();
-
-        Self { nodes, root }
-    }
-
-    pub(crate) fn root(&self) -> Block {
-        self.root
-    }
-
-    pub(crate) fn children<'a>(&'a self, block: Block) -> DomTreeChildIter<'a> {
-        let block = self.nodes[block].children;
-        DomTreeChildIter {
-            domtree: self,
-            block,
-        }
-    }
-}
-
-pub(crate) struct DomTreeChildIter<'a> {
-    domtree: &'a DomTreeWithChildren,
-    block: PackedOption<Block>,
-}
-
-impl<'a> Iterator for DomTreeChildIter<'a> {
-    type Item = Block;
-    fn next(&mut self) -> Option<Block> {
-        self.block.expand().map(|block| {
-            self.block = self.domtree.nodes[block].next;
-            block
-        })
-    }
-}

cranelift/codegen/src/egraph/elaborate.rs

@@ -2,9 +2,8 @@
 //! in CFG nodes.
 
 use super::cost::Cost;
-use super::domtree::DomTreeWithChildren;
 use super::Stats;
-use crate::dominator_tree::DominatorTree;
+use crate::dominator_tree::DominatorTreePreorder;
 use crate::fx::{FxHashMap, FxHashSet};
 use crate::hash_map::Entry as HashEntry;
 use crate::inst_predicates::is_pure_for_egraph;
@@ -18,8 +17,7 @@ use smallvec::{smallvec, SmallVec};
 
 pub(crate) struct Elaborator<'a> {
     func: &'a mut Function,
-    domtree: &'a DominatorTree,
-    domtree_children: &'a DomTreeWithChildren,
+    domtree: &'a DominatorTreePreorder,
     loop_analysis: &'a LoopAnalysis,
     /// Map from Value that is produced by a pure Inst (and was thus
     /// not in the side-effecting skeleton) to the value produced by
@@ -137,8 +135,7 @@ enum BlockStackEntry {
 impl<'a> Elaborator<'a> {
     pub(crate) fn new(
         func: &'a mut Function,
-        domtree: &'a DominatorTree,
-        domtree_children: &'a DomTreeWithChildren,
+        domtree: &'a DominatorTreePreorder,
         loop_analysis: &'a LoopAnalysis,
         remat_values: &'a FxHashSet<Value>,
         stats: &'a mut Stats,
@@ -150,7 +147,6 @@ impl<'a> Elaborator<'a> {
         Self {
             func,
             domtree,
-            domtree_children,
             loop_analysis,
             value_to_elaborated_value: ScopedHashMap::with_capacity(num_values),
             value_to_best_value,
@@ -557,11 +553,7 @@ impl<'a> Elaborator<'a> {
                             let data = &self.loop_stack[loop_hoist_level];
                             // `data.hoist_block` should dominate `before`'s block.
                             let before_block = self.func.layout.inst_block(before).unwrap();
-                            debug_assert!(self.domtree.dominates(
-                                data.hoist_block,
-                                before_block,
-                                &self.func.layout
-                            ));
+                            debug_assert!(self.domtree.dominates(data.hoist_block, before_block));
                             // Determine the instruction at which we
                             // insert in `data.hoist_block`.
                             let before = self.func.layout.last_inst(data.hoist_block).unwrap();
@@ -762,10 +754,9 @@ impl<'a> Elaborator<'a> {
         }
     }
 
-    fn elaborate_domtree(&mut self, domtree: &DomTreeWithChildren) {
-        let root = domtree.root();
+    fn elaborate_domtree(&mut self, domtree: &DominatorTreePreorder) {
         self.block_stack.push(BlockStackEntry::Elaborate {
-            block: root,
+            block: self.func.layout.entry_block().unwrap(),
             idom: None,
         });
 
@@ -808,7 +799,7 @@ impl<'a> Elaborator<'a> {
         self.stats.elaborate_func += 1;
         self.stats.elaborate_func_pre_insts += self.func.dfg.num_insts() as u64;
         self.compute_best_values();
-        self.elaborate_domtree(&self.domtree_children);
+        self.elaborate_domtree(&self.domtree);
         self.stats.elaborate_func_post_insts += self.func.dfg.num_insts() as u64;
     }
 }