1use crate::json;
13use crate::serde_json::Value;
14use std::collections::HashMap;
15use std::hash::{Hash, Hasher};
16
17#[derive(Debug, Clone)]
19pub enum TraverserValue {
20 Vertex(String),
22 Edge {
24 id: String,
25 source: String,
26 target: String,
27 label: String,
28 },
29 Property(String, Value),
31 Path(Path),
33 Map(HashMap<String, Value>),
35 List(Vec<Value>),
37 String(String),
39 Integer(i64),
40 Float(f64),
41 Boolean(bool),
42 Null,
44}
45
46impl TraverserValue {
47 pub fn as_vertex_id(&self) -> Option<&str> {
49 match self {
50 TraverserValue::Vertex(id) => Some(id),
51 TraverserValue::String(s) => Some(s),
52 _ => None,
53 }
54 }
55
56 pub fn as_edge(&self) -> Option<(&str, &str, &str, &str)> {
58 match self {
59 TraverserValue::Edge {
60 id,
61 source,
62 target,
63 label,
64 } => Some((id, source, target, label)),
65 _ => None,
66 }
67 }
68
69 pub fn is_null(&self) -> bool {
71 matches!(self, TraverserValue::Null)
72 }
73
74 pub fn to_json(&self) -> Value {
76 match self {
77 TraverserValue::Vertex(id) => json!({ "type": "vertex", "id": id }),
78 TraverserValue::Edge {
79 id,
80 source,
81 target,
82 label,
83 } => {
84 json!({
85 "type": "edge",
86 "id": id,
87 "source": source,
88 "target": target,
89 "label": label
90 })
91 }
92 TraverserValue::Property(key, value) => {
93 json!({ "key": key, "value": value })
94 }
95 TraverserValue::Path(path) => path.to_json(),
96 TraverserValue::Map(map) => {
97 Value::Object(map.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
98 }
99 TraverserValue::List(list) => Value::Array(list.clone()),
100 TraverserValue::String(s) => Value::String(s.clone()),
101 TraverserValue::Integer(i) => json!(i),
102 TraverserValue::Float(f) => json!(f),
103 TraverserValue::Boolean(b) => Value::Bool(*b),
104 TraverserValue::Null => Value::Null,
105 }
106 }
107}
108
109impl PartialEq for TraverserValue {
110 fn eq(&self, other: &Self) -> bool {
111 match (self, other) {
112 (TraverserValue::Vertex(a), TraverserValue::Vertex(b)) => a == b,
113 (TraverserValue::String(a), TraverserValue::String(b)) => a == b,
114 (TraverserValue::Integer(a), TraverserValue::Integer(b)) => a == b,
115 (TraverserValue::Boolean(a), TraverserValue::Boolean(b)) => a == b,
116 (TraverserValue::Null, TraverserValue::Null) => true,
117 _ => false,
118 }
119 }
120}
121
122impl Eq for TraverserValue {}
123
124impl Hash for TraverserValue {
125 fn hash<H: Hasher>(&self, state: &mut H) {
126 match self {
127 TraverserValue::Vertex(id) => {
128 0u8.hash(state);
129 id.hash(state);
130 }
131 TraverserValue::String(s) => {
132 1u8.hash(state);
133 s.hash(state);
134 }
135 TraverserValue::Integer(i) => {
136 2u8.hash(state);
137 i.hash(state);
138 }
139 TraverserValue::Boolean(b) => {
140 3u8.hash(state);
141 b.hash(state);
142 }
143 TraverserValue::Null => {
144 4u8.hash(state);
145 }
146 _ => {
147 255u8.hash(state);
149 format!("{:?}", self).hash(state);
150 }
151 }
152 }
153}
154
155#[derive(Debug, Clone, Default)]
157pub struct Path {
158 objects: Vec<TraverserValue>,
160 labels: Vec<Vec<String>>,
162}
163
164impl Path {
165 pub fn new() -> Self {
167 Self::default()
168 }
169
170 pub fn start(value: TraverserValue) -> Self {
172 Self {
173 objects: vec![value],
174 labels: vec![vec![]],
175 }
176 }
177
178 pub fn extend(&mut self, value: TraverserValue) {
180 self.objects.push(value);
181 self.labels.push(vec![]);
182 }
183
184 pub fn extend_with_label(&mut self, value: TraverserValue, label: String) {
186 self.objects.push(value);
187 self.labels.push(vec![label]);
188 }
189
190 pub fn add_label(&mut self, label: String) {
192 if let Some(last_labels) = self.labels.last_mut() {
193 if !last_labels.contains(&label) {
194 last_labels.push(label);
195 }
196 }
197 }
198
199 pub fn objects(&self) -> &[TraverserValue] {
201 &self.objects
202 }
203
204 pub fn labels_at(&self, index: usize) -> &[String] {
206 self.labels.get(index).map(|v| v.as_slice()).unwrap_or(&[])
207 }
208
209 pub fn get(&self, label: &str) -> Option<&TraverserValue> {
211 for (i, labels) in self.labels.iter().enumerate() {
212 if labels.contains(&label.to_string()) {
213 return self.objects.get(i);
214 }
215 }
216 None
217 }
218
219 pub fn get_all(&self, label: &str) -> Vec<&TraverserValue> {
221 let mut result = Vec::new();
222 for (i, labels) in self.labels.iter().enumerate() {
223 if labels.contains(&label.to_string()) {
224 if let Some(obj) = self.objects.get(i) {
225 result.push(obj);
226 }
227 }
228 }
229 result
230 }
231
232 pub fn has_label(&self, label: &str) -> bool {
234 self.labels
235 .iter()
236 .any(|labels| labels.contains(&label.to_string()))
237 }
238
239 pub fn len(&self) -> usize {
241 self.objects.len()
242 }
243
244 pub fn is_empty(&self) -> bool {
246 self.objects.is_empty()
247 }
248
249 pub fn head(&self) -> Option<&TraverserValue> {
251 self.objects.last()
252 }
253
254 pub fn clone_and_extend(&self, value: TraverserValue) -> Self {
256 let mut new_path = self.clone();
257 new_path.extend(value);
258 new_path
259 }
260
261 pub fn to_json(&self) -> Value {
263 let objects: Vec<_> = self.objects.iter().map(|o| o.to_json()).collect();
264 let labels: Vec<_> = self
265 .labels
266 .iter()
267 .map(|l| Value::Array(l.iter().map(|s| Value::String(s.clone())).collect()))
268 .collect();
269
270 json!({
271 "objects": objects,
272 "labels": labels
273 })
274 }
275
276 pub fn retract(&mut self, keep_labels: &[String]) {
278 let mut new_objects = Vec::new();
279 let mut new_labels = Vec::new();
280
281 for (i, labels) in self.labels.iter().enumerate() {
282 if labels.iter().any(|l| keep_labels.contains(l)) {
283 if let Some(obj) = self.objects.get(i) {
284 new_objects.push(obj.clone());
285 new_labels.push(labels.clone());
286 }
287 }
288 }
289
290 self.objects = new_objects;
291 self.labels = new_labels;
292 }
293}
294
295#[derive(Debug, Clone, Default)]
297pub struct LoopState {
298 loops: HashMap<String, u32>,
300 current_loop: Option<String>,
302}
303
304impl LoopState {
305 pub fn new() -> Self {
307 Self::default()
308 }
309
310 pub fn init_loop(&mut self, name: &str) {
312 self.loops.insert(name.to_string(), 0);
313 self.current_loop = Some(name.to_string());
314 }
315
316 pub fn incr_loop(&mut self, name: &str) {
318 if let Some(count) = self.loops.get_mut(name) {
319 *count += 1;
320 }
321 }
322
323 pub fn loop_count(&self, name: &str) -> u32 {
325 self.loops.get(name).copied().unwrap_or(0)
326 }
327
328 pub fn current_count(&self) -> u32 {
330 self.current_loop
331 .as_ref()
332 .and_then(|name| self.loops.get(name))
333 .copied()
334 .unwrap_or(0)
335 }
336
337 pub fn reset_loop(&mut self, name: &str) {
339 self.loops.remove(name);
340 if self.current_loop.as_deref() == Some(name) {
341 self.current_loop = None;
342 }
343 }
344}
345
346#[derive(Debug, Clone)]
348pub struct Traverser {
349 value: TraverserValue,
351 bulk: u64,
353 path: Option<Path>,
355 loops: Option<LoopState>,
357 sack: Option<Value>,
359 tags: HashMap<String, String>,
361}
362
363impl Traverser {
364 pub fn new(vertex_id: &str) -> Self {
366 Self {
367 value: TraverserValue::Vertex(vertex_id.to_string()),
368 bulk: 1,
369 path: None,
370 loops: None,
371 sack: None,
372 tags: HashMap::new(),
373 }
374 }
375
376 pub fn with_value(value: TraverserValue) -> Self {
378 Self {
379 value,
380 bulk: 1,
381 path: None,
382 loops: None,
383 sack: None,
384 tags: HashMap::new(),
385 }
386 }
387
388 pub fn value(&self) -> &TraverserValue {
390 &self.value
391 }
392
393 pub fn set_value(&mut self, value: TraverserValue) {
395 self.value = value;
396 }
397
398 pub fn bulk(&self) -> u64 {
400 self.bulk
401 }
402
403 pub fn set_bulk(&mut self, bulk: u64) {
405 self.bulk = bulk;
406 }
407
408 pub fn multiply_bulk(&mut self, factor: u64) {
410 self.bulk *= factor;
411 }
412
413 pub fn path(&self) -> Option<&Path> {
415 self.path.as_ref()
416 }
417
418 pub fn path_mut(&mut self) -> Option<&mut Path> {
420 self.path.as_mut()
421 }
422
423 pub fn enable_path(&mut self) {
425 if self.path.is_none() {
426 self.path = Some(Path::start(self.value.clone()));
427 }
428 }
429
430 pub fn extend_path(&mut self, value: TraverserValue) {
432 if let Some(ref mut path) = self.path {
433 path.extend(value);
434 }
435 }
436
437 pub fn add_path_label(&mut self, label: String) {
439 if let Some(ref mut path) = self.path {
440 path.add_label(label);
441 }
442 }
443
444 pub fn loops(&self) -> Option<&LoopState> {
446 self.loops.as_ref()
447 }
448
449 pub fn loops_mut(&mut self) -> Option<&mut LoopState> {
451 self.loops.as_mut()
452 }
453
454 pub fn enable_loops(&mut self) {
456 if self.loops.is_none() {
457 self.loops = Some(LoopState::new());
458 }
459 }
460
461 pub fn init_loop(&mut self, name: &str) {
463 self.enable_loops();
464 if let Some(ref mut loops) = self.loops {
465 loops.init_loop(name);
466 }
467 }
468
469 pub fn incr_loop(&mut self, name: &str) {
471 if let Some(ref mut loops) = self.loops {
472 loops.incr_loop(name);
473 }
474 }
475
476 pub fn loop_count(&self, name: &str) -> u32 {
478 self.loops.as_ref().map(|l| l.loop_count(name)).unwrap_or(0)
479 }
480
481 pub fn sack(&self) -> Option<&Value> {
483 self.sack.as_ref()
484 }
485
486 pub fn set_sack(&mut self, value: Value) {
488 self.sack = Some(value);
489 }
490
491 pub fn tag(&self, key: &str) -> Option<&String> {
493 self.tags.get(key)
494 }
495
496 pub fn set_tag(&mut self, key: String, value: String) {
498 self.tags.insert(key, value);
499 }
500
501 pub fn split(&self) -> Self {
503 Self {
504 value: self.value.clone(),
505 bulk: 1, path: self.path.clone(),
507 loops: self.loops.clone(),
508 sack: self.sack.clone(),
509 tags: self.tags.clone(),
510 }
511 }
512
513 pub fn clone_with_value(&self, value: TraverserValue) -> Self {
515 let mut new = self.clone();
516 new.value = value;
517 new
518 }
519
520 pub fn merge(&mut self, other: &Traverser) {
522 self.bulk += other.bulk;
523 }
524}
525
526#[derive(Debug, Clone, PartialEq, Eq, Hash)]
528pub enum TraverserRequirement {
529 Bulk,
531 Path,
533 SingleLoop,
535 NestedLoop,
537 Labels,
539 Sack,
541 Barrier,
543 Mutates,
545}
546
547pub struct TraverserGenerator {
549 requirements: Vec<TraverserRequirement>,
550}
551
552impl TraverserGenerator {
553 pub fn new(requirements: Vec<TraverserRequirement>) -> Self {
555 Self { requirements }
556 }
557
558 pub fn generate(&self, vertex_id: &str) -> Traverser {
560 let mut traverser = Traverser::new(vertex_id);
561 self.apply_requirements(&mut traverser);
562 traverser
563 }
564
565 pub fn generate_value(&self, value: TraverserValue) -> Traverser {
567 let mut traverser = Traverser::with_value(value);
568 self.apply_requirements(&mut traverser);
569 traverser
570 }
571
572 pub fn generate_many(&self, vertex_ids: &[String]) -> Vec<Traverser> {
574 vertex_ids.iter().map(|id| self.generate(id)).collect()
575 }
576
577 fn apply_requirements(&self, traverser: &mut Traverser) {
579 for req in &self.requirements {
580 match req {
581 TraverserRequirement::Path => traverser.enable_path(),
582 TraverserRequirement::SingleLoop | TraverserRequirement::NestedLoop => {
583 traverser.enable_loops()
584 }
585 _ => {}
586 }
587 }
588 }
589}
590
591#[cfg(test)]
592mod tests {
593 use super::*;
594
595 #[test]
596 fn test_traverser_new() {
597 let t = Traverser::new("v1");
598 assert!(matches!(t.value(), TraverserValue::Vertex(id) if id == "v1"));
599 assert_eq!(t.bulk(), 1);
600 }
601
602 #[test]
603 fn test_traverser_path() {
604 let mut t = Traverser::new("v1");
605 assert!(t.path().is_none());
606
607 t.enable_path();
608 assert!(t.path().is_some());
609 assert_eq!(t.path().unwrap().len(), 1);
610
611 t.extend_path(TraverserValue::Vertex("v2".to_string()));
612 assert_eq!(t.path().unwrap().len(), 2);
613 }
614
615 #[test]
616 fn test_traverser_loops() {
617 let mut t = Traverser::new("v1");
618 t.init_loop("repeat_0");
619 assert_eq!(t.loop_count("repeat_0"), 0);
620
621 t.incr_loop("repeat_0");
622 assert_eq!(t.loop_count("repeat_0"), 1);
623
624 t.incr_loop("repeat_0");
625 assert_eq!(t.loop_count("repeat_0"), 2);
626 }
627
628 #[test]
629 fn test_traverser_split() {
630 let mut t = Traverser::new("v1");
631 t.set_bulk(5);
632 t.enable_path();
633
634 let split = t.split();
635 assert_eq!(split.bulk(), 1); assert!(split.path().is_some()); }
638
639 #[test]
640 fn test_path_labels() {
641 let mut path = Path::new();
642 path.extend_with_label(TraverserValue::Vertex("v1".to_string()), "a".to_string());
643 path.extend_with_label(TraverserValue::Vertex("v2".to_string()), "b".to_string());
644
645 assert!(path.has_label("a"));
646 assert!(path.has_label("b"));
647 assert!(!path.has_label("c"));
648
649 let v1 = path.get("a").unwrap();
650 assert!(matches!(v1, TraverserValue::Vertex(id) if id == "v1"));
651 }
652
653 #[test]
654 fn test_path_retract() {
655 let mut path = Path::new();
656 path.extend_with_label(TraverserValue::Vertex("v1".to_string()), "a".to_string());
657 path.extend(TraverserValue::Vertex("v2".to_string())); path.extend_with_label(TraverserValue::Vertex("v3".to_string()), "b".to_string());
659
660 path.retract(&["a".to_string(), "b".to_string()]);
661 assert_eq!(path.len(), 2); }
663
664 #[test]
665 fn test_traverser_generator() {
666 let gen = TraverserGenerator::new(vec![
667 TraverserRequirement::Path,
668 TraverserRequirement::SingleLoop,
669 ]);
670
671 let t = gen.generate("v1");
672 assert!(t.path().is_some());
673 assert!(t.loops().is_some());
674 }
675}