turbo_tasks/graph/
visit.rs

1use std::{future::Future, iter::Map};
2
3use anyhow::Result;
4use tracing::Span;
5
6use super::VisitControlFlow;
7
8/// A trait that allows a graph traversal to visit the edges of a node
9/// transitively.
10pub trait Visit<Node, Edge = (), Impl = ()> {
11    type EdgesIntoIter: IntoIterator<Item = (Node, Edge)>;
12    type EdgesFuture: Future<Output = Result<Self::EdgesIntoIter>> + Send;
13
14    /// Visits an edge. Should return a
15    /// [`VisitControlFlow`] that indicates whether to:
16    /// * continue visiting the neighbor node edges;
17    /// * skip visiting the neighbor node's edges;
18    /// * abort the traversal entirely.
19    fn visit(&mut self, node: &Node, edge: Option<&Edge>) -> VisitControlFlow {
20        let _ = node;
21        let _ = edge;
22        VisitControlFlow::Continue
23    }
24
25    /// Returns a future that resolves to the outgoing edges of the given `node`.
26    ///
27    /// Lifetimes:
28    /// - The returned future's lifetime cannot depend on the reference to self because there are
29    ///   multiple `edges` futures created and awaited concurrently.
30    /// - The returned future's lifetime cannot depend on `node` because `GraphStore::insert`
31    ///   returns a node reference that's only valid for the lifetime of its `&mut self` reference.
32    fn edges(&mut self, node: &Node) -> Self::EdgesFuture;
33
34    /// Returns a [Span] for the given `node`, under which all edges are processed.
35    fn span(&mut self, _node: &Node, _edge: Option<&Edge>) -> Span {
36        Span::none()
37    }
38}
39
40// The different `Impl*` here are necessary in order to avoid the `Conflicting
41// implementations of trait` error when implementing `Visit` on different
42// kinds of `FnMut`.
43// See https://users.rust-lang.org/t/conflicting-implementation-when-implementing-traits-for-fn/53359/3
44
45pub struct ImplWithEdgeRef;
46
47impl<Node, Edge, VisitFn, NeighFut, NeighIt> Visit<Node, Edge, ImplWithEdgeRef> for VisitFn
48where
49    VisitFn: FnMut(&Node) -> NeighFut,
50    NeighFut: Future<Output = Result<NeighIt>> + Send,
51    NeighIt: IntoIterator<Item = (Node, Edge)>,
52{
53    type EdgesIntoIter = NeighIt;
54    type EdgesFuture = NeighFut;
55
56    fn edges(&mut self, node: &Node) -> Self::EdgesFuture {
57        (self)(node)
58    }
59}
60
61pub struct ImplWithEdgeValue;
62
63impl<Node, Edge, VisitFn, NeighFut, NeighIt> Visit<Node, Edge, ImplWithEdgeValue> for VisitFn
64where
65    Node: Clone,
66    VisitFn: FnMut(Node) -> NeighFut,
67    NeighFut: Future<Output = Result<NeighIt>> + Send,
68    NeighIt: IntoIterator<Item = (Node, Edge)>,
69{
70    type EdgesIntoIter = NeighIt;
71    type EdgesFuture = NeighFut;
72
73    fn edges(&mut self, node: &Node) -> Self::EdgesFuture {
74        (self)(node.clone())
75    }
76}
77
78pub struct ImplRef;
79
80impl<Node, VisitFn, NeighFut, NeighIt> Visit<Node, (), ImplRef> for VisitFn
81where
82    VisitFn: FnMut(&Node) -> NeighFut,
83    NeighFut: Future<Output = Result<NeighIt>> + Send,
84    NeighIt: IntoIterator<Item = Node>,
85{
86    type EdgesIntoIter = Map<NeighIt::IntoIter, fn(Node) -> (Node, ())>;
87    type EdgesFuture = NoEdgeFuture<NeighFut>;
88
89    fn edges(&mut self, node: &Node) -> Self::EdgesFuture {
90        NoEdgeFuture((self)(node))
91    }
92}
93
94pub struct ImplValue;
95
96impl<Node, VisitFn, NeighFut, NeighIt> Visit<Node, (), ImplValue> for VisitFn
97where
98    Node: Clone,
99    VisitFn: FnMut(Node) -> NeighFut,
100    NeighFut: Future<Output = Result<NeighIt>> + Send,
101    NeighIt: IntoIterator<Item = Node>,
102{
103    type EdgesIntoIter = Map<NeighIt::IntoIter, fn(Node) -> (Node, ())>;
104    type EdgesFuture = NoEdgeFuture<NeighFut>;
105
106    fn edges(&mut self, node: &Node) -> Self::EdgesFuture {
107        NoEdgeFuture((self)(node.clone()))
108    }
109}
110
111pub struct NoEdgeFuture<F>(F);
112
113impl<F, Node, NeighIt> Future for NoEdgeFuture<F>
114where
115    F: Future<Output = Result<NeighIt>> + Send,
116    NeighIt: IntoIterator<Item = Node>,
117{
118    type Output = Result<Map<NeighIt::IntoIter, fn(Node) -> (Node, ())>>;
119
120    fn poll(
121        self: std::pin::Pin<&mut Self>,
122        cx: &mut std::task::Context<'_>,
123    ) -> std::task::Poll<Self::Output> {
124        let future = unsafe { self.map_unchecked_mut(|s| &mut s.0) };
125        future.poll(cx).map(|res| {
126            res.map(|it| {
127                fn map_fn<Node>(node: Node) -> (Node, ()) {
128                    (node, ())
129                }
130                let f: fn(Node) -> (Node, ()) = map_fn;
131                it.into_iter().map(f)
132            })
133        })
134    }
135}
turbo_tasks/graph/visit.rs

turbo_tasks/graph/
visit.rs
