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        | (OffsetStorage::EliasFano(_), NeighborStorage::Direct(_)) => {
48            bfs_flexible(adjacency, workspace);
49        }
50    }
51}
52
53fn bfs_flexible(adjacency: &NeighborCsr, workspace: &mut TraversalCacheWorkspace) {
54    let mut cursor = 0;
55    while cursor < workspace.visited.len() {
56        let node = workspace.visited[cursor];
57        cursor += 1;
58        append_adjacency(adjacency, node, workspace);
59    }
60}
61
62fn bfs_adaptive(
63    offsets: &EliasFano,
64    neighbors: &super::adaptive::AdaptivePackedU32,
65    workspace: &mut TraversalCacheWorkspace,
66) {
67    let mut cursor = 0;
68    while cursor < workspace.visited.len() {
69        let node = workspace.visited[cursor].index();
70        cursor += 1;
71        let start = offsets.get(node) as usize;
72        let end = offsets.get(node + 1) as usize;
73        neighbors.for_each(start, end, |raw| {
74            push_unseen(NodeIndex::new(raw), workspace);
75        });
76    }
77}
78
79fn bfs_direct(offsets: &[u32], neighbors: &[u32], workspace: &mut TraversalCacheWorkspace) {
80    let mut cursor = 0;
81    while cursor < workspace.visited.len() {
82        let node = workspace.visited[cursor].index();
83        cursor += 1;
84        let start = offsets[node] as usize;
85        let end = offsets[node + 1] as usize;
86        for &raw in &neighbors[start..end] {
87            let neighbor = NodeIndex::new(raw);
88            if workspace.mark(neighbor) {
89                workspace.visited.push(neighbor);
90            }
91        }
92    }
93}
94
95fn bfs_packed(offsets: &[u32], neighbors: &PackedU32, workspace: &mut TraversalCacheWorkspace) {
96    let mut cursor = 0;
97    while cursor < workspace.visited.len() {
98        let node = workspace.visited[cursor].index();
99        cursor += 1;
100        let start = offsets[node] as usize;
101        let end = offsets[node + 1] as usize;
102        append_packed(neighbors, start, end, workspace);
103    }
104}
105
106fn bfs_succinct(
107    offsets: &EliasFano,
108    neighbors: &PackedU32,
109    workspace: &mut TraversalCacheWorkspace,
110) {
111    let mut cursor = 0;
112    while cursor < workspace.visited.len() {
113        let node = workspace.visited[cursor].index();
114        cursor += 1;
115        let start = offsets.get(node) as usize;
116        let end = offsets.get(node + 1) as usize;
117        append_packed(neighbors, start, end, workspace);
118    }
119}
120
121fn bfs_both(cache: &TraversalCache, workspace: &mut TraversalCacheWorkspace) {
122    let mut cursor = 0;
123    while cursor < workspace.visited.len() {
124        let node = workspace.visited[cursor];
125        cursor += 1;
126        append_adjacency(&cache.outgoing, node, workspace);
127        append_adjacency(&cache.incoming, node, workspace);
128    }
129}
130
131fn append_adjacency(
132    adjacency: &NeighborCsr,
133    node: NodeIndex,
134    workspace: &mut TraversalCacheWorkspace,
135) {
136    let start = adjacency.offsets.get(node.index()) as usize;
137    let end = adjacency.offsets.get(node.index() + 1) as usize;
138    match &adjacency.neighbors {
139        NeighborStorage::Direct(values) => {
140            for &raw in &values[start..end] {
141                push_unseen(NodeIndex::new(raw), workspace);
142            }
143        }
144        NeighborStorage::Packed(values) => append_packed(values, start, end, workspace),
145        NeighborStorage::Adaptive(values) => values.for_each(start, end, |raw| {
146            push_unseen(NodeIndex::new(raw), workspace);
147        }),
148    }
149}
150
151fn append_packed(
152    neighbors: &PackedU32,
153    start: usize,
154    end: usize,
155    workspace: &mut TraversalCacheWorkspace,
156) {
157    neighbors.for_each(start, end, |raw| {
158        push_unseen(NodeIndex::new(raw), workspace);
159    });
160}
161
162#[inline]
163fn push_unseen(neighbor: NodeIndex, workspace: &mut TraversalCacheWorkspace) {
164    if workspace.mark(neighbor) {
165        workspace.visited.push(neighbor);
166    }
167}
weavatrix_graph/traversal_cache/eager.rs

weavatrix_graph/traversal_cache/
eager.rs
