cranelift_codegen/egraph/
domtree.rs

1//! Extended domtree with various traversal support.
2
3use crate::dominator_tree::DominatorTree;
4use crate::ir::{Block, Function};
5use cranelift_entity::{packed_option::PackedOption, SecondaryMap};
6
7#[derive(Clone, Debug)]
8pub(crate) struct DomTreeWithChildren {
9    nodes: SecondaryMap<Block, DomTreeNode>,
10    root: Block,
11}
12
13#[derive(Clone, Copy, Debug, Default)]
14struct DomTreeNode {
15    children: PackedOption<Block>,
16    next: PackedOption<Block>,
17}
18
19impl DomTreeWithChildren {
20    pub(crate) fn new(func: &Function, domtree: &DominatorTree) -> DomTreeWithChildren {
21        let mut nodes: SecondaryMap<Block, DomTreeNode> =
22            SecondaryMap::with_capacity(func.dfg.num_blocks());
23
24        for block in func.layout.blocks() {
25            let idom_inst = match domtree.idom(block) {
26                Some(idom_inst) => idom_inst,
27                None => continue,
28            };
29            let idom = func
30                .layout
31                .inst_block(idom_inst)
32                .expect("Dominating instruction should be part of a block");
33
34            nodes[block].next = nodes[idom].children;
35            nodes[idom].children = block.into();
36        }
37
38        let root = func.layout.entry_block().unwrap();
39
40        Self { nodes, root }
41    }
42
43    pub(crate) fn root(&self) -> Block {
44        self.root
45    }
46
47    pub(crate) fn children<'a>(&'a self, block: Block) -> DomTreeChildIter<'a> {
48        let block = self.nodes[block].children;
49        DomTreeChildIter {
50            domtree: self,
51            block,
52        }
53    }
54}
55
56pub(crate) struct DomTreeChildIter<'a> {
57    domtree: &'a DomTreeWithChildren,
58    block: PackedOption<Block>,
59}
60
61impl<'a> Iterator for DomTreeChildIter<'a> {
62    type Item = Block;
63    fn next(&mut self) -> Option<Block> {
64        self.block.expand().map(|block| {
65            self.block = self.domtree.nodes[block].next;
66            block
67        })
68    }
69}