turbo_tasks/graph/
graph_store.rs

1use std::hash::Hash;
2
3use rustc_hash::FxHashSet;
4
5use super::VisitedNodes;
6
7/// A graph store is a data structure that will be built up during a graph
8/// traversal. It is used to store the results of the traversal.
9pub trait GraphStore: Send {
10    type Node: Send;
11    type Handle: Clone + Send;
12
13    // TODO(alexkirsz) An `entry(from_handle) -> Entry` API would be more
14    // efficient, as right now we're getting the same key multiple times.
15    /// Inserts a node into the graph store, and returns a handle to it.
16    ///
17    /// If this method returns `None`, the node edges will not be visited.
18    fn insert(
19        &mut self,
20        from_handle: Option<Self::Handle>,
21        node: GraphNode<Self::Node>,
22    ) -> Option<(Self::Handle, &Self::Node)>;
23}
24
25/// Utility type to ensure that GraphStore::insert can only ever be called from
26/// within this module, as a GraphNode can't be constructed outside of it.
27#[derive(Clone, Copy, Eq, PartialEq, Debug, Hash, Ord, PartialOrd)]
28pub struct GraphNode<Node>(pub(super) Node);
29
30impl<Node> GraphNode<Node> {
31    /// Consumes this `GraphNode` and returns the underlying node.
32    pub fn into_node(self) -> Node {
33        self.0
34    }
35
36    /// Returns a reference the underlying node.
37    pub fn node(&self) -> &Node {
38        &self.0
39    }
40
41    /// Returns a mutable reference the underlying node.
42    pub fn node_mut(&mut self) -> &mut Node {
43        &mut self.0
44    }
45}
46
47/// A [`GraphStore`] wrapper that skips nodes that have already been
48/// visited.
49///
50/// This is necessary to avoid repeated work when traversing non-tree
51/// graphs (i.e. where a node can have more than one incoming edge).
52#[derive(Debug)]
53pub struct SkipDuplicates<StoreImpl>
54where
55    StoreImpl: GraphStore,
56{
57    store: StoreImpl,
58    visited: FxHashSet<StoreImpl::Node>,
59}
60
61impl<StoreImpl> SkipDuplicates<StoreImpl>
62where
63    StoreImpl: GraphStore,
64{
65    pub fn new(store: StoreImpl) -> Self {
66        Self {
67            store,
68            visited: FxHashSet::default(),
69        }
70    }
71
72    pub fn new_with_visited_nodes(store: StoreImpl, visited: FxHashSet<StoreImpl::Node>) -> Self {
73        Self { store, visited }
74    }
75}
76
77impl<StoreImpl> GraphStore for SkipDuplicates<StoreImpl>
78where
79    StoreImpl: GraphStore,
80    StoreImpl::Node: Eq + std::hash::Hash + Clone,
81{
82    type Node = StoreImpl::Node;
83    type Handle = StoreImpl::Handle;
84
85    fn insert(
86        &mut self,
87        from_handle: Option<Self::Handle>,
88        node: GraphNode<StoreImpl::Node>,
89    ) -> Option<(Self::Handle, &StoreImpl::Node)> {
90        if !self.visited.contains(node.node()) {
91            self.visited.insert(node.node().clone());
92            self.store.insert(from_handle, node)
93        } else {
94            // Always insert the node into the store, even if we've already
95            // visited it. This is necessary to ensure that the store sees all
96            // edges.
97            self.store.insert(from_handle, node);
98            None
99        }
100    }
101}
102
103impl<StoreImpl> SkipDuplicates<StoreImpl>
104where
105    StoreImpl: GraphStore,
106{
107    /// Consumes the wrapper and returns the underlying store.
108    pub fn into_inner(self) -> StoreImpl {
109        self.store
110    }
111
112    /// Consumes the wrapper and returns the underlying store along with the visited nodes.
113    pub fn into_inner_with_visited(self) -> (StoreImpl, VisitedNodes<StoreImpl::Node>) {
114        (self.store, VisitedNodes(self.visited))
115    }
116}
117
118/// A [`GraphStore`] wrapper that skips nodes that have already been
119/// visited, based on a key extracted from the node.
120///
121/// This is necessary to avoid repeated work when traversing non-tree
122/// graphs (i.e. where a node can have more than one incoming edge).
123#[derive(Debug)]
124pub struct SkipDuplicatesWithKey<StoreImpl, Key, KeyExtractor>
125where
126    StoreImpl: GraphStore,
127    Key: Send + Eq + Hash,
128    KeyExtractor: Send + Fn(&StoreImpl::Node) -> &Key,
129{
130    store: StoreImpl,
131    visited: FxHashSet<Key>,
132    key_extractor: KeyExtractor,
133}
134
135impl<StoreImpl, Key, KeyExtractor> SkipDuplicatesWithKey<StoreImpl, Key, KeyExtractor>
136where
137    StoreImpl: GraphStore,
138    Key: Send + Eq + std::hash::Hash + Clone,
139    KeyExtractor: Send + Fn(&StoreImpl::Node) -> &Key,
140{
141    pub fn new(store: StoreImpl, key_extractor: KeyExtractor) -> Self {
142        Self {
143            store,
144            visited: FxHashSet::default(),
145            key_extractor,
146        }
147    }
148}
149
150impl<StoreImpl, Key, KeyExtractor> GraphStore
151    for SkipDuplicatesWithKey<StoreImpl, Key, KeyExtractor>
152where
153    StoreImpl: GraphStore,
154    StoreImpl::Node: Eq + std::hash::Hash + Clone,
155    Key: Send + Eq + std::hash::Hash + Clone,
156    KeyExtractor: Send + Fn(&StoreImpl::Node) -> &Key,
157{
158    type Node = StoreImpl::Node;
159    type Handle = StoreImpl::Handle;
160
161    fn insert(
162        &mut self,
163        from_handle: Option<Self::Handle>,
164        node: GraphNode<StoreImpl::Node>,
165    ) -> Option<(Self::Handle, &StoreImpl::Node)> {
166        let key = (self.key_extractor)(node.node());
167
168        if !self.visited.contains(key) {
169            self.visited.insert(key.clone());
170            self.store.insert(from_handle, node)
171        } else {
172            // Always insert the node into the store, even if we've already
173            // visited it. This is necessary to ensure that the store sees all
174            // edges.
175            self.store.insert(from_handle, node);
176            None
177        }
178    }
179}
180
181impl<StoreImpl, Key, KeyExtractor> SkipDuplicatesWithKey<StoreImpl, Key, KeyExtractor>
182where
183    StoreImpl: GraphStore,
184    Key: Send + Eq + std::hash::Hash + Clone,
185    KeyExtractor: Send + Fn(&StoreImpl::Node) -> &Key,
186{
187    /// Consumes the wrapper and returns the underlying store.
188    pub fn into_inner(self) -> StoreImpl {
189        self.store
190    }
191
192    /// Consumes the wrapper and returns the underlying store along with the visited nodes.
193    pub fn into_inner_with_visited(self) -> (StoreImpl, VisitedNodes<Key>) {
194        (self.store, VisitedNodes(self.visited))
195    }
196}
turbo_tasks/graph/graph_store.rs

turbo_tasks/graph/
graph_store.rs
