referrerpolicy=no-referrer-when-downgrade

fork_tree/
lib.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! Utility library for managing tree-like ordered data with logic for pruning
19//! the tree while finalizing nodes.
20
21#![warn(missing_docs)]
22
23use codec::{Decode, Encode};
24use std::{cmp::Reverse, fmt};
25
26/// Error occurred when iterating with the tree.
27#[derive(Clone, Debug, PartialEq)]
28pub enum Error<E> {
29	/// Adding duplicate node to tree.
30	Duplicate,
31	/// Finalizing descendent of tree node without finalizing ancestor(s).
32	UnfinalizedAncestor,
33	/// Imported or finalized node that is an ancestor of previously finalized node.
34	Revert,
35	/// Error thrown by user when checking for node ancestry.
36	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/// Result of finalizing a node (that could be a part of the tree or not).
60#[derive(Debug, PartialEq)]
61pub enum FinalizationResult<V> {
62	/// The tree has changed, optionally return the value associated with the finalized node.
63	Changed(Option<V>),
64	/// The tree has not changed.
65	Unchanged,
66}
67
68/// Filtering action.
69#[derive(Debug, PartialEq)]
70pub enum FilterAction {
71	/// Remove the node and its subtree.
72	Remove,
73	/// Maintain the node.
74	KeepNode,
75	/// Maintain the node and its subtree.
76	KeepTree,
77}
78
79/// A tree data structure that stores several nodes across multiple branches.
80///
81/// Top-level branches are called roots. The tree has functionality for
82/// finalizing nodes, which means that node is traversed, and all competing
83/// branches are pruned. It also guarantees that nodes in the tree are finalized
84/// in order. Each node is uniquely identified by its hash but can be ordered by
85/// its number. In order to build the tree an external function must be provided
86/// when interacting with the tree to establish a node's ancestry.
87#[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	/// Create a new empty tree instance.
99	pub fn new() -> ForkTree<H, N, V> {
100		ForkTree { roots: Vec::new(), best_finalized_number: None }
101	}
102
103	/// Rebalance the tree.
104	///
105	/// For each tree level sort child nodes by max branch depth (decreasing).
106	///
107	/// Most operations in the tree are performed with depth-first search
108	/// starting from the leftmost node at every level, since this tree is meant
109	/// to be used in a blockchain context, a good heuristic is that the node
110	/// we'll be looking for at any point will likely be in one of the deepest chains
111	/// (i.e. the longest ones).
112	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	/// Import a new node into the tree.
122	///
123	/// The given function `is_descendent_of` should return `true` if the second
124	/// hash (target) is a descendent of the first hash (base).
125	///
126	/// This method assumes that nodes in the same branch are imported in order.
127	///
128	/// Returns `true` if the imported node is a root.
129	// WARNING: some users of this method (i.e. consensus epoch changes tree) currently silently
130	// rely on a **post-order DFS** traversal. If we are using instead a top-down traversal method
131	// then the `is_descendent_of` closure, when used after a warp-sync, may end up querying the
132	// backend for a block (the one corresponding to the root) that is not present and thus will
133	// return a wrong result.
134	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			// Rebalance may be required only if we've extended the branch depth.
165			self.rebalance();
166		}
167
168		Ok(is_root)
169	}
170
171	/// Iterates over the existing roots in the tree.
172	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		// we need to reverse the order of roots to maintain the expected
178		// ordering since the iterator uses a stack to track state.
179		ForkTreeIterator { stack: self.roots.iter().rev().collect() }
180	}
181
182	/// Iterates the nodes in the tree in pre-order.
183	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	/// Map fork tree into values of new types.
188	///
189	/// Tree traversal technique (e.g. BFS vs DFS) is left as not specified and
190	/// may be subject to change in the future. In other words, your predicates
191	/// should not rely on the observed traversal technique currently in use.
192	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	/// Find a node in the tree that is the deepest ancestor of the given
235	/// block hash and which passes the given predicate.
236	///
237	/// The given function `is_descendent_of` should return `true` if the
238	/// second hash (target) is a descendent of the first hash (base).
239	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	/// Same as [`find_node_where`](ForkTree::find_node_where), but returns mutable reference.
260	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	/// Same as [`find_node_where`](ForkTree::find_node_where), but returns indices.
283	///
284	/// The returned indices represent the full path to reach the matching node starting
285	/// from one of the roots, i.e. the earliest index in the traverse path goes first,
286	/// and the final index in the traverse path goes last.
287	///
288	/// If a node is found that matches the predicate the returned path should always
289	/// contain at least one index, otherwise `None` is returned.
290	// WARNING: some users of this method (i.e. consensus epoch changes tree) currently silently
291	// rely on a **post-order DFS** traversal. If we are using instead a top-down traversal method
292	// then the `is_descendent_of` closure, when used after a warp-sync, will end up querying the
293	// backend for a block (the one corresponding to the root) that is not present and thus will
294	// return a wrong result.
295	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			// The second element in the stack tuple tracks what is the **next** children
318			// index to search into. If we find an ancestor then we stop searching into
319			// alternative branches and we focus on the current path up to the root.
320			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 the element we are looking for is a descendent of the current root
337			// then we can stop the search.
338			if is_descendent {
339				break;
340			}
341			root_idx += 1;
342		}
343
344		Ok(if found {
345			// The path is the root index followed by the indices of all the children
346			// we were processing when we found the element (remember the stack
347			// contains the index of the **next** children to process).
348			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	/// Prune the tree, removing all non-canonical nodes.
357	///
358	/// We find the node in the tree that is the deepest ancestor of the given hash
359	/// and that passes the given predicate. If such a node exists, we re-root the
360	/// tree to this node. Otherwise the tree remains unchanged.
361	///
362	/// The given function `is_descendent_of` should return `true` if the second
363	/// hash (target) is a descendent of the first hash (base).
364	///
365	/// Returns all pruned nodes data.
366	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		// Find and detach the new root from the removed nodes
387		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		// If, because of the `predicate`, the new root is not the deepest ancestor
395		// of `hash` then we can remove all the nodes that are descendants of the new
396		// `root` but not ancestors of `hash`.
397		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				// Now we are positioned just above block identified by `hash`
408				break;
409			};
410			// Preserve only the ancestor node, the siblings are removed
411			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		// Curr now points to our direct ancestor, if necessary remove any node that is
419		// not a descendant of `hash`.
420		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	/// Finalize a root in the tree and return it, return `None` in case no root
437	/// with the given hash exists. All other roots are pruned, and the children
438	/// of the finalized node become the new roots.
439	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	/// Finalize root at given position. See `finalize_root` comment for details.
447	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	/// Finalize a node in the tree. This method will make sure that the node
455	/// being finalized is either an existing root (and return its data), or a
456	/// node from a competing branch (not in the tree), tree pruning is done
457	/// accordingly. The given function `is_descendent_of` should return `true`
458	/// if the second hash (target) is a descendent of the first hash (base).
459	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		// check if one of the current roots is being finalized
476		if let Some(root) = self.finalize_root(hash) {
477			return Ok(FinalizationResult::Changed(Some(root)));
478		}
479
480		// make sure we're not finalizing a descendent of any root
481		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		// we finalized a block earlier than any existing root (or possibly
488		// another fork not part of the tree). make sure to only keep roots that
489		// are part of the finalized branch
490		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	/// Finalize a node in the tree and all its ancestors. The given function
511	/// `is_descendent_of` should return `true` if the second hash (target) is
512	// a descendent of the first hash (base).
513	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		// check if one of the current roots is being finalized
530		if let Some(root) = self.finalize_root(hash) {
531			return Ok(FinalizationResult::Changed(Some(root)));
532		}
533
534		// we need to:
535		// 1) remove all roots that are not ancestors AND not descendants of finalized block;
536		// 2) if node is descendant - just leave it;
537		// 3) if node is ancestor - 'open it'
538		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 we have met finalized root - open it and return
553			if is_finalized {
554				return Ok(FinalizationResult::Changed(Some(self.finalize_root_at(idx))));
555			}
556
557			// if node is descendant of finalized block - just leave it as is
558			if is_descendant {
559				idx += 1;
560				continue;
561			}
562
563			// if node is ancestor of finalized block - remove it and continue with children
564			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			// if node is neither ancestor, nor descendant of the finalized block - remove it
572			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	/// Checks if any node in the tree is finalized by either finalizing the
586	/// node itself or a node's descendent that's not in the tree, guaranteeing
587	/// that the node being finalized isn't a descendent of (or equal to) any of
588	/// the node's children. Returns `Some(true)` if the node being finalized is
589	/// a root, `Some(false)` if the node being finalized is not a root, and
590	/// `None` if no node in the tree is finalized. The given `predicate` is
591	/// checked on the prospective finalized root and must pass for finalization
592	/// to occur. The given function `is_descendent_of` should return `true` if
593	/// the second hash (target) is a descendent of the first hash (base).
594	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		// check if the given hash is equal or a descendent of any node in the
613		// tree, if we find a valid node that passes the predicate then we must
614		// ensure that we're not finalizing past any of its child nodes.
615		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	/// Finalize a root in the tree by either finalizing the node itself or a
634	/// node's descendent that's not in the tree, guaranteeing that the node
635	/// being finalized isn't a descendent of (or equal to) any of the root's
636	/// children. The given `predicate` is checked on the prospective finalized
637	/// root and must pass for finalization to occur. The given function
638	/// `is_descendent_of` should return `true` if the second hash (target) is a
639	/// descendent of the first hash (base).
640	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		// check if the given hash is equal or a a descendent of any root, if we
659		// find a valid root that passes the predicate then we must ensure that
660		// we're not finalizing past any children node.
661		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		// Retain only roots that are descendants of the finalized block (this
686		// happens if the node has been properly finalized) or that are
687		// ancestors (or equal) to the finalized block (in this case the node
688		// wasn't finalized earlier presumably because the predicate didn't
689		// pass).
690		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	/// Remove from the tree some nodes (and their subtrees) using a `filter` predicate.
715	///
716	/// The `filter` is called over tree nodes and returns a filter action:
717	/// - `Remove` if the node and its subtree should be removed;
718	/// - `KeepNode` if we should maintain the node and keep processing the tree.
719	/// - `KeepTree` if we should maintain the node and its entire subtree.
720	///
721	/// An iterator over all the pruned nodes is returned.
722	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
774// Workaround for: https://github.com/rust-lang/rust/issues/34537
775use 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		/// Finds the max depth among all branches descendent from this node.
790		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			// child nodes are stored ordered by max branch height (decreasing),
814			// we want to keep this ordering while iterating but since we're
815			// using a stack for iterator state we need to reverse it.
816			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			// child nodes are stored ordered by max branch height (decreasing),
831			// we want to keep this ordering while iterating but since we're
832			// using a stack for iterator state we need to reverse it.
833			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		//
864		//     +---B-c-C---D---E
865		//     |
866		//     |   +---G
867		//     |   | 
868		// 0---A---F---H---I
869		//     |       |
870		//     |       +---L-m-M---N
871		//     |           |
872		//     |           +---O
873		//     +---J---K
874		//
875		// (where N is not a part of fork tree)
876		//
877		// NOTE: the tree will get automatically rebalance on import and won't be laid out like the
878		// diagram above. the children will be ordered by subtree depth and the longest branches
879		// will be on the leftmost side of the tree.
880		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			// This is a trick to have lowercase blocks be direct parents of their
883			// uppercase correspondent (A excluded)
884			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			// finalizing "A" opens up three possible forks
957			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			// finalizing "B" will progress on its fork and remove any other competing forks
971			tree.finalize_root(&"B");
972
973			assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("C", 30)],);
974
975			// all the other forks have been pruned
976			assert!(tree.roots.len() == 1);
977		}
978
979		{
980			let mut tree = finalize_a();
981
982			// finalizing "J" will progress on its fork and remove any other competing forks
983			tree.finalize_root(&"J");
984
985			assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("K", 30)],);
986
987			// all the other forks have been pruned
988			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		// finalizing a block prior to any in the node doesn't change the tree
999		assert_eq!(tree.finalize(&"0", 0, &is_descendent_of), Ok(FinalizationResult::Unchanged));
1000
1001		assert_eq!(tree.roots, original_roots);
1002
1003		// finalizing "A" opens up three possible forks
1004		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		// finalizing anything lower than what we observed will fail
1015		assert_eq!(tree.best_finalized_number, Some(10));
1016
1017		assert_eq!(tree.finalize(&"Z", 10, &is_descendent_of), Err(Error::Revert));
1018
1019		// trying to finalize a node without finalizing its ancestors first will fail
1020		assert_eq!(tree.finalize(&"H", 30, &is_descendent_of), Err(Error::UnfinalizedAncestor));
1021
1022		// after finalizing "F" we can finalize "H"
1023		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		// finalizing a node from another fork that isn't part of the tree clears the tree
1039		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		// finalizing a block prior to any in the node doesn't change the tree
1054		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		// finalizing "A" opens up three possible forks
1062		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		// finalizing H:
1073		// 1) removes roots that are not ancestors/descendants of H (B, J)
1074		// 2) opens root that is ancestor of H (F -> G+H)
1075		// 3) finalizes the just opened root H (H -> I + L)
1076		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		// finalizing N (which is not a part of the tree):
1089		// 1) removes roots that are not ancestors/descendants of N (I)
1090		// 2) opens root that is ancestor of N (L -> M+O)
1091		// 3) removes roots that are not ancestors/descendants of N (O)
1092		// 4) opens root that is ancestor of N (M -> {})
1093		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				// A0 #1 - (B #2) - (C #5) - D #10 - E #15 - (F #100)
1115				//                            \
1116				//                             - (G #100)
1117				//
1118				// A1 #1
1119				//
1120				// Nodes B, C, F and G  are not part of the tree.
1121				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		// finalizing "D" is not allowed since it is not a root.
1156		assert_eq!(
1157			tree.finalize_with_descendent_if(&"D", 10, &is_descendent_of, |c| c.effective <= 10),
1158			Err(Error::UnfinalizedAncestor)
1159		);
1160
1161		// finalizing "D" will finalize a block from the tree, but it can't be applied yet
1162		// since it is not a root change.
1163		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		// finalizing "E" is not allowed since there are not finalized ancestors.
1170		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		// finalizing "B" doesn't finalize "A0" since the predicate doesn't pass,
1177		// although it will clear out "A1" from the tree
1178		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		// finalizing "C" will finalize the node "A0" and prune it out of the tree
1186		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		// finalizing "F" will fail since it would finalize past "E" without finalizing "D" first
1204		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		// it will work with "G" though since it is not in the same branch as "E"
1211		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		// "E" will be pruned out
1223		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			// Deep tree where we want to call `finalizes_any_with_descendent_if`. The
1263			// search for the node should first check the predicate (which is cheaper) and
1264			// only then call `is_descendent_of`
1265			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			// "L" is a descendent of "K", but the predicate will only pass for "K",
1273			// therefore only one call to `is_descendent_of` should be made
1274			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			// Multiple roots in the tree where we want to call `finalize_with_descendent_if`.
1286			// The search for the root node should first check the predicate (which is cheaper)
1287			// and only then call `is_descendent_of`
1288			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			// "L" is a descendent of "K", but the predicate will only pass for "K",
1296			// therefore only one call to `is_descendent_of` should be made
1297			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		// Extend the single root fork-tree to also exercise the roots order during map.
1311		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		// Check content and order
1321		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		// This is to re-root the tree not at the immediate ancestor, but the one just before.
1379		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		// A - B
1415		//  \
1416		//   — C
1417		//
1418		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		// when searching the tree we reach node `C`, but the
1432		// predicate doesn't pass. we should backtrack to `B`, but not to `A`,
1433		// since "B" fulfills the predicate.
1434		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		// the tree is automatically rebalanced on import, therefore we should iterate in preorder
1444		// exploring the longest forks first. check the ascii art above to understand the expected
1445		// output below.
1446		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's add a block "P" which is a descendent of block "O"
1452		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		// this should re-order the tree, since the branch "A -> B -> C -> D -> E" is no longer tied
1462		// with 5 blocks depth. additionally "O" should be visited before "M" now, since it has one
1463		// descendent "P" which makes that branch 6 blocks long.
1464		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		// A(t) --- B(f) --- C(t) --- D(f)
1529		//      \-- E(t) --- F(f)
1530		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		// Test for the post-order DFS traversal requirement as specified by the
1585		// `find_node_index_where` and `import` comments.
1586		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		// Post order traversal requirement for `find_node_index_where`
1593		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		// Post order traversal requirement for `import`
1600		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}