cargo_modules/graph/
walker.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4
5use std::collections::HashSet;
6
7use petgraph::{Direction, graph::NodeIndex, visit::EdgeRef};
8
9use crate::graph::{Edge, Graph, Node};
10
11pub(crate) struct GraphWalker {
12    direction: Direction,
13    pub(crate) nodes_visited: HashSet<NodeIndex>,
14}
15
16impl GraphWalker {
17    pub(crate) fn new(direction: Direction) -> Self {
18        let nodes_visited = HashSet::default();
19
20        Self {
21            direction,
22            nodes_visited,
23        }
24    }
25
26    pub(crate) fn walk_graph<F>(
27        &mut self,
28        graph: &Graph<Node, Edge>,
29        origin_node_idx: NodeIndex,
30        predicate: F,
31    ) where
32        F: Fn(&Edge, &Node, usize) -> bool,
33    {
34        self.visit_node_recursively(graph, origin_node_idx, 0, &predicate);
35    }
36
37    pub(crate) fn visit_node_recursively<F>(
38        &mut self,
39        graph: &Graph<Node, Edge>,
40        node_idx: NodeIndex,
41        depth: usize,
42        predicate: &F,
43    ) where
44        F: Fn(&Edge, &Node, usize) -> bool,
45    {
46        if self.nodes_visited.contains(&node_idx) {
47            return;
48        }
49
50        self.nodes_visited.insert(node_idx);
51
52        let edges_directed = graph.edges_directed(node_idx, self.direction);
53
54        for edge_ref in edges_directed {
55            let node_idx = match self.direction {
56                Direction::Outgoing => edge_ref.target(),
57                Direction::Incoming => edge_ref.source(),
58            };
59
60            if !predicate(edge_ref.weight(), &graph[node_idx], depth + 1) {
61                continue;
62            }
63
64            self.visit_node_recursively(graph, node_idx, depth + 1, predicate);
65        }
66    }
67}