use std::collections::{BTreeMap, BTreeSet};
use brink_format::DefinitionId;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CallGraph {
pub nodes: BTreeSet<DefinitionId>,
pub edges: BTreeMap<DefinitionId, BTreeSet<DefinitionId>>,
}
impl CallGraph {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn add_node(&mut self, def: DefinitionId) {
self.nodes.insert(def);
self.edges.entry(def).or_default();
}
pub fn add_edge(&mut self, from: DefinitionId, to: DefinitionId) {
self.add_node(from);
self.add_node(to);
self.edges.entry(from).or_default().insert(to);
}
fn reachable_forward(&self, start: DefinitionId) -> BTreeSet<DefinitionId> {
let mut seen = BTreeSet::new();
let mut stack = vec![start];
while let Some(n) = stack.pop() {
if seen.insert(n)
&& let Some(callees) = self.edges.get(&n)
{
stack.extend(callees.iter().copied());
}
}
seen
}
fn reachable_backward(&self, target: DefinitionId) -> BTreeSet<DefinitionId> {
let mut seen = BTreeSet::new();
let mut stack = vec![target];
while let Some(n) = stack.pop() {
if seen.insert(n) {
for (caller, callees) in &self.edges {
if callees.contains(&n) {
stack.push(*caller);
}
}
}
}
seen
}
}
#[must_use]
pub fn strongly_connected_components(graph: &CallGraph) -> Vec<BTreeSet<DefinitionId>> {
let mut assigned: BTreeSet<DefinitionId> = BTreeSet::new();
let mut components: Vec<BTreeSet<DefinitionId>> = Vec::new();
for &node in &graph.nodes {
if assigned.contains(&node) {
continue;
}
let fwd = graph.reachable_forward(node);
let bwd = graph.reachable_backward(node);
let component: BTreeSet<DefinitionId> = fwd.intersection(&bwd).copied().collect();
assigned.extend(component.iter().copied());
components.push(component);
}
components.sort_by_key(|c| c.iter().next().copied());
components
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct SccGraph {
pub order: Vec<BTreeSet<DefinitionId>>,
pub depends_on: BTreeMap<DefinitionId, BTreeSet<DefinitionId>>,
pub member_of: BTreeMap<DefinitionId, DefinitionId>,
}
#[must_use]
pub fn scc_graph(graph: &CallGraph) -> SccGraph {
let components = strongly_connected_components(graph);
if components.is_empty() {
return SccGraph::default();
}
let mut owner: BTreeMap<DefinitionId, usize> = BTreeMap::new();
for (idx, comp) in components.iter().enumerate() {
for &n in comp {
owner.insert(n, idx);
}
}
let component_key = |idx: usize| components[idx].iter().next().copied();
let mut depends_on_idx: BTreeMap<usize, BTreeSet<usize>> = BTreeMap::new();
let mut dependents_idx: BTreeMap<usize, BTreeSet<usize>> = BTreeMap::new();
for idx in 0..components.len() {
depends_on_idx.insert(idx, BTreeSet::new());
dependents_idx.insert(idx, BTreeSet::new());
}
for (caller, callees) in &graph.edges {
let Some(&from) = owner.get(caller) else {
continue;
};
for callee in callees {
let Some(&to) = owner.get(callee) else {
continue;
};
if from != to {
depends_on_idx.entry(from).or_default().insert(to);
dependents_idx.entry(to).or_default().insert(from);
}
}
}
let mut remaining = depends_on_idx.clone();
let mut order_idx: Vec<usize> = Vec::new();
loop {
let mut ready: Vec<usize> = remaining
.iter()
.filter(|(_, deps)| deps.is_empty())
.map(|(&idx, _)| idx)
.collect();
if ready.is_empty() {
break;
}
ready.sort_by_key(|&idx| component_key(idx));
for idx in ready {
if remaining.remove(&idx).is_none() {
continue;
}
order_idx.push(idx);
if let Some(deps) = dependents_idx.get(&idx) {
for &dep in deps {
if let Some(set) = remaining.get_mut(&dep) {
set.remove(&idx);
}
}
}
}
}
let mut leftover: Vec<usize> = remaining.keys().copied().collect();
leftover.sort_by_key(|&idx| component_key(idx));
order_idx.extend(leftover);
let order: Vec<BTreeSet<DefinitionId>> = order_idx
.into_iter()
.map(|idx| components[idx].clone())
.collect();
let mut depends_on: BTreeMap<DefinitionId, BTreeSet<DefinitionId>> = BTreeMap::new();
let mut member_of: BTreeMap<DefinitionId, DefinitionId> = BTreeMap::new();
for (idx, comp) in components.iter().enumerate() {
let Some(comp_id) = component_key(idx) else {
continue;
};
for &member in comp {
member_of.insert(member, comp_id);
}
let deps: BTreeSet<DefinitionId> = depends_on_idx
.get(&idx)
.into_iter()
.flatten()
.filter_map(|&dep_idx| component_key(dep_idx))
.collect();
depends_on.insert(comp_id, deps);
}
SccGraph {
order,
depends_on,
member_of,
}
}
#[must_use]
pub fn topo_order(graph: &CallGraph) -> Vec<BTreeSet<DefinitionId>> {
scc_graph(graph).order
}
#[cfg(test)]
mod tests {
use super::*;
use brink_format::DefinitionTag;
fn def(n: u64) -> DefinitionId {
DefinitionId::new(DefinitionTag::Address, n)
}
#[test]
fn empty_graph_has_no_components() {
let g = CallGraph::new();
assert!(topo_order(&g).is_empty());
}
#[test]
fn isolated_node_is_its_own_component() {
let mut g = CallGraph::new();
g.add_node(def(1));
let sccs = strongly_connected_components(&g);
assert_eq!(sccs, vec![BTreeSet::from([def(1)])]);
}
#[test]
fn direct_recursion_is_a_singleton_component() {
let mut g = CallGraph::new();
g.add_edge(def(1), def(1));
let sccs = strongly_connected_components(&g);
assert_eq!(sccs, vec![BTreeSet::from([def(1)])]);
}
#[test]
fn linear_chain_orders_callee_before_caller() {
let mut g = CallGraph::new();
g.add_edge(def(1), def(2));
g.add_edge(def(2), def(3));
let order = topo_order(&g);
let flat: Vec<DefinitionId> = order.into_iter().flatten().collect();
let pos = |d: DefinitionId| flat.iter().position(|&x| x == d).expect("present");
assert!(pos(def(3)) < pos(def(2)), "callee c before caller b");
assert!(pos(def(2)) < pos(def(1)), "callee b before caller a");
}
#[test]
fn mutual_recursion_is_one_component() {
let mut g = CallGraph::new();
g.add_edge(def(1), def(2));
g.add_edge(def(2), def(1));
g.add_edge(def(3), def(1));
let sccs = strongly_connected_components(&g);
let ab = sccs
.iter()
.find(|c| c.contains(&def(1)))
.expect("component containing a");
assert_eq!(
ab,
&BTreeSet::from([def(1), def(2)]),
"a and b fold into one SCC"
);
let order = topo_order(&g);
let ab_idx = order
.iter()
.position(|c| c.contains(&def(1)))
.expect("ab component present");
let c_idx = order
.iter()
.position(|c| c.contains(&def(3)))
.expect("c component present");
assert!(
ab_idx < c_idx,
"the mutually-recursive pair solves before its caller"
);
}
#[test]
fn disconnected_components_both_appear() {
let mut g = CallGraph::new();
g.add_edge(def(1), def(2));
g.add_edge(def(10), def(20));
let order = topo_order(&g);
let flat: BTreeSet<DefinitionId> = order.into_iter().flatten().collect();
assert_eq!(flat, BTreeSet::from([def(1), def(2), def(10), def(20)]));
}
#[test]
fn scc_graph_order_matches_topo_order() {
let mut g = CallGraph::new();
g.add_edge(def(1), def(2));
g.add_edge(def(2), def(3));
g.add_edge(def(4), def(1));
assert_eq!(scc_graph(&g).order, topo_order(&g));
}
#[test]
fn scc_graph_member_of_maps_every_node_to_its_component_id() {
let mut g = CallGraph::new();
g.add_edge(def(1), def(2));
g.add_edge(def(2), def(1));
g.add_edge(def(3), def(1));
let sg = scc_graph(&g);
assert_eq!(sg.member_of.get(&def(1)), Some(&def(1)));
assert_eq!(sg.member_of.get(&def(2)), Some(&def(1)));
assert_eq!(sg.member_of.get(&def(3)), Some(&def(3)));
}
#[test]
fn scc_graph_depends_on_is_the_condensation_adjacency() {
let mut g = CallGraph::new();
g.add_edge(def(1), def(2));
g.add_edge(def(2), def(3));
let sg = scc_graph(&g);
assert_eq!(sg.depends_on.get(&def(1)), Some(&BTreeSet::from([def(2)])));
assert_eq!(sg.depends_on.get(&def(2)), Some(&BTreeSet::from([def(3)])));
assert_eq!(sg.depends_on.get(&def(3)), Some(&BTreeSet::new()));
}
#[test]
fn scc_graph_depends_on_never_names_the_components_own_id() {
let mut g = CallGraph::new();
g.add_edge(def(1), def(2));
g.add_edge(def(2), def(1));
let sg = scc_graph(&g);
assert_eq!(sg.depends_on.get(&def(1)), Some(&BTreeSet::new()));
}
#[test]
fn scc_graph_empty_graph_is_empty() {
let g = CallGraph::new();
assert_eq!(scc_graph(&g), SccGraph::default());
}
}