guppy/petgraph_support/
dfs.rs

1// Copyright (c) The cargo-guppy Contributors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use petgraph::{
5    prelude::*,
6    visit::{
7        DfsPostOrder, IntoEdgeReferences, IntoEdges, IntoEdgesDirected, Reversed,
8        ReversedEdgeReference, VisitMap,
9    },
10};
11
12/// `DfsPostOrder::next`, adapted for a buffered filter that's also FnMut.
13pub fn dfs_next_buffered_filter<N, VM, G>(
14    dfs: &mut DfsPostOrder<N, VM>,
15    graph: G,
16    mut buffered_filter: impl BufferedEdgeFilter<G>,
17) -> Option<N>
18where
19    N: Copy + PartialEq,
20    VM: VisitMap<N>,
21    G: IntoEdges<NodeId = N>,
22{
23    // Adapted from DfsPostOrder::next in petgraph 0.5.0.
24    while let Some(&nx) = dfs.stack.last() {
25        if dfs.discovered.visit(nx) {
26            // First time visiting `nx`: Push neighbors, don't pop `nx`
27            let neighbors = graph.edges(nx).flat_map(|edge| {
28                buffered_filter
29                    .filter(edge)
30                    .into_iter()
31                    .map(|edge| edge.target())
32            });
33            for succ in neighbors {
34                if !dfs.discovered.is_visited(&succ) {
35                    dfs.stack.push(succ);
36                }
37            }
38        } else {
39            dfs.stack.pop();
40            if dfs.finished.visit(nx) {
41                // Second time: All reachable nodes must have been finished
42                return Some(nx);
43            }
44        }
45    }
46    None
47}
48
49/// A buffered filter is a graph traversal edge filter that can buffer up some edges to be
50/// returned later.
51pub trait BufferedEdgeFilter<G>
52where
53    G: IntoEdges,
54{
55    /// Returns a list of edge references to follow.
56    type Iter: IntoIterator<Item = G::EdgeRef>;
57
58    fn filter(&mut self, edge: G::EdgeRef) -> Self::Iter;
59}
60
61impl<G, T> BufferedEdgeFilter<G> for &mut T
62where
63    T: BufferedEdgeFilter<G>,
64    G: IntoEdges,
65{
66    /// Returns a list of node IDs to follow.
67    type Iter = T::Iter;
68
69    #[inline]
70    fn filter(&mut self, edge: G::EdgeRef) -> Self::Iter {
71        (*self).filter(edge)
72    }
73}
74
75#[derive(Debug)]
76pub struct SimpleEdgeFilterFn<F>(pub F);
77
78impl<F, G> BufferedEdgeFilter<G> for SimpleEdgeFilterFn<F>
79where
80    F: FnMut(G::EdgeRef) -> bool,
81    G: IntoEdges,
82{
83    type Iter = Option<G::EdgeRef>;
84
85    #[inline]
86    fn filter(&mut self, edge: G::EdgeRef) -> Self::Iter {
87        if (self.0)(edge) {
88            Some(edge)
89        } else {
90            None
91        }
92    }
93}
94
95#[derive(Debug)]
96pub struct BufferedEdgeFilterFn<F>(pub F);
97
98impl<F, G, I> BufferedEdgeFilter<G> for BufferedEdgeFilterFn<F>
99where
100    F: FnMut(G::EdgeRef) -> I,
101    G: IntoEdges,
102    I: IntoIterator<Item = G::EdgeRef>,
103{
104    type Iter = I;
105
106    #[inline]
107    fn filter(&mut self, edge: G::EdgeRef) -> Self::Iter {
108        (self.0)(edge)
109    }
110}
111
112#[derive(Debug)]
113pub struct ReversedBufferedFilter<T>(pub T);
114
115impl<T, G> BufferedEdgeFilter<Reversed<G>> for ReversedBufferedFilter<T>
116where
117    T: BufferedEdgeFilter<G>,
118    G: IntoEdgesDirected,
119{
120    type Iter =
121        ReversedEdgeReferences<<<T as BufferedEdgeFilter<G>>::Iter as IntoIterator>::IntoIter>;
122
123    fn filter(&mut self, edge: <Reversed<G> as IntoEdgeReferences>::EdgeRef) -> Self::Iter {
124        ReversedEdgeReferences {
125            iter: self.0.filter(edge.into_unreversed()).into_iter(),
126        }
127    }
128}
129
130// TODO: replace with upstream impl
131pub struct ReversedEdgeReferences<I> {
132    iter: I,
133}
134
135impl<I> Iterator for ReversedEdgeReferences<I>
136where
137    I: Iterator,
138{
139    type Item = ReversedEdgeReference<I::Item>;
140
141    #[inline]
142    fn next(&mut self) -> Option<Self::Item> {
143        // Ugh this sucks! Should be supported upstream.
144        // SAFETY: this is just a newtype. This is SUCH a horrible hack.
145        let item = self.iter.next()?;
146        Some(unsafe { horrible_transmute::<I::Item, ReversedEdgeReference<I::Item>>(item) })
147    }
148
149    #[inline]
150    fn size_hint(&self) -> (usize, Option<usize>) {
151        self.iter.size_hint()
152    }
153}
154
155unsafe fn horrible_transmute<A, B>(a: A) -> B {
156    let b = ::core::ptr::read(&a as *const A as *const B);
157    ::core::mem::forget(a);
158    b
159}
160
161// Check that the above transmute is actually safe.
162static_assertions::assert_eq_size!((), ReversedEdgeReference<()>);