petgraph/visit/
traversal.rs

1use alloc::{collections::VecDeque, vec::Vec};
2
3use super::{
4    GraphRef, IntoNeighbors, IntoNeighborsDirected, IntoNodeIdentifiers, Reversed, VisitMap,
5    Visitable,
6};
7use crate::Incoming;
8
9/// Visit nodes of a graph in a depth-first-search (DFS) emitting nodes in
10/// preorder (when they are first discovered).
11///
12/// The traversal starts at a given node and only traverses nodes reachable
13/// from it.
14///
15/// `Dfs` is not recursive.
16///
17/// `Dfs` does not itself borrow the graph, and because of this you can run
18/// a traversal over a graph while still retaining mutable access to it, if you
19/// use it like the following example:
20///
21/// ```
22/// use petgraph::Graph;
23/// use petgraph::visit::Dfs;
24///
25/// let mut graph = Graph::<_,()>::new();
26/// let a = graph.add_node(0);
27///
28/// let mut dfs = Dfs::new(&graph, a);
29/// while let Some(nx) = dfs.next(&graph) {
30///     // we can access `graph` mutably here still
31///     graph[nx] += 1;
32/// }
33///
34/// assert_eq!(graph[a], 1);
35/// ```
36///
37/// **Note:** The algorithm may not behave correctly if nodes are removed
38/// during iteration. It may not necessarily visit added nodes or edges.
39#[derive(Clone, Debug)]
40pub struct Dfs<N, VM> {
41    /// The stack of nodes to visit
42    pub stack: Vec<N>,
43    /// The map of discovered nodes
44    pub discovered: VM,
45}
46
47impl<N, VM> Default for Dfs<N, VM>
48where
49    VM: Default,
50{
51    fn default() -> Self {
52        Dfs {
53            stack: Vec::new(),
54            discovered: VM::default(),
55        }
56    }
57}
58
59impl<N, VM> Dfs<N, VM>
60where
61    N: Copy + PartialEq,
62    VM: VisitMap<N>,
63{
64    /// Create a new **Dfs**, using the graph's visitor map, and put **start**
65    /// in the stack of nodes to visit.
66    pub fn new<G>(graph: G, start: N) -> Self
67    where
68        G: GraphRef + Visitable<NodeId = N, Map = VM>,
69    {
70        let mut dfs = Dfs::empty(graph);
71        dfs.move_to(start);
72        dfs
73    }
74
75    /// Create a `Dfs` from a vector and a visit map
76    pub fn from_parts(stack: Vec<N>, discovered: VM) -> Self {
77        Dfs { stack, discovered }
78    }
79
80    /// Clear the visit state
81    pub fn reset<G>(&mut self, graph: G)
82    where
83        G: GraphRef + Visitable<NodeId = N, Map = VM>,
84    {
85        graph.reset_map(&mut self.discovered);
86        self.stack.clear();
87    }
88
89    /// Create a new **Dfs** using the graph's visitor map, and no stack.
90    pub fn empty<G>(graph: G) -> Self
91    where
92        G: GraphRef + Visitable<NodeId = N, Map = VM>,
93    {
94        Dfs {
95            stack: Vec::new(),
96            discovered: graph.visit_map(),
97        }
98    }
99
100    /// Keep the discovered map, but clear the visit stack and restart
101    /// the dfs from a particular node.
102    pub fn move_to(&mut self, start: N) {
103        self.stack.clear();
104        self.stack.push(start);
105    }
106
107    /// Return the next node in the dfs, or **None** if the traversal is done.
108    pub fn next<G>(&mut self, graph: G) -> Option<N>
109    where
110        G: IntoNeighbors<NodeId = N>,
111    {
112        while let Some(node) = self.stack.pop() {
113            if self.discovered.visit(node) {
114                for succ in graph.neighbors(node) {
115                    if !self.discovered.is_visited(&succ) {
116                        self.stack.push(succ);
117                    }
118                }
119                return Some(node);
120            }
121        }
122        None
123    }
124}
125
126/// Visit nodes in a depth-first-search (DFS) emitting nodes in postorder
127/// (each node after all its descendants have been emitted).
128///
129/// `DfsPostOrder` is not recursive.
130///
131/// The traversal starts at a given node and only traverses nodes reachable
132/// from it.
133#[derive(Clone, Debug)]
134pub struct DfsPostOrder<N, VM> {
135    /// The stack of nodes to visit
136    pub stack: Vec<N>,
137    /// The map of discovered nodes
138    pub discovered: VM,
139    /// The map of finished nodes
140    pub finished: VM,
141}
142
143impl<N, VM> Default for DfsPostOrder<N, VM>
144where
145    VM: Default,
146{
147    fn default() -> Self {
148        DfsPostOrder {
149            stack: Vec::new(),
150            discovered: VM::default(),
151            finished: VM::default(),
152        }
153    }
154}
155
156impl<N, VM> DfsPostOrder<N, VM>
157where
158    N: Copy + PartialEq,
159    VM: VisitMap<N>,
160{
161    /// Create a new `DfsPostOrder` using the graph's visitor map, and put
162    /// `start` in the stack of nodes to visit.
163    pub fn new<G>(graph: G, start: N) -> Self
164    where
165        G: GraphRef + Visitable<NodeId = N, Map = VM>,
166    {
167        let mut dfs = Self::empty(graph);
168        dfs.move_to(start);
169        dfs
170    }
171
172    /// Create a new `DfsPostOrder` using the graph's visitor map, and no stack.
173    pub fn empty<G>(graph: G) -> Self
174    where
175        G: GraphRef + Visitable<NodeId = N, Map = VM>,
176    {
177        DfsPostOrder {
178            stack: Vec::new(),
179            discovered: graph.visit_map(),
180            finished: graph.visit_map(),
181        }
182    }
183
184    /// Clear the visit state
185    pub fn reset<G>(&mut self, graph: G)
186    where
187        G: GraphRef + Visitable<NodeId = N, Map = VM>,
188    {
189        graph.reset_map(&mut self.discovered);
190        graph.reset_map(&mut self.finished);
191        self.stack.clear();
192    }
193
194    /// Keep the discovered and finished map, but clear the visit stack and restart
195    /// the dfs from a particular node.
196    pub fn move_to(&mut self, start: N) {
197        self.stack.clear();
198        self.stack.push(start);
199    }
200
201    /// Return the next node in the traversal, or `None` if the traversal is done.
202    pub fn next<G>(&mut self, graph: G) -> Option<N>
203    where
204        G: IntoNeighbors<NodeId = N>,
205    {
206        while let Some(&nx) = self.stack.last() {
207            if self.discovered.visit(nx) {
208                // First time visiting `nx`: Push neighbors, don't pop `nx`
209                for succ in graph.neighbors(nx) {
210                    if !self.discovered.is_visited(&succ) {
211                        self.stack.push(succ);
212                    }
213                }
214            } else {
215                self.stack.pop();
216                if self.finished.visit(nx) {
217                    // Second time: All reachable nodes must have been finished
218                    return Some(nx);
219                }
220            }
221        }
222        None
223    }
224}
225
226/// A breadth first search (BFS) of a graph.
227///
228/// The traversal starts at a given node and only traverses nodes reachable
229/// from it.
230///
231/// `Bfs` is not recursive.
232///
233/// `Bfs` does not itself borrow the graph, and because of this you can run
234/// a traversal over a graph while still retaining mutable access to it, if you
235/// use it like the following example:
236///
237/// ```
238/// use petgraph::Graph;
239/// use petgraph::visit::Bfs;
240///
241/// let mut graph = Graph::<_,()>::new();
242/// let a = graph.add_node(0);
243///
244/// let mut bfs = Bfs::new(&graph, a);
245/// while let Some(nx) = bfs.next(&graph) {
246///     // we can access `graph` mutably here still
247///     graph[nx] += 1;
248/// }
249///
250/// assert_eq!(graph[a], 1);
251/// ```
252///
253/// **Note:** The algorithm may not behave correctly if nodes are removed
254/// during iteration. It may not necessarily visit added nodes or edges.
255#[derive(Clone)]
256pub struct Bfs<N, VM> {
257    /// The queue of nodes to visit
258    pub stack: VecDeque<N>,
259    /// The map of discovered nodes
260    pub discovered: VM,
261}
262
263impl<N, VM> Default for Bfs<N, VM>
264where
265    VM: Default,
266{
267    fn default() -> Self {
268        Bfs {
269            stack: VecDeque::new(),
270            discovered: VM::default(),
271        }
272    }
273}
274
275impl<N, VM> Bfs<N, VM>
276where
277    N: Copy + PartialEq,
278    VM: VisitMap<N>,
279{
280    /// Create a new **Bfs**, using the graph's visitor map, and put **start**
281    /// in the stack of nodes to visit.
282    pub fn new<G>(graph: G, start: N) -> Self
283    where
284        G: GraphRef + Visitable<NodeId = N, Map = VM>,
285    {
286        let mut discovered = graph.visit_map();
287        discovered.visit(start);
288        let mut stack = VecDeque::new();
289        stack.push_front(start);
290        Bfs { stack, discovered }
291    }
292
293    /// Return the next node in the bfs, or **None** if the traversal is done.
294    pub fn next<G>(&mut self, graph: G) -> Option<N>
295    where
296        G: IntoNeighbors<NodeId = N>,
297    {
298        if let Some(node) = self.stack.pop_front() {
299            for succ in graph.neighbors(node) {
300                if self.discovered.visit(succ) {
301                    self.stack.push_back(succ);
302                }
303            }
304
305            return Some(node);
306        }
307        None
308    }
309}
310
311/// A topological order traversal for a graph.
312///
313/// **Note** that `Topo` only visits nodes that are not part of cycles,
314/// i.e. nodes in a true DAG. Use other visitors like `DfsPostOrder` or
315/// algorithms like kosaraju_scc to handle graphs with possible cycles.
316#[derive(Clone)]
317pub struct Topo<N, VM> {
318    tovisit: Vec<N>,
319    ordered: VM,
320}
321
322impl<N, VM> Default for Topo<N, VM>
323where
324    VM: Default,
325{
326    fn default() -> Self {
327        Topo {
328            tovisit: Vec::new(),
329            ordered: VM::default(),
330        }
331    }
332}
333
334impl<N, VM> Topo<N, VM>
335where
336    N: Copy + PartialEq,
337    VM: VisitMap<N>,
338{
339    /// Create a new `Topo`, using the graph's visitor map, and put all
340    /// initial nodes in the to visit list.
341    pub fn new<G>(graph: G) -> Self
342    where
343        G: IntoNodeIdentifiers + IntoNeighborsDirected + Visitable<NodeId = N, Map = VM>,
344    {
345        let mut topo = Self::empty(graph);
346        topo.extend_with_initials(graph);
347        topo
348    }
349
350    /// Create a new `Topo` with initial nodes.
351    ///
352    /// Nodes with incoming edges are ignored.
353    pub fn with_initials<G, I>(graph: G, initials: I) -> Self
354    where
355        G: IntoNeighborsDirected + Visitable<NodeId = N, Map = VM>,
356        I: IntoIterator<Item = N>,
357    {
358        Topo {
359            tovisit: initials
360                .into_iter()
361                .filter(|&n| graph.neighbors_directed(n, Incoming).next().is_none())
362                .collect(),
363            ordered: graph.visit_map(),
364        }
365    }
366
367    fn extend_with_initials<G>(&mut self, g: G)
368    where
369        G: IntoNodeIdentifiers + IntoNeighborsDirected<NodeId = N>,
370    {
371        // find all initial nodes (nodes without incoming edges)
372        self.tovisit.extend(
373            g.node_identifiers()
374                .filter(move |&a| g.neighbors_directed(a, Incoming).next().is_none()),
375        );
376    }
377
378    /* Private until it has a use */
379    /// Create a new `Topo`, using the graph's visitor map with *no* starting
380    /// index specified.
381    fn empty<G>(graph: G) -> Self
382    where
383        G: GraphRef + Visitable<NodeId = N, Map = VM>,
384    {
385        Topo {
386            ordered: graph.visit_map(),
387            tovisit: Vec::new(),
388        }
389    }
390
391    /// Clear visited state, and put all initial nodes in the to visit list.
392    pub fn reset<G>(&mut self, graph: G)
393    where
394        G: IntoNodeIdentifiers + IntoNeighborsDirected + Visitable<NodeId = N, Map = VM>,
395    {
396        graph.reset_map(&mut self.ordered);
397        self.tovisit.clear();
398        self.extend_with_initials(graph);
399    }
400
401    /// Return the next node in the current topological order traversal, or
402    /// `None` if the traversal is at the end.
403    ///
404    /// *Note:* The graph may not have a complete topological order, and the only
405    /// way to know is to run the whole traversal and make sure it visits every node.
406    pub fn next<G>(&mut self, g: G) -> Option<N>
407    where
408        G: IntoNeighborsDirected + Visitable<NodeId = N, Map = VM>,
409    {
410        // Take an unvisited element and find which of its neighbors are next
411        while let Some(nix) = self.tovisit.pop() {
412            if self.ordered.is_visited(&nix) {
413                continue;
414            }
415            self.ordered.visit(nix);
416            for neigh in g.neighbors(nix) {
417                // Look at each neighbor, and those that only have incoming edges
418                // from the already ordered list, they are the next to visit.
419                if Reversed(g)
420                    .neighbors(neigh)
421                    .all(|b| self.ordered.is_visited(&b))
422                {
423                    self.tovisit.push(neigh);
424                }
425            }
426            return Some(nix);
427        }
428        None
429    }
430}
431
432/// A walker is a traversal state, but where part of the traversal
433/// information is supplied manually to each next call.
434///
435/// This for example allows graph traversals that don't hold a borrow of the
436/// graph they are traversing.
437pub trait Walker<Context> {
438    type Item;
439    /// Advance to the next item
440    fn walk_next(&mut self, context: Context) -> Option<Self::Item>;
441
442    /// Create an iterator out of the walker and given `context`.
443    fn iter(self, context: Context) -> WalkerIter<Self, Context>
444    where
445        Self: Sized,
446        Context: Clone,
447    {
448        WalkerIter {
449            walker: self,
450            context,
451        }
452    }
453}
454
455/// A walker and its context wrapped into an iterator.
456#[derive(Clone, Debug)]
457pub struct WalkerIter<W, C> {
458    walker: W,
459    context: C,
460}
461
462impl<W, C> WalkerIter<W, C>
463where
464    W: Walker<C>,
465    C: Clone,
466{
467    pub fn context(&self) -> C {
468        self.context.clone()
469    }
470
471    pub fn inner_ref(&self) -> &W {
472        &self.walker
473    }
474
475    pub fn inner_mut(&mut self) -> &mut W {
476        &mut self.walker
477    }
478}
479
480impl<W, C> Iterator for WalkerIter<W, C>
481where
482    W: Walker<C>,
483    C: Clone,
484{
485    type Item = W::Item;
486    fn next(&mut self) -> Option<Self::Item> {
487        self.walker.walk_next(self.context.clone())
488    }
489}
490
491impl<C, W: ?Sized> Walker<C> for &mut W
492where
493    W: Walker<C>,
494{
495    type Item = W::Item;
496    fn walk_next(&mut self, context: C) -> Option<Self::Item> {
497        (**self).walk_next(context)
498    }
499}
500
501impl<G> Walker<G> for Dfs<G::NodeId, G::Map>
502where
503    G: IntoNeighbors + Visitable,
504{
505    type Item = G::NodeId;
506    fn walk_next(&mut self, context: G) -> Option<Self::Item> {
507        self.next(context)
508    }
509}
510
511impl<G> Walker<G> for DfsPostOrder<G::NodeId, G::Map>
512where
513    G: IntoNeighbors + Visitable,
514{
515    type Item = G::NodeId;
516    fn walk_next(&mut self, context: G) -> Option<Self::Item> {
517        self.next(context)
518    }
519}
520
521impl<G> Walker<G> for Bfs<G::NodeId, G::Map>
522where
523    G: IntoNeighbors + Visitable,
524{
525    type Item = G::NodeId;
526    fn walk_next(&mut self, context: G) -> Option<Self::Item> {
527        self.next(context)
528    }
529}
530
531impl<G> Walker<G> for Topo<G::NodeId, G::Map>
532where
533    G: IntoNeighborsDirected + Visitable,
534{
535    type Item = G::NodeId;
536    fn walk_next(&mut self, context: G) -> Option<Self::Item> {
537        self.next(context)
538    }
539}