use crate::data_structures::BitSet;
use crate::ir::cfg::{BasicBlock, BasicBlockId, ControlFlowGraph, Successors};
pub struct Postorder {
visited: BitSet<BasicBlockId>,
visit_stack: Vec<(BasicBlockId, Successors)>,
root_is_start_block: bool,
}
impl<'lt> Postorder {
pub fn new(cfg: &ControlFlowGraph, root: BasicBlockId) -> Postorder {
let mut po = Postorder {
visited: BitSet::new_empty(cfg.blocks.len_idx()),
visit_stack: Vec::new(),
root_is_start_block: root == cfg.start(),
};
po.visited.insert(root);
po.visit_stack.push((root, cfg.successors(root)));
po.traverse_successor(cfg);
po
}
fn traverse_successor(&mut self, cfg: &ControlFlowGraph) {
while let Some(bb) = self
.visit_stack
.last_mut()
.and_then(|(_, iter)| iter.next())
{
if !self.visited.put(bb) {
self.visit_stack.push((bb, cfg[bb].terminator.successors()));
}
}
}
}
pub struct PostorderIter<'lt> {
pub base: Postorder,
pub cfg: &'lt ControlFlowGraph,
}
impl<'lt> PostorderIter<'lt> {
pub fn new(cfg: &'lt ControlFlowGraph, root: BasicBlockId) -> Self {
Self {
base: Postorder::new(cfg, root),
cfg,
}
}
}
impl<'lt> Iterator for PostorderIter<'lt> {
type Item = (BasicBlockId, &'lt BasicBlock);
fn next(&mut self) -> Option<(BasicBlockId, &'lt BasicBlock)> {
let next = self.base.visit_stack.pop();
if next.is_some() {
self.base.traverse_successor(self.cfg);
}
next.map(|(bb, _)| (bb, &self.cfg[bb]))
}
fn size_hint(&self) -> (usize, Option<usize>) {
let upper = self.cfg.blocks.len() - self.base.visited.ones().count();
let lower = if self.base.root_is_start_block {
upper
} else {
self.base.visit_stack.len()
};
(lower, Some(upper))
}
}
pub struct PostorderIterMut<'lt> {
base: Postorder,
cfg: &'lt mut ControlFlowGraph,
}
impl<'lt> PostorderIterMut<'lt> {
pub fn new(cfg: &'lt mut ControlFlowGraph, root: BasicBlockId) -> Self {
Self {
base: Postorder::new(cfg, root),
cfg,
}
}
}
impl<'lt> Iterator for PostorderIterMut<'lt> {
type Item = (BasicBlockId, &'lt mut BasicBlock);
fn next(&mut self) -> Option<(BasicBlockId, &'lt mut BasicBlock)> {
let next = self.base.visit_stack.pop();
if next.is_some() {
self.base.traverse_successor(self.cfg);
}
next.map(|(bb, _)| (bb, unsafe { std::mem::transmute(&mut self.cfg[bb]) }))
}
fn size_hint(&self) -> (usize, Option<usize>) {
let upper = self.cfg.blocks.len() - self.base.visited.ones().count();
let lower = if self.base.root_is_start_block {
upper
} else {
self.base.visit_stack.len()
};
(lower, Some(upper))
}
}
#[derive(Clone, Debug)]
pub struct ReversePostorder {
blocks: Vec<BasicBlockId>,
idx: usize,
}
impl ReversePostorder {
pub fn new(cfg: &ControlFlowGraph, root: BasicBlockId) -> Self {
let blocks: Vec<_> = PostorderIter::new(cfg, root).map(|(bb, _)| bb).collect();
let len = blocks.len();
ReversePostorder { blocks, idx: len }
}
pub fn reset(&mut self) {
self.idx = self.blocks.len();
}
}
impl Iterator for ReversePostorder {
type Item = BasicBlockId;
fn next(&mut self) -> Option<BasicBlockId> {
if self.idx == 0 {
return None;
}
self.idx -= 1;
self.blocks.get(self.idx).copied()
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.idx, Some(self.idx))
}
}
impl ExactSizeIterator for ReversePostorder {}
pub struct ReversePostorderIter<'lt> {
cfg: &'lt ControlFlowGraph,
base: ReversePostorder,
}
impl<'lt> ReversePostorderIter<'lt> {
pub fn new(cfg: &'lt ControlFlowGraph, root: BasicBlockId) -> Self {
Self {
base: ReversePostorder::new(cfg, root),
cfg,
}
}
}
impl<'lt> Iterator for ReversePostorderIter<'lt> {
type Item = (BasicBlockId, &'lt BasicBlock);
fn next(&mut self) -> Option<(BasicBlockId, &'lt BasicBlock)> {
self.base.next().map(|bb| (bb, &self.cfg[bb]))
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.base.size_hint()
}
}
pub struct ReversePostorderIterMut<'lt> {
cfg: &'lt mut ControlFlowGraph,
base: ReversePostorder,
}
impl<'lt> ReversePostorderIterMut<'lt> {
pub fn new(cfg: &'lt mut ControlFlowGraph, root: BasicBlockId) -> Self {
Self {
base: ReversePostorder::new(cfg, root),
cfg,
}
}
}
impl<'lt> Iterator for ReversePostorderIterMut<'lt> {
type Item = (BasicBlockId, &'lt mut BasicBlock);
fn next(&mut self) -> Option<(BasicBlockId, &'lt mut BasicBlock)> {
self.base
.next()
.map(|bb| (bb, unsafe { std::mem::transmute(&mut self.cfg[bb]) }))
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.base.size_hint()
}
}