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>(
297 &self,
298 hash: &H,
299 number: &N,
300 is_descendent_of: &F,
301 predicate: &P,
302 ) -> Result<Option<Vec<usize>>, Error<E>>
303 where
304 E: std::error::Error,
305 F: Fn(&H, &H) -> Result<bool, E>,
306 P: Fn(&V) -> bool,
307 {
308 let mut stack = vec![];
309 let mut root_idx = 0;
310 let mut found = false;
311 let mut is_descendent = false;
312
313 while root_idx < self.roots.len() {
314 if *number <= self.roots[root_idx].number {
315 root_idx += 1;
316 continue
317 }
318 stack.push((&self.roots[root_idx], 0));
322 while let Some((node, i)) = stack.pop() {
323 if i < node.children.len() && !is_descendent {
324 stack.push((node, i + 1));
325 if node.children[i].number < *number {
326 stack.push((&node.children[i], 0));
327 }
328 } else if is_descendent || is_descendent_of(&node.hash, hash)? {
329 is_descendent = true;
330 if predicate(&node.data) {
331 found = true;
332 break
333 }
334 }
335 }
336
337 if is_descendent {
340 break
341 }
342 root_idx += 1;
343 }
344
345 Ok(if found {
346 let path: Vec<_> =
350 std::iter::once(root_idx).chain(stack.iter().map(|(_, i)| *i - 1)).collect();
351 Some(path)
352 } else {
353 None
354 })
355 }
356
357 pub fn prune<F, E, P>(
368 &mut self,
369 hash: &H,
370 number: &N,
371 is_descendent_of: &F,
372 predicate: &P,
373 ) -> Result<impl Iterator<Item = (H, N, V)>, Error<E>>
374 where
375 E: std::error::Error,
376 F: Fn(&H, &H) -> Result<bool, E>,
377 P: Fn(&V) -> bool,
378 {
379 let new_root_path =
380 match self.find_node_index_where(hash, number, is_descendent_of, predicate)? {
381 Some(path) => path,
382 None => return Ok(RemovedIterator { stack: Vec::new() }),
383 };
384
385 let mut removed = std::mem::take(&mut self.roots);
386
387 let root_siblings = new_root_path
389 .iter()
390 .take(new_root_path.len() - 1)
391 .fold(&mut removed, |curr, idx| &mut curr[*idx].children);
392 let root = root_siblings.remove(new_root_path[new_root_path.len() - 1]);
393 self.roots = vec![root];
394
395 let mut curr = &mut self.roots[0];
399 loop {
400 let mut maybe_ancestor_idx = None;
401 for (idx, child) in curr.children.iter().enumerate() {
402 if child.number < *number && is_descendent_of(&child.hash, hash)? {
403 maybe_ancestor_idx = Some(idx);
404 break
405 }
406 }
407 let Some(ancestor_idx) = maybe_ancestor_idx else {
408 break
410 };
411 let mut next_siblings = std::mem::take(&mut curr.children);
413 let next = next_siblings.remove(ancestor_idx);
414 curr.children = vec![next];
415 removed.append(&mut next_siblings);
416 curr = &mut curr.children[0];
417 }
418
419 let children = std::mem::take(&mut curr.children);
422 for child in children {
423 if child.number == *number && child.hash == *hash ||
424 *number < child.number && is_descendent_of(hash, &child.hash)?
425 {
426 curr.children.push(child);
427 } else {
428 removed.push(child);
429 }
430 }
431
432 self.rebalance();
433
434 Ok(RemovedIterator { stack: removed })
435 }
436
437 pub fn finalize_root(&mut self, hash: &H) -> Option<V> {
441 self.roots
442 .iter()
443 .position(|node| node.hash == *hash)
444 .map(|position| self.finalize_root_at(position))
445 }
446
447 fn finalize_root_at(&mut self, position: usize) -> V {
449 let node = self.roots.swap_remove(position);
450 self.roots = node.children;
451 self.best_finalized_number = Some(node.number);
452 node.data
453 }
454
455 pub fn finalize<F, E>(
461 &mut self,
462 hash: &H,
463 number: N,
464 is_descendent_of: &F,
465 ) -> Result<FinalizationResult<V>, Error<E>>
466 where
467 E: std::error::Error,
468 F: Fn(&H, &H) -> Result<bool, E>,
469 {
470 if let Some(ref best_finalized_number) = self.best_finalized_number {
471 if number <= *best_finalized_number {
472 return Err(Error::Revert)
473 }
474 }
475
476 if let Some(root) = self.finalize_root(hash) {
478 return Ok(FinalizationResult::Changed(Some(root)))
479 }
480
481 for root in self.roots.iter() {
483 if number > root.number && is_descendent_of(&root.hash, hash)? {
484 return Err(Error::UnfinalizedAncestor)
485 }
486 }
487
488 let mut changed = false;
492 let roots = std::mem::take(&mut self.roots);
493
494 for root in roots {
495 if root.number > number && is_descendent_of(hash, &root.hash)? {
496 self.roots.push(root);
497 } else {
498 changed = true;
499 }
500 }
501
502 self.best_finalized_number = Some(number);
503
504 if changed {
505 Ok(FinalizationResult::Changed(None))
506 } else {
507 Ok(FinalizationResult::Unchanged)
508 }
509 }
510
511 pub fn finalize_with_ancestors<F, E>(
515 &mut self,
516 hash: &H,
517 number: N,
518 is_descendent_of: &F,
519 ) -> Result<FinalizationResult<V>, Error<E>>
520 where
521 E: std::error::Error,
522 F: Fn(&H, &H) -> Result<bool, E>,
523 {
524 if let Some(ref best_finalized_number) = self.best_finalized_number {
525 if number <= *best_finalized_number {
526 return Err(Error::Revert)
527 }
528 }
529
530 if let Some(root) = self.finalize_root(hash) {
532 return Ok(FinalizationResult::Changed(Some(root)))
533 }
534
535 let mut changed = false;
540 let mut idx = 0;
541 while idx != self.roots.len() {
542 let (is_finalized, is_descendant, is_ancestor) = {
543 let root = &self.roots[idx];
544 let is_finalized = root.hash == *hash;
545 let is_descendant =
546 !is_finalized && root.number > number && is_descendent_of(hash, &root.hash)?;
547 let is_ancestor = !is_finalized &&
548 !is_descendant && root.number < number &&
549 is_descendent_of(&root.hash, hash)?;
550 (is_finalized, is_descendant, is_ancestor)
551 };
552
553 if is_finalized {
555 return Ok(FinalizationResult::Changed(Some(self.finalize_root_at(idx))))
556 }
557
558 if is_descendant {
560 idx += 1;
561 continue
562 }
563
564 if is_ancestor {
566 let root = self.roots.swap_remove(idx);
567 self.roots.extend(root.children);
568 changed = true;
569 continue
570 }
571
572 self.roots.swap_remove(idx);
574 changed = true;
575 }
576
577 self.best_finalized_number = Some(number);
578
579 if changed {
580 Ok(FinalizationResult::Changed(None))
581 } else {
582 Ok(FinalizationResult::Unchanged)
583 }
584 }
585
586 pub fn finalizes_any_with_descendent_if<F, P, E>(
596 &self,
597 hash: &H,
598 number: N,
599 is_descendent_of: &F,
600 predicate: P,
601 ) -> Result<Option<bool>, Error<E>>
602 where
603 E: std::error::Error,
604 F: Fn(&H, &H) -> Result<bool, E>,
605 P: Fn(&V) -> bool,
606 {
607 if let Some(ref best_finalized_number) = self.best_finalized_number {
608 if number <= *best_finalized_number {
609 return Err(Error::Revert)
610 }
611 }
612
613 for node in self.node_iter() {
617 if predicate(&node.data) && (node.hash == *hash || is_descendent_of(&node.hash, hash)?)
618 {
619 for child in node.children.iter() {
620 if child.number <= number &&
621 (child.hash == *hash || is_descendent_of(&child.hash, hash)?)
622 {
623 return Err(Error::UnfinalizedAncestor)
624 }
625 }
626
627 return Ok(Some(self.roots.iter().any(|root| root.hash == node.hash)))
628 }
629 }
630
631 Ok(None)
632 }
633
634 pub fn finalize_with_descendent_if<F, P, E>(
642 &mut self,
643 hash: &H,
644 number: N,
645 is_descendent_of: &F,
646 predicate: P,
647 ) -> Result<FinalizationResult<V>, Error<E>>
648 where
649 E: std::error::Error,
650 F: Fn(&H, &H) -> Result<bool, E>,
651 P: Fn(&V) -> bool,
652 {
653 if let Some(ref best_finalized_number) = self.best_finalized_number {
654 if number <= *best_finalized_number {
655 return Err(Error::Revert)
656 }
657 }
658
659 let mut position = None;
663 for (i, root) in self.roots.iter().enumerate() {
664 if predicate(&root.data) && (root.hash == *hash || is_descendent_of(&root.hash, hash)?)
665 {
666 for child in root.children.iter() {
667 if child.number <= number &&
668 (child.hash == *hash || is_descendent_of(&child.hash, hash)?)
669 {
670 return Err(Error::UnfinalizedAncestor)
671 }
672 }
673
674 position = Some(i);
675 break
676 }
677 }
678
679 let node_data = position.map(|i| {
680 let node = self.roots.swap_remove(i);
681 self.roots = node.children;
682 self.best_finalized_number = Some(node.number);
683 node.data
684 });
685
686 let mut changed = false;
692 let roots = std::mem::take(&mut self.roots);
693
694 for root in roots {
695 let retain = root.number > number && is_descendent_of(hash, &root.hash)? ||
696 root.number == number && root.hash == *hash ||
697 is_descendent_of(&root.hash, hash)?;
698
699 if retain {
700 self.roots.push(root);
701 } else {
702 changed = true;
703 }
704 }
705
706 self.best_finalized_number = Some(number);
707
708 match (node_data, changed) {
709 (Some(data), _) => Ok(FinalizationResult::Changed(Some(data))),
710 (None, true) => Ok(FinalizationResult::Changed(None)),
711 (None, false) => Ok(FinalizationResult::Unchanged),
712 }
713 }
714
715 pub fn drain_filter<F>(&mut self, filter: F) -> impl Iterator<Item = (H, N, V)>
724 where
725 F: Fn(&H, &N, &V) -> FilterAction,
726 {
727 let mut removed = vec![];
728 let mut retained = Vec::new();
729
730 let mut queue: Vec<_> = std::mem::take(&mut self.roots)
731 .into_iter()
732 .rev()
733 .map(|node| (usize::MAX, node))
734 .collect();
735 let mut next_queue = Vec::new();
736
737 while !queue.is_empty() {
738 for (parent_idx, mut node) in queue.drain(..) {
739 match filter(&node.hash, &node.number, &node.data) {
740 FilterAction::KeepNode => {
741 let node_idx = retained.len();
742 let children = std::mem::take(&mut node.children);
743 retained.push((parent_idx, node));
744 for child in children.into_iter().rev() {
745 next_queue.push((node_idx, child));
746 }
747 },
748 FilterAction::KeepTree => {
749 retained.push((parent_idx, node));
750 },
751 FilterAction::Remove => {
752 removed.push(node);
753 },
754 }
755 }
756
757 std::mem::swap(&mut queue, &mut next_queue);
758 }
759
760 while let Some((parent_idx, node)) = retained.pop() {
761 if parent_idx == usize::MAX {
762 self.roots.push(node);
763 } else {
764 retained[parent_idx].1.children.push(node);
765 }
766 }
767
768 if !removed.is_empty() {
769 self.rebalance();
770 }
771 RemovedIterator { stack: removed }
772 }
773}
774
775use node_implementation::Node;
777
778mod node_implementation {
779 use super::*;
780
781 #[derive(Clone, Debug, Decode, Encode, PartialEq)]
782 pub struct Node<H, N, V> {
783 pub hash: H,
784 pub number: N,
785 pub data: V,
786 pub children: Vec<Node<H, N, V>>,
787 }
788
789 impl<H: PartialEq, N: Ord, V> Node<H, N, V> {
790 pub fn max_depth(&self) -> usize {
792 let mut max: usize = 0;
793 let mut stack = vec![(self, 0)];
794 while let Some((node, height)) = stack.pop() {
795 if height > max {
796 max = height;
797 }
798 node.children.iter().for_each(|n| stack.push((n, height + 1)));
799 }
800 max
801 }
802 }
803}
804
805struct ForkTreeIterator<'a, H, N, V> {
806 stack: Vec<&'a Node<H, N, V>>,
807}
808
809impl<'a, H, N, V> Iterator for ForkTreeIterator<'a, H, N, V> {
810 type Item = &'a Node<H, N, V>;
811
812 fn next(&mut self) -> Option<Self::Item> {
813 self.stack.pop().inspect(|node| {
814 self.stack.extend(node.children.iter().rev());
818 })
819 }
820}
821
822struct RemovedIterator<H, N, V> {
823 stack: Vec<Node<H, N, V>>,
824}
825
826impl<H, N, V> Iterator for RemovedIterator<H, N, V> {
827 type Item = (H, N, V);
828
829 fn next(&mut self) -> Option<Self::Item> {
830 self.stack.pop().map(|mut node| {
831 let children = std::mem::take(&mut node.children);
835
836 self.stack.extend(children.into_iter().rev());
837 (node.hash, node.number, node.data)
838 })
839 }
840}
841
842#[cfg(test)]
843mod test {
844 use crate::FilterAction;
845
846 use super::{Error, FinalizationResult, ForkTree};
847
848 #[derive(Debug, PartialEq)]
849 struct TestError;
850
851 impl std::fmt::Display for TestError {
852 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
853 write!(f, "TestError")
854 }
855 }
856
857 impl std::error::Error for TestError {}
858
859 fn test_fork_tree<'a>(
860 ) -> (ForkTree<&'a str, u64, u32>, impl Fn(&&str, &&str) -> Result<bool, TestError>) {
861 let mut tree = ForkTree::new();
862
863 #[rustfmt::skip]
864 let is_descendent_of = |base: &&str, block: &&str| -> Result<bool, TestError> {
882 let letters = vec!["B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O"];
883 let block = block.to_uppercase();
886 match (*base, block) {
887 ("A", b) => Ok(letters.into_iter().any(|n| n == b)),
888 ("B" | "c", b) => Ok(b == "C" || b == "D" || b == "E"),
889 ("C", b) => Ok(b == "D" || b == "E"),
890 ("D", b) => Ok(b == "E"),
891 ("E", _) => Ok(false),
892 ("F", b) =>
893 Ok(b == "G" || b == "H" || b == "I" || b == "L" || b == "M" || b == "N" || b == "O"),
894 ("G", _) => Ok(false),
895 ("H", b) => Ok(b == "I" || b == "L" || b == "M" || b == "N" || b == "O"),
896 ("I", _) => Ok(false),
897 ("J", b) => Ok(b == "K"),
898 ("K", _) => Ok(false),
899 ("L", b) => Ok(b == "M" || b == "N" || b == "O"),
900 ("m", b) => Ok(b == "M" || b == "N"),
901 ("M", b) => Ok(b == "N"),
902 ("O", _) => Ok(false),
903 ("0", _) => Ok(true),
904 _ => Ok(false),
905 }
906 };
907
908 tree.import("A", 10, 1, &is_descendent_of).unwrap();
909 tree.import("B", 20, 2, &is_descendent_of).unwrap();
910 tree.import("C", 30, 3, &is_descendent_of).unwrap();
911 tree.import("D", 40, 4, &is_descendent_of).unwrap();
912 tree.import("E", 50, 5, &is_descendent_of).unwrap();
913 tree.import("F", 20, 2, &is_descendent_of).unwrap();
914 tree.import("G", 30, 3, &is_descendent_of).unwrap();
915 tree.import("H", 30, 3, &is_descendent_of).unwrap();
916 tree.import("I", 40, 4, &is_descendent_of).unwrap();
917 tree.import("L", 40, 4, &is_descendent_of).unwrap();
918 tree.import("M", 50, 5, &is_descendent_of).unwrap();
919 tree.import("O", 50, 5, &is_descendent_of).unwrap();
920 tree.import("J", 20, 2, &is_descendent_of).unwrap();
921 tree.import("K", 30, 3, &is_descendent_of).unwrap();
922
923 (tree, is_descendent_of)
924 }
925
926 #[test]
927 fn import_doesnt_revert() {
928 let (mut tree, is_descendent_of) = test_fork_tree();
929
930 tree.finalize_root(&"A");
931
932 assert_eq!(tree.best_finalized_number, Some(10));
933
934 assert_eq!(tree.import("A", 10, 1, &is_descendent_of), Err(Error::Revert));
935 }
936
937 #[test]
938 fn import_doesnt_add_duplicates() {
939 let (mut tree, is_descendent_of) = test_fork_tree();
940
941 assert_eq!(tree.import("A", 10, 1, &is_descendent_of), Err(Error::Duplicate));
942
943 assert_eq!(tree.import("I", 40, 4, &is_descendent_of), Err(Error::Duplicate));
944
945 assert_eq!(tree.import("G", 30, 3, &is_descendent_of), Err(Error::Duplicate));
946
947 assert_eq!(tree.import("K", 30, 3, &is_descendent_of), Err(Error::Duplicate));
948 }
949
950 #[test]
951 fn finalize_root_works() {
952 let finalize_a = || {
953 let (mut tree, ..) = test_fork_tree();
954
955 assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("A", 10)]);
956
957 tree.finalize_root(&"A");
959
960 assert_eq!(
961 tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(),
962 vec![("B", 20), ("F", 20), ("J", 20)],
963 );
964
965 tree
966 };
967
968 {
969 let mut tree = finalize_a();
970
971 tree.finalize_root(&"B");
973
974 assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("C", 30)],);
975
976 assert!(tree.roots.len() == 1);
978 }
979
980 {
981 let mut tree = finalize_a();
982
983 tree.finalize_root(&"J");
985
986 assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("K", 30)],);
987
988 assert!(tree.roots.len() == 1);
990 }
991 }
992
993 #[test]
994 fn finalize_works() {
995 let (mut tree, is_descendent_of) = test_fork_tree();
996
997 let original_roots = tree.roots.clone();
998
999 assert_eq!(tree.finalize(&"0", 0, &is_descendent_of), Ok(FinalizationResult::Unchanged));
1001
1002 assert_eq!(tree.roots, original_roots);
1003
1004 assert_eq!(
1006 tree.finalize(&"A", 10, &is_descendent_of),
1007 Ok(FinalizationResult::Changed(Some(1))),
1008 );
1009
1010 assert_eq!(
1011 tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(),
1012 vec![("B", 20), ("F", 20), ("J", 20)],
1013 );
1014
1015 assert_eq!(tree.best_finalized_number, Some(10));
1017
1018 assert_eq!(tree.finalize(&"Z", 10, &is_descendent_of), Err(Error::Revert));
1019
1020 assert_eq!(tree.finalize(&"H", 30, &is_descendent_of), Err(Error::UnfinalizedAncestor));
1022
1023 assert_eq!(
1025 tree.finalize(&"F", 20, &is_descendent_of),
1026 Ok(FinalizationResult::Changed(Some(2))),
1027 );
1028
1029 assert_eq!(
1030 tree.finalize(&"H", 30, &is_descendent_of),
1031 Ok(FinalizationResult::Changed(Some(3))),
1032 );
1033
1034 assert_eq!(
1035 tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(),
1036 vec![("L", 40), ("I", 40)],
1037 );
1038
1039 assert_eq!(
1041 tree.finalize(&"Z", 50, &is_descendent_of),
1042 Ok(FinalizationResult::Changed(None)),
1043 );
1044
1045 assert!(tree.roots.is_empty());
1046 }
1047
1048 #[test]
1049 fn finalize_with_ancestor_works() {
1050 let (mut tree, is_descendent_of) = test_fork_tree();
1051
1052 let original_roots = tree.roots.clone();
1053
1054 assert_eq!(
1056 tree.finalize_with_ancestors(&"0", 0, &is_descendent_of),
1057 Ok(FinalizationResult::Unchanged),
1058 );
1059
1060 assert_eq!(tree.roots, original_roots);
1061
1062 assert_eq!(
1064 tree.finalize_with_ancestors(&"A", 10, &is_descendent_of),
1065 Ok(FinalizationResult::Changed(Some(1))),
1066 );
1067
1068 assert_eq!(
1069 tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(),
1070 vec![("B", 20), ("F", 20), ("J", 20)],
1071 );
1072
1073 assert_eq!(
1078 tree.finalize_with_ancestors(&"H", 30, &is_descendent_of),
1079 Ok(FinalizationResult::Changed(Some(3))),
1080 );
1081
1082 assert_eq!(
1083 tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(),
1084 vec![("L", 40), ("I", 40)],
1085 );
1086
1087 assert_eq!(tree.best_finalized_number, Some(30));
1088
1089 assert_eq!(
1095 tree.finalize_with_ancestors(&"N", 60, &is_descendent_of),
1096 Ok(FinalizationResult::Changed(None)),
1097 );
1098
1099 assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![],);
1100
1101 assert_eq!(tree.best_finalized_number, Some(60));
1102 }
1103
1104 #[test]
1105 fn finalize_with_descendent_works() {
1106 #[derive(Debug, PartialEq)]
1107 struct Change {
1108 effective: u64,
1109 }
1110
1111 let (mut tree, is_descendent_of) = {
1112 let mut tree = ForkTree::new();
1113
1114 let is_descendent_of = |base: &&str, block: &&str| -> Result<bool, TestError> {
1115 match (*base, *block) {
1123 ("A0", b) => Ok(b == "B" || b == "C" || b == "D" || b == "E" || b == "G"),
1124 ("A1", _) => Ok(false),
1125 ("C", b) => Ok(b == "D"),
1126 ("D", b) => Ok(b == "E" || b == "F" || b == "G"),
1127 ("E", b) => Ok(b == "F"),
1128 _ => Ok(false),
1129 }
1130 };
1131
1132 let is_root = tree.import("A0", 1, Change { effective: 5 }, &is_descendent_of).unwrap();
1133 assert!(is_root);
1134 let is_root = tree.import("A1", 1, Change { effective: 5 }, &is_descendent_of).unwrap();
1135 assert!(is_root);
1136 let is_root =
1137 tree.import("D", 10, Change { effective: 10 }, &is_descendent_of).unwrap();
1138 assert!(!is_root);
1139 let is_root =
1140 tree.import("E", 15, Change { effective: 50 }, &is_descendent_of).unwrap();
1141 assert!(!is_root);
1142
1143 (tree, is_descendent_of)
1144 };
1145
1146 assert_eq!(
1147 tree.finalizes_any_with_descendent_if(
1148 &"B",
1149 2,
1150 &is_descendent_of,
1151 |c| c.effective <= 2,
1152 ),
1153 Ok(None),
1154 );
1155
1156 assert_eq!(
1158 tree.finalize_with_descendent_if(&"D", 10, &is_descendent_of, |c| c.effective <= 10),
1159 Err(Error::UnfinalizedAncestor)
1160 );
1161
1162 assert_eq!(
1165 tree.finalizes_any_with_descendent_if(&"D", 10, &is_descendent_of, |c| c.effective ==
1166 10),
1167 Ok(Some(false)),
1168 );
1169
1170 assert_eq!(
1172 tree.finalizes_any_with_descendent_if(&"E", 15, &is_descendent_of, |c| c.effective ==
1173 10),
1174 Err(Error::UnfinalizedAncestor)
1175 );
1176
1177 assert_eq!(
1180 tree.finalize_with_descendent_if(&"B", 2, &is_descendent_of, |c| c.effective <= 2),
1181 Ok(FinalizationResult::Changed(None)),
1182 );
1183
1184 assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("A0", 1)],);
1185
1186 assert_eq!(
1188 tree.finalizes_any_with_descendent_if(
1189 &"C",
1190 5,
1191 &is_descendent_of,
1192 |c| c.effective <= 5,
1193 ),
1194 Ok(Some(true)),
1195 );
1196
1197 assert_eq!(
1198 tree.finalize_with_descendent_if(&"C", 5, &is_descendent_of, |c| c.effective <= 5),
1199 Ok(FinalizationResult::Changed(Some(Change { effective: 5 }))),
1200 );
1201
1202 assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("D", 10)],);
1203
1204 assert_eq!(
1206 tree.finalizes_any_with_descendent_if(&"F", 100, &is_descendent_of, |c| c.effective <=
1207 100,),
1208 Err(Error::UnfinalizedAncestor),
1209 );
1210
1211 assert_eq!(
1213 tree.finalizes_any_with_descendent_if(&"G", 100, &is_descendent_of, |c| c.effective <=
1214 100),
1215 Ok(Some(true)),
1216 );
1217
1218 assert_eq!(
1219 tree.finalize_with_descendent_if(&"G", 100, &is_descendent_of, |c| c.effective <= 100),
1220 Ok(FinalizationResult::Changed(Some(Change { effective: 10 }))),
1221 );
1222
1223 assert_eq!(tree.roots().count(), 0);
1225 }
1226
1227 #[test]
1228 fn iter_iterates_in_preorder() {
1229 let (tree, ..) = test_fork_tree();
1230 assert_eq!(
1231 tree.iter().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(),
1232 vec![
1233 ("A", 10),
1234 ("B", 20),
1235 ("C", 30),
1236 ("D", 40),
1237 ("E", 50),
1238 ("F", 20),
1239 ("H", 30),
1240 ("L", 40),
1241 ("M", 50),
1242 ("O", 50),
1243 ("I", 40),
1244 ("G", 30),
1245 ("J", 20),
1246 ("K", 30),
1247 ],
1248 );
1249 }
1250
1251 #[test]
1252 fn minimizes_calls_to_is_descendent_of() {
1253 use std::sync::atomic::{AtomicUsize, Ordering};
1254
1255 let n_is_descendent_of_calls = AtomicUsize::new(0);
1256
1257 let is_descendent_of = |_: &&str, _: &&str| -> Result<bool, TestError> {
1258 n_is_descendent_of_calls.fetch_add(1, Ordering::SeqCst);
1259 Ok(true)
1260 };
1261
1262 {
1263 let mut tree = ForkTree::new();
1267 let letters = vec!["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K"];
1268
1269 for (i, letter) in letters.iter().enumerate() {
1270 tree.import::<_, TestError>(*letter, i, i, &|_, _| Ok(true)).unwrap();
1271 }
1272
1273 assert_eq!(
1276 tree.finalizes_any_with_descendent_if(&"L", 11, &is_descendent_of, |i| *i == 10,),
1277 Ok(Some(false)),
1278 );
1279
1280 assert_eq!(n_is_descendent_of_calls.load(Ordering::SeqCst), 1);
1281 }
1282
1283 n_is_descendent_of_calls.store(0, Ordering::SeqCst);
1284
1285 {
1286 let mut tree = ForkTree::new();
1290 let letters = vec!["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K"];
1291
1292 for (i, letter) in letters.iter().enumerate() {
1293 tree.import::<_, TestError>(*letter, i, i, &|_, _| Ok(false)).unwrap();
1294 }
1295
1296 assert_eq!(
1299 tree.finalize_with_descendent_if(&"L", 11, &is_descendent_of, |i| *i == 10,),
1300 Ok(FinalizationResult::Changed(Some(10))),
1301 );
1302
1303 assert_eq!(n_is_descendent_of_calls.load(Ordering::SeqCst), 1);
1304 }
1305 }
1306
1307 #[test]
1308 fn map_works() {
1309 let (mut tree, _) = test_fork_tree();
1310
1311 let is_descendent_of = |_: &&str, _: &&str| -> Result<bool, TestError> { Ok(false) };
1313 let is_root = tree.import("A1", 10, 1, &is_descendent_of).unwrap();
1314 assert!(is_root);
1315 let is_root = tree.import("A2", 10, 1, &is_descendent_of).unwrap();
1316 assert!(is_root);
1317
1318 let old_tree = tree.clone();
1319 let new_tree = tree.map(&mut |hash, _, _| hash.to_owned());
1320
1321 assert!(new_tree.iter().all(|(hash, _, data)| hash == data));
1323 assert_eq!(
1324 old_tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(),
1325 new_tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(),
1326 );
1327 }
1328
1329 #[test]
1330 fn prune_works_for_in_tree_hashes() {
1331 let (mut tree, is_descendent_of) = test_fork_tree();
1332
1333 let removed = tree.prune(&"C", &30, &is_descendent_of, &|_| true).unwrap();
1334
1335 assert_eq!(tree.roots.iter().map(|node| node.hash).collect::<Vec<_>>(), vec!["B"]);
1336
1337 assert_eq!(
1338 tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(),
1339 vec!["B", "C", "D", "E"],
1340 );
1341
1342 assert_eq!(
1343 removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(),
1344 vec!["A", "F", "H", "L", "M", "O", "I", "G", "J", "K"]
1345 );
1346
1347 let removed = tree.prune(&"E", &50, &is_descendent_of, &|_| true).unwrap();
1348
1349 assert_eq!(tree.roots.iter().map(|node| node.hash).collect::<Vec<_>>(), vec!["D"]);
1350
1351 assert_eq!(tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(), vec!["D", "E"]);
1352
1353 assert_eq!(removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(), vec!["B", "C"]);
1354 }
1355
1356 #[test]
1357 fn prune_works_for_out_of_tree_hashes() {
1358 let (mut tree, is_descendent_of) = test_fork_tree();
1359
1360 let removed = tree.prune(&"c", &25, &is_descendent_of, &|_| true).unwrap();
1361
1362 assert_eq!(tree.roots.iter().map(|node| node.hash).collect::<Vec<_>>(), vec!["B"]);
1363
1364 assert_eq!(
1365 tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(),
1366 vec!["B", "C", "D", "E"],
1367 );
1368
1369 assert_eq!(
1370 removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(),
1371 vec!["A", "F", "H", "L", "M", "O", "I", "G", "J", "K"]
1372 );
1373 }
1374
1375 #[test]
1376 fn prune_works_for_not_direct_ancestor() {
1377 let (mut tree, is_descendent_of) = test_fork_tree();
1378
1379 let removed = tree.prune(&"m", &45, &is_descendent_of, &|height| *height == 3).unwrap();
1381
1382 assert_eq!(tree.roots.iter().map(|node| node.hash).collect::<Vec<_>>(), vec!["H"]);
1383
1384 assert_eq!(tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(), vec!["H", "L", "M"],);
1385
1386 assert_eq!(
1387 removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(),
1388 vec!["O", "I", "A", "B", "C", "D", "E", "F", "G", "J", "K"]
1389 );
1390 }
1391
1392 #[test]
1393 fn prune_works_for_far_away_ancestor() {
1394 let (mut tree, is_descendent_of) = test_fork_tree();
1395
1396 let removed = tree.prune(&"m", &45, &is_descendent_of, &|height| *height == 2).unwrap();
1397
1398 assert_eq!(tree.roots.iter().map(|node| node.hash).collect::<Vec<_>>(), vec!["F"]);
1399
1400 assert_eq!(
1401 tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(),
1402 vec!["F", "H", "L", "M"],
1403 );
1404
1405 assert_eq!(
1406 removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(),
1407 vec!["O", "I", "G", "A", "B", "C", "D", "E", "J", "K"]
1408 );
1409 }
1410
1411 #[test]
1412 fn find_node_backtracks_after_finding_highest_descending_node() {
1413 let mut tree = ForkTree::new();
1414
1415 let is_descendent_of = |base: &&str, block: &&str| -> Result<bool, TestError> {
1420 match (*base, *block) {
1421 ("A", b) => Ok(b == "B" || b == "C" || b == "D"),
1422 ("B", b) | ("C", b) => Ok(b == "D"),
1423 ("0", _) => Ok(true),
1424 _ => Ok(false),
1425 }
1426 };
1427
1428 tree.import("A", 1, 1, &is_descendent_of).unwrap();
1429 tree.import("B", 2, 2, &is_descendent_of).unwrap();
1430 tree.import("C", 2, 4, &is_descendent_of).unwrap();
1431
1432 let node = tree.find_node_where(&"D", &3, &is_descendent_of, &|data| *data < 3).unwrap();
1436
1437 assert_eq!(node.unwrap().hash, "B");
1438 }
1439
1440 #[test]
1441 fn rebalance_works() {
1442 let (mut tree, _) = test_fork_tree();
1443
1444 assert_eq!(
1448 tree.iter().map(|(h, _, _)| *h).collect::<Vec<_>>(),
1449 vec!["A", "B", "C", "D", "E", "F", "H", "L", "M", "O", "I", "G", "J", "K"],
1450 );
1451
1452 let is_descendent_of = |base: &&str, block: &&str| -> Result<bool, TestError> {
1454 match (*base, *block) {
1455 (b, "P") => Ok(vec!["A", "F", "H", "L", "O"].into_iter().any(|n| n == b)),
1456 _ => Ok(false),
1457 }
1458 };
1459
1460 tree.import("P", 60, 6, &is_descendent_of).unwrap();
1461
1462 assert_eq!(
1466 tree.iter().map(|(h, _, _)| *h).collect::<Vec<_>>(),
1467 ["A", "F", "H", "L", "O", "P", "M", "I", "G", "B", "C", "D", "E", "J", "K"]
1468 );
1469 }
1470
1471 #[test]
1472 fn drain_filter_works() {
1473 let (mut tree, _) = test_fork_tree();
1474
1475 let filter = |h: &&str, _: &u64, _: &u32| match *h {
1476 "A" | "B" | "F" | "G" => FilterAction::KeepNode,
1477 "C" => FilterAction::KeepTree,
1478 "H" | "J" => FilterAction::Remove,
1479 _ => panic!("Unexpected filtering for node: {}", *h),
1480 };
1481
1482 let removed = tree.drain_filter(filter);
1483
1484 assert_eq!(
1485 tree.iter().map(|(h, _, _)| *h).collect::<Vec<_>>(),
1486 ["A", "B", "C", "D", "E", "F", "G"]
1487 );
1488
1489 assert_eq!(
1490 removed.map(|(h, _, _)| h).collect::<Vec<_>>(),
1491 ["H", "L", "M", "O", "I", "J", "K"]
1492 );
1493 }
1494
1495 #[test]
1496 fn find_node_index_works() {
1497 let (tree, is_descendent_of) = test_fork_tree();
1498
1499 let path = tree
1500 .find_node_index_where(&"D", &40, &is_descendent_of, &|_| true)
1501 .unwrap()
1502 .unwrap();
1503 assert_eq!(path, [0, 0, 0]);
1504
1505 let path = tree
1506 .find_node_index_where(&"O", &50, &is_descendent_of, &|_| true)
1507 .unwrap()
1508 .unwrap();
1509 assert_eq!(path, [0, 1, 0, 0]);
1510
1511 let path = tree
1512 .find_node_index_where(&"N", &60, &is_descendent_of, &|_| true)
1513 .unwrap()
1514 .unwrap();
1515 assert_eq!(path, [0, 1, 0, 0, 0]);
1516 }
1517
1518 #[test]
1519 fn find_node_index_with_predicate_works() {
1520 let is_descendent_of = |parent: &char, child: &char| match *parent {
1521 'A' => Ok(['B', 'C', 'D', 'E', 'F'].contains(child)),
1522 'B' => Ok(['C', 'D'].contains(child)),
1523 'C' => Ok(['D'].contains(child)),
1524 'E' => Ok(['F'].contains(child)),
1525 'D' | 'F' => Ok(false),
1526 _ => Err(TestError),
1527 };
1528
1529 let mut tree: ForkTree<char, u8, bool> = ForkTree::new();
1532 tree.import('A', 1, true, &is_descendent_of).unwrap();
1533 tree.import('B', 2, false, &is_descendent_of).unwrap();
1534 tree.import('C', 3, true, &is_descendent_of).unwrap();
1535 tree.import('D', 4, false, &is_descendent_of).unwrap();
1536
1537 tree.import('E', 2, true, &is_descendent_of).unwrap();
1538 tree.import('F', 3, false, &is_descendent_of).unwrap();
1539
1540 let path = tree
1541 .find_node_index_where(&'D', &4, &is_descendent_of, &|&value| !value)
1542 .unwrap()
1543 .unwrap();
1544 assert_eq!(path, [0, 0]);
1545
1546 let path = tree
1547 .find_node_index_where(&'D', &4, &is_descendent_of, &|&value| value)
1548 .unwrap()
1549 .unwrap();
1550 assert_eq!(path, [0, 0, 0]);
1551
1552 let path = tree
1553 .find_node_index_where(&'F', &3, &is_descendent_of, &|&value| !value)
1554 .unwrap();
1555 assert_eq!(path, None);
1556
1557 let path = tree
1558 .find_node_index_where(&'F', &3, &is_descendent_of, &|&value| value)
1559 .unwrap()
1560 .unwrap();
1561 assert_eq!(path, [0, 1]);
1562 }
1563
1564 #[test]
1565 fn find_node_works() {
1566 let (tree, is_descendent_of) = test_fork_tree();
1567
1568 let node = tree.find_node_where(&"B", &20, &is_descendent_of, &|_| true).unwrap().unwrap();
1569 assert_eq!((node.hash, node.number), ("A", 10));
1570
1571 let node = tree.find_node_where(&"D", &40, &is_descendent_of, &|_| true).unwrap().unwrap();
1572 assert_eq!((node.hash, node.number), ("C", 30));
1573
1574 let node = tree.find_node_where(&"O", &50, &is_descendent_of, &|_| true).unwrap().unwrap();
1575 assert_eq!((node.hash, node.number), ("L", 40));
1576
1577 let node = tree.find_node_where(&"N", &60, &is_descendent_of, &|_| true).unwrap().unwrap();
1578 assert_eq!((node.hash, node.number), ("M", 50));
1579 }
1580
1581 #[test]
1582 fn post_order_traversal_requirement() {
1583 let (mut tree, is_descendent_of) = test_fork_tree();
1584
1585 let is_descendent_of_for_post_order = |parent: &&str, child: &&str| match *parent {
1588 "A" => Err(TestError),
1589 "K" if *child == "Z" => Ok(true),
1590 _ => is_descendent_of(parent, child),
1591 };
1592
1593 let path = tree
1595 .find_node_index_where(&"N", &60, &is_descendent_of_for_post_order, &|_| true)
1596 .unwrap()
1597 .unwrap();
1598 assert_eq!(path, [0, 1, 0, 0, 0]);
1599
1600 let res = tree.import(&"Z", 100, 10, &is_descendent_of_for_post_order);
1602 assert_eq!(res, Ok(false));
1603 assert_eq!(
1604 tree.iter().map(|node| *node.0).collect::<Vec<_>>(),
1605 vec!["A", "B", "C", "D", "E", "F", "H", "L", "M", "O", "I", "G", "J", "K", "Z"],
1606 );
1607 }
1608}