turbo_tasks/graph/
adjacency_map.rs

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

turbo_tasks/graph/
adjacency_map.rs
