petgraph/visit/
filter.rs

1use crate::prelude::*;
2
3use fixedbitset::FixedBitSet;
4use std::collections::HashSet;
5use std::marker::PhantomData;
6
7use crate::data::DataMap;
8use crate::visit::{Data, NodeCompactIndexable, NodeCount};
9use crate::visit::{
10    EdgeIndexable, GraphBase, GraphProp, IntoEdgeReferences, IntoEdges, IntoEdgesDirected,
11    IntoNeighbors, IntoNeighborsDirected, IntoNodeIdentifiers, IntoNodeReferences, NodeIndexable,
12    NodeRef, VisitMap, Visitable,
13};
14
15/// A graph filter for nodes.
16pub trait FilterNode<N> {
17    /// Return true to have the node be part of the graph
18    fn include_node(&self, node: N) -> bool;
19}
20
21impl<F, N> FilterNode<N> for F
22where
23    F: Fn(N) -> bool,
24{
25    fn include_node(&self, n: N) -> bool {
26        (*self)(n)
27    }
28}
29
30/// This filter includes the nodes that are contained in the set.
31impl<N> FilterNode<N> for FixedBitSet
32where
33    FixedBitSet: VisitMap<N>,
34{
35    fn include_node(&self, n: N) -> bool {
36        self.is_visited(&n)
37    }
38}
39
40/// This filter includes the nodes that are contained in the set.
41impl<N, S> FilterNode<N> for HashSet<N, S>
42where
43    HashSet<N, S>: VisitMap<N>,
44{
45    fn include_node(&self, n: N) -> bool {
46        self.is_visited(&n)
47    }
48}
49
50// Can't express these as a generic impl over all references since that would conflict with the
51// impl for Fn.
52impl<N> FilterNode<N> for &FixedBitSet
53where
54    FixedBitSet: VisitMap<N>,
55{
56    fn include_node(&self, n: N) -> bool {
57        self.is_visited(&n)
58    }
59}
60
61impl<N, S> FilterNode<N> for &HashSet<N, S>
62where
63    HashSet<N, S>: VisitMap<N>,
64{
65    fn include_node(&self, n: N) -> bool {
66        self.is_visited(&n)
67    }
68}
69
70/// A node-filtering graph adaptor.
71#[derive(Copy, Clone, Debug)]
72pub struct NodeFiltered<G, F>(pub G, pub F);
73
74impl<F, G> NodeFiltered<G, F>
75where
76    G: GraphBase,
77    F: Fn(G::NodeId) -> bool,
78{
79    /// Create an `NodeFiltered` adaptor from the closure `filter`.
80    pub fn from_fn(graph: G, filter: F) -> Self {
81        NodeFiltered(graph, filter)
82    }
83}
84
85impl<G, F> GraphBase for NodeFiltered<G, F>
86where
87    G: GraphBase,
88{
89    type NodeId = G::NodeId;
90    type EdgeId = G::EdgeId;
91}
92
93impl<'a, G, F> IntoNeighbors for &'a NodeFiltered<G, F>
94where
95    G: IntoNeighbors,
96    F: FilterNode<G::NodeId>,
97{
98    type Neighbors = NodeFilteredNeighbors<'a, G::Neighbors, F>;
99    fn neighbors(self, n: G::NodeId) -> Self::Neighbors {
100        NodeFilteredNeighbors {
101            include_source: self.1.include_node(n),
102            iter: self.0.neighbors(n),
103            f: &self.1,
104        }
105    }
106}
107
108/// A filtered neighbors iterator.
109#[derive(Debug, Clone)]
110pub struct NodeFilteredNeighbors<'a, I, F: 'a> {
111    include_source: bool,
112    iter: I,
113    f: &'a F,
114}
115
116impl<I, F> Iterator for NodeFilteredNeighbors<'_, I, F>
117where
118    I: Iterator,
119    I::Item: Copy,
120    F: FilterNode<I::Item>,
121{
122    type Item = I::Item;
123    fn next(&mut self) -> Option<Self::Item> {
124        let f = self.f;
125        if !self.include_source {
126            None
127        } else {
128            self.iter.find(move |&target| f.include_node(target))
129        }
130    }
131    fn size_hint(&self) -> (usize, Option<usize>) {
132        let (_, upper) = self.iter.size_hint();
133        (0, upper)
134    }
135}
136
137impl<'a, G, F> IntoNeighborsDirected for &'a NodeFiltered<G, F>
138where
139    G: IntoNeighborsDirected,
140    F: FilterNode<G::NodeId>,
141{
142    type NeighborsDirected = NodeFilteredNeighbors<'a, G::NeighborsDirected, F>;
143    fn neighbors_directed(self, n: G::NodeId, dir: Direction) -> Self::NeighborsDirected {
144        NodeFilteredNeighbors {
145            include_source: self.1.include_node(n),
146            iter: self.0.neighbors_directed(n, dir),
147            f: &self.1,
148        }
149    }
150}
151
152impl<'a, G, F> IntoNodeIdentifiers for &'a NodeFiltered<G, F>
153where
154    G: IntoNodeIdentifiers,
155    F: FilterNode<G::NodeId>,
156{
157    type NodeIdentifiers = NodeFilteredNeighbors<'a, G::NodeIdentifiers, F>;
158    fn node_identifiers(self) -> Self::NodeIdentifiers {
159        NodeFilteredNeighbors {
160            include_source: true,
161            iter: self.0.node_identifiers(),
162            f: &self.1,
163        }
164    }
165}
166
167impl<'a, G, F> IntoNodeReferences for &'a NodeFiltered<G, F>
168where
169    G: IntoNodeReferences,
170    F: FilterNode<G::NodeId>,
171{
172    type NodeRef = G::NodeRef;
173    type NodeReferences = NodeFilteredNodes<'a, G::NodeReferences, F>;
174    fn node_references(self) -> Self::NodeReferences {
175        NodeFilteredNodes {
176            include_source: true,
177            iter: self.0.node_references(),
178            f: &self.1,
179        }
180    }
181}
182
183/// A filtered node references iterator.
184#[derive(Debug, Clone)]
185pub struct NodeFilteredNodes<'a, I, F: 'a> {
186    include_source: bool,
187    iter: I,
188    f: &'a F,
189}
190
191impl<I, F> Iterator for NodeFilteredNodes<'_, I, F>
192where
193    I: Iterator,
194    I::Item: Copy + NodeRef,
195    F: FilterNode<<I::Item as NodeRef>::NodeId>,
196{
197    type Item = I::Item;
198    fn next(&mut self) -> Option<Self::Item> {
199        let f = self.f;
200        if !self.include_source {
201            None
202        } else {
203            self.iter.find(move |&target| f.include_node(target.id()))
204        }
205    }
206    fn size_hint(&self) -> (usize, Option<usize>) {
207        let (_, upper) = self.iter.size_hint();
208        (0, upper)
209    }
210}
211
212impl<'a, G, F> IntoEdgeReferences for &'a NodeFiltered<G, F>
213where
214    G: IntoEdgeReferences,
215    F: FilterNode<G::NodeId>,
216{
217    type EdgeRef = G::EdgeRef;
218    type EdgeReferences = NodeFilteredEdgeReferences<'a, G, G::EdgeReferences, F>;
219    fn edge_references(self) -> Self::EdgeReferences {
220        NodeFilteredEdgeReferences {
221            graph: PhantomData,
222            iter: self.0.edge_references(),
223            f: &self.1,
224        }
225    }
226}
227
228/// A filtered edges iterator.
229#[derive(Debug, Clone)]
230pub struct NodeFilteredEdgeReferences<'a, G, I, F: 'a> {
231    graph: PhantomData<G>,
232    iter: I,
233    f: &'a F,
234}
235
236impl<G, I, F> Iterator for NodeFilteredEdgeReferences<'_, G, I, F>
237where
238    F: FilterNode<G::NodeId>,
239    G: IntoEdgeReferences,
240    I: Iterator<Item = G::EdgeRef>,
241{
242    type Item = I::Item;
243    fn next(&mut self) -> Option<Self::Item> {
244        let f = self.f;
245        self.iter
246            .find(move |&edge| f.include_node(edge.source()) && f.include_node(edge.target()))
247    }
248    fn size_hint(&self) -> (usize, Option<usize>) {
249        let (_, upper) = self.iter.size_hint();
250        (0, upper)
251    }
252}
253
254impl<'a, G, F> IntoEdges for &'a NodeFiltered<G, F>
255where
256    G: IntoEdges,
257    F: FilterNode<G::NodeId>,
258{
259    type Edges = NodeFilteredEdges<'a, G, G::Edges, F>;
260    fn edges(self, a: G::NodeId) -> Self::Edges {
261        NodeFilteredEdges {
262            graph: PhantomData,
263            include_source: self.1.include_node(a),
264            iter: self.0.edges(a),
265            f: &self.1,
266            dir: Direction::Outgoing,
267        }
268    }
269}
270
271impl<'a, G, F> IntoEdgesDirected for &'a NodeFiltered<G, F>
272where
273    G: IntoEdgesDirected,
274    F: FilterNode<G::NodeId>,
275{
276    type EdgesDirected = NodeFilteredEdges<'a, G, G::EdgesDirected, F>;
277    fn edges_directed(self, a: G::NodeId, dir: Direction) -> Self::EdgesDirected {
278        NodeFilteredEdges {
279            graph: PhantomData,
280            include_source: self.1.include_node(a),
281            iter: self.0.edges_directed(a, dir),
282            f: &self.1,
283            dir,
284        }
285    }
286}
287
288/// A filtered edges iterator.
289#[derive(Debug, Clone)]
290pub struct NodeFilteredEdges<'a, G, I, F: 'a> {
291    graph: PhantomData<G>,
292    include_source: bool,
293    iter: I,
294    f: &'a F,
295    dir: Direction,
296}
297
298impl<G, I, F> Iterator for NodeFilteredEdges<'_, G, I, F>
299where
300    F: FilterNode<G::NodeId>,
301    G: IntoEdges,
302    I: Iterator<Item = G::EdgeRef>,
303{
304    type Item = I::Item;
305    fn next(&mut self) -> Option<Self::Item> {
306        if !self.include_source {
307            None
308        } else {
309            let dir = self.dir;
310            let f = self.f;
311            self.iter.find(move |&edge| {
312                f.include_node(match dir {
313                    Direction::Outgoing => edge.target(),
314                    Direction::Incoming => edge.source(),
315                })
316            })
317        }
318    }
319    fn size_hint(&self) -> (usize, Option<usize>) {
320        let (_, upper) = self.iter.size_hint();
321        (0, upper)
322    }
323}
324
325impl<G, F> DataMap for NodeFiltered<G, F>
326where
327    G: DataMap,
328    F: FilterNode<G::NodeId>,
329{
330    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight> {
331        if self.1.include_node(id) {
332            self.0.node_weight(id)
333        } else {
334            None
335        }
336    }
337
338    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight> {
339        self.0.edge_weight(id)
340    }
341}
342
343macro_rules! access0 {
344    ($e:expr) => {
345        $e.0
346    };
347}
348
349Data! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
350NodeIndexable! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
351EdgeIndexable! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
352GraphProp! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
353Visitable! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
354
355/// A graph filter for edges
356pub trait FilterEdge<Edge> {
357    /// Return true to have the edge be part of the graph
358    fn include_edge(&self, edge: Edge) -> bool;
359}
360
361impl<F, N> FilterEdge<N> for F
362where
363    F: Fn(N) -> bool,
364{
365    fn include_edge(&self, n: N) -> bool {
366        (*self)(n)
367    }
368}
369
370/// An edge-filtering graph adaptor.
371///
372/// The adaptor may filter out edges. The filter implements the trait
373/// `FilterEdge`. Closures of type `Fn(G::EdgeRef) -> bool` already
374/// implement this trait.
375///
376/// The filter may use edge source, target, id, and weight to select whether to
377/// include the edge or not.
378#[derive(Copy, Clone, Debug)]
379pub struct EdgeFiltered<G, F>(pub G, pub F);
380
381impl<F, G> EdgeFiltered<G, F>
382where
383    G: IntoEdgeReferences,
384    F: Fn(G::EdgeRef) -> bool,
385{
386    /// Create an `EdgeFiltered` adaptor from the closure `filter`.
387    pub fn from_fn(graph: G, filter: F) -> Self {
388        EdgeFiltered(graph, filter)
389    }
390}
391
392impl<G, F> GraphBase for EdgeFiltered<G, F>
393where
394    G: GraphBase,
395{
396    type NodeId = G::NodeId;
397    type EdgeId = G::EdgeId;
398}
399
400impl<'a, G, F> IntoNeighbors for &'a EdgeFiltered<G, F>
401where
402    G: IntoEdges,
403    F: FilterEdge<G::EdgeRef>,
404{
405    type Neighbors = EdgeFilteredNeighbors<'a, G, F>;
406    fn neighbors(self, n: G::NodeId) -> Self::Neighbors {
407        EdgeFilteredNeighbors {
408            iter: self.0.edges(n),
409            f: &self.1,
410        }
411    }
412}
413
414impl<'a, G, F> IntoNeighborsDirected for &'a EdgeFiltered<G, F>
415where
416    G: IntoEdgesDirected,
417    F: FilterEdge<G::EdgeRef>,
418{
419    type NeighborsDirected = EdgeFilteredNeighborsDirected<'a, G, F>;
420    fn neighbors_directed(self, n: G::NodeId, dir: Direction) -> Self::NeighborsDirected {
421        EdgeFilteredNeighborsDirected {
422            iter: self.0.edges_directed(n, dir),
423            f: &self.1,
424            from: n,
425        }
426    }
427}
428
429/// A filtered neighbors iterator.
430#[derive(Debug, Clone)]
431pub struct EdgeFilteredNeighbors<'a, G, F: 'a>
432where
433    G: IntoEdges,
434{
435    iter: G::Edges,
436    f: &'a F,
437}
438
439impl<G, F> Iterator for EdgeFilteredNeighbors<'_, G, F>
440where
441    F: FilterEdge<G::EdgeRef>,
442    G: IntoEdges,
443{
444    type Item = G::NodeId;
445    fn next(&mut self) -> Option<Self::Item> {
446        let f = self.f;
447        (&mut self.iter)
448            .filter_map(move |edge| {
449                if f.include_edge(edge) {
450                    Some(edge.target())
451                } else {
452                    None
453                }
454            })
455            .next()
456    }
457    fn size_hint(&self) -> (usize, Option<usize>) {
458        let (_, upper) = self.iter.size_hint();
459        (0, upper)
460    }
461}
462
463impl<'a, G, F> IntoEdgeReferences for &'a EdgeFiltered<G, F>
464where
465    G: IntoEdgeReferences,
466    F: FilterEdge<G::EdgeRef>,
467{
468    type EdgeRef = G::EdgeRef;
469    type EdgeReferences = EdgeFilteredEdges<'a, G, G::EdgeReferences, F>;
470    fn edge_references(self) -> Self::EdgeReferences {
471        EdgeFilteredEdges {
472            graph: PhantomData,
473            iter: self.0.edge_references(),
474            f: &self.1,
475        }
476    }
477}
478
479impl<'a, G, F> IntoEdges for &'a EdgeFiltered<G, F>
480where
481    G: IntoEdges,
482    F: FilterEdge<G::EdgeRef>,
483{
484    type Edges = EdgeFilteredEdges<'a, G, G::Edges, F>;
485    fn edges(self, n: G::NodeId) -> Self::Edges {
486        EdgeFilteredEdges {
487            graph: PhantomData,
488            iter: self.0.edges(n),
489            f: &self.1,
490        }
491    }
492}
493
494impl<'a, G, F> IntoEdgesDirected for &'a EdgeFiltered<G, F>
495where
496    G: IntoEdgesDirected,
497    F: FilterEdge<G::EdgeRef>,
498{
499    type EdgesDirected = EdgeFilteredEdges<'a, G, G::EdgesDirected, F>;
500
501    fn edges_directed(self, n: G::NodeId, dir: Direction) -> Self::EdgesDirected {
502        EdgeFilteredEdges {
503            graph: PhantomData,
504            iter: self.0.edges_directed(n, dir),
505            f: &self.1,
506        }
507    }
508}
509
510/// A filtered edges iterator.
511#[derive(Debug, Clone)]
512pub struct EdgeFilteredEdges<'a, G, I, F: 'a> {
513    graph: PhantomData<G>,
514    iter: I,
515    f: &'a F,
516}
517
518impl<G, I, F> Iterator for EdgeFilteredEdges<'_, G, I, F>
519where
520    F: FilterEdge<G::EdgeRef>,
521    G: IntoEdgeReferences,
522    I: Iterator<Item = G::EdgeRef>,
523{
524    type Item = I::Item;
525    fn next(&mut self) -> Option<Self::Item> {
526        let f = self.f;
527        self.iter.find(move |&edge| f.include_edge(edge))
528    }
529    fn size_hint(&self) -> (usize, Option<usize>) {
530        let (_, upper) = self.iter.size_hint();
531        (0, upper)
532    }
533}
534
535/// A filtered neighbors-directed iterator.
536#[derive(Debug, Clone)]
537pub struct EdgeFilteredNeighborsDirected<'a, G, F: 'a>
538where
539    G: IntoEdgesDirected,
540{
541    iter: G::EdgesDirected,
542    f: &'a F,
543    from: G::NodeId,
544}
545
546impl<G, F> Iterator for EdgeFilteredNeighborsDirected<'_, G, F>
547where
548    F: FilterEdge<G::EdgeRef>,
549    G: IntoEdgesDirected,
550{
551    type Item = G::NodeId;
552    fn next(&mut self) -> Option<Self::Item> {
553        let f = self.f;
554        let from = self.from;
555        (&mut self.iter)
556            .filter_map(move |edge| {
557                if f.include_edge(edge) {
558                    if edge.source() != from {
559                        Some(edge.source())
560                    } else {
561                        Some(edge.target()) // includes case where from == source == target
562                    }
563                } else {
564                    None
565                }
566            })
567            .next()
568    }
569    fn size_hint(&self) -> (usize, Option<usize>) {
570        let (_, upper) = self.iter.size_hint();
571        (0, upper)
572    }
573}
574
575Data! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
576GraphProp! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
577IntoNodeIdentifiers! {delegate_impl [['a, G, F], G, &'a EdgeFiltered<G, F>, access0]}
578IntoNodeReferences! {delegate_impl [['a, G, F], G, &'a EdgeFiltered<G, F>, access0]}
579NodeCompactIndexable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
580NodeCount! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
581NodeIndexable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
582EdgeIndexable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
583Visitable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
petgraph/visit/filter.rs

petgraph/visit/
filter.rs
