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}