use alloc::collections::BTreeMap;
use alloc::vec::Vec;
use crate::stack::StackNode;
use crate::NodeId;
#[derive(Clone, PartialEq, Eq, Debug, Default)]
pub struct StackGraph {
pub nodes: BTreeMap<NodeId, StackNode>,
pub edges: BTreeMap<NodeId, BTreeMap<NodeId, i64>>,
}
impl StackGraph {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn add_node(&mut self, id: NodeId, node: StackNode) {
self.nodes.insert(id, node);
}
pub fn add_edge(&mut self, source: NodeId, sink: NodeId, precedence: i64) {
self.edges
.entry(source)
.or_default()
.entry(sink)
.and_modify(|existing| *existing = (*existing).max(precedence))
.or_insert(precedence);
}
#[must_use]
pub fn successors(&self, id: NodeId) -> Vec<(NodeId, i64)> {
self.edges
.get(&id)
.map(|adjacent| {
adjacent
.iter()
.map(|(sink, precedence)| (*sink, *precedence))
.collect()
})
.unwrap_or_default()
}
#[must_use]
pub fn predecessors(&self, id: NodeId) -> Vec<NodeId> {
self.edges
.iter()
.filter(|(_, adjacent)| adjacent.contains_key(&id))
.map(|(source, _)| *source)
.collect()
}
}
#[cfg(test)]
mod tests {
use super::StackGraph;
use crate::stack::{StackNode, StackScope};
use crate::NodeId;
use alloc::vec;
#[test]
fn add_edge_keeps_the_precedence_of_a_new_edge() {
let mut graph = StackGraph::new();
graph.add_edge(NodeId(1), NodeId(2), 1);
assert_eq!(graph.successors(NodeId(1)), vec![(NodeId(2), 1)]);
}
#[test]
fn add_edge_keeps_the_higher_precedence_of_a_repeated_edge() {
let mut graph = StackGraph::new();
graph.add_edge(NodeId(1), NodeId(2), 1);
graph.add_edge(NodeId(1), NodeId(2), 0);
graph.add_edge(NodeId(3), NodeId(4), 0);
graph.add_edge(NodeId(3), NodeId(4), 2);
assert_eq!(graph.successors(NodeId(1)), vec![(NodeId(2), 1)]);
assert_eq!(graph.successors(NodeId(3)), vec![(NodeId(4), 2)]);
}
#[test]
fn successors_follow_sink_order() {
let mut graph = StackGraph::new();
graph.add_edge(NodeId(1), NodeId(9), 0);
graph.add_edge(NodeId(1), NodeId(3), 1);
assert_eq!(
graph.successors(NodeId(1)),
vec![(NodeId(3), 1), (NodeId(9), 0)]
);
assert_eq!(graph.successors(NodeId(9)), vec![]);
}
#[test]
fn predecessors_find_every_source_of_a_sink() {
let mut graph = StackGraph::new();
graph.add_edge(NodeId(5), NodeId(1), 0);
graph.add_edge(NodeId(2), NodeId(1), 1);
graph.add_edge(NodeId(2), NodeId(7), 0);
assert_eq!(graph.predecessors(NodeId(1)), vec![NodeId(2), NodeId(5)]);
assert_eq!(graph.predecessors(NodeId(5)), vec![]);
}
#[test]
fn add_node_replaces_the_node_under_an_id() {
let mut graph = StackGraph::new();
graph.add_node(NodeId(1), StackNode::Root);
graph.add_node(NodeId(1), StackNode::Scope(StackScope::default()));
assert_eq!(
graph.nodes.get(&NodeId(1)),
Some(&StackNode::Scope(StackScope::default()))
);
assert_eq!(graph.nodes.len(), 1);
}
}