fluidattacks_blends_domain/stack/
graph.rs1use alloc::collections::BTreeMap;
2use alloc::vec::Vec;
3
4use crate::stack::StackNode;
5use crate::NodeId;
6
7#[derive(Clone, PartialEq, Eq, Debug, Default)]
9#[non_exhaustive]
10pub struct StackGraph {
11 pub nodes: BTreeMap<NodeId, StackNode>,
13 pub edges: BTreeMap<NodeId, BTreeMap<NodeId, i64>>,
15}
16
17impl StackGraph {
18 #[must_use]
20 pub fn new() -> Self {
21 Self::default()
22 }
23
24 pub fn add_node(&mut self, id: NodeId, node: StackNode) {
26 self.nodes.insert(id, node);
27 }
28
29 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 #[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 #[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}