turbo_tasks/graph/
adjacency_map.rs

1use std::{
2    collections::{VecDeque, hash_map::Entry},
3    hash::Hash,
4};
5
6use rustc_hash::{FxHashMap, FxHashSet};
7use turbo_tasks_macros::{TraceRawVcs, ValueDebugFormat};
8
9use crate::{self as turbo_tasks, NonLocalValue, graph::graph_store::GraphStore};
10
11/// A graph traversal that builds an adjacency map
12#[derive(Debug, Clone, TraceRawVcs, ValueDebugFormat)]
13pub struct AdjacencyMap<T, E> {
14    adjacency_map: FxHashMap<T, Vec<(T, E)>>,
15    roots: Vec<T>,
16}
17
18unsafe impl<T, E> NonLocalValue for AdjacencyMap<T, E>
19where
20    T: NonLocalValue,
21    E: NonLocalValue,
22{
23}
24impl<T, E> PartialEq for AdjacencyMap<T, E>
25where
26    T: Eq + Hash,
27    E: Eq,
28{
29    fn eq(&self, other: &Self) -> bool {
30        self.adjacency_map == other.adjacency_map && self.roots == other.roots
31    }
32}
33
34impl<T, E> Eq for AdjacencyMap<T, E>
35where
36    T: Eq + Hash,
37    E: Eq,
38{
39}
40
41impl<T, E> Default for AdjacencyMap<T, E>
42where
43    T: Eq + Hash + Clone,
44{
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50impl<T, E> AdjacencyMap<T, E>
51where
52    T: Eq + Hash + Clone,
53{
54    /// Creates a new adjacency map
55    pub fn new() -> Self {
56        Self {
57            adjacency_map: FxHashMap::default(),
58            roots: Vec::new(),
59        }
60    }
61
62    /// Returns an iterator over the root nodes of the graph
63    pub fn roots(&self) -> impl Iterator<Item = &T> {
64        self.roots.iter()
65    }
66
67    /// Returns an iterator over the children of the given node
68    pub fn get(&self, node: &T) -> Option<impl Iterator<Item = &(T, E)>> {
69        self.adjacency_map.get(node).map(|vec| vec.iter())
70    }
71
72    /// Returns the number of nodes in the graph
73    pub fn len(&self) -> usize {
74        self.adjacency_map.len()
75    }
76
77    /// Returns true if the graph is empty
78    pub fn is_empty(&self) -> bool {
79        self.adjacency_map.is_empty()
80    }
81
82    pub fn reversed(&self) -> AdjacencyMap<T, &E> {
83        let mut reversed = AdjacencyMap::new();
84
85        for root in &self.roots {
86            reversed.roots.push(root.clone());
87        }
88
89        for (from, edges) in &self.adjacency_map {
90            for (to, edge) in edges {
91                let vec = reversed
92                    .adjacency_map
93                    .entry(to.clone())
94                    .or_insert_with(|| Vec::with_capacity(1));
95                vec.push((from.clone(), edge));
96            }
97        }
98
99        reversed
100    }
101}
102
103impl<T, E> GraphStore for AdjacencyMap<T, E>
104where
105    T: Eq + Hash + Clone + Send,
106    E: Send,
107{
108    type Node = T;
109    type Edge = E;
110    type Handle = T;
111
112    fn insert(&mut self, from: Option<(&T, E)>, node: T) {
113        if let Some((from_node, edge)) = from {
114            let vec = self
115                .adjacency_map
116                .entry(from_node.clone())
117                .or_insert_with(|| Vec::with_capacity(1));
118            vec.push((node, edge));
119        } else {
120            self.roots.push(node);
121        };
122    }
123
124    fn try_enter(&mut self, node: &T) -> Option<T> {
125        match self.adjacency_map.entry(node.clone()) {
126            Entry::Occupied(_) => None,
127            Entry::Vacant(e) => {
128                e.insert(Vec::new());
129                Some(node.clone())
130            }
131        }
132    }
133}
134
135impl<T, E> AdjacencyMap<T, E>
136where
137    T: Eq + Hash + Clone,
138{
139    /// Returns an owned iterator over the nodes in postorder topological order,
140    /// starting from the roots.
141    pub fn into_postorder_topological(self) -> IntoPostorderTopologicalIter<T, E> {
142        IntoPostorderTopologicalIter {
143            adjacency_map: self.adjacency_map,
144            stack: self
145                .roots
146                .into_iter()
147                .rev()
148                .map(|root| (ReverseTopologicalPass::Pre, root))
149                .collect(),
150            visited: FxHashSet::default(),
151        }
152    }
153
154    /// Returns an owned iterator over all edges (node pairs) in reverse breadth first order,
155    /// starting from the roots.
156    pub fn into_breadth_first_edges(self) -> IntoBreadthFirstEdges<T, E> {
157        IntoBreadthFirstEdges {
158            adjacency_map: self.adjacency_map,
159            queue: self.roots.into_iter().map(|root| (None, root)).collect(),
160            expanded: FxHashSet::default(),
161        }
162    }
163
164    /// Returns an iterator over the nodes in postorder topological order,
165    /// starting from the roots.
166    pub fn postorder_topological(&self) -> PostorderTopologicalIter<'_, T, E> {
167        PostorderTopologicalIter {
168            adjacency_map: &self.adjacency_map,
169            stack: self
170                .roots
171                .iter()
172                .rev()
173                .map(|root| (ReverseTopologicalPass::Pre, root))
174                .collect(),
175            visited: FxHashSet::default(),
176        }
177    }
178}
179
180#[derive(Debug)]
181enum ReverseTopologicalPass {
182    Pre,
183    Post,
184}
185
186/// An iterator over the nodes of a graph in postorder topological order, starting
187/// from the roots.
188pub struct IntoPostorderTopologicalIter<T, E>
189where
190    T: Eq + Hash + Clone,
191{
192    adjacency_map: FxHashMap<T, Vec<(T, E)>>,
193    stack: Vec<(ReverseTopologicalPass, T)>,
194    visited: FxHashSet<T>,
195}
196
197impl<T, E> Iterator for IntoPostorderTopologicalIter<T, E>
198where
199    T: Eq + Hash + Clone,
200    E: Clone,
201{
202    type Item = T;
203
204    fn next(&mut self) -> Option<Self::Item> {
205        let current = loop {
206            let (pass, current) = self.stack.pop()?;
207
208            match pass {
209                ReverseTopologicalPass::Post => {
210                    break current;
211                }
212                ReverseTopologicalPass::Pre => {
213                    if self.visited.contains(&current) {
214                        continue;
215                    }
216
217                    self.visited.insert(current.clone());
218
219                    let Some(neighbors) = self.adjacency_map.get(&current) else {
220                        break current;
221                    };
222
223                    self.stack.push((ReverseTopologicalPass::Post, current));
224                    self.stack.extend(
225                        neighbors
226                            .iter()
227                            .rev()
228                            .map(|(neighbor, _)| (ReverseTopologicalPass::Pre, neighbor.clone())),
229                    );
230                }
231            }
232        };
233
234        Some(current)
235    }
236}
237
238pub struct IntoBreadthFirstEdges<T, E>
239where
240    T: Eq + std::hash::Hash + Clone,
241{
242    adjacency_map: FxHashMap<T, Vec<(T, E)>>,
243    queue: VecDeque<(Option<(T, E)>, T)>,
244    expanded: FxHashSet<T>,
245}
246
247impl<T, E> Iterator for IntoBreadthFirstEdges<T, E>
248where
249    T: Eq + std::hash::Hash + Clone,
250    E: Clone,
251{
252    type Item = (Option<(T, E)>, T);
253
254    fn next(&mut self) -> Option<Self::Item> {
255        let (parent, current) = self.queue.pop_front()?;
256
257        let Some(neighbors) = self.adjacency_map.get(&current) else {
258            return Some((parent, current));
259        };
260
261        if self.expanded.insert(current.clone()) {
262            self.queue.extend(
263                neighbors.iter().map(|(neighbor, edge)| {
264                    (Some((current.clone(), edge.clone())), neighbor.clone())
265                }),
266            );
267        }
268
269        Some((parent, current))
270    }
271}
272
273/// An iterator over the nodes of a graph in postorder topological order, starting
274/// from the roots.
275pub struct PostorderTopologicalIter<'graph, T, E>
276where
277    T: Eq + Hash + Clone,
278{
279    adjacency_map: &'graph FxHashMap<T, Vec<(T, E)>>,
280    stack: Vec<(ReverseTopologicalPass, &'graph T)>,
281    visited: FxHashSet<&'graph T>,
282}
283
284impl<'graph, T, E> Iterator for PostorderTopologicalIter<'graph, T, E>
285where
286    T: Eq + Hash + Clone,
287{
288    type Item = &'graph T;
289
290    fn next(&mut self) -> Option<Self::Item> {
291        let current = loop {
292            let (pass, current) = self.stack.pop()?;
293
294            match pass {
295                ReverseTopologicalPass::Post => {
296                    break current;
297                }
298                ReverseTopologicalPass::Pre => {
299                    if self.visited.contains(current) {
300                        continue;
301                    }
302
303                    self.visited.insert(current);
304
305                    let Some(neighbors) = self.adjacency_map.get(current) else {
306                        break current;
307                    };
308
309                    self.stack.push((ReverseTopologicalPass::Post, current));
310                    self.stack.extend(
311                        neighbors
312                            .iter()
313                            .rev()
314                            .map(|(neighbor, _)| (ReverseTopologicalPass::Pre, neighbor)),
315                    );
316                }
317            }
318        };
319
320        Some(current)
321    }
322}
turbo_tasks/graph/adjacency_map.rs

turbo_tasks/graph/
adjacency_map.rs
