turbo_tasks/graph/
adjacency_map.rs

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

turbo_tasks/graph/
adjacency_map.rs
