use crate::cfg::analysis::find_entry;
use crate::cfg::{BlockId, Cfg};
use petgraph::algo::dominators::simple_fast;
use petgraph::graph::NodeIndex;
use std::collections::HashMap;
#[derive(Debug, Clone)]
pub struct DominatorTree {
root: NodeIndex,
immediate_dominator: HashMap<NodeIndex, Option<NodeIndex>>,
children: HashMap<NodeIndex, Vec<NodeIndex>>,
}
impl DominatorTree {
pub fn new(cfg: &Cfg) -> Option<Self> {
let entry = find_entry(cfg)?;
let dominators = simple_fast(cfg, entry);
let mut immediate_dominator = HashMap::new();
let mut children: HashMap<NodeIndex, Vec<NodeIndex>> = HashMap::new();
for node in cfg.node_indices() {
let idom = dominators.immediate_dominator(node);
immediate_dominator.insert(node, idom);
if let Some(parent) = idom {
children.entry(parent).or_default().push(node);
}
}
Some(Self {
root: entry,
immediate_dominator,
children,
})
}
pub fn root(&self) -> NodeIndex {
self.root
}
pub fn immediate_dominator(&self, node: NodeIndex) -> Option<NodeIndex> {
self.immediate_dominator.get(&node).copied().flatten()
}
pub fn dominates(&self, a: NodeIndex, b: NodeIndex) -> bool {
if a == b {
return true; }
let mut current = b;
while let Some(idom) = self.immediate_dominator(current) {
if idom == a {
return true;
}
current = idom;
}
false
}
pub fn children(&self, node: NodeIndex) -> &[NodeIndex] {
self.children.get(&node).map_or(&[], |v| v.as_slice())
}
pub fn strictly_dominates(&self, a: NodeIndex, b: NodeIndex) -> bool {
a != b && self.dominates(a, b)
}
pub fn dominators(&self, node: NodeIndex) -> Dominators<'_> {
Dominators {
tree: self,
current: Some(node),
}
}
pub fn common_dominator(&self, a: NodeIndex, b: NodeIndex) -> Option<NodeIndex> {
let a_doms: std::collections::HashSet<NodeIndex> = self.dominators(a).collect();
for dom in self.dominators(b) {
if a_doms.contains(&dom) {
return Some(dom);
}
}
None
}
pub fn depth(&self, node: NodeIndex) -> usize {
let mut depth = 0;
let mut current = node;
while let Some(idom) = self.immediate_dominator(current) {
depth += 1;
current = idom;
}
depth
}
pub fn apply_dominator_depths(&self, cfg: &mut Cfg) {
for node in cfg.node_indices() {
let depth = self.depth(node) as i64;
if let Some(block) = cfg.node_weight_mut(node) {
block.coord_x = depth;
}
}
}
pub fn new_with_depths(cfg: &mut Cfg) -> Option<Self> {
let dom_tree = Self::new(cfg)?;
dom_tree.apply_dominator_depths(cfg);
Some(dom_tree)
}
pub(crate) fn from_parts(
root: NodeIndex,
immediate_dominator: HashMap<NodeIndex, Option<NodeIndex>>,
children: HashMap<NodeIndex, Vec<NodeIndex>>,
) -> Self {
Self {
root,
immediate_dominator,
children,
}
}
}
pub struct Dominators<'a> {
tree: &'a DominatorTree,
current: Option<NodeIndex>,
}
impl<'a> Iterator for Dominators<'a> {
type Item = NodeIndex;
fn next(&mut self) -> Option<Self::Item> {
let node = self.current?;
self.current = self.tree.immediate_dominator(node);
Some(node)
}
}
pub fn compute_dominator_tree(cfg: &Cfg) -> Option<DominatorTree> {
DominatorTree::new(cfg)
}
pub fn immediate_dominator_id(
tree: &DominatorTree,
block_id: BlockId,
cfg: &Cfg,
) -> Option<BlockId> {
let node = node_from_id(cfg, block_id)?;
let idom_node = tree.immediate_dominator(node)?;
Some(cfg[idom_node].id)
}
fn node_from_id(cfg: &Cfg, block_id: BlockId) -> Option<NodeIndex> {
cfg.node_indices().find(|&n| cfg[n].id == block_id)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cfg::{BasicBlock, BlockKind, EdgeType, Terminator};
use petgraph::graph::DiGraph;
fn create_diamond_cfg() -> Cfg {
let mut g = DiGraph::new();
let b0 = g.add_node(BasicBlock {
id: 0,
db_id: None,
kind: BlockKind::Entry,
statements: vec![],
terminator: Terminator::SwitchInt {
targets: vec![1],
otherwise: 2,
},
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b1 = g.add_node(BasicBlock {
id: 1,
db_id: None,
kind: BlockKind::Normal,
statements: vec!["branch 1".to_string()],
terminator: Terminator::Goto { target: 3 },
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b2 = g.add_node(BasicBlock {
id: 2,
db_id: None,
kind: BlockKind::Normal,
statements: vec!["branch 2".to_string()],
terminator: Terminator::Goto { target: 3 },
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b3 = g.add_node(BasicBlock {
id: 3,
db_id: None,
kind: BlockKind::Exit,
statements: vec![],
terminator: Terminator::Return,
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
g.add_edge(b0, b1, EdgeType::TrueBranch);
g.add_edge(b0, b2, EdgeType::FalseBranch);
g.add_edge(b1, b3, EdgeType::Fallthrough);
g.add_edge(b2, b3, EdgeType::Fallthrough);
g
}
#[test]
fn test_dominator_tree_construction() {
let cfg = create_diamond_cfg();
let dom_tree = DominatorTree::new(&cfg).expect("CFG has entry");
assert_eq!(dom_tree.root(), NodeIndex::new(0));
assert_eq!(dom_tree.immediate_dominator(NodeIndex::new(0)), None);
assert_eq!(
dom_tree.immediate_dominator(NodeIndex::new(1)),
Some(NodeIndex::new(0))
);
assert_eq!(
dom_tree.immediate_dominator(NodeIndex::new(2)),
Some(NodeIndex::new(0))
);
assert_eq!(
dom_tree.immediate_dominator(NodeIndex::new(3)),
Some(NodeIndex::new(0))
);
}
#[test]
fn test_dominates() {
let cfg = create_diamond_cfg();
let dom_tree = DominatorTree::new(&cfg).expect("CFG has entry");
let entry = NodeIndex::new(0);
let node1 = NodeIndex::new(1);
let node3 = NodeIndex::new(3);
assert!(dom_tree.dominates(entry, entry));
assert!(dom_tree.dominates(entry, node1));
assert!(dom_tree.dominates(entry, node3));
assert!(!dom_tree.dominates(node1, entry));
assert!(dom_tree.dominates(node1, node1));
assert!(dom_tree.dominates(node3, node3));
}
#[test]
fn test_children() {
let cfg = create_diamond_cfg();
let dom_tree = DominatorTree::new(&cfg).expect("CFG has entry");
let entry = NodeIndex::new(0);
let children = dom_tree.children(entry);
assert_eq!(children.len(), 3);
assert!(children.contains(&NodeIndex::new(1)));
assert!(children.contains(&NodeIndex::new(2)));
assert!(children.contains(&NodeIndex::new(3)));
}
#[test]
fn test_strictly_dominates() {
let cfg = create_diamond_cfg();
let dom_tree = DominatorTree::new(&cfg).expect("CFG has entry");
let entry = NodeIndex::new(0);
let node1 = NodeIndex::new(1);
assert!(dom_tree.strictly_dominates(entry, node1));
assert!(!dom_tree.strictly_dominates(entry, entry));
}
#[test]
fn test_dominators_iterator() {
let cfg = create_diamond_cfg();
let dom_tree = DominatorTree::new(&cfg).expect("CFG has entry");
let node3 = NodeIndex::new(3);
let doms: Vec<_> = dom_tree.dominators(node3).collect();
assert_eq!(doms.len(), 2);
assert_eq!(doms[0], node3);
assert_eq!(doms[1], NodeIndex::new(0));
}
#[test]
fn test_common_dominator() {
let cfg = create_diamond_cfg();
let dom_tree = DominatorTree::new(&cfg).expect("CFG has entry");
let node1 = NodeIndex::new(1);
let node2 = NodeIndex::new(2);
let entry = NodeIndex::new(0);
assert_eq!(dom_tree.common_dominator(node1, node2), Some(entry));
assert_eq!(dom_tree.common_dominator(node1, node1), Some(node1));
}
#[test]
fn test_depth() {
let cfg = create_diamond_cfg();
let dom_tree = DominatorTree::new(&cfg).expect("CFG has entry");
assert_eq!(dom_tree.depth(NodeIndex::new(0)), 0);
assert_eq!(dom_tree.depth(NodeIndex::new(1)), 1);
assert_eq!(dom_tree.depth(NodeIndex::new(2)), 1);
assert_eq!(dom_tree.depth(NodeIndex::new(3)), 1);
}
#[test]
fn test_empty_cfg() {
let cfg: Cfg = DiGraph::new();
assert!(DominatorTree::new(&cfg).is_none());
}
#[test]
fn test_linear_cfg() {
let mut g = DiGraph::new();
let b0 = g.add_node(BasicBlock {
id: 0,
db_id: None,
kind: BlockKind::Entry,
statements: vec![],
terminator: Terminator::Goto { target: 1 },
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b1 = g.add_node(BasicBlock {
id: 1,
db_id: None,
kind: BlockKind::Normal,
statements: vec![],
terminator: Terminator::Goto { target: 2 },
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b2 = g.add_node(BasicBlock {
id: 2,
db_id: None,
kind: BlockKind::Normal,
statements: vec![],
terminator: Terminator::Goto { target: 3 },
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b3 = g.add_node(BasicBlock {
id: 3,
db_id: None,
kind: BlockKind::Exit,
statements: vec![],
terminator: Terminator::Return,
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
g.add_edge(b0, b1, EdgeType::Fallthrough);
g.add_edge(b1, b2, EdgeType::Fallthrough);
g.add_edge(b2, b3, EdgeType::Fallthrough);
let dom_tree = DominatorTree::new(&g).expect("CFG has entry");
assert_eq!(dom_tree.immediate_dominator(b0), None);
assert_eq!(dom_tree.immediate_dominator(b1), Some(b0));
assert_eq!(dom_tree.immediate_dominator(b2), Some(b1));
assert_eq!(dom_tree.immediate_dominator(b3), Some(b2));
}
#[test]
fn test_apply_dominator_depths_diamond_cfg() {
let mut cfg = create_diamond_cfg();
let dom_tree = DominatorTree::new(&cfg).expect("CFG has entry");
dom_tree.apply_dominator_depths(&mut cfg);
assert_eq!(
cfg[NodeIndex::new(0)].coord_x,
0,
"Entry should have depth 0"
);
assert_eq!(
cfg[NodeIndex::new(1)].coord_x,
1,
"Node 1 should have depth 1"
);
assert_eq!(
cfg[NodeIndex::new(2)].coord_x,
1,
"Node 2 should have depth 1"
);
assert_eq!(
cfg[NodeIndex::new(3)].coord_x,
1,
"Node 3 should have depth 1"
);
}
#[test]
fn test_apply_dominator_depths_linear_cfg() {
let mut g = DiGraph::new();
let b0 = g.add_node(BasicBlock {
id: 0,
db_id: None,
kind: BlockKind::Entry,
statements: vec![],
terminator: Terminator::Goto { target: 1 },
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b1 = g.add_node(BasicBlock {
id: 1,
db_id: None,
kind: BlockKind::Normal,
statements: vec![],
terminator: Terminator::Goto { target: 2 },
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b2 = g.add_node(BasicBlock {
id: 2,
db_id: None,
kind: BlockKind::Normal,
statements: vec![],
terminator: Terminator::Goto { target: 3 },
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b3 = g.add_node(BasicBlock {
id: 3,
db_id: None,
kind: BlockKind::Exit,
statements: vec![],
terminator: Terminator::Return,
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
g.add_edge(b0, b1, EdgeType::Fallthrough);
g.add_edge(b1, b2, EdgeType::Fallthrough);
g.add_edge(b2, b3, EdgeType::Fallthrough);
let dom_tree = DominatorTree::new(&g).expect("CFG has entry");
dom_tree.apply_dominator_depths(&mut g);
assert_eq!(g[b0].coord_x, 0, "Entry should have depth 0");
assert_eq!(g[b1].coord_x, 1, "Node 1 should have depth 1");
assert_eq!(g[b2].coord_x, 2, "Node 2 should have depth 2");
assert_eq!(g[b3].coord_x, 3, "Node 3 should have depth 3");
}
#[test]
fn test_new_with_depths_diamond_cfg() {
let mut cfg = create_diamond_cfg();
let dom_tree = DominatorTree::new_with_depths(&mut cfg).expect("CFG has entry");
assert_eq!(dom_tree.root(), NodeIndex::new(0));
assert_eq!(
cfg[NodeIndex::new(0)].coord_x,
0,
"Entry should have depth 0"
);
assert_eq!(
cfg[NodeIndex::new(1)].coord_x,
1,
"Node 1 should have depth 1"
);
assert_eq!(
cfg[NodeIndex::new(2)].coord_x,
1,
"Node 2 should have depth 1"
);
assert_eq!(
cfg[NodeIndex::new(3)].coord_x,
1,
"Node 3 should have depth 1"
);
}
#[test]
fn test_depth_method_matches_coord_x() {
let mut cfg = create_diamond_cfg();
let dom_tree = DominatorTree::new(&cfg).expect("CFG has entry");
dom_tree.apply_dominator_depths(&mut cfg);
for node in cfg.node_indices() {
let expected_depth = dom_tree.depth(node) as i64;
let actual_coord_x = cfg[node].coord_x;
assert_eq!(
actual_coord_x, expected_depth,
"Node {:?} coord_x should match depth()",
node
);
}
}
#[test]
fn test_complex_nested_dominator_structure() {
let mut g = DiGraph::new();
let b0 = g.add_node(BasicBlock {
id: 0,
db_id: None,
kind: BlockKind::Entry,
statements: vec![],
terminator: Terminator::Goto { target: 1 },
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b1 = g.add_node(BasicBlock {
id: 1,
db_id: None,
kind: BlockKind::Normal,
statements: vec![],
terminator: Terminator::Goto { target: 2 },
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b2 = g.add_node(BasicBlock {
id: 2,
db_id: None,
kind: BlockKind::Normal,
statements: vec![],
terminator: Terminator::SwitchInt {
targets: vec![3],
otherwise: 4,
},
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b3 = g.add_node(BasicBlock {
id: 3,
db_id: None,
kind: BlockKind::Normal,
statements: vec![],
terminator: Terminator::Goto { target: 5 },
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b4 = g.add_node(BasicBlock {
id: 4,
db_id: None,
kind: BlockKind::Normal,
statements: vec![],
terminator: Terminator::Goto { target: 5 },
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
let b5 = g.add_node(BasicBlock {
id: 5,
db_id: None,
kind: BlockKind::Exit,
statements: vec![],
terminator: Terminator::Return,
source_location: None,
coord_x: 0,
coord_y: 0,
coord_z: 0,
});
g.add_edge(b0, b1, EdgeType::Fallthrough);
g.add_edge(b1, b2, EdgeType::Fallthrough);
g.add_edge(b2, b3, EdgeType::TrueBranch);
g.add_edge(b2, b4, EdgeType::FalseBranch);
g.add_edge(b3, b5, EdgeType::Fallthrough);
g.add_edge(b4, b5, EdgeType::Fallthrough);
let dom_tree = DominatorTree::new(&g).expect("CFG has entry");
dom_tree.apply_dominator_depths(&mut g);
assert_eq!(g[b0].coord_x, 0, "Entry should have depth 0");
assert_eq!(g[b1].coord_x, 1, "Node 1 should have depth 1");
assert_eq!(g[b2].coord_x, 2, "Node 2 should have depth 2");
assert_eq!(g[b3].coord_x, 3, "Node 3 should have depth 3");
assert_eq!(g[b4].coord_x, 3, "Node 4 should have depth 3");
assert_eq!(g[b5].coord_x, 3, "Node 5 should have depth 3");
}
}