Skip to main content

fluidattacks_blends_domain/stack/
graph.rs

1use alloc::collections::BTreeMap;
2use alloc::vec::Vec;
3
4use crate::stack::StackNode;
5use crate::NodeId;
6
7/// The stack graph: every push/pop/scope node plus the paths between them.
8#[derive(Clone, PartialEq, Eq, Debug, Default)]
9#[non_exhaustive]
10pub struct StackGraph {
11    /// Every node, keyed by id.
12    pub nodes: BTreeMap<NodeId, StackNode>,
13    /// Adjacency: source id -> sink id -> the edge's precedence.
14    pub edges: BTreeMap<NodeId, BTreeMap<NodeId, i64>>,
15}
16
17impl StackGraph {
18    /// An empty stack graph.
19    #[must_use]
20    pub fn new() -> Self {
21        Self::default()
22    }
23
24    /// Inserts `node` under `id`, replacing any node already there.
25    pub fn add_node(&mut self, id: NodeId, node: StackNode) {
26        self.nodes.insert(id, node);
27    }
28
29    /// Connects `source` to `sink`, keeping the higher precedence on repeat.
30    pub fn add_edge(&mut self, source: NodeId, sink: NodeId, precedence: i64) {
31        self.edges
32            .entry(source)
33            .or_default()
34            .entry(sink)
35            .and_modify(|existing| *existing = (*existing).max(precedence))
36            .or_insert(precedence);
37    }
38
39    /// `id`'s outgoing edges, as (sink, precedence) pairs.
40    #[must_use]
41    pub fn successors(&self, id: NodeId) -> Vec<(NodeId, i64)> {
42        self.edges
43            .get(&id)
44            .map(|adjacent| {
45                adjacent
46                    .iter()
47                    .map(|(sink, precedence)| (*sink, *precedence))
48                    .collect()
49            })
50            .unwrap_or_default()
51    }
52
53    /// Ids of every node with an edge into `id`.
54    #[must_use]
55    pub fn predecessors(&self, id: NodeId) -> Vec<NodeId> {
56        self.edges
57            .iter()
58            .filter(|(_, adjacent)| adjacent.contains_key(&id))
59            .map(|(source, _)| *source)
60            .collect()
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::StackGraph;
67    use crate::stack::{StackNode, StackScope};
68    use crate::NodeId;
69    use alloc::vec;
70
71    #[test]
72    fn add_edge_keeps_the_precedence_of_a_new_edge() {
73        let mut graph = StackGraph::new();
74        graph.add_edge(NodeId(1), NodeId(2), 1);
75
76        assert_eq!(graph.successors(NodeId(1)), vec![(NodeId(2), 1)]);
77    }
78
79    #[test]
80    fn add_edge_keeps_the_higher_precedence_of_a_repeated_edge() {
81        let mut graph = StackGraph::new();
82        graph.add_edge(NodeId(1), NodeId(2), 1);
83        graph.add_edge(NodeId(1), NodeId(2), 0);
84        graph.add_edge(NodeId(3), NodeId(4), 0);
85        graph.add_edge(NodeId(3), NodeId(4), 2);
86
87        assert_eq!(graph.successors(NodeId(1)), vec![(NodeId(2), 1)]);
88        assert_eq!(graph.successors(NodeId(3)), vec![(NodeId(4), 2)]);
89    }
90
91    #[test]
92    fn successors_follow_sink_order() {
93        let mut graph = StackGraph::new();
94        graph.add_edge(NodeId(1), NodeId(9), 0);
95        graph.add_edge(NodeId(1), NodeId(3), 1);
96
97        assert_eq!(
98            graph.successors(NodeId(1)),
99            vec![(NodeId(3), 1), (NodeId(9), 0)]
100        );
101        assert_eq!(graph.successors(NodeId(9)), vec![]);
102    }
103
104    #[test]
105    fn predecessors_find_every_source_of_a_sink() {
106        let mut graph = StackGraph::new();
107        graph.add_edge(NodeId(5), NodeId(1), 0);
108        graph.add_edge(NodeId(2), NodeId(1), 1);
109        graph.add_edge(NodeId(2), NodeId(7), 0);
110
111        assert_eq!(graph.predecessors(NodeId(1)), vec![NodeId(2), NodeId(5)]);
112        assert_eq!(graph.predecessors(NodeId(5)), vec![]);
113    }
114
115    #[test]
116    fn add_node_replaces_the_node_under_an_id() {
117        let mut graph = StackGraph::new();
118        graph.add_node(NodeId(1), StackNode::Root);
119        graph.add_node(NodeId(1), StackNode::Scope(StackScope::default()));
120
121        assert_eq!(
122            graph.nodes.get(&NodeId(1)),
123            Some(&StackNode::Scope(StackScope::default()))
124        );
125        assert_eq!(graph.nodes.len(), 1);
126    }
127}