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