wasmtime/cranelift/codegen/src/egraph/domtree.rs

//! 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)]
pub(crate) struct DomTreeWithChildren {
    nodes: SecondaryMap<Block, DomTreeNode>,
    root: Block,
}

#[derive(Clone, Copy, Debug, Default)]
struct DomTreeNode {
    children: PackedOption<Block>,
    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 idom_inst = match domtree.idom(block) {
                Some(idom_inst) => idom_inst,
                None => continue,
            };
            let idom = func
                .layout
                .inst_block(idom_inst)
                .expect("Dominating instruction should be part of a block");

            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
        })
    }
}