1use std::{ops::Index, sync::Arc, num::NonZeroUsize};
2use std::hash::{Hash, Hasher};
3use std::collections::hash_map::DefaultHasher;
4use im::Vector;
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use once_cell::sync::Lazy;
8use parking_lot::RwLock;
9use lru::LruCache;
10use std::fmt::{self, Debug};
11use crate::error::PoolResult;
12use crate::node_type::NodeEnum;
13use crate::{
14 error::error_helpers,
15 mark::Mark,
16 node::Node,
17 ops::{AttrsRef, MarkRef, NodeRef},
18 types::NodeId,
19};
20
21static SHARD_INDEX_CACHE: Lazy<RwLock<LruCache<String, usize>>> =
23 Lazy::new(|| RwLock::new(LruCache::new(NonZeroUsize::new(10000).unwrap())));
24
25#[derive(Clone, PartialEq, Serialize, Deserialize)]
26pub struct Tree {
27 pub root_id: NodeId,
28 pub nodes: Vector<im::HashMap<NodeId, Arc<Node>>>, pub parent_map: im::HashMap<NodeId, NodeId>,
30 #[serde(skip)]
31 num_shards: usize, }
33impl Debug for Tree {
34 fn fmt(
35 &self,
36 f: &mut fmt::Formatter<'_>,
37 ) -> fmt::Result {
38 let nodes = self
40 .nodes
41 .iter()
42 .filter(|node| !node.is_empty())
43 .collect::<Vec<_>>();
44 f.debug_struct("Tree")
45 .field("root_id", &self.root_id)
46 .field("nodes", &nodes)
47 .field("parent_map", &self.parent_map)
48 .field("num_shards", &self.num_shards)
49 .finish()
50 }
51}
52
53impl Tree {
54 #[inline]
55 pub fn get_shard_index(
56 &self,
57 id: &NodeId,
58 ) -> usize {
59 {
61 let cache = SHARD_INDEX_CACHE.read();
62 if let Some(&index) = cache.peek(id) {
63 return index;
64 }
65 }
66
67 let mut hasher = DefaultHasher::new();
69 id.hash(&mut hasher);
70 let index = (hasher.finish() as usize) % self.num_shards;
71
72 {
74 let mut cache = SHARD_INDEX_CACHE.write();
75 cache.put(id.clone(), index);
76 }
77
78 index
79 }
80
81 #[inline]
82 pub fn get_shard_indices(
83 &self,
84 ids: &[&NodeId],
85 ) -> Vec<usize> {
86 ids.iter().map(|id| self.get_shard_index(id)).collect()
87 }
88
89 #[inline]
91 pub fn get_shard_index_batch<'a>(
92 &self,
93 ids: &'a [&'a NodeId],
94 ) -> Vec<(usize, &'a NodeId)> {
95 let mut results = Vec::with_capacity(ids.len());
96 let mut cache_misses = Vec::new();
97
98 {
100 let cache = SHARD_INDEX_CACHE.read();
101 for &id in ids {
102 if let Some(&index) = cache.peek(id) {
103 results.push((index, id));
104 } else {
105 cache_misses.push(id);
106 }
107 }
108 }
109
110 if !cache_misses.is_empty() {
112 let mut cache = SHARD_INDEX_CACHE.write();
113 for &id in &cache_misses {
114 let mut hasher = DefaultHasher::new();
115 id.hash(&mut hasher);
116 let index = (hasher.finish() as usize) % self.num_shards;
117 cache.put(id.clone(), index);
118 results.push((index, id));
119 }
120 }
121
122 results
123 }
124
125 pub fn clear_shard_cache() {
127 let mut cache = SHARD_INDEX_CACHE.write();
128 cache.clear();
129 }
130
131 pub fn contains_node(
132 &self,
133 id: &NodeId,
134 ) -> bool {
135 let shard_index = self.get_shard_index(id);
136 self.nodes[shard_index].contains_key(id)
137 }
138
139 pub fn get_node(
140 &self,
141 id: &NodeId,
142 ) -> Option<Arc<Node>> {
143 let shard_index = self.get_shard_index(id);
144 self.nodes[shard_index].get(id).cloned()
145 }
146
147 pub fn get_parent_node(
148 &self,
149 id: &NodeId,
150 ) -> Option<Arc<Node>> {
151 self.parent_map.get(id).and_then(|parent_id| {
152 let shard_index = self.get_shard_index(parent_id);
153 self.nodes[shard_index].get(parent_id).cloned()
154 })
155 }
156 pub fn from(nodes: NodeEnum) -> Self {
157 let num_shards = std::cmp::max(
158 std::thread::available_parallelism()
159 .map(NonZeroUsize::get)
160 .unwrap_or(2),
161 2,
162 );
163 let mut shards = Vector::from(vec![im::HashMap::new(); num_shards]);
164 let mut parent_map = im::HashMap::new();
165 let (root_node, children) = nodes.into_parts();
166 let root_id = root_node.id.clone();
167
168 let mut hasher = DefaultHasher::new();
169 root_id.hash(&mut hasher);
170 let shard_index = (hasher.finish() as usize) % num_shards;
171 shards[shard_index] =
172 shards[shard_index].update(root_id.clone(), Arc::new(root_node));
173
174 fn process_children(
175 children: Vec<NodeEnum>,
176 parent_id: &NodeId,
177 shards: &mut Vector<im::HashMap<NodeId, Arc<Node>>>,
178 parent_map: &mut im::HashMap<NodeId, NodeId>,
179 num_shards: usize,
180 ) {
181 for child in children {
182 let (node, grand_children) = child.into_parts();
183 let node_id = node.id.clone();
184 let mut hasher = DefaultHasher::new();
185 node_id.hash(&mut hasher);
186 let shard_index = (hasher.finish() as usize) % num_shards;
187 shards[shard_index] =
188 shards[shard_index].update(node_id.clone(), Arc::new(node));
189 parent_map.insert(node_id.clone(), parent_id.clone());
190
191 process_children(
193 grand_children,
194 &node_id,
195 shards,
196 parent_map,
197 num_shards,
198 );
199 }
200 }
201
202 process_children(
203 children,
204 &root_id,
205 &mut shards,
206 &mut parent_map,
207 num_shards,
208 );
209
210 Self { root_id, nodes: shards, parent_map, num_shards }
211 }
212
213 pub fn new(root: Node) -> Self {
214 let num_shards = std::cmp::max(
215 std::thread::available_parallelism()
216 .map(NonZeroUsize::get)
217 .unwrap_or(2),
218 2,
219 );
220 let mut nodes = Vector::from(vec![im::HashMap::new(); num_shards]);
221 let root_id = root.id.clone();
222 let mut hasher = DefaultHasher::new();
223 root_id.hash(&mut hasher);
224 let shard_index = (hasher.finish() as usize) % num_shards;
225 nodes[shard_index] =
226 nodes[shard_index].update(root_id.clone(), Arc::new(root));
227 Self { root_id, nodes, parent_map: im::HashMap::new(), num_shards }
228 }
229
230 pub fn update_attr(
231 &mut self,
232 id: &NodeId,
233 new_values: im::HashMap<String, Value>,
234 ) -> PoolResult<()> {
235 let shard_index = self.get_shard_index(id);
236 let node = self.nodes[shard_index]
237 .get(id)
238 .ok_or(error_helpers::node_not_found(id.clone()))?;
239 let old_values = node.attrs.clone();
240 let mut new_node = node.as_ref().clone();
241 let new_attrs = old_values.update(new_values);
242 new_node.attrs = new_attrs.clone();
243 self.nodes[shard_index] =
244 self.nodes[shard_index].update(id.clone(), Arc::new(new_node));
245 Ok(())
246 }
247 pub fn update_node(
248 &mut self,
249 node: Node,
250 ) -> PoolResult<()> {
251 let shard_index = self.get_shard_index(&node.id);
252 self.nodes[shard_index] =
253 self.nodes[shard_index].update(node.id.clone(), Arc::new(node));
254 Ok(())
255 }
256
257 pub fn add(
268 &mut self,
269 parent_id: &NodeId,
270 nodes: Vec<NodeEnum>,
271 ) -> PoolResult<()> {
272 let parent_shard_index = self.get_shard_index(&parent_id);
274 let parent_node = self.nodes[parent_shard_index]
275 .get(parent_id)
276 .ok_or(error_helpers::parent_not_found(parent_id.clone()))?;
277 let mut new_parent = parent_node.as_ref().clone();
278
279 let zenliang: Vector<String> =
281 nodes.iter().map(|n| n.0.id.clone()).collect();
282 new_parent.content.extend(zenliang);
283
284 self.nodes[parent_shard_index] = self.nodes[parent_shard_index]
286 .update(parent_id.clone(), Arc::new(new_parent));
287
288 let mut node_queue = Vec::new();
290 node_queue.push((nodes, parent_id.clone()));
291 while let Some((current_children, current_parent_id)) = node_queue.pop()
292 {
293 for child in current_children {
294 let (mut child_node, grand_children) = child.into_parts();
296 let current_node_id = child_node.id.clone();
297
298 let grand_children_ids: Vector<String> =
300 grand_children.iter().map(|n| n.0.id.clone()).collect();
301 child_node.content.extend(grand_children_ids);
302
303 let shard_index = self.get_shard_index(¤t_node_id);
305 self.nodes[shard_index] = self.nodes[shard_index]
306 .update(current_node_id.clone(), Arc::new(child_node));
307
308 self.parent_map
310 .insert(current_node_id.clone(), current_parent_id.clone());
311
312 node_queue.push((grand_children, current_node_id.clone()));
314 }
315 }
316 Ok(())
317 }
318 pub fn add_at_index(
320 &mut self,
321 parent_id: &NodeId,
322 index: usize,
323 node: &Node,
324 ) -> PoolResult<()> {
325 let parent_shard_index = self.get_shard_index(parent_id);
327 let parent = self.nodes[parent_shard_index]
328 .get(parent_id)
329 .ok_or(error_helpers::parent_not_found(parent_id.clone()))?;
330 let mut new_parent = parent.as_ref().clone();
331 new_parent.content.insert(index, node.id.clone());
332 self.nodes[parent_shard_index] = self.nodes[parent_shard_index]
333 .update(parent_id.clone(), Arc::new(new_parent));
334 self.parent_map.insert(node.id.clone(), parent_id.clone());
335 Ok(())
336 }
337 pub fn add_node(
338 &mut self,
339 parent_id: &NodeId,
340 nodes: &Vec<Node>,
341 ) -> PoolResult<()> {
342 let parent_shard_index = self.get_shard_index(parent_id);
343 let parent = self.nodes[parent_shard_index]
344 .get(parent_id)
345 .ok_or(error_helpers::parent_not_found(parent_id.clone()))?;
346 let mut new_parent = parent.as_ref().clone();
347 new_parent.content.push_back(nodes[0].id.clone());
348 self.nodes[parent_shard_index] = self.nodes[parent_shard_index]
349 .update(parent_id.clone(), Arc::new(new_parent));
350 self.parent_map.insert(nodes[0].id.clone(), parent_id.clone());
351 for node in nodes {
352 let shard_index = self.get_shard_index(&node.id);
353 for child_id in &node.content {
354 self.parent_map.insert(child_id.clone(), node.id.clone());
355 }
356 self.nodes[shard_index] = self.nodes[shard_index]
357 .update(node.id.clone(), Arc::new(node.clone()));
358 }
359 Ok(())
360 }
361
362 pub fn node(
363 &mut self,
364 key: &str,
365 ) -> NodeRef<'_> {
366 NodeRef::new(self, key.to_string())
367 }
368 pub fn mark(
369 &mut self,
370 key: &str,
371 ) -> MarkRef<'_> {
372 MarkRef::new(self, key.to_string())
373 }
374 pub fn attrs(
375 &mut self,
376 key: &str,
377 ) -> AttrsRef<'_> {
378 AttrsRef::new(self, key.to_string())
379 }
380
381 pub fn children(
382 &self,
383 parent_id: &NodeId,
384 ) -> Option<im::Vector<NodeId>> {
385 self.get_node(parent_id).map(|n| n.content.clone())
386 }
387
388 pub fn children_node(
389 &self,
390 parent_id: &NodeId,
391 ) -> Option<im::Vector<Arc<Node>>> {
392 self.children(parent_id)
393 .map(|ids| ids.iter().filter_map(|id| self.get_node(id)).collect())
394 }
395 pub fn all_children(
397 &self,
398 parent_id: &NodeId,
399 filter: Option<&dyn Fn(&Node) -> bool>,
400 ) -> Option<NodeEnum> {
401 if let Some(node) = self.get_node(parent_id) {
402 let mut child_enums = Vec::new();
403 for child_id in &node.content {
404 if let Some(child_node) = self.get_node(child_id) {
405 if let Some(filter_fn) = filter {
407 if !filter_fn(child_node.as_ref()) {
408 continue; }
410 }
411 if let Some(child_enum) =
413 self.all_children(child_id, filter)
414 {
415 child_enums.push(child_enum);
416 }
417 }
418 }
419 Some(NodeEnum(node.as_ref().clone(), child_enums))
420 } else {
421 None
422 }
423 }
424
425 pub fn children_count(
426 &self,
427 parent_id: &NodeId,
428 ) -> usize {
429 self.get_node(parent_id).map(|n| n.content.len()).unwrap_or(0)
430 }
431 pub fn remove_mark_by_name(
432 &mut self,
433 id: &NodeId,
434 mark_name: &str,
435 ) -> PoolResult<()> {
436 let shard_index = self.get_shard_index(id);
437 let node = self.nodes[shard_index]
438 .get(id)
439 .ok_or(error_helpers::node_not_found(id.clone()))?;
440 let mut new_node = node.as_ref().clone();
441 new_node.marks = new_node
442 .marks
443 .iter()
444 .filter(|&m| m.r#type != mark_name)
445 .cloned()
446 .collect();
447 self.nodes[shard_index] =
448 self.nodes[shard_index].update(id.clone(), Arc::new(new_node));
449 Ok(())
450 }
451 pub fn get_marks(
452 &self,
453 id: &NodeId,
454 ) -> Option<im::Vector<Mark>> {
455 self.get_node(id).map(|n| n.marks.clone())
456 }
457
458 pub fn remove_mark(
459 &mut self,
460 id: &NodeId,
461 mark: Mark,
462 ) -> PoolResult<()> {
463 let shard_index = self.get_shard_index(id);
464 let node = self.nodes[shard_index]
465 .get(id)
466 .ok_or(error_helpers::node_not_found(id.clone()))?;
467 let mut new_node = node.as_ref().clone();
468 new_node.marks =
469 new_node.marks.iter().filter(|&m| !m.eq(&mark)).cloned().collect();
470 self.nodes[shard_index] =
471 self.nodes[shard_index].update(id.clone(), Arc::new(new_node));
472 Ok(())
473 }
474
475 pub fn add_mark(
476 &mut self,
477 id: &NodeId,
478 marks: &Vec<Mark>,
479 ) -> PoolResult<()> {
480 let shard_index = self.get_shard_index(id);
481 let node = self.nodes[shard_index]
482 .get(id)
483 .ok_or(error_helpers::node_not_found(id.clone()))?;
484 let mut new_node = node.as_ref().clone();
485 new_node.marks.extend(marks.clone());
486 self.nodes[shard_index] =
487 self.nodes[shard_index].update(id.clone(), Arc::new(new_node));
488 Ok(())
489 }
490
491 pub fn move_node(
492 &mut self,
493 source_parent_id: &NodeId,
494 target_parent_id: &NodeId,
495 node_id: &NodeId,
496 position: Option<usize>,
497 ) -> PoolResult<()> {
498 let source_shard_index = self.get_shard_index(source_parent_id);
499 let target_shard_index = self.get_shard_index(target_parent_id);
500 let node_shard_index = self.get_shard_index(node_id);
501 let source_parent = self.nodes[source_shard_index]
502 .get(source_parent_id)
503 .ok_or(error_helpers::parent_not_found(source_parent_id.clone()))?;
504 let target_parent = self.nodes[target_shard_index]
505 .get(target_parent_id)
506 .ok_or(error_helpers::parent_not_found(target_parent_id.clone()))?;
507 let _node = self.nodes[node_shard_index]
508 .get(node_id)
509 .ok_or(error_helpers::node_not_found(node_id.clone()))?;
510 if !source_parent.content.contains(node_id) {
511 return Err(error_helpers::invalid_parenting(
512 node_id.clone(),
513 source_parent_id.clone(),
514 ));
515 }
516 let mut new_source_parent = source_parent.as_ref().clone();
517 new_source_parent.content = new_source_parent
518 .content
519 .iter()
520 .filter(|&id| id != node_id)
521 .cloned()
522 .collect();
523 let mut new_target_parent = target_parent.as_ref().clone();
524 if let Some(pos) = position {
525 let insert_pos = pos.min(new_target_parent.content.len());
527
528 if insert_pos == new_target_parent.content.len() {
529 new_target_parent.content.push_back(node_id.clone());
531 } else {
532 let mut new_content = im::Vector::new();
534
535 for (i, child_id) in
537 new_target_parent.content.iter().enumerate()
538 {
539 if i == insert_pos {
540 new_content.push_back(node_id.clone());
542 }
543 new_content.push_back(child_id.clone());
544 }
545
546 new_target_parent.content = new_content;
547 }
548 } else {
549 new_target_parent.content.push_back(node_id.clone());
551 }
552 self.nodes[source_shard_index] = self.nodes[source_shard_index]
553 .update(source_parent_id.clone(), Arc::new(new_source_parent));
554 self.nodes[target_shard_index] = self.nodes[target_shard_index]
555 .update(target_parent_id.clone(), Arc::new(new_target_parent));
556 self.parent_map.insert(node_id.clone(), target_parent_id.clone());
557 Ok(())
558 }
559
560 pub fn remove_node(
561 &mut self,
562 parent_id: &NodeId,
563 nodes: Vec<NodeId>,
564 ) -> PoolResult<()> {
565 let parent_shard_index = self.get_shard_index(parent_id);
566 let parent = self.nodes[parent_shard_index]
567 .get(parent_id)
568 .ok_or(error_helpers::parent_not_found(parent_id.clone()))?;
569 if nodes.contains(&self.root_id) {
570 return Err(error_helpers::cannot_remove_root());
571 }
572 for node_id in &nodes {
573 if !parent.content.contains(node_id) {
574 return Err(error_helpers::invalid_parenting(
575 node_id.clone(),
576 parent_id.clone(),
577 ));
578 }
579 }
580 let nodes_to_remove: std::collections::HashSet<_> =
581 nodes.iter().collect();
582 let filtered_children: im::Vector<NodeId> = parent
583 .as_ref()
584 .content
585 .iter()
586 .filter(|&id| !nodes_to_remove.contains(id))
587 .cloned()
588 .collect();
589 let mut parent_node = parent.as_ref().clone();
590 parent_node.content = filtered_children;
591 self.nodes[parent_shard_index] = self.nodes[parent_shard_index]
592 .update(parent_id.clone(), Arc::new(parent_node));
593 let mut remove_nodes = Vec::new();
594 for node_id in nodes {
595 self.remove_subtree(&node_id, &mut remove_nodes)?;
596 }
597 Ok(())
598 }
599 pub fn remove_node_by_id(
601 &mut self,
602 node_id: &NodeId,
603 ) -> PoolResult<()> {
604 if node_id == &self.root_id {
606 return Err(error_helpers::cannot_remove_root());
607 }
608
609 let shard_index = self.get_shard_index(node_id);
610 let _ = self.nodes[shard_index]
611 .get(node_id)
612 .ok_or(error_helpers::node_not_found(node_id.clone()))?;
613
614 if let Some(parent_id) = self.parent_map.get(node_id).cloned() {
616 let parent_shard_index = self.get_shard_index(&parent_id);
617 if let Some(parent_node) =
618 self.nodes[parent_shard_index].get(&parent_id)
619 {
620 let mut new_parent = parent_node.as_ref().clone();
621 new_parent.content = new_parent
622 .content
623 .iter()
624 .filter(|&id| id != node_id)
625 .cloned()
626 .collect();
627 self.nodes[parent_shard_index] = self.nodes[parent_shard_index]
628 .update(parent_id.clone(), Arc::new(new_parent));
629 }
630 }
631
632 let mut remove_nodes = Vec::new();
634 self.remove_subtree(node_id, &mut remove_nodes)?;
635
636 Ok(())
638 }
639
640 pub fn remove_node_by_index(
642 &mut self,
643 parent_id: &NodeId,
644 index: usize,
645 ) -> PoolResult<()> {
646 let shard_index = self.get_shard_index(parent_id);
647 let parent = self.nodes[shard_index]
648 .get(parent_id)
649 .ok_or(error_helpers::parent_not_found(parent_id.clone()))?;
650 let mut new_parent = parent.as_ref().clone();
651 let remove_node_id = new_parent.content.remove(index);
652 self.nodes[shard_index] = self.nodes[shard_index]
653 .update(parent_id.clone(), Arc::new(new_parent));
654 let mut remove_nodes = Vec::new();
655 self.remove_subtree(&remove_node_id, &mut remove_nodes)?;
656 for node in remove_nodes {
657 let shard_index = self.get_shard_index(&node.id);
658 self.nodes[shard_index].remove(&node.id);
659 self.parent_map.remove(&node.id);
660 }
661 Ok(())
662 }
663
664 fn remove_subtree(
666 &mut self,
667 node_id: &NodeId,
668 remove_nodes: &mut Vec<Node>,
669 ) -> PoolResult<()> {
670 if node_id == &self.root_id {
671 return Err(error_helpers::cannot_remove_root());
672 }
673 let shard_index = self.get_shard_index(node_id);
674 let _ = self.nodes[shard_index]
675 .get(node_id)
676 .ok_or(error_helpers::node_not_found(node_id.clone()))?;
677 if let Some(children) = self.children(node_id) {
678 for child_id in children {
679 self.remove_subtree(&child_id, remove_nodes)?;
680 }
681 }
682 self.parent_map.remove(node_id);
683 if let Some(remove_node) = self.nodes[shard_index].remove(node_id) {
684 remove_nodes.push(remove_node.as_ref().clone());
685 }
686 Ok(())
687 }
688}
689
690impl Index<&NodeId> for Tree {
691 type Output = Arc<Node>;
692 fn index(
693 &self,
694 index: &NodeId,
695 ) -> &Self::Output {
696 let shard_index = self.get_shard_index(index);
697 self.nodes[shard_index].get(index).expect("Node not found")
698 }
699}
700
701impl Index<&str> for Tree {
702 type Output = Arc<Node>;
703 fn index(
704 &self,
705 index: &str,
706 ) -> &Self::Output {
707 let node_id = NodeId::from(index);
708 let shard_index = self.get_shard_index(&node_id);
709 self.nodes[shard_index].get(&node_id).expect("Node not found")
710 }
711}
712
713#[cfg(test)]
714mod tests {
715 use super::*;
716 use crate::node::Node;
717 use crate::node_type::NodeEnum;
718 use crate::types::NodeId;
719 use crate::attrs::Attrs;
720 use crate::mark::Mark;
721 use im::HashMap;
722 use serde_json::json;
723
724 fn create_test_node(id: &str) -> Node {
725 Node::new(id, "test".to_string(), Attrs::default(), vec![], vec![])
726 }
727
728 #[test]
729 fn test_tree_creation() {
730 let root = create_test_node("root");
731 let tree = Tree::new(root.clone());
732 assert_eq!(tree.root_id, root.id);
733 assert!(tree.contains_node(&root.id));
734 }
735
736 #[test]
737 fn test_add_node() {
738 let root = create_test_node("root");
739 let mut tree = Tree::new(root.clone());
740
741 let child = create_test_node("child");
742 let nodes = vec![child.clone()];
743
744 tree.add_node(&root.id, &nodes).unwrap();
745 dbg!(&tree);
746 assert!(tree.contains_node(&child.id));
747 assert_eq!(tree.children(&root.id).unwrap().len(), 1);
748 }
749
750 #[test]
751 fn test_remove_node() {
752 let root = create_test_node("root");
753 let mut tree = Tree::new(root.clone());
754
755 let child = create_test_node("child");
756 let nodes = vec![child.clone()];
757
758 tree.add_node(&root.id, &nodes).unwrap();
759 dbg!(&tree);
760 tree.remove_node(&root.id, vec![child.id.clone()]).unwrap();
761 dbg!(&tree);
762 assert!(!tree.contains_node(&child.id));
763 assert_eq!(tree.children(&root.id).unwrap().len(), 0);
764 }
765
766 #[test]
767 fn test_move_node() {
768 let parent1 = create_test_node("parent1");
770 let parent2 = create_test_node("parent2");
771 let mut tree = Tree::new(parent1.clone());
772
773 tree.add_node(&parent1.id, &vec![parent2.clone()]).unwrap();
775
776 let child1 = create_test_node("child1");
778 let child2 = create_test_node("child2");
779 let child3 = create_test_node("child3");
780
781 tree.add_node(&parent1.id, &vec![child1.clone()]).unwrap();
783 tree.add_node(&parent1.id, &vec![child2.clone()]).unwrap();
784 tree.add_node(&parent1.id, &vec![child3.clone()]).unwrap();
785
786 let parent1_children = tree.children(&parent1.id).unwrap();
788 assert_eq!(parent1_children.len(), 4); assert_eq!(parent1_children[0], parent2.id);
790 assert_eq!(parent1_children[1], child1.id);
791 assert_eq!(parent1_children[2], child2.id);
792 assert_eq!(parent1_children[3], child3.id);
793
794 tree.move_node(&parent1.id, &parent2.id, &child1.id, None).unwrap();
796
797 let parent1_children = tree.children(&parent1.id).unwrap();
799 let parent2_children = tree.children(&parent2.id).unwrap();
800 assert_eq!(parent1_children.len(), 3); assert_eq!(parent2_children.len(), 1); assert_eq!(parent2_children[0], child1.id);
803
804 tree.move_node(&parent1.id, &parent2.id, &child2.id, Some(1)).unwrap();
806
807 let parent1_children = tree.children(&parent1.id).unwrap();
809 let parent2_children = tree.children(&parent2.id).unwrap();
810 assert_eq!(parent1_children.len(), 2); assert_eq!(parent2_children.len(), 2); assert_eq!(parent2_children[0], child1.id);
813 assert_eq!(parent2_children[1], child2.id);
814
815 let child1_parent = tree.get_parent_node(&child1.id).unwrap();
817 let child2_parent = tree.get_parent_node(&child2.id).unwrap();
818 assert_eq!(child1_parent.id, parent2.id);
819 assert_eq!(child2_parent.id, parent2.id);
820 }
821
822 #[test]
823 fn test_update_attr() {
824 let root = create_test_node("root");
825 let mut tree = Tree::new(root.clone());
826
827 let mut attrs = HashMap::new();
828 attrs.insert("key".to_string(), json!("value"));
829
830 tree.update_attr(&root.id, attrs).unwrap();
831
832 let node = tree.get_node(&root.id).unwrap();
833 dbg!(&node);
834 assert_eq!(node.attrs.get("key").unwrap(), &json!("value"));
835 }
836
837 #[test]
838 fn test_add_mark() {
839 let root = create_test_node("root");
840 let mut tree = Tree::new(root.clone());
841
842 let mark = Mark { r#type: "test".to_string(), attrs: Attrs::default() };
843 tree.add_mark(&root.id, &vec![mark.clone()]).unwrap();
844 dbg!(&tree);
845 let node = tree.get_node(&root.id).unwrap();
846 assert!(node.marks.contains(&mark));
847 }
848
849 #[test]
850 fn test_remove_mark() {
851 let root = create_test_node("root");
852 let mut tree = Tree::new(root.clone());
853
854 let mark = Mark { r#type: "test".to_string(), attrs: Attrs::default() };
855 tree.add_mark(&root.id, &vec![mark.clone()]).unwrap();
856 dbg!(&tree);
857 tree.remove_mark(&root.id, mark.clone()).unwrap();
858 dbg!(&tree);
859 let node = tree.get_node(&root.id).unwrap();
860 assert!(!node.marks.contains(&mark));
861 }
862
863 #[test]
864 fn test_all_children() {
865 let root = create_test_node("root");
866 let mut tree = Tree::new(root.clone());
867
868 let child1 = create_test_node("child1");
869 let child2 = create_test_node("child2");
870
871 tree.add_node(&root.id, &vec![child1.clone()]).unwrap();
872 tree.add_node(&root.id, &vec![child2.clone()]).unwrap();
873 dbg!(&tree);
874 let all_children = tree.all_children(&root.id, None).unwrap();
875 assert_eq!(all_children.1.len(), 2);
876 }
877
878 #[test]
879 fn test_children_count() {
880 let root = create_test_node("root");
881 let mut tree = Tree::new(root.clone());
882
883 let child1 = create_test_node("child1");
884 let child2 = create_test_node("child2");
885
886 tree.add_node(&root.id, &vec![child1.clone()]).unwrap();
887 tree.add_node(&root.id, &vec![child2.clone()]).unwrap();
888
889 assert_eq!(tree.children_count(&root.id), 2);
890 }
891
892 #[test]
893 fn test_remove_node_by_id_updates_parent() {
894 let root = create_test_node("root");
895 let mut tree = Tree::new(root.clone());
896
897 let child = create_test_node("child");
898 tree.add_node(&root.id, &vec![child.clone()]).unwrap();
899
900 assert_eq!(tree.children_count(&root.id), 1);
902 assert!(tree.contains_node(&child.id));
903
904 tree.remove_node_by_id(&child.id).unwrap();
906
907 assert_eq!(tree.children_count(&root.id), 0);
909 assert!(!tree.contains_node(&child.id));
910 }
911
912 #[test]
913 fn test_move_node_position_edge_cases() {
914 let root = create_test_node("root");
915 let mut tree = Tree::new(root.clone());
916
917 let container = create_test_node("container");
918 tree.add_node(&root.id, &vec![container.clone()]).unwrap();
919
920 let child1 = create_test_node("child1");
921 let child2 = create_test_node("child2");
922 let child3 = create_test_node("child3");
923
924 tree.add_node(&root.id, &vec![child1.clone()]).unwrap();
925 tree.add_node(&root.id, &vec![child2.clone()]).unwrap();
926 tree.add_node(&root.id, &vec![child3.clone()]).unwrap();
927
928 tree.move_node(&root.id, &container.id, &child1.id, Some(100)).unwrap();
930
931 let container_children = tree.children(&container.id).unwrap();
932 assert_eq!(container_children.len(), 1);
933 assert_eq!(container_children[0], child1.id);
934
935 tree.move_node(&root.id, &container.id, &child2.id, Some(0)).unwrap();
937
938 let container_children = tree.children(&container.id).unwrap();
939 assert_eq!(container_children.len(), 2);
940 assert_eq!(container_children[0], child2.id);
941 assert_eq!(container_children[1], child1.id);
942 }
943
944 #[test]
945 fn test_cannot_remove_root_node() {
946 let root = create_test_node("root");
947 let mut tree = Tree::new(root.clone());
948
949 let result = tree.remove_node_by_id(&root.id);
951 assert!(result.is_err());
952 }
953
954 #[test]
955 fn test_get_parent_node() {
956 let root = create_test_node("root");
957 let mut tree = Tree::new(root.clone());
958
959 let child = create_test_node("child");
960 tree.add_node(&root.id, &vec![child.clone()]).unwrap();
961
962 let parent = tree.get_parent_node(&child.id).unwrap();
963 assert_eq!(parent.id, root.id);
964 }
965}