Skip to main content

weavatrix_graph/algo/walk/
events.rs

1use crate::IndexGraphView;
2use crate::Vec;
3use crate::algo::traversal::{Direction, for_each_adjacent};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6pub enum TraversalControl {
7    Continue,
8    Break,
9}
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum DfsEvent<Node, Edge> {
13    Discover(Node),
14    TreeEdge {
15        edge: Edge,
16        source: Node,
17        target: Node,
18    },
19    BackEdge {
20        edge: Edge,
21        source: Node,
22        target: Node,
23    },
24    CrossForwardEdge {
25        edge: Edge,
26        source: Node,
27        target: Node,
28    },
29    Finish(Node),
30}
31
32#[derive(Debug, Clone)]
33struct Frame<Node, Edge> {
34    node: Node,
35    adjacent: Vec<(Edge, Node)>,
36    next: usize,
37}
38
39/// Reusable color map and iterative stack for DFS event traversal.
40#[derive(Debug, Clone)]
41pub struct DfsEventWorkspace<Node, Edge> {
42    colors: Vec<u8>,
43    stack: Vec<Frame<Node, Edge>>,
44}
45
46impl<Node, Edge> DfsEventWorkspace<Node, Edge> {
47    #[must_use]
48    pub const fn new() -> Self {
49        Self {
50            colors: Vec::new(),
51            stack: Vec::new(),
52        }
53    }
54
55    fn begin(&mut self, node_bound: usize) {
56        self.colors.resize(node_bound, 0);
57        self.colors.fill(0);
58        self.stack.clear();
59    }
60}
61
62impl<Node, Edge> Default for DfsEventWorkspace<Node, Edge> {
63    fn default() -> Self {
64        Self::new()
65    }
66}
67
68pub fn depth_first_search<G, I, V>(
69    graph: &G,
70    starts: I,
71    workspace: &mut DfsEventWorkspace<G::Node, G::Edge>,
72    visitor: V,
73) -> bool
74where
75    G: IndexGraphView,
76    I: IntoIterator<Item = G::Node>,
77    V: FnMut(DfsEvent<G::Node, G::Edge>) -> TraversalControl,
78{
79    depth_first_search_filtered(
80        graph,
81        starts,
82        Direction::Outgoing,
83        workspace,
84        |_| true,
85        visitor,
86    )
87}
88
89pub fn depth_first_search_filtered<G, I, F, V>(
90    graph: &G,
91    starts: I,
92    direction: Direction,
93    workspace: &mut DfsEventWorkspace<G::Node, G::Edge>,
94    mut keep_edge: F,
95    mut visitor: V,
96) -> bool
97where
98    G: IndexGraphView,
99    I: IntoIterator<Item = G::Node>,
100    F: FnMut(G::Edge) -> bool,
101    V: FnMut(DfsEvent<G::Node, G::Edge>) -> TraversalControl,
102{
103    workspace.begin(graph.node_bound());
104    for root in starts {
105        if !graph.contains_node(root) || color::<G>(workspace, root) != 0 {
106            continue;
107        }
108        if !discover(
109            graph,
110            root,
111            direction,
112            workspace,
113            &mut keep_edge,
114            &mut visitor,
115        ) {
116            return false;
117        }
118        while !workspace.stack.is_empty() {
119            let frame_index = workspace.stack.len() - 1;
120            let next = workspace.stack[frame_index].next;
121            if let Some(&(edge, target)) = workspace.stack[frame_index].adjacent.get(next) {
122                let source = workspace.stack[frame_index].node;
123                workspace.stack[frame_index].next += 1;
124                let event = match color::<G>(workspace, target) {
125                    0 => DfsEvent::TreeEdge {
126                        edge,
127                        source,
128                        target,
129                    },
130                    1 => DfsEvent::BackEdge {
131                        edge,
132                        source,
133                        target,
134                    },
135                    _ => DfsEvent::CrossForwardEdge {
136                        edge,
137                        source,
138                        target,
139                    },
140                };
141                if visitor(event) == TraversalControl::Break {
142                    return false;
143                }
144                if color::<G>(workspace, target) == 0
145                    && !discover(
146                        graph,
147                        target,
148                        direction,
149                        workspace,
150                        &mut keep_edge,
151                        &mut visitor,
152                    )
153                {
154                    return false;
155                }
156            } else {
157                let Some(frame) = workspace.stack.pop() else {
158                    break;
159                };
160                workspace.colors[G::node_slot(frame.node)] = 2;
161                if visitor(DfsEvent::Finish(frame.node)) == TraversalControl::Break {
162                    return false;
163                }
164            }
165        }
166    }
167    true
168}
169
170fn discover<G, F, V>(
171    graph: &G,
172    node: G::Node,
173    direction: Direction,
174    workspace: &mut DfsEventWorkspace<G::Node, G::Edge>,
175    keep_edge: &mut F,
176    visitor: &mut V,
177) -> bool
178where
179    G: IndexGraphView,
180    F: FnMut(G::Edge) -> bool,
181    V: FnMut(DfsEvent<G::Node, G::Edge>) -> TraversalControl,
182{
183    workspace.colors[G::node_slot(node)] = 1;
184    if visitor(DfsEvent::Discover(node)) == TraversalControl::Break {
185        return false;
186    }
187    let mut adjacent = Vec::new();
188    for_each_adjacent(graph, node, direction, keep_edge, |edge, target| {
189        adjacent.push((edge, target));
190    });
191    workspace.stack.push(Frame {
192        node,
193        adjacent,
194        next: 0,
195    });
196    true
197}
198
199fn color<G>(workspace: &DfsEventWorkspace<G::Node, G::Edge>, node: G::Node) -> u8
200where
201    G: IndexGraphView,
202{
203    workspace
204        .colors
205        .get(G::node_slot(node))
206        .copied()
207        .unwrap_or(2)
208}