diff --git a/src/cfg/mod.rs b/src/cfg/mod.rs index 59c9aa0..8214103 100644 --- a/src/cfg/mod.rs +++ b/src/cfg/mod.rs @@ -3,7 +3,7 @@ // Borrowed from regalloc2's cfg.rs, which is also Apache-2.0 with // LLVM exception. -use crate::entity::PerEntity; +use crate::entity::{EntityRef, PerEntity}; use crate::ir::{Block, FunctionBody, Terminator}; use smallvec::SmallVec; @@ -26,6 +26,32 @@ pub struct CFGInfo { pub postorder_pos: PerEntity>, /// Domtree parents, indexed by block. pub domtree: PerEntity, + /// Domtree children. + pub domtree_children: PerEntity, +} + +#[derive(Clone, Debug, Default)] +pub struct DomtreeChildren { + pub child: Block, + pub next: Block, +} + +pub struct DomtreeChildIter<'a> { + domtree_children: &'a PerEntity, + block: Block, +} + +impl<'a> Iterator for DomtreeChildIter<'a> { + type Item = Block; + fn next(&mut self) -> Option { + if self.block.is_invalid() { + None + } else { + let block = self.block; + self.block = self.domtree_children[block].next; + Some(block) + } + } } impl CFGInfo { @@ -53,11 +79,17 @@ impl CFGInfo { postorder_pos[*block] = Some(i); } - let domtree = domtree::calculate( - |block| &&block_preds[block], - &postorder[..], - f.entry, - ); + let domtree = domtree::calculate(|block| &&block_preds[block], &postorder[..], f.entry); + + let mut domtree_children: PerEntity = PerEntity::default(); + for block in f.blocks.iter().rev() { + let idom = domtree[block]; + if idom.is_valid() { + let next = domtree_children[idom].child; + domtree_children[block].next = next; + domtree_children[idom].child = block; + } + } CFGInfo { entry: f.entry, @@ -67,6 +99,7 @@ impl CFGInfo { postorder, postorder_pos, domtree, + domtree_children, } } @@ -74,6 +107,13 @@ impl CFGInfo { domtree::dominates(&self.domtree, a, b) } + pub fn dom_children<'a>(&'a self, block: Block) -> DomtreeChildIter<'a> { + DomtreeChildIter { + domtree_children: &self.domtree_children, + block: self.domtree_children[block].child, + } + } + pub fn succs(&self, block: Block) -> &[Block] { &self.block_succs[block] } diff --git a/src/entity.rs b/src/entity.rs index 6189a39..f098594 100644 --- a/src/entity.rs +++ b/src/entity.rs @@ -90,7 +90,7 @@ impl EntityVec { self.0.len() } - pub fn iter(&self) -> impl Iterator { + pub fn iter(&self) -> impl DoubleEndedIterator { (0..self.0.len()).map(|index| Idx::new(index)) }