1#![warn(missing_docs)]
22
23use codec::{Decode, Encode};
24use std::{cmp::Reverse, fmt};
25
26#[derive(Clone, Debug, PartialEq)]
28pub enum Error<E> {
29 Duplicate,
31 UnfinalizedAncestor,
33 Revert,
35 Client(E),
37}
38
39impl<E: std::error::Error> fmt::Display for Error<E> {
40 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
41 let message = match *self {
42 Error::Duplicate => "Hash already exists in Tree".into(),
43 Error::UnfinalizedAncestor => "Finalized descendent of Tree node without finalizing its ancestor(s) first".into(),
44 Error::Revert => "Tried to import or finalize node that is an ancestor of a previously finalized node".into(),
45 Error::Client(ref err) => format!("Client error: {}", err),
46 };
47 write!(f, "{}", message)
48 }
49}
50
51impl<E: std::error::Error> std::error::Error for Error<E> {}
52
53impl<E: std::error::Error> From<E> for Error<E> {
54 fn from(err: E) -> Error<E> {
55 Error::Client(err)
56 }
57}
58
59#[derive(Debug, PartialEq)]
61pub enum FinalizationResult<V> {
62 Changed(Option<V>),
64 Unchanged,
66}
67
68#[derive(Debug, PartialEq)]
70pub enum FilterAction {
71 Remove,
73 KeepNode,
75 KeepTree,
77}
78
79#[derive(Clone, Debug, Decode, Encode, PartialEq)]
88pub struct ForkTree<H, N, V> {
89 roots: Vec<Node<H, N, V>>,
90 best_finalized_number: Option<N>,
91}
92
93impl<H, N, V> ForkTree<H, N, V>
94where
95 H: PartialEq,
96 N: Ord,
97{
98 pub fn new() -> ForkTree<H, N, V> {
100 ForkTree { roots: Vec::new(), best_finalized_number: None }
101 }
102
103 pub fn rebalance(&mut self) {
113 self.roots.sort_by_key(|n| Reverse(n.max_depth()));
114 let mut stack: Vec<_> = self.roots.iter_mut().collect();
115 while let Some(node) = stack.pop() {
116 node.children.sort_by_key(|n| Reverse(n.max_depth()));
117 stack.extend(node.children.iter_mut());
118 }
119 }
120
121 pub fn import<F, E>(
135 &mut self,
136 hash: H,
137 number: N,
138 data: V,
139 is_descendent_of: &F,
140 ) -> Result<bool, Error<E>>
141 where
142 E: std::error::Error,
143 F: Fn(&H, &H) -> Result<bool, E>,
144 {
145 if let Some(ref best_finalized_number) = self.best_finalized_number {
146 if number <= *best_finalized_number {
147 return Err(Error::Revert);
148 }
149 }
150
151 let (children, is_root) =
152 match self.find_node_where_mut(&hash, &number, is_descendent_of, &|_| true)? {
153 Some(parent) => (&mut parent.children, false),
154 None => (&mut self.roots, true),
155 };
156
157 if children.iter().any(|elem| elem.hash == hash) {
158 return Err(Error::Duplicate);
159 }
160
161 children.push(Node { data, hash, number, children: Default::default() });
162
163 if children.len() == 1 {
164 self.rebalance();
166 }
167
168 Ok(is_root)
169 }
170
171 pub fn roots(&self) -> impl Iterator<Item = (&H, &N, &V)> {
173 self.roots.iter().map(|node| (&node.hash, &node.number, &node.data))
174 }
175
176 fn node_iter(&self) -> impl Iterator<Item = &Node<H, N, V>> {
177 ForkTreeIterator { stack: self.roots.iter().rev().collect() }
180 }
181
182 pub fn iter(&self) -> impl Iterator<Item = (&H, &N, &V)> {
184 self.node_iter().map(|node| (&node.hash, &node.number, &node.data))
185 }
186
187 pub fn map<VT, F>(self, f: &mut F) -> ForkTree<H, N, VT>
193 where
194 F: FnMut(&H, &N, V) -> VT,
195 {
196 let mut queue: Vec<_> =
197 self.roots.into_iter().rev().map(|node| (usize::MAX, node)).collect();
198 let mut next_queue = Vec::new();
199 let mut output = Vec::new();
200
201 while !queue.is_empty() {
202 for (parent_index, node) in queue.drain(..) {
203 let new_data = f(&node.hash, &node.number, node.data);
204 let new_node = Node {
205 hash: node.hash,
206 number: node.number,
207 data: new_data,
208 children: Vec::with_capacity(node.children.len()),
209 };
210
211 let node_id = output.len();
212 output.push((parent_index, new_node));
213
214 for child in node.children.into_iter().rev() {
215 next_queue.push((node_id, child));
216 }
217 }
218
219 std::mem::swap(&mut queue, &mut next_queue);
220 }
221
222 let mut roots = Vec::new();
223 while let Some((parent_index, new_node)) = output.pop() {
224 if parent_index == usize::MAX {
225 roots.push(new_node);
226 } else {
227 output[parent_index].1.children.push(new_node);
228 }
229 }
230
231 ForkTree { roots, best_finalized_number: self.best_finalized_number }
232 }
233
234 pub fn find_node_where<F, E, P>(
240 &self,
241 hash: &H,
242 number: &N,
243 is_descendent_of: &F,
244 predicate: &P,
245 ) -> Result<Option<&Node<H, N, V>>, Error<E>>
246 where
247 E: std::error::Error,
248 F: Fn(&H, &H) -> Result<bool, E>,
249 P: Fn(&V) -> bool,
250 {
251 let maybe_path = self.find_node_index_where(hash, number, is_descendent_of, predicate)?;
252 Ok(maybe_path.map(|path| {
253 let children =
254 path.iter().take(path.len() - 1).fold(&self.roots, |curr, &i| &curr[i].children);
255 &children[path[path.len() - 1]]
256 }))
257 }
258
259 pub fn find_node_where_mut<F, E, P>(
261 &mut self,
262 hash: &H,
263 number: &N,
264 is_descendent_of: &F,
265 predicate: &P,
266 ) -> Result<Option<&mut Node<H, N, V>>, Error<E>>
267 where
268 E: std::error::Error,
269 F: Fn(&H, &H) -> Result<bool, E>,
270 P: Fn(&V) -> bool,
271 {
272 let maybe_path = self.find_node_index_where(hash, number, is_descendent_of, predicate)?;
273 Ok(maybe_path.map(|path| {
274 let children = path
275 .iter()
276 .take(path.len() - 1)
277 .fold(&mut self.roots, |curr, &i| &mut curr[i].children);
278 &mut children[path[path.len() - 1]]
279 }))
280 }
281
282 pub fn find_node_index_where<F, E, P>(
296 &self,
297 hash: &H,
298 number: &N,
299 is_descendent_of: &F,
300 predicate: &P,
301 ) -> Result<Option<Vec<usize>>, Error<E>>
302 where
303 E: std::error::Error,
304 F: Fn(&H, &H) -> Result<bool, E>,
305 P: Fn(&V) -> bool,
306 {
307 let mut stack = vec![];
308 let mut root_idx = 0;
309 let mut found = false;
310 let mut is_descendent = false;
311
312 while root_idx < self.roots.len() {
313 if *number <= self.roots[root_idx].number {
314 root_idx += 1;
315 continue;
316 }
317 stack.push((&self.roots[root_idx], 0));
321 while let Some((node, i)) = stack.pop() {
322 if i < node.children.len() && !is_descendent {
323 stack.push((node, i + 1));
324 if node.children[i].number < *number {
325 stack.push((&node.children[i], 0));
326 }
327 } else if is_descendent || is_descendent_of(&node.hash, hash)? {
328 is_descendent = true;
329 if predicate(&node.data) {
330 found = true;
331 break;
332 }
333 }
334 }
335
336 if is_descendent {
339 break;
340 }
341 root_idx += 1;
342 }
343
344 Ok(if found {
345 let path: Vec<_> =
349 std::iter::once(root_idx).chain(stack.iter().map(|(_, i)| *i - 1)).collect();
350 Some(path)
351 } else {
352 None
353 })
354 }
355
356 pub fn prune<F, E, P>(
367 &mut self,
368 hash: &H,
369 number: &N,
370 is_descendent_of: &F,
371 predicate: &P,
372 ) -> Result<impl Iterator<Item = (H, N, V)>, Error<E>>
373 where
374 E: std::error::Error,
375 F: Fn(&H, &H) -> Result<bool, E>,
376 P: Fn(&V) -> bool,
377 {
378 let new_root_path =
379 match self.find_node_index_where(hash, number, is_descendent_of, predicate)? {
380 Some(path) => path,
381 None => return Ok(RemovedIterator { stack: Vec::new() }),
382 };
383
384 let mut removed = std::mem::take(&mut self.roots);
385
386 let root_siblings = new_root_path
388 .iter()
389 .take(new_root_path.len() - 1)
390 .fold(&mut removed, |curr, idx| &mut curr[*idx].children);
391 let root = root_siblings.remove(new_root_path[new_root_path.len() - 1]);
392 self.roots = vec![root];
393
394 let mut curr = &mut self.roots[0];
398 loop {
399 let mut maybe_ancestor_idx = None;
400 for (idx, child) in curr.children.iter().enumerate() {
401 if child.number < *number && is_descendent_of(&child.hash, hash)? {
402 maybe_ancestor_idx = Some(idx);
403 break;
404 }
405 }
406 let Some(ancestor_idx) = maybe_ancestor_idx else {
407 break;
409 };
410 let mut next_siblings = std::mem::take(&mut curr.children);
412 let next = next_siblings.remove(ancestor_idx);
413 curr.children = vec![next];
414 removed.append(&mut next_siblings);
415 curr = &mut curr.children[0];
416 }
417
418 let children = std::mem::take(&mut curr.children);
421 for child in children {
422 if child.number == *number && child.hash == *hash ||
423 *number < child.number && is_descendent_of(hash, &child.hash)?
424 {
425 curr.children.push(child);
426 } else {
427 removed.push(child);
428 }
429 }
430
431 self.rebalance();
432
433 Ok(RemovedIterator { stack: removed })
434 }
435
436 pub fn finalize_root(&mut self, hash: &H) -> Option<V> {
440 self.roots
441 .iter()
442 .position(|node| node.hash == *hash)
443 .map(|position| self.finalize_root_at(position))
444 }
445
446 fn finalize_root_at(&mut self, position: usize) -> V {
448 let node = self.roots.swap_remove(position);
449 self.roots = node.children;
450 self.best_finalized_number = Some(node.number);
451 node.data
452 }
453
454 pub fn finalize<F, E>(
460 &mut self,
461 hash: &H,
462 number: N,
463 is_descendent_of: &F,
464 ) -> Result<FinalizationResult<V>, Error<E>>
465 where
466 E: std::error::Error,
467 F: Fn(&H, &H) -> Result<bool, E>,
468 {
469 if let Some(ref best_finalized_number) = self.best_finalized_number {
470 if number <= *best_finalized_number {
471 return Err(Error::Revert);
472 }
473 }
474
475 if let Some(root) = self.finalize_root(hash) {
477 return Ok(FinalizationResult::Changed(Some(root)));
478 }
479
480 for root in self.roots.iter() {
482 if number > root.number && is_descendent_of(&root.hash, hash)? {
483 return Err(Error::UnfinalizedAncestor);
484 }
485 }
486
487 let mut changed = false;
491 let roots = std::mem::take(&mut self.roots);
492
493 for root in roots {
494 if root.number > number && is_descendent_of(hash, &root.hash)? {
495 self.roots.push(root);
496 } else {
497 changed = true;
498 }
499 }
500
501 self.best_finalized_number = Some(number);
502
503 if changed {
504 Ok(FinalizationResult::Changed(None))
505 } else {
506 Ok(FinalizationResult::Unchanged)
507 }
508 }
509
510 pub fn finalize_with_ancestors<F, E>(
514 &mut self,
515 hash: &H,
516 number: N,
517 is_descendent_of: &F,
518 ) -> Result<FinalizationResult<V>, Error<E>>
519 where
520 E: std::error::Error,
521 F: Fn(&H, &H) -> Result<bool, E>,
522 {
523 if let Some(ref best_finalized_number) = self.best_finalized_number {
524 if number <= *best_finalized_number {
525 return Err(Error::Revert);
526 }
527 }
528
529 if let Some(root) = self.finalize_root(hash) {
531 return Ok(FinalizationResult::Changed(Some(root)));
532 }
533
534 let mut changed = false;
539 let mut idx = 0;
540 while idx != self.roots.len() {
541 let (is_finalized, is_descendant, is_ancestor) = {
542 let root = &self.roots[idx];
543 let is_finalized = root.hash == *hash;
544 let is_descendant =
545 !is_finalized && root.number > number && is_descendent_of(hash, &root.hash)?;
546 let is_ancestor = !is_finalized &&
547 !is_descendant && root.number < number &&
548 is_descendent_of(&root.hash, hash)?;
549 (is_finalized, is_descendant, is_ancestor)
550 };
551
552 if is_finalized {
554 return Ok(FinalizationResult::Changed(Some(self.finalize_root_at(idx))));
555 }
556
557 if is_descendant {
559 idx += 1;
560 continue;
561 }
562
563 if is_ancestor {
565 let root = self.roots.swap_remove(idx);
566 self.roots.extend(root.children);
567 changed = true;
568 continue;
569 }
570
571 self.roots.swap_remove(idx);
573 changed = true;
574 }
575
576 self.best_finalized_number = Some(number);
577
578 if changed {
579 Ok(FinalizationResult::Changed(None))
580 } else {
581 Ok(FinalizationResult::Unchanged)
582 }
583 }
584
585 pub fn finalizes_any_with_descendent_if<F, P, E>(
595 &self,
596 hash: &H,
597 number: N,
598 is_descendent_of: &F,
599 predicate: P,
600 ) -> Result<Option<bool>, Error<E>>
601 where
602 E: std::error::Error,
603 F: Fn(&H, &H) -> Result<bool, E>,
604 P: Fn(&V) -> bool,
605 {
606 if let Some(ref best_finalized_number) = self.best_finalized_number {
607 if number <= *best_finalized_number {
608 return Err(Error::Revert);
609 }
610 }
611
612 for node in self.node_iter() {
616 if predicate(&node.data) && (node.hash == *hash || is_descendent_of(&node.hash, hash)?)
617 {
618 for child in node.children.iter() {
619 if child.number <= number &&
620 (child.hash == *hash || is_descendent_of(&child.hash, hash)?)
621 {
622 return Err(Error::UnfinalizedAncestor);
623 }
624 }
625
626 return Ok(Some(self.roots.iter().any(|root| root.hash == node.hash)));
627 }
628 }
629
630 Ok(None)
631 }
632
633 pub fn finalize_with_descendent_if<F, P, E>(
641 &mut self,
642 hash: &H,
643 number: N,
644 is_descendent_of: &F,
645 predicate: P,
646 ) -> Result<FinalizationResult<V>, Error<E>>
647 where
648 E: std::error::Error,
649 F: Fn(&H, &H) -> Result<bool, E>,
650 P: Fn(&V) -> bool,
651 {
652 if let Some(ref best_finalized_number) = self.best_finalized_number {
653 if number <= *best_finalized_number {
654 return Err(Error::Revert);
655 }
656 }
657
658 let mut position = None;
662 for (i, root) in self.roots.iter().enumerate() {
663 if predicate(&root.data) && (root.hash == *hash || is_descendent_of(&root.hash, hash)?)
664 {
665 for child in root.children.iter() {
666 if child.number <= number &&
667 (child.hash == *hash || is_descendent_of(&child.hash, hash)?)
668 {
669 return Err(Error::UnfinalizedAncestor);
670 }
671 }
672
673 position = Some(i);
674 break;
675 }
676 }
677
678 let node_data = position.map(|i| {
679 let node = self.roots.swap_remove(i);
680 self.roots = node.children;
681 self.best_finalized_number = Some(node.number);
682 node.data
683 });
684
685 let mut changed = false;
691 let roots = std::mem::take(&mut self.roots);
692
693 for root in roots {
694 let retain = root.number > number && is_descendent_of(hash, &root.hash)? ||
695 root.number == number && root.hash == *hash ||
696 is_descendent_of(&root.hash, hash)?;
697
698 if retain {
699 self.roots.push(root);
700 } else {
701 changed = true;
702 }
703 }
704
705 self.best_finalized_number = Some(number);
706
707 match (node_data, changed) {
708 (Some(data), _) => Ok(FinalizationResult::Changed(Some(data))),
709 (None, true) => Ok(FinalizationResult::Changed(None)),
710 (None, false) => Ok(FinalizationResult::Unchanged),
711 }
712 }
713
714 pub fn drain_filter<F>(&mut self, filter: F) -> impl Iterator<Item = (H, N, V)>
723 where
724 F: Fn(&H, &N, &V) -> FilterAction,
725 {
726 let mut removed = vec![];
727 let mut retained = Vec::new();
728
729 let mut queue: Vec<_> = std::mem::take(&mut self.roots)
730 .into_iter()
731 .rev()
732 .map(|node| (usize::MAX, node))
733 .collect();
734 let mut next_queue = Vec::new();
735
736 while !queue.is_empty() {
737 for (parent_idx, mut node) in queue.drain(..) {
738 match filter(&node.hash, &node.number, &node.data) {
739 FilterAction::KeepNode => {
740 let node_idx = retained.len();
741 let children = std::mem::take(&mut node.children);
742 retained.push((parent_idx, node));
743 for child in children.into_iter().rev() {
744 next_queue.push((node_idx, child));
745 }
746 },
747 FilterAction::KeepTree => {
748 retained.push((parent_idx, node));
749 },
750 FilterAction::Remove => {
751 removed.push(node);
752 },
753 }
754 }
755
756 std::mem::swap(&mut queue, &mut next_queue);
757 }
758
759 while let Some((parent_idx, node)) = retained.pop() {
760 if parent_idx == usize::MAX {
761 self.roots.push(node);
762 } else {
763 retained[parent_idx].1.children.push(node);
764 }
765 }
766
767 if !removed.is_empty() {
768 self.rebalance();
769 }
770 RemovedIterator { stack: removed }
771 }
772}
773
774use node_implementation::Node;
776
777mod node_implementation {
778 use super::*;
779
780 #[derive(Clone, Debug, Decode, Encode, PartialEq)]
781 pub struct Node<H, N, V> {
782 pub hash: H,
783 pub number: N,
784 pub data: V,
785 pub children: Vec<Node<H, N, V>>,
786 }
787
788 impl<H: PartialEq, N: Ord, V> Node<H, N, V> {
789 pub fn max_depth(&self) -> usize {
791 let mut max: usize = 0;
792 let mut stack = vec![(self, 0)];
793 while let Some((node, height)) = stack.pop() {
794 if height > max {
795 max = height;
796 }
797 node.children.iter().for_each(|n| stack.push((n, height + 1)));
798 }
799 max
800 }
801 }
802}
803
804struct ForkTreeIterator<'a, H, N, V> {
805 stack: Vec<&'a Node<H, N, V>>,
806}
807
808impl<'a, H, N, V> Iterator for ForkTreeIterator<'a, H, N, V> {
809 type Item = &'a Node<H, N, V>;
810
811 fn next(&mut self) -> Option<Self::Item> {
812 self.stack.pop().inspect(|node| {
813 self.stack.extend(node.children.iter().rev());
817 })
818 }
819}
820
821struct RemovedIterator<H, N, V> {
822 stack: Vec<Node<H, N, V>>,
823}
824
825impl<H, N, V> Iterator for RemovedIterator<H, N, V> {
826 type Item = (H, N, V);
827
828 fn next(&mut self) -> Option<Self::Item> {
829 self.stack.pop().map(|mut node| {
830 let children = std::mem::take(&mut node.children);
834
835 self.stack.extend(children.into_iter().rev());
836 (node.hash, node.number, node.data)
837 })
838 }
839}
840
841#[cfg(test)]
842mod test {
843 use crate::FilterAction;
844
845 use super::{Error, FinalizationResult, ForkTree};
846
847 #[derive(Debug, PartialEq)]
848 struct TestError;
849
850 impl std::fmt::Display for TestError {
851 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
852 write!(f, "TestError")
853 }
854 }
855
856 impl std::error::Error for TestError {}
857
858 fn test_fork_tree<'a>(
859 ) -> (ForkTree<&'a str, u64, u32>, impl Fn(&&str, &&str) -> Result<bool, TestError>) {
860 let mut tree = ForkTree::new();
861
862 #[rustfmt::skip]
863 let is_descendent_of = |base: &&str, block: &&str| -> Result<bool, TestError> {
881 let letters = vec!["B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O"];
882 let block = block.to_uppercase();
885 match (*base, block) {
886 ("A", b) => Ok(letters.into_iter().any(|n| n == b)),
887 ("B" | "c", b) => Ok(b == "C" || b == "D" || b == "E"),
888 ("C", b) => Ok(b == "D" || b == "E"),
889 ("D", b) => Ok(b == "E"),
890 ("E", _) => Ok(false),
891 ("F", b) =>
892 Ok(b == "G" || b == "H" || b == "I" || b == "L" || b == "M" || b == "N" || b == "O"),
893 ("G", _) => Ok(false),
894 ("H", b) => Ok(b == "I" || b == "L" || b == "M" || b == "N" || b == "O"),
895 ("I", _) => Ok(false),
896 ("J", b) => Ok(b == "K"),
897 ("K", _) => Ok(false),
898 ("L", b) => Ok(b == "M" || b == "N" || b == "O"),
899 ("m", b) => Ok(b == "M" || b == "N"),
900 ("M", b) => Ok(b == "N"),
901 ("O", _) => Ok(false),
902 ("0", _) => Ok(true),
903 _ => Ok(false),
904 }
905 };
906
907 tree.import("A", 10, 1, &is_descendent_of).unwrap();
908 tree.import("B", 20, 2, &is_descendent_of).unwrap();
909 tree.import("C", 30, 3, &is_descendent_of).unwrap();
910 tree.import("D", 40, 4, &is_descendent_of).unwrap();
911 tree.import("E", 50, 5, &is_descendent_of).unwrap();
912 tree.import("F", 20, 2, &is_descendent_of).unwrap();
913 tree.import("G", 30, 3, &is_descendent_of).unwrap();
914 tree.import("H", 30, 3, &is_descendent_of).unwrap();
915 tree.import("I", 40, 4, &is_descendent_of).unwrap();
916 tree.import("L", 40, 4, &is_descendent_of).unwrap();
917 tree.import("M", 50, 5, &is_descendent_of).unwrap();
918 tree.import("O", 50, 5, &is_descendent_of).unwrap();
919 tree.import("J", 20, 2, &is_descendent_of).unwrap();
920 tree.import("K", 30, 3, &is_descendent_of).unwrap();
921
922 (tree, is_descendent_of)
923 }
924
925 #[test]
926 fn import_doesnt_revert() {
927 let (mut tree, is_descendent_of) = test_fork_tree();
928
929 tree.finalize_root(&"A");
930
931 assert_eq!(tree.best_finalized_number, Some(10));
932
933 assert_eq!(tree.import("A", 10, 1, &is_descendent_of), Err(Error::Revert));
934 }
935
936 #[test]
937 fn import_doesnt_add_duplicates() {
938 let (mut tree, is_descendent_of) = test_fork_tree();
939
940 assert_eq!(tree.import("A", 10, 1, &is_descendent_of), Err(Error::Duplicate));
941
942 assert_eq!(tree.import("I", 40, 4, &is_descendent_of), Err(Error::Duplicate));
943
944 assert_eq!(tree.import("G", 30, 3, &is_descendent_of), Err(Error::Duplicate));
945
946 assert_eq!(tree.import("K", 30, 3, &is_descendent_of), Err(Error::Duplicate));
947 }
948
949 #[test]
950 fn finalize_root_works() {
951 let finalize_a = || {
952 let (mut tree, ..) = test_fork_tree();
953
954 assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("A", 10)]);
955
956 tree.finalize_root(&"A");
958
959 assert_eq!(
960 tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(),
961 vec![("B", 20), ("F", 20), ("J", 20)],
962 );
963
964 tree
965 };
966
967 {
968 let mut tree = finalize_a();
969
970 tree.finalize_root(&"B");
972
973 assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("C", 30)],);
974
975 assert!(tree.roots.len() == 1);
977 }
978
979 {
980 let mut tree = finalize_a();
981
982 tree.finalize_root(&"J");
984
985 assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("K", 30)],);
986
987 assert!(tree.roots.len() == 1);
989 }
990 }
991
992 #[test]
993 fn finalize_works() {
994 let (mut tree, is_descendent_of) = test_fork_tree();
995
996 let original_roots = tree.roots.clone();
997
998 assert_eq!(tree.finalize(&"0", 0, &is_descendent_of), Ok(FinalizationResult::Unchanged));
1000
1001 assert_eq!(tree.roots, original_roots);
1002
1003 assert_eq!(
1005 tree.finalize(&"A", 10, &is_descendent_of),
1006 Ok(FinalizationResult::Changed(Some(1))),
1007 );
1008
1009 assert_eq!(
1010 tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(),
1011 vec![("B", 20), ("F", 20), ("J", 20)],
1012 );
1013
1014 assert_eq!(tree.best_finalized_number, Some(10));
1016
1017 assert_eq!(tree.finalize(&"Z", 10, &is_descendent_of), Err(Error::Revert));
1018
1019 assert_eq!(tree.finalize(&"H", 30, &is_descendent_of), Err(Error::UnfinalizedAncestor));
1021
1022 assert_eq!(
1024 tree.finalize(&"F", 20, &is_descendent_of),
1025 Ok(FinalizationResult::Changed(Some(2))),
1026 );
1027
1028 assert_eq!(
1029 tree.finalize(&"H", 30, &is_descendent_of),
1030 Ok(FinalizationResult::Changed(Some(3))),
1031 );
1032
1033 assert_eq!(
1034 tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(),
1035 vec![("L", 40), ("I", 40)],
1036 );
1037
1038 assert_eq!(
1040 tree.finalize(&"Z", 50, &is_descendent_of),
1041 Ok(FinalizationResult::Changed(None)),
1042 );
1043
1044 assert!(tree.roots.is_empty());
1045 }
1046
1047 #[test]
1048 fn finalize_with_ancestor_works() {
1049 let (mut tree, is_descendent_of) = test_fork_tree();
1050
1051 let original_roots = tree.roots.clone();
1052
1053 assert_eq!(
1055 tree.finalize_with_ancestors(&"0", 0, &is_descendent_of),
1056 Ok(FinalizationResult::Unchanged),
1057 );
1058
1059 assert_eq!(tree.roots, original_roots);
1060
1061 assert_eq!(
1063 tree.finalize_with_ancestors(&"A", 10, &is_descendent_of),
1064 Ok(FinalizationResult::Changed(Some(1))),
1065 );
1066
1067 assert_eq!(
1068 tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(),
1069 vec![("B", 20), ("F", 20), ("J", 20)],
1070 );
1071
1072 assert_eq!(
1077 tree.finalize_with_ancestors(&"H", 30, &is_descendent_of),
1078 Ok(FinalizationResult::Changed(Some(3))),
1079 );
1080
1081 assert_eq!(
1082 tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(),
1083 vec![("L", 40), ("I", 40)],
1084 );
1085
1086 assert_eq!(tree.best_finalized_number, Some(30));
1087
1088 assert_eq!(
1094 tree.finalize_with_ancestors(&"N", 60, &is_descendent_of),
1095 Ok(FinalizationResult::Changed(None)),
1096 );
1097
1098 assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![],);
1099
1100 assert_eq!(tree.best_finalized_number, Some(60));
1101 }
1102
1103 #[test]
1104 fn finalize_with_descendent_works() {
1105 #[derive(Debug, PartialEq)]
1106 struct Change {
1107 effective: u64,
1108 }
1109
1110 let (mut tree, is_descendent_of) = {
1111 let mut tree = ForkTree::new();
1112
1113 let is_descendent_of = |base: &&str, block: &&str| -> Result<bool, TestError> {
1114 match (*base, *block) {
1122 ("A0", b) => Ok(b == "B" || b == "C" || b == "D" || b == "E" || b == "G"),
1123 ("A1", _) => Ok(false),
1124 ("C", b) => Ok(b == "D"),
1125 ("D", b) => Ok(b == "E" || b == "F" || b == "G"),
1126 ("E", b) => Ok(b == "F"),
1127 _ => Ok(false),
1128 }
1129 };
1130
1131 let is_root = tree.import("A0", 1, Change { effective: 5 }, &is_descendent_of).unwrap();
1132 assert!(is_root);
1133 let is_root = tree.import("A1", 1, Change { effective: 5 }, &is_descendent_of).unwrap();
1134 assert!(is_root);
1135 let is_root =
1136 tree.import("D", 10, Change { effective: 10 }, &is_descendent_of).unwrap();
1137 assert!(!is_root);
1138 let is_root =
1139 tree.import("E", 15, Change { effective: 50 }, &is_descendent_of).unwrap();
1140 assert!(!is_root);
1141
1142 (tree, is_descendent_of)
1143 };
1144
1145 assert_eq!(
1146 tree.finalizes_any_with_descendent_if(
1147 &"B",
1148 2,
1149 &is_descendent_of,
1150 |c| c.effective <= 2,
1151 ),
1152 Ok(None),
1153 );
1154
1155 assert_eq!(
1157 tree.finalize_with_descendent_if(&"D", 10, &is_descendent_of, |c| c.effective <= 10),
1158 Err(Error::UnfinalizedAncestor)
1159 );
1160
1161 assert_eq!(
1164 tree.finalizes_any_with_descendent_if(&"D", 10, &is_descendent_of, |c| c.effective ==
1165 10),
1166 Ok(Some(false)),
1167 );
1168
1169 assert_eq!(
1171 tree.finalizes_any_with_descendent_if(&"E", 15, &is_descendent_of, |c| c.effective ==
1172 10),
1173 Err(Error::UnfinalizedAncestor)
1174 );
1175
1176 assert_eq!(
1179 tree.finalize_with_descendent_if(&"B", 2, &is_descendent_of, |c| c.effective <= 2),
1180 Ok(FinalizationResult::Changed(None)),
1181 );
1182
1183 assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("A0", 1)],);
1184
1185 assert_eq!(
1187 tree.finalizes_any_with_descendent_if(
1188 &"C",
1189 5,
1190 &is_descendent_of,
1191 |c| c.effective <= 5,
1192 ),
1193 Ok(Some(true)),
1194 );
1195
1196 assert_eq!(
1197 tree.finalize_with_descendent_if(&"C", 5, &is_descendent_of, |c| c.effective <= 5),
1198 Ok(FinalizationResult::Changed(Some(Change { effective: 5 }))),
1199 );
1200
1201 assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("D", 10)],);
1202
1203 assert_eq!(
1205 tree.finalizes_any_with_descendent_if(&"F", 100, &is_descendent_of, |c| c.effective <=
1206 100,),
1207 Err(Error::UnfinalizedAncestor),
1208 );
1209
1210 assert_eq!(
1212 tree.finalizes_any_with_descendent_if(&"G", 100, &is_descendent_of, |c| c.effective <=
1213 100),
1214 Ok(Some(true)),
1215 );
1216
1217 assert_eq!(
1218 tree.finalize_with_descendent_if(&"G", 100, &is_descendent_of, |c| c.effective <= 100),
1219 Ok(FinalizationResult::Changed(Some(Change { effective: 10 }))),
1220 );
1221
1222 assert_eq!(tree.roots().count(), 0);
1224 }
1225
1226 #[test]
1227 fn iter_iterates_in_preorder() {
1228 let (tree, ..) = test_fork_tree();
1229 assert_eq!(
1230 tree.iter().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(),
1231 vec![
1232 ("A", 10),
1233 ("B", 20),
1234 ("C", 30),
1235 ("D", 40),
1236 ("E", 50),
1237 ("F", 20),
1238 ("H", 30),
1239 ("L", 40),
1240 ("M", 50),
1241 ("O", 50),
1242 ("I", 40),
1243 ("G", 30),
1244 ("J", 20),
1245 ("K", 30),
1246 ],
1247 );
1248 }
1249
1250 #[test]
1251 fn minimizes_calls_to_is_descendent_of() {
1252 use std::sync::atomic::{AtomicUsize, Ordering};
1253
1254 let n_is_descendent_of_calls = AtomicUsize::new(0);
1255
1256 let is_descendent_of = |_: &&str, _: &&str| -> Result<bool, TestError> {
1257 n_is_descendent_of_calls.fetch_add(1, Ordering::SeqCst);
1258 Ok(true)
1259 };
1260
1261 {
1262 let mut tree = ForkTree::new();
1266 let letters = vec!["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K"];
1267
1268 for (i, letter) in letters.iter().enumerate() {
1269 tree.import::<_, TestError>(*letter, i, i, &|_, _| Ok(true)).unwrap();
1270 }
1271
1272 assert_eq!(
1275 tree.finalizes_any_with_descendent_if(&"L", 11, &is_descendent_of, |i| *i == 10,),
1276 Ok(Some(false)),
1277 );
1278
1279 assert_eq!(n_is_descendent_of_calls.load(Ordering::SeqCst), 1);
1280 }
1281
1282 n_is_descendent_of_calls.store(0, Ordering::SeqCst);
1283
1284 {
1285 let mut tree = ForkTree::new();
1289 let letters = vec!["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K"];
1290
1291 for (i, letter) in letters.iter().enumerate() {
1292 tree.import::<_, TestError>(*letter, i, i, &|_, _| Ok(false)).unwrap();
1293 }
1294
1295 assert_eq!(
1298 tree.finalize_with_descendent_if(&"L", 11, &is_descendent_of, |i| *i == 10,),
1299 Ok(FinalizationResult::Changed(Some(10))),
1300 );
1301
1302 assert_eq!(n_is_descendent_of_calls.load(Ordering::SeqCst), 1);
1303 }
1304 }
1305
1306 #[test]
1307 fn map_works() {
1308 let (mut tree, _) = test_fork_tree();
1309
1310 let is_descendent_of = |_: &&str, _: &&str| -> Result<bool, TestError> { Ok(false) };
1312 let is_root = tree.import("A1", 10, 1, &is_descendent_of).unwrap();
1313 assert!(is_root);
1314 let is_root = tree.import("A2", 10, 1, &is_descendent_of).unwrap();
1315 assert!(is_root);
1316
1317 let old_tree = tree.clone();
1318 let new_tree = tree.map(&mut |hash, _, _| hash.to_owned());
1319
1320 assert!(new_tree.iter().all(|(hash, _, data)| hash == data));
1322 assert_eq!(
1323 old_tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(),
1324 new_tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(),
1325 );
1326 }
1327
1328 #[test]
1329 fn prune_works_for_in_tree_hashes() {
1330 let (mut tree, is_descendent_of) = test_fork_tree();
1331
1332 let removed = tree.prune(&"C", &30, &is_descendent_of, &|_| true).unwrap();
1333
1334 assert_eq!(tree.roots.iter().map(|node| node.hash).collect::<Vec<_>>(), vec!["B"]);
1335
1336 assert_eq!(
1337 tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(),
1338 vec!["B", "C", "D", "E"],
1339 );
1340
1341 assert_eq!(
1342 removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(),
1343 vec!["A", "F", "H", "L", "M", "O", "I", "G", "J", "K"]
1344 );
1345
1346 let removed = tree.prune(&"E", &50, &is_descendent_of, &|_| true).unwrap();
1347
1348 assert_eq!(tree.roots.iter().map(|node| node.hash).collect::<Vec<_>>(), vec!["D"]);
1349
1350 assert_eq!(tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(), vec!["D", "E"]);
1351
1352 assert_eq!(removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(), vec!["B", "C"]);
1353 }
1354
1355 #[test]
1356 fn prune_works_for_out_of_tree_hashes() {
1357 let (mut tree, is_descendent_of) = test_fork_tree();
1358
1359 let removed = tree.prune(&"c", &25, &is_descendent_of, &|_| true).unwrap();
1360
1361 assert_eq!(tree.roots.iter().map(|node| node.hash).collect::<Vec<_>>(), vec!["B"]);
1362
1363 assert_eq!(
1364 tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(),
1365 vec!["B", "C", "D", "E"],
1366 );
1367
1368 assert_eq!(
1369 removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(),
1370 vec!["A", "F", "H", "L", "M", "O", "I", "G", "J", "K"]
1371 );
1372 }
1373
1374 #[test]
1375 fn prune_works_for_not_direct_ancestor() {
1376 let (mut tree, is_descendent_of) = test_fork_tree();
1377
1378 let removed = tree.prune(&"m", &45, &is_descendent_of, &|height| *height == 3).unwrap();
1380
1381 assert_eq!(tree.roots.iter().map(|node| node.hash).collect::<Vec<_>>(), vec!["H"]);
1382
1383 assert_eq!(tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(), vec!["H", "L", "M"],);
1384
1385 assert_eq!(
1386 removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(),
1387 vec!["O", "I", "A", "B", "C", "D", "E", "F", "G", "J", "K"]
1388 );
1389 }
1390
1391 #[test]
1392 fn prune_works_for_far_away_ancestor() {
1393 let (mut tree, is_descendent_of) = test_fork_tree();
1394
1395 let removed = tree.prune(&"m", &45, &is_descendent_of, &|height| *height == 2).unwrap();
1396
1397 assert_eq!(tree.roots.iter().map(|node| node.hash).collect::<Vec<_>>(), vec!["F"]);
1398
1399 assert_eq!(
1400 tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(),
1401 vec!["F", "H", "L", "M"],
1402 );
1403
1404 assert_eq!(
1405 removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(),
1406 vec!["O", "I", "G", "A", "B", "C", "D", "E", "J", "K"]
1407 );
1408 }
1409
1410 #[test]
1411 fn find_node_backtracks_after_finding_highest_descending_node() {
1412 let mut tree = ForkTree::new();
1413
1414 let is_descendent_of = |base: &&str, block: &&str| -> Result<bool, TestError> {
1419 match (*base, *block) {
1420 ("A", b) => Ok(b == "B" || b == "C" || b == "D"),
1421 ("B", b) | ("C", b) => Ok(b == "D"),
1422 ("0", _) => Ok(true),
1423 _ => Ok(false),
1424 }
1425 };
1426
1427 tree.import("A", 1, 1, &is_descendent_of).unwrap();
1428 tree.import("B", 2, 2, &is_descendent_of).unwrap();
1429 tree.import("C", 2, 4, &is_descendent_of).unwrap();
1430
1431 let node = tree.find_node_where(&"D", &3, &is_descendent_of, &|data| *data < 3).unwrap();
1435
1436 assert_eq!(node.unwrap().hash, "B");
1437 }
1438
1439 #[test]
1440 fn rebalance_works() {
1441 let (mut tree, _) = test_fork_tree();
1442
1443 assert_eq!(
1447 tree.iter().map(|(h, _, _)| *h).collect::<Vec<_>>(),
1448 vec!["A", "B", "C", "D", "E", "F", "H", "L", "M", "O", "I", "G", "J", "K"],
1449 );
1450
1451 let is_descendent_of = |base: &&str, block: &&str| -> Result<bool, TestError> {
1453 match (*base, *block) {
1454 (b, "P") => Ok(vec!["A", "F", "H", "L", "O"].into_iter().any(|n| n == b)),
1455 _ => Ok(false),
1456 }
1457 };
1458
1459 tree.import("P", 60, 6, &is_descendent_of).unwrap();
1460
1461 assert_eq!(
1465 tree.iter().map(|(h, _, _)| *h).collect::<Vec<_>>(),
1466 ["A", "F", "H", "L", "O", "P", "M", "I", "G", "B", "C", "D", "E", "J", "K"]
1467 );
1468 }
1469
1470 #[test]
1471 fn drain_filter_works() {
1472 let (mut tree, _) = test_fork_tree();
1473
1474 let filter = |h: &&str, _: &u64, _: &u32| match *h {
1475 "A" | "B" | "F" | "G" => FilterAction::KeepNode,
1476 "C" => FilterAction::KeepTree,
1477 "H" | "J" => FilterAction::Remove,
1478 _ => panic!("Unexpected filtering for node: {}", *h),
1479 };
1480
1481 let removed = tree.drain_filter(filter);
1482
1483 assert_eq!(
1484 tree.iter().map(|(h, _, _)| *h).collect::<Vec<_>>(),
1485 ["A", "B", "C", "D", "E", "F", "G"]
1486 );
1487
1488 assert_eq!(
1489 removed.map(|(h, _, _)| h).collect::<Vec<_>>(),
1490 ["H", "L", "M", "O", "I", "J", "K"]
1491 );
1492 }
1493
1494 #[test]
1495 fn find_node_index_works() {
1496 let (tree, is_descendent_of) = test_fork_tree();
1497
1498 let path = tree
1499 .find_node_index_where(&"D", &40, &is_descendent_of, &|_| true)
1500 .unwrap()
1501 .unwrap();
1502 assert_eq!(path, [0, 0, 0]);
1503
1504 let path = tree
1505 .find_node_index_where(&"O", &50, &is_descendent_of, &|_| true)
1506 .unwrap()
1507 .unwrap();
1508 assert_eq!(path, [0, 1, 0, 0]);
1509
1510 let path = tree
1511 .find_node_index_where(&"N", &60, &is_descendent_of, &|_| true)
1512 .unwrap()
1513 .unwrap();
1514 assert_eq!(path, [0, 1, 0, 0, 0]);
1515 }
1516
1517 #[test]
1518 fn find_node_index_with_predicate_works() {
1519 let is_descendent_of = |parent: &char, child: &char| match *parent {
1520 'A' => Ok(['B', 'C', 'D', 'E', 'F'].contains(child)),
1521 'B' => Ok(['C', 'D'].contains(child)),
1522 'C' => Ok(['D'].contains(child)),
1523 'E' => Ok(['F'].contains(child)),
1524 'D' | 'F' => Ok(false),
1525 _ => Err(TestError),
1526 };
1527
1528 let mut tree: ForkTree<char, u8, bool> = ForkTree::new();
1531 tree.import('A', 1, true, &is_descendent_of).unwrap();
1532 tree.import('B', 2, false, &is_descendent_of).unwrap();
1533 tree.import('C', 3, true, &is_descendent_of).unwrap();
1534 tree.import('D', 4, false, &is_descendent_of).unwrap();
1535
1536 tree.import('E', 2, true, &is_descendent_of).unwrap();
1537 tree.import('F', 3, false, &is_descendent_of).unwrap();
1538
1539 let path = tree
1540 .find_node_index_where(&'D', &4, &is_descendent_of, &|&value| !value)
1541 .unwrap()
1542 .unwrap();
1543 assert_eq!(path, [0, 0]);
1544
1545 let path = tree
1546 .find_node_index_where(&'D', &4, &is_descendent_of, &|&value| value)
1547 .unwrap()
1548 .unwrap();
1549 assert_eq!(path, [0, 0, 0]);
1550
1551 let path = tree
1552 .find_node_index_where(&'F', &3, &is_descendent_of, &|&value| !value)
1553 .unwrap();
1554 assert_eq!(path, None);
1555
1556 let path = tree
1557 .find_node_index_where(&'F', &3, &is_descendent_of, &|&value| value)
1558 .unwrap()
1559 .unwrap();
1560 assert_eq!(path, [0, 1]);
1561 }
1562
1563 #[test]
1564 fn find_node_works() {
1565 let (tree, is_descendent_of) = test_fork_tree();
1566
1567 let node = tree.find_node_where(&"B", &20, &is_descendent_of, &|_| true).unwrap().unwrap();
1568 assert_eq!((node.hash, node.number), ("A", 10));
1569
1570 let node = tree.find_node_where(&"D", &40, &is_descendent_of, &|_| true).unwrap().unwrap();
1571 assert_eq!((node.hash, node.number), ("C", 30));
1572
1573 let node = tree.find_node_where(&"O", &50, &is_descendent_of, &|_| true).unwrap().unwrap();
1574 assert_eq!((node.hash, node.number), ("L", 40));
1575
1576 let node = tree.find_node_where(&"N", &60, &is_descendent_of, &|_| true).unwrap().unwrap();
1577 assert_eq!((node.hash, node.number), ("M", 50));
1578 }
1579
1580 #[test]
1581 fn post_order_traversal_requirement() {
1582 let (mut tree, is_descendent_of) = test_fork_tree();
1583
1584 let is_descendent_of_for_post_order = |parent: &&str, child: &&str| match *parent {
1587 "A" => Err(TestError),
1588 "K" if *child == "Z" => Ok(true),
1589 _ => is_descendent_of(parent, child),
1590 };
1591
1592 let path = tree
1594 .find_node_index_where(&"N", &60, &is_descendent_of_for_post_order, &|_| true)
1595 .unwrap()
1596 .unwrap();
1597 assert_eq!(path, [0, 1, 0, 0, 0]);
1598
1599 let res = tree.import(&"Z", 100, 10, &is_descendent_of_for_post_order);
1601 assert_eq!(res, Ok(false));
1602 assert_eq!(
1603 tree.iter().map(|node| *node.0).collect::<Vec<_>>(),
1604 vec!["A", "B", "C", "D", "E", "F", "H", "L", "M", "O", "I", "G", "J", "K", "Z"],
1605 );
1606 }
1607}