Skip to main content

weavatrix_graph/traversal_cache/
eager.rs

1use super::TraversalCache;
2use super::core::{NeighborCsr, NeighborStorage, OffsetStorage};
3use super::elias_fano::EliasFano;
4use super::packed::PackedU32;
5use super::walk::TraversalCacheWorkspace;
6use crate::{Direction, NodeIndex};
7
8impl TraversalCache {
9    /// Runs a materialized BFS without allocating after the workspace has grown.
10    #[must_use]
11    pub fn bfs_with_workspace<'workspace>(
12        &self,
13        start: NodeIndex,
14        direction: Direction,
15        workspace: &'workspace mut TraversalCacheWorkspace,
16    ) -> &'workspace [NodeIndex] {
17        workspace.begin(self.node_count());
18        if !self.contains(start) {
19            return &workspace.visited;
20        }
21        workspace.mark(start);
22        workspace.visited.push(start);
23        match direction {
24            Direction::Outgoing => bfs_adjacency(&self.outgoing, workspace),
25            Direction::Incoming => bfs_adjacency(&self.incoming, workspace),
26            Direction::Both => bfs_both(self, workspace),
27        }
28        &workspace.visited
29    }
30}
31
32fn bfs_adjacency(adjacency: &NeighborCsr, workspace: &mut TraversalCacheWorkspace) {
33    match (&adjacency.offsets, &adjacency.neighbors) {
34        (OffsetStorage::Direct(offsets), NeighborStorage::Direct(neighbors)) => {
35            bfs_direct(offsets, neighbors, workspace);
36        }
37        (OffsetStorage::Direct(offsets), NeighborStorage::Packed(neighbors)) => {
38            bfs_packed(offsets, neighbors, workspace);
39        }
40        (OffsetStorage::EliasFano(offsets), NeighborStorage::Packed(neighbors)) => {
41            bfs_succinct(offsets, neighbors, workspace);
42        }
43        (OffsetStorage::EliasFano(offsets), NeighborStorage::Adaptive(neighbors)) => {
44            bfs_adaptive(offsets, neighbors, workspace);
45        }
46        (OffsetStorage::Direct(_), NeighborStorage::Adaptive(_)) => {
47            unreachable!("adaptive neighbors are reserved for compact storage");
48        }
49        (OffsetStorage::EliasFano(_), NeighborStorage::Direct(_)) => {
50            unreachable!("direct neighbors always use direct offsets");
51        }
52    }
53}
54
55fn bfs_adaptive(
56    offsets: &EliasFano,
57    neighbors: &super::adaptive::AdaptivePackedU32,
58    workspace: &mut TraversalCacheWorkspace,
59) {
60    let mut cursor = 0;
61    while cursor < workspace.visited.len() {
62        let node = workspace.visited[cursor].index();
63        cursor += 1;
64        let start = offsets.get(node) as usize;
65        let end = offsets.get(node + 1) as usize;
66        neighbors.for_each(start, end, |raw| {
67            push_unseen(NodeIndex::new(raw), workspace);
68        });
69    }
70}
71
72fn bfs_direct(offsets: &[u32], neighbors: &[u32], workspace: &mut TraversalCacheWorkspace) {
73    let mut cursor = 0;
74    while cursor < workspace.visited.len() {
75        let node = workspace.visited[cursor].index();
76        cursor += 1;
77        let start = offsets[node] as usize;
78        let end = offsets[node + 1] as usize;
79        for &raw in &neighbors[start..end] {
80            let neighbor = NodeIndex::new(raw);
81            if workspace.mark(neighbor) {
82                workspace.visited.push(neighbor);
83            }
84        }
85    }
86}
87
88fn bfs_packed(offsets: &[u32], neighbors: &PackedU32, workspace: &mut TraversalCacheWorkspace) {
89    let mut cursor = 0;
90    while cursor < workspace.visited.len() {
91        let node = workspace.visited[cursor].index();
92        cursor += 1;
93        let start = offsets[node] as usize;
94        let end = offsets[node + 1] as usize;
95        append_packed(neighbors, start, end, workspace);
96    }
97}
98
99fn bfs_succinct(
100    offsets: &EliasFano,
101    neighbors: &PackedU32,
102    workspace: &mut TraversalCacheWorkspace,
103) {
104    let mut cursor = 0;
105    while cursor < workspace.visited.len() {
106        let node = workspace.visited[cursor].index();
107        cursor += 1;
108        let start = offsets.get(node) as usize;
109        let end = offsets.get(node + 1) as usize;
110        append_packed(neighbors, start, end, workspace);
111    }
112}
113
114fn bfs_both(cache: &TraversalCache, workspace: &mut TraversalCacheWorkspace) {
115    let mut cursor = 0;
116    while cursor < workspace.visited.len() {
117        let node = workspace.visited[cursor];
118        cursor += 1;
119        append_adjacency(&cache.outgoing, node, workspace);
120        append_adjacency(&cache.incoming, node, workspace);
121    }
122}
123
124fn append_adjacency(
125    adjacency: &NeighborCsr,
126    node: NodeIndex,
127    workspace: &mut TraversalCacheWorkspace,
128) {
129    let start = adjacency.offsets.get(node.index()) as usize;
130    let end = adjacency.offsets.get(node.index() + 1) as usize;
131    match &adjacency.neighbors {
132        NeighborStorage::Direct(values) => {
133            for &raw in &values[start..end] {
134                push_unseen(NodeIndex::new(raw), workspace);
135            }
136        }
137        NeighborStorage::Packed(values) => append_packed(values, start, end, workspace),
138        NeighborStorage::Adaptive(values) => values.for_each(start, end, |raw| {
139            push_unseen(NodeIndex::new(raw), workspace);
140        }),
141    }
142}
143
144fn append_packed(
145    neighbors: &PackedU32,
146    start: usize,
147    end: usize,
148    workspace: &mut TraversalCacheWorkspace,
149) {
150    neighbors.for_each(start, end, |raw| {
151        push_unseen(NodeIndex::new(raw), workspace);
152    });
153}
154
155#[inline]
156fn push_unseen(neighbor: NodeIndex, workspace: &mut TraversalCacheWorkspace) {
157    if workspace.mark(neighbor) {
158        workspace.visited.push(neighbor);
159    }
160}