1 use crate::prelude::*;
2 
3 use fixedbitset::FixedBitSet;
4 use std::collections::HashSet;
5 use std::marker::PhantomData;
6 
7 use crate::data::DataMap;
8 use crate::visit::{Data, NodeCompactIndexable, NodeCount};
9 use 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.
16 pub trait FilterNode<N> {
17     /// Return true to have the node be part of the graph
include_node(&self, node: N) -> bool18     fn include_node(&self, node: N) -> bool;
19 }
20 
21 impl<F, N> FilterNode<N> for F
22 where
23     F: Fn(N) -> bool,
24 {
include_node(&self, n: N) -> bool25     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.
31 impl<N> FilterNode<N> for FixedBitSet
32 where
33     FixedBitSet: VisitMap<N>,
34 {
include_node(&self, n: N) -> bool35     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.
41 impl<N, S> FilterNode<N> for HashSet<N, S>
42 where
43     HashSet<N, S>: VisitMap<N>,
44 {
include_node(&self, n: N) -> bool45     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.
52 impl<N> FilterNode<N> for &FixedBitSet
53 where
54     FixedBitSet: VisitMap<N>,
55 {
include_node(&self, n: N) -> bool56     fn include_node(&self, n: N) -> bool {
57         self.is_visited(&n)
58     }
59 }
60 
61 impl<N, S> FilterNode<N> for &HashSet<N, S>
62 where
63     HashSet<N, S>: VisitMap<N>,
64 {
include_node(&self, n: N) -> bool65     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)]
72 pub struct NodeFiltered<G, F>(pub G, pub F);
73 
74 impl<F, G> NodeFiltered<G, F>
75 where
76     G: GraphBase,
77     F: Fn(G::NodeId) -> bool,
78 {
79     /// Create an `NodeFiltered` adaptor from the closure `filter`.
from_fn(graph: G, filter: F) -> Self80     pub fn from_fn(graph: G, filter: F) -> Self {
81         NodeFiltered(graph, filter)
82     }
83 }
84 
85 impl<G, F> GraphBase for NodeFiltered<G, F>
86 where
87     G: GraphBase,
88 {
89     type NodeId = G::NodeId;
90     type EdgeId = G::EdgeId;
91 }
92 
93 impl<'a, G, F> IntoNeighbors for &'a NodeFiltered<G, F>
94 where
95     G: IntoNeighbors,
96     F: FilterNode<G::NodeId>,
97 {
98     type Neighbors = NodeFilteredNeighbors<'a, G::Neighbors, F>;
neighbors(self, n: G::NodeId) -> Self::Neighbors99     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)]
110 pub struct NodeFilteredNeighbors<'a, I, F: 'a> {
111     include_source: bool,
112     iter: I,
113     f: &'a F,
114 }
115 
116 impl<'a, I, F> Iterator for NodeFilteredNeighbors<'a, I, F>
117 where
118     I: Iterator,
119     I::Item: Copy,
120     F: FilterNode<I::Item>,
121 {
122     type Item = I::Item;
next(&mut self) -> Option<Self::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     }
size_hint(&self) -> (usize, Option<usize>)131     fn size_hint(&self) -> (usize, Option<usize>) {
132         let (_, upper) = self.iter.size_hint();
133         (0, upper)
134     }
135 }
136 
137 impl<'a, G, F> IntoNeighborsDirected for &'a NodeFiltered<G, F>
138 where
139     G: IntoNeighborsDirected,
140     F: FilterNode<G::NodeId>,
141 {
142     type NeighborsDirected = NodeFilteredNeighbors<'a, G::NeighborsDirected, F>;
neighbors_directed(self, n: G::NodeId, dir: Direction) -> Self::NeighborsDirected143     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 
152 impl<'a, G, F> IntoNodeIdentifiers for &'a NodeFiltered<G, F>
153 where
154     G: IntoNodeIdentifiers,
155     F: FilterNode<G::NodeId>,
156 {
157     type NodeIdentifiers = NodeFilteredNeighbors<'a, G::NodeIdentifiers, F>;
node_identifiers(self) -> Self::NodeIdentifiers158     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 
167 impl<'a, G, F> IntoNodeReferences for &'a NodeFiltered<G, F>
168 where
169     G: IntoNodeReferences,
170     F: FilterNode<G::NodeId>,
171 {
172     type NodeRef = G::NodeRef;
173     type NodeReferences = NodeFilteredNodes<'a, G::NodeReferences, F>;
node_references(self) -> Self::NodeReferences174     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)]
185 pub struct NodeFilteredNodes<'a, I, F: 'a> {
186     include_source: bool,
187     iter: I,
188     f: &'a F,
189 }
190 
191 impl<'a, I, F> Iterator for NodeFilteredNodes<'a, I, F>
192 where
193     I: Iterator,
194     I::Item: Copy + NodeRef,
195     F: FilterNode<<I::Item as NodeRef>::NodeId>,
196 {
197     type Item = I::Item;
next(&mut self) -> Option<Self::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     }
size_hint(&self) -> (usize, Option<usize>)206     fn size_hint(&self) -> (usize, Option<usize>) {
207         let (_, upper) = self.iter.size_hint();
208         (0, upper)
209     }
210 }
211 
212 impl<'a, G, F> IntoEdgeReferences for &'a NodeFiltered<G, F>
213 where
214     G: IntoEdgeReferences,
215     F: FilterNode<G::NodeId>,
216 {
217     type EdgeRef = G::EdgeRef;
218     type EdgeReferences = NodeFilteredEdgeReferences<'a, G, G::EdgeReferences, F>;
edge_references(self) -> Self::EdgeReferences219     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)]
230 pub struct NodeFilteredEdgeReferences<'a, G, I, F: 'a> {
231     graph: PhantomData<G>,
232     iter: I,
233     f: &'a F,
234 }
235 
236 impl<'a, G, I, F> Iterator for NodeFilteredEdgeReferences<'a, G, I, F>
237 where
238     F: FilterNode<G::NodeId>,
239     G: IntoEdgeReferences,
240     I: Iterator<Item = G::EdgeRef>,
241 {
242     type Item = I::Item;
next(&mut self) -> Option<Self::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     }
size_hint(&self) -> (usize, Option<usize>)248     fn size_hint(&self) -> (usize, Option<usize>) {
249         let (_, upper) = self.iter.size_hint();
250         (0, upper)
251     }
252 }
253 
254 impl<'a, G, F> IntoEdges for &'a NodeFiltered<G, F>
255 where
256     G: IntoEdges,
257     F: FilterNode<G::NodeId>,
258 {
259     type Edges = NodeFilteredEdges<'a, G, G::Edges, F>;
edges(self, a: G::NodeId) -> Self::Edges260     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 
271 impl<'a, G, F> IntoEdgesDirected for &'a NodeFiltered<G, F>
272 where
273     G: IntoEdgesDirected,
274     F: FilterNode<G::NodeId>,
275 {
276     type EdgesDirected = NodeFilteredEdges<'a, G, G::EdgesDirected, F>;
edges_directed(self, a: G::NodeId, dir: Direction) -> Self::EdgesDirected277     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)]
290 pub 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 
298 impl<'a, G, I, F> Iterator for NodeFilteredEdges<'a, G, I, F>
299 where
300     F: FilterNode<G::NodeId>,
301     G: IntoEdges,
302     I: Iterator<Item = G::EdgeRef>,
303 {
304     type Item = I::Item;
next(&mut self) -> Option<Self::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     }
size_hint(&self) -> (usize, Option<usize>)319     fn size_hint(&self) -> (usize, Option<usize>) {
320         let (_, upper) = self.iter.size_hint();
321         (0, upper)
322     }
323 }
324 
325 impl<G, F> DataMap for NodeFiltered<G, F>
326 where
327     G: DataMap,
328     F: FilterNode<G::NodeId>,
329 {
node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight>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 
edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight>338     fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight> {
339         self.0.edge_weight(id)
340     }
341 }
342 
343 macro_rules! access0 {
344     ($e:expr) => {
345         $e.0
346     };
347 }
348 
349 Data! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
350 NodeIndexable! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
351 EdgeIndexable! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
352 GraphProp! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
353 Visitable! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
354 
355 /// A graph filter for edges
356 pub trait FilterEdge<Edge> {
357     /// Return true to have the edge be part of the graph
include_edge(&self, edge: Edge) -> bool358     fn include_edge(&self, edge: Edge) -> bool;
359 }
360 
361 impl<F, N> FilterEdge<N> for F
362 where
363     F: Fn(N) -> bool,
364 {
include_edge(&self, n: N) -> bool365     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)]
379 pub struct EdgeFiltered<G, F>(pub G, pub F);
380 
381 impl<F, G> EdgeFiltered<G, F>
382 where
383     G: IntoEdgeReferences,
384     F: Fn(G::EdgeRef) -> bool,
385 {
386     /// Create an `EdgeFiltered` adaptor from the closure `filter`.
from_fn(graph: G, filter: F) -> Self387     pub fn from_fn(graph: G, filter: F) -> Self {
388         EdgeFiltered(graph, filter)
389     }
390 }
391 
392 impl<G, F> GraphBase for EdgeFiltered<G, F>
393 where
394     G: GraphBase,
395 {
396     type NodeId = G::NodeId;
397     type EdgeId = G::EdgeId;
398 }
399 
400 impl<'a, G, F> IntoNeighbors for &'a EdgeFiltered<G, F>
401 where
402     G: IntoEdges,
403     F: FilterEdge<G::EdgeRef>,
404 {
405     type Neighbors = EdgeFilteredNeighbors<'a, G, F>;
neighbors(self, n: G::NodeId) -> Self::Neighbors406     fn neighbors(self, n: G::NodeId) -> Self::Neighbors {
407         EdgeFilteredNeighbors {
408             iter: self.0.edges(n),
409             f: &self.1,
410         }
411     }
412 }
413 
414 impl<'a, G, F> IntoNeighborsDirected for &'a EdgeFiltered<G, F>
415 where
416     G: IntoEdgesDirected,
417     F: FilterEdge<G::EdgeRef>,
418 {
419     type NeighborsDirected = EdgeFilteredNeighborsDirected<'a, G, F>;
neighbors_directed(self, n: G::NodeId, dir: Direction) -> Self::NeighborsDirected420     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)]
431 pub struct EdgeFilteredNeighbors<'a, G, F: 'a>
432 where
433     G: IntoEdges,
434 {
435     iter: G::Edges,
436     f: &'a F,
437 }
438 
439 impl<'a, G, F> Iterator for EdgeFilteredNeighbors<'a, G, F>
440 where
441     F: FilterEdge<G::EdgeRef>,
442     G: IntoEdges,
443 {
444     type Item = G::NodeId;
next(&mut self) -> Option<Self::Item>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     }
size_hint(&self) -> (usize, Option<usize>)457     fn size_hint(&self) -> (usize, Option<usize>) {
458         let (_, upper) = self.iter.size_hint();
459         (0, upper)
460     }
461 }
462 
463 impl<'a, G, F> IntoEdgeReferences for &'a EdgeFiltered<G, F>
464 where
465     G: IntoEdgeReferences,
466     F: FilterEdge<G::EdgeRef>,
467 {
468     type EdgeRef = G::EdgeRef;
469     type EdgeReferences = EdgeFilteredEdges<'a, G, G::EdgeReferences, F>;
edge_references(self) -> Self::EdgeReferences470     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 
479 impl<'a, G, F> IntoEdges for &'a EdgeFiltered<G, F>
480 where
481     G: IntoEdges,
482     F: FilterEdge<G::EdgeRef>,
483 {
484     type Edges = EdgeFilteredEdges<'a, G, G::Edges, F>;
edges(self, n: G::NodeId) -> Self::Edges485     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 
494 impl<'a, G, F> IntoEdgesDirected for &'a EdgeFiltered<G, F>
495 where
496     G: IntoEdgesDirected,
497     F: FilterEdge<G::EdgeRef>,
498 {
499     type EdgesDirected = EdgeFilteredEdges<'a, G, G::EdgesDirected, F>;
500 
edges_directed(self, n: G::NodeId, dir: Direction) -> Self::EdgesDirected501     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)]
512 pub struct EdgeFilteredEdges<'a, G, I, F: 'a> {
513     graph: PhantomData<G>,
514     iter: I,
515     f: &'a F,
516 }
517 
518 impl<'a, G, I, F> Iterator for EdgeFilteredEdges<'a, G, I, F>
519 where
520     F: FilterEdge<G::EdgeRef>,
521     G: IntoEdgeReferences,
522     I: Iterator<Item = G::EdgeRef>,
523 {
524     type Item = I::Item;
next(&mut self) -> Option<Self::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     }
size_hint(&self) -> (usize, Option<usize>)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)]
537 pub struct EdgeFilteredNeighborsDirected<'a, G, F: 'a>
538 where
539     G: IntoEdgesDirected,
540 {
541     iter: G::EdgesDirected,
542     f: &'a F,
543     from: G::NodeId,
544 }
545 
546 impl<'a, G, F> Iterator for EdgeFilteredNeighborsDirected<'a, G, F>
547 where
548     F: FilterEdge<G::EdgeRef>,
549     G: IntoEdgesDirected,
550 {
551     type Item = G::NodeId;
next(&mut self) -> Option<Self::Item>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     }
size_hint(&self) -> (usize, Option<usize>)569     fn size_hint(&self) -> (usize, Option<usize>) {
570         let (_, upper) = self.iter.size_hint();
571         (0, upper)
572     }
573 }
574 
575 Data! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
576 GraphProp! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
577 IntoNodeIdentifiers! {delegate_impl [['a, G, F], G, &'a EdgeFiltered<G, F>, access0]}
578 IntoNodeReferences! {delegate_impl [['a, G, F], G, &'a EdgeFiltered<G, F>, access0]}
579 NodeCompactIndexable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
580 NodeCount! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
581 NodeIndexable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
582 EdgeIndexable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
583 Visitable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
584