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
83impl<T, E> GraphStore for AdjacencyMap<T, E>
84where
85    T: Eq + Hash + Clone + Send,
86    E: Send,
87{
88    type Node = T;
89    type Edge = E;
90    type Handle = T;
91
92    fn insert(&mut self, from: Option<(&T, E)>, node: T) {
93        if let Some((from_node, edge)) = from {
94            let vec = self
95                .adjacency_map
96                .entry(from_node.clone())
97                .or_insert_with(|| Vec::with_capacity(1));
98            vec.push((node, edge));
99        } else {
100            self.roots.push(node);
101        };
102    }
103
104    fn try_enter(&mut self, node: &T) -> Option<T> {
105        match self.adjacency_map.entry(node.clone()) {
106            Entry::Occupied(_) => None,
107            Entry::Vacant(e) => {
108                e.insert(Vec::new());
109                Some(node.clone())
110            }
111        }
112    }
113}
114
115impl<T, E> AdjacencyMap<T, E>
116where
117    T: Eq + Hash + Clone,
118{
119    /// Returns an owned iterator over the nodes in postorder topological order,
120    /// starting from the roots.
121    pub fn into_postorder_topological(self) -> IntoPostorderTopologicalIter<T, E> {
122        IntoPostorderTopologicalIter {
123            adjacency_map: self.adjacency_map,
124            stack: self
125                .roots
126                .into_iter()
127                .rev()
128                .map(|root| (ReverseTopologicalPass::Pre, root))
129                .collect(),
130            visited: FxHashSet::default(),
131        }
132    }
133
134    /// Returns an owned iterator over all edges (node pairs) in reverse breadth first order,
135    /// starting from the roots.
136    pub fn into_breadth_first_edges(self) -> IntoBreadthFirstEdges<T, E> {
137        IntoBreadthFirstEdges {
138            adjacency_map: self.adjacency_map,
139            queue: self
140                .roots
141                .into_iter()
142                .rev()
143                .map(|root| (None, root))
144                .collect(),
145            expanded: FxHashSet::default(),
146        }
147    }
148
149    /// Returns an iterator over the nodes in postorder topological order,
150    /// starting from the roots.
151    pub fn postorder_topological(&self) -> PostorderTopologicalIter<'_, T, E> {
152        PostorderTopologicalIter {
153            adjacency_map: &self.adjacency_map,
154            stack: self
155                .roots
156                .iter()
157                .rev()
158                .map(|root| (ReverseTopologicalPass::Pre, root))
159                .collect(),
160            visited: FxHashSet::default(),
161        }
162    }
163
164    /// Returns an iterator over the nodes in postorder topological order,
165    /// starting from the given node.
166    pub fn postorder_topological_from_node<'graph>(
167        &'graph self,
168        node: &'graph T,
169    ) -> PostorderTopologicalIter<'graph, T, E> {
170        PostorderTopologicalIter {
171            adjacency_map: &self.adjacency_map,
172            stack: vec![(ReverseTopologicalPass::Pre, node)],
173            visited: FxHashSet::default(),
174        }
175    }
176}
177
178#[derive(Debug)]
179enum ReverseTopologicalPass {
180    Pre,
181    Post,
182}
183
184/// An iterator over the nodes of a graph in postorder topological order, starting
185/// from the roots.
186pub struct IntoPostorderTopologicalIter<T, E>
187where
188    T: Eq + Hash + Clone,
189{
190    adjacency_map: FxHashMap<T, Vec<(T, E)>>,
191    stack: Vec<(ReverseTopologicalPass, T)>,
192    visited: FxHashSet<T>,
193}
194
195impl<T, E> Iterator for IntoPostorderTopologicalIter<T, E>
196where
197    T: Eq + Hash + Clone,
198    E: Clone,
199{
200    type Item = T;
201
202    fn next(&mut self) -> Option<Self::Item> {
203        let current = loop {
204            let (pass, current) = self.stack.pop()?;
205
206            match pass {
207                ReverseTopologicalPass::Post => {
208                    break current;
209                }
210                ReverseTopologicalPass::Pre => {
211                    if self.visited.contains(&current) {
212                        continue;
213                    }
214
215                    self.visited.insert(current.clone());
216
217                    let Some(neighbors) = self.adjacency_map.get(&current) else {
218                        break current;
219                    };
220
221                    self.stack.push((ReverseTopologicalPass::Post, current));
222                    self.stack.extend(
223                        neighbors
224                            .iter()
225                            .rev()
226                            .map(|(neighbor, _)| (ReverseTopologicalPass::Pre, neighbor.clone())),
227                    );
228                }
229            }
230        };
231
232        Some(current)
233    }
234}
235
236pub struct IntoBreadthFirstEdges<T, E>
237where
238    T: Eq + std::hash::Hash + Clone,
239{
240    adjacency_map: FxHashMap<T, Vec<(T, E)>>,
241    queue: VecDeque<(Option<(T, E)>, T)>,
242    expanded: FxHashSet<T>,
243}
244
245impl<T, E> Iterator for IntoBreadthFirstEdges<T, E>
246where
247    T: Eq + std::hash::Hash + Clone,
248    E: Clone,
249{
250    type Item = (Option<(T, E)>, T);
251
252    fn next(&mut self) -> Option<Self::Item> {
253        let (parent, current) = self.queue.pop_front()?;
254
255        let Some(neighbors) = self.adjacency_map.get(&current) else {
256            return Some((parent, current));
257        };
258
259        if self.expanded.insert(current.clone()) {
260            self.queue.extend(
261                neighbors.iter().map(|(neighbor, edge)| {
262                    (Some((current.clone(), edge.clone())), neighbor.clone())
263                }),
264            );
265        }
266
267        Some((parent, current))
268    }
269}
270
271/// An iterator over the nodes of a graph in postorder topological order, starting
272/// from the roots.
273pub struct PostorderTopologicalIter<'graph, T, E>
274where
275    T: Eq + Hash + Clone,
276{
277    adjacency_map: &'graph FxHashMap<T, Vec<(T, E)>>,
278    stack: Vec<(ReverseTopologicalPass, &'graph T)>,
279    visited: FxHashSet<&'graph T>,
280}
281
282impl<'graph, T, E> Iterator for PostorderTopologicalIter<'graph, T, E>
283where
284    T: Eq + Hash + Clone,
285{
286    type Item = &'graph T;
287
288    fn next(&mut self) -> Option<Self::Item> {
289        let current = loop {
290            let (pass, current) = self.stack.pop()?;
291
292            match pass {
293                ReverseTopologicalPass::Post => {
294                    break current;
295                }
296                ReverseTopologicalPass::Pre => {
297                    if self.visited.contains(current) {
298                        continue;
299                    }
300
301                    self.visited.insert(current);
302
303                    let Some(neighbors) = self.adjacency_map.get(current) else {
304                        break current;
305                    };
306
307                    self.stack.push((ReverseTopologicalPass::Post, current));
308                    self.stack.extend(
309                        neighbors
310                            .iter()
311                            .rev()
312                            .map(|(neighbor, _)| (ReverseTopologicalPass::Pre, neighbor)),
313                    );
314                }
315            }
316        };
317
318        Some(current)
319    }
320}
turbo_tasks/graph/adjacency_map.rs

turbo_tasks/graph/
adjacency_map.rs
