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(_) => {
165 format!("{:?}", value)
166 }
167 PropertyValue::Map(_) => format!("{:?}", value),
168 }
169 }
170}
171
172impl Default for PropertyIndex {
173 fn default() -> Self {
174 Self::new()
175 }
176}
177
178#[derive(Debug, Clone)]
180pub struct EdgeTypeIndex {
181 index: Arc<DashMap<String, HashSet<EdgeId>>>,
183}
184
185impl EdgeTypeIndex {
186 pub fn new() -> Self {
188 Self {
189 index: Arc::new(DashMap::new()),
190 }
191 }
192
193 pub fn add_edge(&self, edge: &Edge) {
195 self.index
196 .entry(edge.edge_type.clone())
197 .or_insert_with(HashSet::new)
198 .insert(edge.id.clone());
199 }
200
201 pub fn remove_edge(&self, edge: &Edge) {
203 if let Some(mut set) = self.index.get_mut(&edge.edge_type) {
204 set.remove(&edge.id);
205 }
206 }
207
208 pub fn get_edges_by_type(&self, edge_type: &str) -> Vec<EdgeId> {
210 self.index
211 .get(edge_type)
212 .map(|set| set.iter().cloned().collect())
213 .unwrap_or_default()
214 }
215
216 pub fn all_edge_types(&self) -> Vec<String> {
218 self.index.iter().map(|entry| entry.key().clone()).collect()
219 }
220
221 pub fn count_by_type(&self, edge_type: &str) -> usize {
223 self.index.get(edge_type).map(|set| set.len()).unwrap_or(0)
224 }
225
226 pub fn clear(&self) {
228 self.index.clear();
229 }
230}
231
232impl Default for EdgeTypeIndex {
233 fn default() -> Self {
234 Self::new()
235 }
236}
237
238#[derive(Debug, Clone)]
240pub struct AdjacencyIndex {
241 outgoing: Arc<DashMap<NodeId, HashSet<EdgeId>>>,
243 incoming: Arc<DashMap<NodeId, HashSet<EdgeId>>>,
245}
246
247impl AdjacencyIndex {
248 pub fn new() -> Self {
250 Self {
251 outgoing: Arc::new(DashMap::new()),
252 incoming: Arc::new(DashMap::new()),
253 }
254 }
255
256 pub fn add_edge(&self, edge: &Edge) {
258 self.outgoing
259 .entry(edge.from.clone())
260 .or_insert_with(HashSet::new)
261 .insert(edge.id.clone());
262
263 self.incoming
264 .entry(edge.to.clone())
265 .or_insert_with(HashSet::new)
266 .insert(edge.id.clone());
267 }
268
269 pub fn remove_edge(&self, edge: &Edge) {
271 if let Some(mut set) = self.outgoing.get_mut(&edge.from) {
272 set.remove(&edge.id);
273 }
274 if let Some(mut set) = self.incoming.get_mut(&edge.to) {
275 set.remove(&edge.id);
276 }
277 }
278
279 pub fn get_outgoing_edges(&self, node_id: &NodeId) -> Vec<EdgeId> {
281 self.outgoing
282 .get(node_id)
283 .map(|set| set.iter().cloned().collect())
284 .unwrap_or_default()
285 }
286
287 pub fn get_incoming_edges(&self, node_id: &NodeId) -> Vec<EdgeId> {
289 self.incoming
290 .get(node_id)
291 .map(|set| set.iter().cloned().collect())
292 .unwrap_or_default()
293 }
294
295 pub fn get_all_edges(&self, node_id: &NodeId) -> Vec<EdgeId> {
297 let mut edges = self.get_outgoing_edges(node_id);
298 edges.extend(self.get_incoming_edges(node_id));
299 edges
300 }
301
302 pub fn get_edges_for_nodes(&self, node_ids: &[NodeId]) -> Vec<EdgeId> {
304 let mut result = Vec::with_capacity(node_ids.len() * 4);
305 for node_id in node_ids {
306 if let Some(guard) = self.outgoing.get(node_id) {
307 result.extend(guard.iter().cloned());
308 }
309 }
310
311 result
312 }
313
314 pub fn for_each_outgoing_edge<F: FnMut(&EdgeId)>(&self, node_ids: &[NodeId], mut f: F) {
316 for node_id in node_ids {
317 if let Some(guard) = self.outgoing.get(node_id) {
318 for edge_id in guard.iter() {
319 f(edge_id);
320 }
321 }
322 }
323 }
324
325 pub fn out_degree(&self, node_id: &NodeId) -> usize {
327 self.outgoing.get(node_id).map(|set| set.len()).unwrap_or(0)
328 }
329
330 pub fn in_degree(&self, node_id: &NodeId) -> usize {
332 self.incoming.get(node_id).map(|set| set.len()).unwrap_or(0)
333 }
334
335 pub fn clear(&self) {
337 self.outgoing.clear();
338 self.incoming.clear();
339 }
340}
341
342impl Default for AdjacencyIndex {
343 fn default() -> Self {
344 Self::new()
345 }
346}
347
348#[derive(Debug, Clone)]
350pub struct HyperedgeNodeIndex {
351 index: Arc<DashMap<NodeId, HashSet<HyperedgeId>>>,
353}
354
355impl HyperedgeNodeIndex {
356 pub fn new() -> Self {
358 Self {
359 index: Arc::new(DashMap::new()),
360 }
361 }
362
363 pub fn add_hyperedge(&self, hyperedge: &Hyperedge) {
365 for node_id in &hyperedge.nodes {
366 self.index
367 .entry(node_id.clone())
368 .or_insert_with(HashSet::new)
369 .insert(hyperedge.id.clone());
370 }
371 }
372
373 pub fn remove_hyperedge(&self, hyperedge: &Hyperedge) {
375 for node_id in &hyperedge.nodes {
376 if let Some(mut set) = self.index.get_mut(node_id) {
377 set.remove(&hyperedge.id);
378 }
379 }
380 }
381
382 pub fn get_hyperedges_by_node(&self, node_id: &NodeId) -> Vec<HyperedgeId> {
384 self.index
385 .get(node_id)
386 .map(|set| set.iter().cloned().collect())
387 .unwrap_or_default()
388 }
389
390 pub fn clear(&self) {
392 self.index.clear();
393 }
394}
395
396impl Default for HyperedgeNodeIndex {
397 fn default() -> Self {
398 Self::new()
399 }
400}
401
402#[cfg(test)]
403mod tests {
404 use super::*;
405 use crate::node::NodeBuilder;
406
407 #[test]
408 fn test_label_index() {
409 let index = LabelIndex::new();
410
411 let node1 = NodeBuilder::new().label("Person").label("User").build();
412
413 let node2 = NodeBuilder::new().label("Person").build();
414
415 index.add_node(&node1);
416 index.add_node(&node2);
417
418 let people = index.get_nodes_by_label("Person");
419 assert_eq!(people.len(), 2);
420
421 let users = index.get_nodes_by_label("User");
422 assert_eq!(users.len(), 1);
423
424 assert_eq!(index.count_by_label("Person"), 2);
425 }
426
427 #[test]
428 fn test_property_index() {
429 let index = PropertyIndex::new();
430
431 let node1 = NodeBuilder::new()
432 .property("name", "Alice")
433 .property("age", 30i64)
434 .build();
435
436 let node2 = NodeBuilder::new()
437 .property("name", "Bob")
438 .property("age", 30i64)
439 .build();
440
441 index.add_node(&node1);
442 index.add_node(&node2);
443
444 let alice =
445 index.get_nodes_by_property("name", &PropertyValue::String("Alice".to_string()));
446 assert_eq!(alice.len(), 1);
447
448 let age_30 = index.get_nodes_by_property("age", &PropertyValue::Integer(30));
449 assert_eq!(age_30.len(), 2);
450
451 let with_age = index.get_nodes_with_property("age");
452 assert_eq!(with_age.len(), 2);
453 }
454
455 #[test]
456 fn test_edge_type_index() {
457 let index = EdgeTypeIndex::new();
458
459 let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
460 let edge2 = Edge::create("n2".to_string(), "n3".to_string(), "KNOWS");
461 let edge3 = Edge::create("n1".to_string(), "n3".to_string(), "WORKS_WITH");
462
463 index.add_edge(&edge1);
464 index.add_edge(&edge2);
465 index.add_edge(&edge3);
466
467 let knows_edges = index.get_edges_by_type("KNOWS");
468 assert_eq!(knows_edges.len(), 2);
469
470 let works_with_edges = index.get_edges_by_type("WORKS_WITH");
471 assert_eq!(works_with_edges.len(), 1);
472
473 assert_eq!(index.all_edge_types().len(), 2);
474 }
475
476 #[test]
477 fn test_adjacency_index() {
478 let index = AdjacencyIndex::new();
479
480 let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
481 let edge2 = Edge::create("n1".to_string(), "n3".to_string(), "KNOWS");
482 let edge3 = Edge::create("n2".to_string(), "n1".to_string(), "KNOWS");
483
484 index.add_edge(&edge1);
485 index.add_edge(&edge2);
486 index.add_edge(&edge3);
487
488 assert_eq!(index.out_degree(&"n1".to_string()), 2);
489 assert_eq!(index.in_degree(&"n1".to_string()), 1);
490
491 let outgoing = index.get_outgoing_edges(&"n1".to_string());
492 assert_eq!(outgoing.len(), 2);
493
494 let incoming = index.get_incoming_edges(&"n1".to_string());
495 assert_eq!(incoming.len(), 1);
496 }
497
498 #[test]
499 fn test_get_edges_for_nodes() {
500 let index = AdjacencyIndex::new();
501
502 let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
503 let edge2 = Edge::create("n1".to_string(), "n3".to_string(), "KNOWS");
504 let edge3 = Edge::create("n2".to_string(), "n1".to_string(), "KNOWS");
505 let edge4 = Edge::create("n3".to_string(), "n4".to_string(), "KNOWS");
506
507 index.add_edge(&edge1);
508 index.add_edge(&edge2);
509 index.add_edge(&edge3);
510 index.add_edge(&edge4);
511
512 let result = index.get_edges_for_nodes(&["n1".to_string(), "n2".to_string()]);
513 assert_eq!(result.len(), 3);
514
515 let result_single = index.get_edges_for_nodes(&["n3".to_string()]);
516 assert_eq!(result_single.len(), 1);
517
518 let result_empty = index.get_edges_for_nodes(&["n5".to_string()]);
519 assert_eq!(result_empty.len(), 0);
520
521 let result_multi_empty = index.get_edges_for_nodes(&[]);
522 assert_eq!(result_multi_empty.len(), 0);
523 }
524
525 #[test]
526 fn test_for_each_outgoing_edge() {
527 let index = AdjacencyIndex::new();
528
529 let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
530 let edge2 = Edge::create("n1".to_string(), "n3".to_string(), "KNOWS");
531 let edge3 = Edge::create("n2".to_string(), "n1".to_string(), "KNOWS");
532 let edge4 = Edge::create("n3".to_string(), "n4".to_string(), "KNOWS");
533
534 index.add_edge(&edge1);
535 index.add_edge(&edge2);
536 index.add_edge(&edge3);
537 index.add_edge(&edge4);
538
539 let mut collected = Vec::new();
540 index.for_each_outgoing_edge(&["n1".to_string(), "n2".to_string()], |id| {
541 collected.push(id.clone());
542 });
543 assert_eq!(collected.len(), 3);
544
545 let mut single = Vec::new();
546 index.for_each_outgoing_edge(&["n3".to_string()], |id| {
547 single.push(id.clone());
548 });
549 assert_eq!(single.len(), 1);
550
551 let mut empty = Vec::new();
552 index.for_each_outgoing_edge(&["n5".to_string()], |id| {
553 empty.push(id.clone());
554 });
555 assert_eq!(empty.len(), 0);
556
557 let mut none_called = Vec::new();
558 index.for_each_outgoing_edge(&[], |id| {
559 none_called.push(id.clone());
560 });
561 assert_eq!(none_called.len(), 0);
562 }
563}