1use std::collections::HashSet;
6
7#[cfg(test)]
8use super::DEFAULT_MAX_VISITED;
9use crate::csr::{CsrIndex, Direction};
10use crate::overlay_delta::GraphOverlayDelta;
11use crate::traversal_overlay::{OverlaySubgraphParams, subgraph_overlay};
12
13impl CsrIndex {
14 pub fn subgraph(
33 &self,
34 start_nodes: &[&str],
35 label_filter: &[&str],
36 direction: Direction,
37 max_depth: usize,
38 max_visited: usize,
39 overlay: Option<&GraphOverlayDelta>,
40 ) -> Vec<(String, String, String)> {
41 Self::subgraph_on(
42 Some(self),
43 start_nodes,
44 label_filter,
45 direction,
46 max_depth,
47 max_visited,
48 overlay,
49 )
50 }
51
52 pub fn subgraph_on(
56 partition: Option<&Self>,
57 start_nodes: &[&str],
58 label_filter: &[&str],
59 direction: Direction,
60 max_depth: usize,
61 max_visited: usize,
62 overlay: Option<&GraphOverlayDelta>,
63 ) -> Vec<(String, String, String)> {
64 match (partition, overlay) {
65 (_, Some(ov)) if !ov.is_empty() => subgraph_overlay(
66 partition,
67 OverlaySubgraphParams {
68 start_nodes,
69 label_filter,
70 direction,
71 max_depth,
72 max_visited,
73 },
74 ov,
75 ),
76 (Some(csr), _) => {
77 csr.subgraph_dense(start_nodes, label_filter, direction, max_depth, max_visited)
78 }
79 (None, _) => Vec::new(),
80 }
81 }
82
83 fn subgraph_dense(
85 &self,
86 start_nodes: &[&str],
87 label_filter: &[&str],
88 direction: Direction,
89 max_depth: usize,
90 max_visited: usize,
91 ) -> Vec<(String, String, String)> {
92 let labels = self.label_filter(label_filter);
93 let mut visited: HashSet<u32> = HashSet::new();
94 let mut frontier: Vec<u32> = Vec::new();
95 let mut edges = Vec::new();
96 let mut seen: HashSet<(u32, u32, u32)> = HashSet::new();
99
100 for &node in start_nodes {
101 if let Some(&id) = self.node_to_id.get(node)
102 && visited.insert(id)
103 {
104 frontier.push(id);
105 }
106 }
107
108 for _depth in 0..max_depth {
111 if frontier.is_empty() || visited.len() >= max_visited {
112 break;
113 }
114 let mut candidates: Vec<u32> = Vec::new();
115 for &node_id in &frontier {
116 self.record_access(node_id);
117 if matches!(direction, Direction::Out | Direction::Both) {
118 for (lid, dst) in self.dense_iter_out(node_id) {
119 if labels.keeps(lid) {
120 if seen.insert((node_id, lid, dst)) {
121 edges.push((
122 self.id_to_node[node_id as usize].clone(),
123 self.label_name(lid).to_string(),
124 self.id_to_node[dst as usize].clone(),
125 ));
126 }
127 if !visited.contains(&dst) {
128 candidates.push(dst);
129 }
130 }
131 }
132 }
133 if matches!(direction, Direction::In | Direction::Both) {
134 for (lid, src) in self.dense_iter_in(node_id) {
135 if labels.keeps(lid) {
136 if seen.insert((src, lid, node_id)) {
137 edges.push((
138 self.id_to_node[src as usize].clone(),
139 self.label_name(lid).to_string(),
140 self.id_to_node[node_id as usize].clone(),
141 ));
142 }
143 if !visited.contains(&src) {
144 candidates.push(src);
145 }
146 }
147 }
148 }
149 }
150 frontier = self.admit_by_name(candidates, &mut visited, max_visited);
151 }
152
153 for &node_id in &frontier {
156 if matches!(direction, Direction::Out | Direction::Both) {
157 for (lid, dst) in self.dense_iter_out(node_id) {
158 if labels.keeps(lid)
159 && visited.contains(&dst)
160 && seen.insert((node_id, lid, dst))
161 {
162 edges.push((
163 self.id_to_node[node_id as usize].clone(),
164 self.label_name(lid).to_string(),
165 self.id_to_node[dst as usize].clone(),
166 ));
167 }
168 }
169 }
170 if matches!(direction, Direction::In | Direction::Both) {
171 for (lid, src) in self.dense_iter_in(node_id) {
172 if labels.keeps(lid)
173 && visited.contains(&src)
174 && seen.insert((src, lid, node_id))
175 {
176 edges.push((
177 self.id_to_node[src as usize].clone(),
178 self.label_name(lid).to_string(),
179 self.id_to_node[node_id as usize].clone(),
180 ));
181 }
182 }
183 }
184 }
185
186 edges
187 }
188}
189
190#[cfg(test)]
191mod tests {
192 use super::*;
193 use crate::test_support::{chain_csr, test_memory};
194
195 #[test]
196 fn subgraph_materialization() {
197 let csr = chain_csr();
198 let edges = csr.subgraph(&["a"], &[], Direction::Out, 2, DEFAULT_MAX_VISITED, None);
199 assert_eq!(edges.len(), 3);
200 assert!(edges.contains(&("a".into(), "KNOWS".into(), "b".into())));
201 assert!(edges.contains(&("a".into(), "WORKS".into(), "e".into())));
202 assert!(edges.contains(&("b".into(), "KNOWS".into(), "c".into())));
203 }
204 #[test]
205 fn a_capped_subgraph_expands_only_admitted_levels() {
206 let mut csr = CsrIndex::new(test_memory());
208 csr.add_edge("a", "L", "z").unwrap();
209 csr.add_edge("a", "L", "m").unwrap();
210 csr.add_edge("a", "L", "b").unwrap();
211 csr.add_edge("z", "L", "y").unwrap();
212 csr.add_edge("b", "L", "c").unwrap();
213 let mut edges = csr.subgraph(&["a"], &[], Direction::Out, 3, 3, None);
214 edges.sort();
215 let expected: Vec<(String, String, String)> = ["b", "m", "z"]
217 .iter()
218 .map(|dst| ("a".to_string(), "L".to_string(), dst.to_string()))
219 .collect();
220 assert_eq!(edges, expected);
221 }
222
223 #[test]
224 fn both_directions_return_each_physical_edge_once() {
225 let mut csr = CsrIndex::new(test_memory());
226 csr.add_edge("a", "L", "b").unwrap();
227 csr.add_edge("b", "L", "a").unwrap();
228 csr.add_edge("a", "SELF", "a").unwrap();
229 csr.add_edge_in_collection("a", "L", "c", "first").unwrap();
230 csr.add_edge_in_collection("a", "L", "c", "second").unwrap();
231 let mut edges = csr.subgraph(&["a"], &[], Direction::Both, 3, DEFAULT_MAX_VISITED, None);
232 edges.sort();
233 let expected: Vec<(String, String, String)> = [
234 ("a", "L", "b"),
235 ("a", "L", "c"),
236 ("a", "SELF", "a"),
237 ("b", "L", "a"),
238 ]
239 .iter()
240 .map(|(s, l, d)| (s.to_string(), l.to_string(), d.to_string()))
241 .collect();
242 assert_eq!(edges, expected);
243 }
244
245 fn triples(edges: &[(&str, &str, &str)]) -> Vec<(String, String, String)> {
246 let mut out: Vec<(String, String, String)> = edges
247 .iter()
248 .map(|(s, l, d)| (s.to_string(), l.to_string(), d.to_string()))
249 .collect();
250 out.sort();
251 out
252 }
253
254 fn boundary_csr() -> CsrIndex {
255 let mut csr = CsrIndex::new(test_memory());
256 for (src, dst) in [
257 ("a", "b"),
258 ("a", "c"),
259 ("b", "c"),
260 ("c", "b"),
261 ("c", "a"),
262 ("c", "x"),
263 ] {
264 csr.add_edge(src, "L", dst).unwrap();
265 }
266 csr
267 }
268
269 #[test]
270 fn boundary_nodes_record_edges_among_admitted_nodes() {
271 let csr = boundary_csr();
272 let mut staged = GraphOverlayDelta::new();
274 staged.stage_tombstone("q", "L", "r");
275 for overlay in [None, Some(&staged)] {
276 let mut edges =
277 csr.subgraph(&["a"], &[], Direction::Out, 1, DEFAULT_MAX_VISITED, overlay);
278 edges.sort();
279 assert_eq!(
280 edges,
281 triples(&[
282 ("a", "L", "b"),
283 ("a", "L", "c"),
284 ("b", "L", "c"),
285 ("c", "L", "a"),
286 ("c", "L", "b"),
287 ]),
288 "c -> x leaves the admitted set"
289 );
290 }
291 }
292
293 #[test]
294 fn boundary_edges_follow_the_walk_direction() {
295 let mut csr = CsrIndex::new(test_memory());
296 for (src, dst) in [("b", "a"), ("c", "a"), ("b", "c"), ("a", "z")] {
297 csr.add_edge(src, "L", dst).unwrap();
298 }
299 let mut edges = csr.subgraph(&["a"], &[], Direction::In, 1, DEFAULT_MAX_VISITED, None);
300 edges.sort();
301 assert_eq!(
302 edges,
303 triples(&[("b", "L", "a"), ("b", "L", "c"), ("c", "L", "a")])
304 );
305 }
306
307 #[test]
308 fn depth_zero_records_only_the_start_nodes_edges_among_themselves() {
309 let mut csr = CsrIndex::new(test_memory());
310 csr.add_edge("a", "L", "a").unwrap();
311 csr.add_edge("a", "L", "b").unwrap();
312 let edges = csr.subgraph(&["a"], &[], Direction::Out, 0, DEFAULT_MAX_VISITED, None);
313 assert_eq!(edges, triples(&[("a", "L", "a")]));
314 }
315
316 #[test]
317 fn a_label_set_returns_parallel_edges_of_both_labels() {
318 let mut csr = CsrIndex::new(test_memory());
319 csr.add_edge("a", "K", "b").unwrap();
320 csr.add_edge("a", "W", "b").unwrap();
321 csr.add_edge("a", "X", "c").unwrap();
322 let mut edges = csr.subgraph(
323 &["a"],
324 &["K", "W"],
325 Direction::Out,
326 2,
327 DEFAULT_MAX_VISITED,
328 None,
329 );
330 edges.sort();
331 let expected: Vec<(String, String, String)> = [("a", "K", "b"), ("a", "W", "b")]
332 .iter()
333 .map(|(s, l, d)| (s.to_string(), l.to_string(), d.to_string()))
334 .collect();
335 assert_eq!(edges, expected);
336 }
337}