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	//
291	// WARNING: some users of this method (i.e. consensus epoch changes tree) currently silently
292	// rely on a **post-order DFS** traversal. If we are using instead a top-down traversal method
293	// then the `is_descendent_of` closure, when used after a warp-sync, will end up querying the
294	// backend for a block (the one corresponding to the root) that is not present and thus will
295	// return a wrong result.
296	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			// The second element in the stack tuple tracks what is the **next** children
319			// index to search into. If we find an ancestor then we stop searching into
320			// alternative branches and we focus on the current path up to the root.
321			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 the element we are looking for is a descendent of the current root
338			// then we can stop the search.
339			if is_descendent {
340				break
341			}
342			root_idx += 1;
343		}
344
345		Ok(if found {
346			// The path is the root index followed by the indices of all the children
347			// we were processing when we found the element (remember the stack
348			// contains the index of the **next** children to process).
349			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	/// Prune the tree, removing all non-canonical nodes.
358	///
359	/// We find the node in the tree that is the deepest ancestor of the given hash
360	/// and that passes the given predicate. If such a node exists, we re-root the
361	/// tree to this node. Otherwise the tree remains unchanged.
362	///
363	/// The given function `is_descendent_of` should return `true` if the second
364	/// hash (target) is a descendent of the first hash (base).
365	///
366	/// Returns all pruned nodes data.
367	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		// Find and detach the new root from the removed nodes
388		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		// If, because of the `predicate`, the new root is not the deepest ancestor
396		// of `hash` then we can remove all the nodes that are descendants of the new
397		// `root` but not ancestors of `hash`.
398		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				// Now we are positioned just above block identified by `hash`
409				break
410			};
411			// Preserve only the ancestor node, the siblings are removed
412			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		// Curr now points to our direct ancestor, if necessary remove any node that is
420		// not a descendant of `hash`.
421		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	/// Finalize a root in the tree and return it, return `None` in case no root
438	/// with the given hash exists. All other roots are pruned, and the children
439	/// of the finalized node become the new roots.
440	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	/// Finalize root at given position. See `finalize_root` comment for details.
448	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	/// Finalize a node in the tree. This method will make sure that the node
456	/// being finalized is either an existing root (and return its data), or a
457	/// node from a competing branch (not in the tree), tree pruning is done
458	/// accordingly. The given function `is_descendent_of` should return `true`
459	/// if the second hash (target) is a descendent of the first hash (base).
460	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		// check if one of the current roots is being finalized
477		if let Some(root) = self.finalize_root(hash) {
478			return Ok(FinalizationResult::Changed(Some(root)))
479		}
480
481		// make sure we're not finalizing a descendent of any root
482		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		// we finalized a block earlier than any existing root (or possibly
489		// another fork not part of the tree). make sure to only keep roots that
490		// are part of the finalized branch
491		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	/// Finalize a node in the tree and all its ancestors. The given function
512	/// `is_descendent_of` should return `true` if the second hash (target) is
513	// a descendent of the first hash (base).
514	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		// check if one of the current roots is being finalized
531		if let Some(root) = self.finalize_root(hash) {
532			return Ok(FinalizationResult::Changed(Some(root)))
533		}
534
535		// we need to:
536		// 1) remove all roots that are not ancestors AND not descendants of finalized block;
537		// 2) if node is descendant - just leave it;
538		// 3) if node is ancestor - 'open it'
539		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 we have met finalized root - open it and return
554			if is_finalized {
555				return Ok(FinalizationResult::Changed(Some(self.finalize_root_at(idx))))
556			}
557
558			// if node is descendant of finalized block - just leave it as is
559			if is_descendant {
560				idx += 1;
561				continue
562			}
563
564			// if node is ancestor of finalized block - remove it and continue with children
565			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			// if node is neither ancestor, nor descendant of the finalized block - remove it
573			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	/// Checks if any node in the tree is finalized by either finalizing the
587	/// node itself or a node's descendent that's not in the tree, guaranteeing
588	/// that the node being finalized isn't a descendent of (or equal to) any of
589	/// the node's children. Returns `Some(true)` if the node being finalized is
590	/// a root, `Some(false)` if the node being finalized is not a root, and
591	/// `None` if no node in the tree is finalized. The given `predicate` is
592	/// checked on the prospective finalized root and must pass for finalization
593	/// to occur. The given function `is_descendent_of` should return `true` if
594	/// the second hash (target) is a descendent of the first hash (base).
595	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		// check if the given hash is equal or a descendent of any node in the
614		// tree, if we find a valid node that passes the predicate then we must
615		// ensure that we're not finalizing past any of its child nodes.
616		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	/// Finalize a root in the tree by either finalizing the node itself or a
635	/// node's descendent that's not in the tree, guaranteeing that the node
636	/// being finalized isn't a descendent of (or equal to) any of the root's
637	/// children. The given `predicate` is checked on the prospective finalized
638	/// root and must pass for finalization to occur. The given function
639	/// `is_descendent_of` should return `true` if the second hash (target) is a
640	/// descendent of the first hash (base).
641	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		// check if the given hash is equal or a a descendent of any root, if we
660		// find a valid root that passes the predicate then we must ensure that
661		// we're not finalizing past any children node.
662		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		// Retain only roots that are descendants of the finalized block (this
687		// happens if the node has been properly finalized) or that are
688		// ancestors (or equal) to the finalized block (in this case the node
689		// wasn't finalized earlier presumably because the predicate didn't
690		// pass).
691		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	/// Remove from the tree some nodes (and their subtrees) using a `filter` predicate.
716	///
717	/// The `filter` is called over tree nodes and returns a filter action:
718	/// - `Remove` if the node and its subtree should be removed;
719	/// - `KeepNode` if we should maintain the node and keep processing the tree.
720	/// - `KeepTree` if we should maintain the node and its entire subtree.
721	///
722	/// An iterator over all the pruned nodes is returned.
723	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
775// Workaround for: https://github.com/rust-lang/rust/issues/34537
776use 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		/// Finds the max depth among all branches descendent from this node.
791		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			// child nodes are stored ordered by max branch height (decreasing),
815			// we want to keep this ordering while iterating but since we're
816			// using a stack for iterator state we need to reverse it.
817			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			// child nodes are stored ordered by max branch height (decreasing),
832			// we want to keep this ordering while iterating but since we're
833			// using a stack for iterator state we need to reverse it.
834			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		//
865		//     +---B-c-C---D---E
866		//     |
867		//     |   +---G
868		//     |   | 
869		// 0---A---F---H---I
870		//     |       |
871		//     |       +---L-m-M---N
872		//     |           |
873		//     |           +---O
874		//     +---J---K
875		//
876		// (where N is not a part of fork tree)
877		//
878		// NOTE: the tree will get automatically rebalance on import and won't be laid out like the
879		// diagram above. the children will be ordered by subtree depth and the longest branches
880		// will be on the leftmost side of the tree.
881		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			// This is a trick to have lowercase blocks be direct parents of their
884			// uppercase correspondent (A excluded)
885			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			// finalizing "A" opens up three possible forks
958			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			// finalizing "B" will progress on its fork and remove any other competing forks
972			tree.finalize_root(&"B");
973
974			assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("C", 30)],);
975
976			// all the other forks have been pruned
977			assert!(tree.roots.len() == 1);
978		}
979
980		{
981			let mut tree = finalize_a();
982
983			// finalizing "J" will progress on its fork and remove any other competing forks
984			tree.finalize_root(&"J");
985
986			assert_eq!(tree.roots().map(|(h, n, _)| (*h, *n)).collect::<Vec<_>>(), vec![("K", 30)],);
987
988			// all the other forks have been pruned
989			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		// finalizing a block prior to any in the node doesn't change the tree
1000		assert_eq!(tree.finalize(&"0", 0, &is_descendent_of), Ok(FinalizationResult::Unchanged));
1001
1002		assert_eq!(tree.roots, original_roots);
1003
1004		// finalizing "A" opens up three possible forks
1005		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		// finalizing anything lower than what we observed will fail
1016		assert_eq!(tree.best_finalized_number, Some(10));
1017
1018		assert_eq!(tree.finalize(&"Z", 10, &is_descendent_of), Err(Error::Revert));
1019
1020		// trying to finalize a node without finalizing its ancestors first will fail
1021		assert_eq!(tree.finalize(&"H", 30, &is_descendent_of), Err(Error::UnfinalizedAncestor));
1022
1023		// after finalizing "F" we can finalize "H"
1024		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		// finalizing a node from another fork that isn't part of the tree clears the tree
1040		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		// finalizing a block prior to any in the node doesn't change the tree
1055		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		// finalizing "A" opens up three possible forks
1063		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		// finalizing H:
1074		// 1) removes roots that are not ancestors/descendants of H (B, J)
1075		// 2) opens root that is ancestor of H (F -> G+H)
1076		// 3) finalizes the just opened root H (H -> I + L)
1077		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		// finalizing N (which is not a part of the tree):
1090		// 1) removes roots that are not ancestors/descendants of N (I)
1091		// 2) opens root that is ancestor of N (L -> M+O)
1092		// 3) removes roots that are not ancestors/descendants of N (O)
1093		// 4) opens root that is ancestor of N (M -> {})
1094		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				// A0 #1 - (B #2) - (C #5) - D #10 - E #15 - (F #100)
1116				//                            \
1117				//                             - (G #100)
1118				//
1119				// A1 #1
1120				//
1121				// Nodes B, C, F and G  are not part of the tree.
1122				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		// finalizing "D" is not allowed since it is not a root.
1157		assert_eq!(
1158			tree.finalize_with_descendent_if(&"D", 10, &is_descendent_of, |c| c.effective <= 10),
1159			Err(Error::UnfinalizedAncestor)
1160		);
1161
1162		// finalizing "D" will finalize a block from the tree, but it can't be applied yet
1163		// since it is not a root change.
1164		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		// finalizing "E" is not allowed since there are not finalized ancestors.
1171		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		// finalizing "B" doesn't finalize "A0" since the predicate doesn't pass,
1178		// although it will clear out "A1" from the tree
1179		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		// finalizing "C" will finalize the node "A0" and prune it out of the tree
1187		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		// finalizing "F" will fail since it would finalize past "E" without finalizing "D" first
1205		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		// it will work with "G" though since it is not in the same branch as "E"
1212		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		// "E" will be pruned out
1224		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			// Deep tree where we want to call `finalizes_any_with_descendent_if`. The
1264			// search for the node should first check the predicate (which is cheaper) and
1265			// only then call `is_descendent_of`
1266			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			// "L" is a descendent of "K", but the predicate will only pass for "K",
1274			// therefore only one call to `is_descendent_of` should be made
1275			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			// Multiple roots in the tree where we want to call `finalize_with_descendent_if`.
1287			// The search for the root node should first check the predicate (which is cheaper)
1288			// and only then call `is_descendent_of`
1289			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			// "L" is a descendent of "K", but the predicate will only pass for "K",
1297			// therefore only one call to `is_descendent_of` should be made
1298			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		// Extend the single root fork-tree to also exercise the roots order during map.
1312		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		// Check content and order
1322		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		// This is to re-root the tree not at the immediate ancestor, but the one just before.
1380		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		// A - B
1416		//  \
1417		//   — C
1418		//
1419		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		// when searching the tree we reach node `C`, but the
1433		// predicate doesn't pass. we should backtrack to `B`, but not to `A`,
1434		// since "B" fulfills the predicate.
1435		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		// the tree is automatically rebalanced on import, therefore we should iterate in preorder
1445		// exploring the longest forks first. check the ascii art above to understand the expected
1446		// output below.
1447		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's add a block "P" which is a descendent of block "O"
1453		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		// this should re-order the tree, since the branch "A -> B -> C -> D -> E" is no longer tied
1463		// with 5 blocks depth. additionally "O" should be visited before "M" now, since it has one
1464		// descendent "P" which makes that branch 6 blocks long.
1465		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		// A(t) --- B(f) --- C(t) --- D(f)
1530		//      \-- E(t) --- F(f)
1531		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		// Test for the post-order DFS traversal requirement as specified by the
1586		// `find_node_index_where` and `import` comments.
1587		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		// Post order traversal requirement for `find_node_index_where`
1594		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		// Post order traversal requirement for `import`
1601		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}