1use crate::edge::Edge;
6use crate::hyperedge::{Hyperedge, HyperedgeId};
7use crate::node::Node;
8use crate::types::{EdgeId, NodeId, PropertyValue};
9use dashmap::DashMap;
10use std::collections::HashSet;
11use std::sync::Arc;
12
13#[derive(Debug, Clone)]
15pub struct LabelIndex {
16 index: Arc<DashMap<String, HashSet<NodeId>>>,
18}
19
20impl LabelIndex {
21 pub fn new() -> Self {
23 Self {
24 index: Arc::new(DashMap::new()),
25 }
26 }
27
28 pub fn add_node(&self, node: &Node) {
30 for label in &node.labels {
31 self.index
32 .entry(label.name.clone())
33 .or_insert_with(HashSet::new)
34 .insert(node.id.clone());
35 }
36 }
37
38 pub fn remove_node(&self, node: &Node) {
40 for label in &node.labels {
41 if let Some(mut set) = self.index.get_mut(&label.name) {
42 set.remove(&node.id);
43 }
44 }
45 }
46
47 pub fn get_nodes_by_label(&self, label: &str) -> Vec<NodeId> {
49 self.index
50 .get(label)
51 .map(|set| set.iter().cloned().collect())
52 .unwrap_or_default()
53 }
54
55 pub fn all_labels(&self) -> Vec<String> {
57 self.index.iter().map(|entry| entry.key().clone()).collect()
58 }
59
60 pub fn count_by_label(&self, label: &str) -> usize {
62 self.index.get(label).map(|set| set.len()).unwrap_or(0)
63 }
64
65 pub fn clear(&self) {
67 self.index.clear();
68 }
69}
70
71impl Default for LabelIndex {
72 fn default() -> Self {
73 Self::new()
74 }
75}
76
77#[derive(Debug, Clone)]
79pub struct PropertyIndex {
80 index: Arc<DashMap<String, DashMap<String, HashSet<NodeId>>>>,
82}
83
84impl PropertyIndex {
85 pub fn new() -> Self {
87 Self {
88 index: Arc::new(DashMap::new()),
89 }
90 }
91
92 pub fn add_node(&self, node: &Node) {
94 for (key, value) in &node.properties {
95 let value_str = self.property_value_to_string(value);
96 self.index
97 .entry(key.clone())
98 .or_insert_with(DashMap::new)
99 .entry(value_str)
100 .or_insert_with(HashSet::new)
101 .insert(node.id.clone());
102 }
103 }
104
105 pub fn remove_node(&self, node: &Node) {
107 for (key, value) in &node.properties {
108 let value_str = self.property_value_to_string(value);
109 if let Some(value_map) = self.index.get(key) {
110 if let Some(mut set) = value_map.get_mut(&value_str) {
111 set.remove(&node.id);
112 }
113 }
114 }
115 }
116
117 pub fn get_nodes_by_property(&self, key: &str, value: &PropertyValue) -> Vec<NodeId> {
119 let value_str = self.property_value_to_string(value);
120 self.index
121 .get(key)
122 .and_then(|value_map| {
123 value_map
124 .get(&value_str)
125 .map(|set| set.iter().cloned().collect())
126 })
127 .unwrap_or_default()
128 }
129
130 pub fn get_nodes_with_property(&self, key: &str) -> Vec<NodeId> {
132 self.index
133 .get(key)
134 .map(|value_map| {
135 let mut result = HashSet::new();
136 for entry in value_map.iter() {
137 for id in entry.value().iter() {
138 result.insert(id.clone());
139 }
140 }
141 result.into_iter().collect()
142 })
143 .unwrap_or_default()
144 }
145
146 pub fn all_property_keys(&self) -> Vec<String> {
148 self.index.iter().map(|entry| entry.key().clone()).collect()
149 }
150
151 pub fn clear(&self) {
153 self.index.clear();
154 }
155
156 fn property_value_to_string(&self, value: &PropertyValue) -> String {
158 match value {
159 PropertyValue::Null => "null".to_string(),
160 PropertyValue::Boolean(b) => b.to_string(),
161 PropertyValue::Integer(i) => i.to_string(),
162 PropertyValue::Float(f) => f.to_string(),
163 PropertyValue::String(s) => s.clone(),
164 PropertyValue::FloatArray(_) | PropertyValue::Array(_) | PropertyValue::List(_) => format!("{:?}", value),
165 PropertyValue::Map(_) => format!("{:?}", value),
166 }
167 }
168}
169
170impl Default for PropertyIndex {
171 fn default() -> Self {
172 Self::new()
173 }
174}
175
176#[derive(Debug, Clone)]
178pub struct EdgeTypeIndex {
179 index: Arc<DashMap<String, HashSet<EdgeId>>>,
181}
182
183impl EdgeTypeIndex {
184 pub fn new() -> Self {
186 Self {
187 index: Arc::new(DashMap::new()),
188 }
189 }
190
191 pub fn add_edge(&self, edge: &Edge) {
193 self.index
194 .entry(edge.edge_type.clone())
195 .or_insert_with(HashSet::new)
196 .insert(edge.id.clone());
197 }
198
199 pub fn remove_edge(&self, edge: &Edge) {
201 if let Some(mut set) = self.index.get_mut(&edge.edge_type) {
202 set.remove(&edge.id);
203 }
204 }
205
206 pub fn get_edges_by_type(&self, edge_type: &str) -> Vec<EdgeId> {
208 self.index
209 .get(edge_type)
210 .map(|set| set.iter().cloned().collect())
211 .unwrap_or_default()
212 }
213
214 pub fn all_edge_types(&self) -> Vec<String> {
216 self.index.iter().map(|entry| entry.key().clone()).collect()
217 }
218
219 pub fn count_by_type(&self, edge_type: &str) -> usize {
221 self.index.get(edge_type).map(|set| set.len()).unwrap_or(0)
222 }
223
224 pub fn clear(&self) {
226 self.index.clear();
227 }
228}
229
230impl Default for EdgeTypeIndex {
231 fn default() -> Self {
232 Self::new()
233 }
234}
235
236#[derive(Debug, Clone)]
238pub struct AdjacencyIndex {
239 outgoing: Arc<DashMap<NodeId, HashSet<EdgeId>>>,
241 incoming: Arc<DashMap<NodeId, HashSet<EdgeId>>>,
243}
244
245impl AdjacencyIndex {
246 pub fn new() -> Self {
248 Self {
249 outgoing: Arc::new(DashMap::new()),
250 incoming: Arc::new(DashMap::new()),
251 }
252 }
253
254 pub fn add_edge(&self, edge: &Edge) {
256 self.outgoing
257 .entry(edge.from.clone())
258 .or_insert_with(HashSet::new)
259 .insert(edge.id.clone());
260
261 self.incoming
262 .entry(edge.to.clone())
263 .or_insert_with(HashSet::new)
264 .insert(edge.id.clone());
265 }
266
267 pub fn remove_edge(&self, edge: &Edge) {
269 if let Some(mut set) = self.outgoing.get_mut(&edge.from) {
270 set.remove(&edge.id);
271 }
272 if let Some(mut set) = self.incoming.get_mut(&edge.to) {
273 set.remove(&edge.id);
274 }
275 }
276
277 pub fn get_outgoing_edges(&self, node_id: &NodeId) -> Vec<EdgeId> {
279 self.outgoing
280 .get(node_id)
281 .map(|set| set.iter().cloned().collect())
282 .unwrap_or_default()
283 }
284
285 pub fn get_incoming_edges(&self, node_id: &NodeId) -> Vec<EdgeId> {
287 self.incoming
288 .get(node_id)
289 .map(|set| set.iter().cloned().collect())
290 .unwrap_or_default()
291 }
292
293 pub fn get_all_edges(&self, node_id: &NodeId) -> Vec<EdgeId> {
295 let mut edges = self.get_outgoing_edges(node_id);
296 edges.extend(self.get_incoming_edges(node_id));
297 edges
298 }
299
300 pub fn get_edges_for_nodes(&self, node_ids: &[NodeId]) -> Vec<EdgeId> {
302 let mut result = Vec::with_capacity(node_ids.len() * 4);
303 for node_id in node_ids {
304 if let Some(guard) = self.outgoing.get(node_id) {
305 result.extend(guard.iter().cloned());
306 }
307 }
308
309 result
310 }
311
312 pub fn for_each_outgoing_edge<F: FnMut(&EdgeId)>(&self, node_ids: &[NodeId], mut f: F) {
314 for node_id in node_ids {
315 if let Some(guard) = self.outgoing.get(node_id) {
316 for edge_id in guard.iter() {
317 f(edge_id);
318 }
319 }
320 }
321 }
322
323 pub fn out_degree(&self, node_id: &NodeId) -> usize {
325 self.outgoing.get(node_id).map(|set| set.len()).unwrap_or(0)
326 }
327
328 pub fn in_degree(&self, node_id: &NodeId) -> usize {
330 self.incoming.get(node_id).map(|set| set.len()).unwrap_or(0)
331 }
332
333 pub fn clear(&self) {
335 self.outgoing.clear();
336 self.incoming.clear();
337 }
338}
339
340impl Default for AdjacencyIndex {
341 fn default() -> Self {
342 Self::new()
343 }
344}
345
346#[derive(Debug, Clone)]
348pub struct HyperedgeNodeIndex {
349 index: Arc<DashMap<NodeId, HashSet<HyperedgeId>>>,
351}
352
353impl HyperedgeNodeIndex {
354 pub fn new() -> Self {
356 Self {
357 index: Arc::new(DashMap::new()),
358 }
359 }
360
361 pub fn add_hyperedge(&self, hyperedge: &Hyperedge) {
363 for node_id in &hyperedge.nodes {
364 self.index
365 .entry(node_id.clone())
366 .or_insert_with(HashSet::new)
367 .insert(hyperedge.id.clone());
368 }
369 }
370
371 pub fn remove_hyperedge(&self, hyperedge: &Hyperedge) {
373 for node_id in &hyperedge.nodes {
374 if let Some(mut set) = self.index.get_mut(node_id) {
375 set.remove(&hyperedge.id);
376 }
377 }
378 }
379
380 pub fn get_hyperedges_by_node(&self, node_id: &NodeId) -> Vec<HyperedgeId> {
382 self.index
383 .get(node_id)
384 .map(|set| set.iter().cloned().collect())
385 .unwrap_or_default()
386 }
387
388 pub fn clear(&self) {
390 self.index.clear();
391 }
392}
393
394impl Default for HyperedgeNodeIndex {
395 fn default() -> Self {
396 Self::new()
397 }
398}
399
400#[cfg(test)]
401mod tests {
402 use super::*;
403 use crate::node::NodeBuilder;
404
405 #[test]
406 fn test_label_index() {
407 let index = LabelIndex::new();
408
409 let node1 = NodeBuilder::new().label("Person").label("User").build();
410
411 let node2 = NodeBuilder::new().label("Person").build();
412
413 index.add_node(&node1);
414 index.add_node(&node2);
415
416 let people = index.get_nodes_by_label("Person");
417 assert_eq!(people.len(), 2);
418
419 let users = index.get_nodes_by_label("User");
420 assert_eq!(users.len(), 1);
421
422 assert_eq!(index.count_by_label("Person"), 2);
423 }
424
425 #[test]
426 fn test_property_index() {
427 let index = PropertyIndex::new();
428
429 let node1 = NodeBuilder::new()
430 .property("name", "Alice")
431 .property("age", 30i64)
432 .build();
433
434 let node2 = NodeBuilder::new()
435 .property("name", "Bob")
436 .property("age", 30i64)
437 .build();
438
439 index.add_node(&node1);
440 index.add_node(&node2);
441
442 let alice =
443 index.get_nodes_by_property("name", &PropertyValue::String("Alice".to_string()));
444 assert_eq!(alice.len(), 1);
445
446 let age_30 = index.get_nodes_by_property("age", &PropertyValue::Integer(30));
447 assert_eq!(age_30.len(), 2);
448
449 let with_age = index.get_nodes_with_property("age");
450 assert_eq!(with_age.len(), 2);
451 }
452
453 #[test]
454 fn test_edge_type_index() {
455 let index = EdgeTypeIndex::new();
456
457 let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
458 let edge2 = Edge::create("n2".to_string(), "n3".to_string(), "KNOWS");
459 let edge3 = Edge::create("n1".to_string(), "n3".to_string(), "WORKS_WITH");
460
461 index.add_edge(&edge1);
462 index.add_edge(&edge2);
463 index.add_edge(&edge3);
464
465 let knows_edges = index.get_edges_by_type("KNOWS");
466 assert_eq!(knows_edges.len(), 2);
467
468 let works_with_edges = index.get_edges_by_type("WORKS_WITH");
469 assert_eq!(works_with_edges.len(), 1);
470
471 assert_eq!(index.all_edge_types().len(), 2);
472 }
473
474 #[test]
475 fn test_adjacency_index() {
476 let index = AdjacencyIndex::new();
477
478 let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
479 let edge2 = Edge::create("n1".to_string(), "n3".to_string(), "KNOWS");
480 let edge3 = Edge::create("n2".to_string(), "n1".to_string(), "KNOWS");
481
482 index.add_edge(&edge1);
483 index.add_edge(&edge2);
484 index.add_edge(&edge3);
485
486 assert_eq!(index.out_degree(&"n1".to_string()), 2);
487 assert_eq!(index.in_degree(&"n1".to_string()), 1);
488
489 let outgoing = index.get_outgoing_edges(&"n1".to_string());
490 assert_eq!(outgoing.len(), 2);
491
492 let incoming = index.get_incoming_edges(&"n1".to_string());
493 assert_eq!(incoming.len(), 1);
494 }
495
496 #[test]
497 fn test_get_edges_for_nodes() {
498 let index = AdjacencyIndex::new();
499
500 let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
501 let edge2 = Edge::create("n1".to_string(), "n3".to_string(), "KNOWS");
502 let edge3 = Edge::create("n2".to_string(), "n1".to_string(), "KNOWS");
503 let edge4 = Edge::create("n3".to_string(), "n4".to_string(), "KNOWS");
504
505 index.add_edge(&edge1);
506 index.add_edge(&edge2);
507 index.add_edge(&edge3);
508 index.add_edge(&edge4);
509
510 let result = index.get_edges_for_nodes(&["n1".to_string(), "n2".to_string()]);
511 assert_eq!(result.len(), 3);
512
513 let result_single = index.get_edges_for_nodes(&["n3".to_string()]);
514 assert_eq!(result_single.len(), 1);
515
516 let result_empty = index.get_edges_for_nodes(&["n5".to_string()]);
517 assert_eq!(result_empty.len(), 0);
518
519 let result_multi_empty = index.get_edges_for_nodes(&[]);
520 assert_eq!(result_multi_empty.len(), 0);
521 }
522
523 #[test]
524 fn test_for_each_outgoing_edge() {
525 let index = AdjacencyIndex::new();
526
527 let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
528 let edge2 = Edge::create("n1".to_string(), "n3".to_string(), "KNOWS");
529 let edge3 = Edge::create("n2".to_string(), "n1".to_string(), "KNOWS");
530 let edge4 = Edge::create("n3".to_string(), "n4".to_string(), "KNOWS");
531
532 index.add_edge(&edge1);
533 index.add_edge(&edge2);
534 index.add_edge(&edge3);
535 index.add_edge(&edge4);
536
537 let mut collected = Vec::new();
538 index.for_each_outgoing_edge(&["n1".to_string(), "n2".to_string()], |id| {
539 collected.push(id.clone());
540 });
541 assert_eq!(collected.len(), 3);
542
543 let mut single = Vec::new();
544 index.for_each_outgoing_edge(&["n3".to_string()], |id| {
545 single.push(id.clone());
546 });
547 assert_eq!(single.len(), 1);
548
549 let mut empty = Vec::new();
550 index.for_each_outgoing_edge(&["n5".to_string()], |id| {
551 empty.push(id.clone());
552 });
553 assert_eq!(empty.len(), 0);
554
555 let mut none_called = Vec::new();
556 index.for_each_outgoing_edge(&[], |id| {
557 none_called.push(id.clone());
558 });
559 assert_eq!(none_called.len(), 0);
560 }
561}