// Copyright 2019-2020 Parity Technologies (UK) Ltd.
// This file is part of Substrate.

// Substrate is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Substrate is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.

//! Utility library for managing tree-like ordered data with logic for pruning
//! the tree while finalizing nodes.

#![warn(missing_docs)]

use std::cmp::Reverse;
use std::fmt;
use codec::{Decode, Encode};

/// Error occurred when iterating with the tree.
#[derive(Clone, Debug, PartialEq)]
pub enum Error<E> {
	/// Adding duplicate node to tree.
	Duplicate,
	/// Finalizing descendent of tree node without finalizing ancestor(s).
	UnfinalizedAncestor,
	/// Imported or finalized node that is an ancestor of previously finalized node.
	Revert,
	/// Error throw by client when checking for node ancestry.
	Client(E),
}

impl<E: std::error::Error> fmt::Display for Error<E> {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		let message = match *self {
			Error::Duplicate => "Hash already exists in Tree".into(),
			Error::UnfinalizedAncestor => "Finalized descendent of Tree node without finalizing its ancestor(s) first".into(),
			Error::Revert => "Tried to import or finalize node that is an ancestor of a previously finalized node".into(),
			Error::Client(ref err) => format!("Client error: {}", err),
		};
		write!(f, "{}", message)
	}
}

impl<E: std::error::Error> std::error::Error for Error<E> {
	fn cause(&self) -> Option<&dyn std::error::Error> {
		None
	}
}

impl<E: std::error::Error> From<E> for Error<E> {
	fn from(err: E) -> Error<E> {
		Error::Client(err)
	}
}

/// Result of finalizing a node (that could be a part of the tree or not).
#[derive(Debug, PartialEq)]
pub enum FinalizationResult<V> {
	/// The tree has changed, optionally return the value associated with the finalized node.
	Changed(Option<V>),
	/// The tree has not changed.
	Unchanged,
}

/// A tree data structure that stores several nodes across multiple branches.
/// Top-level branches are called roots. The tree has functionality for
/// finalizing nodes, which means that that node is traversed, and all competing
/// branches are pruned. It also guarantees that nodes in the tree are finalized
/// in order. Each node is uniquely identified by its hash but can be ordered by
/// its number. In order to build the tree an external function must be provided
/// when interacting with the tree to establish a node's ancestry.
#[derive(Clone, Debug, Decode, Encode, PartialEq)]
pub struct ForkTree<H, N, V> {
	roots: Vec<Node<H, N, V>>,
	best_finalized_number: Option<N>,
}

impl<H, N, V> ForkTree<H, N, V> where
	H: PartialEq + Clone,
	N: Ord + Clone,
	V: Clone,
{
	/// Prune the tree, removing all non-canonical nodes. We find the node in the
	/// tree that is the deepest ancestor of the given hash and that passes the
	/// given predicate. If such a node exists, we re-root the tree to this
	/// node. Otherwise the tree remains unchanged. The given function
	/// `is_descendent_of` should return `true` if the second hash (target) is a
	/// descendent of the first hash (base).
	///
	/// Returns all pruned node data.
	pub fn prune<F, E, P>(
		&mut self,
		hash: &H,
		number: &N,
		is_descendent_of: &F,
		predicate: &P,
	) -> Result<impl Iterator<Item=(H, N, V)>, Error<E>>
		where E: std::error::Error,
			  F: Fn(&H, &H) -> Result<bool, E>,
			  P: Fn(&V) -> bool,
	{
		let new_root_index = self.find_node_index_where(
			hash,
			number,
			is_descendent_of,
			predicate,
		)?;

		let removed = if let Some(mut root_index) = new_root_index {
			let mut old_roots = std::mem::replace(&mut self.roots, Vec::new());

			let mut root = None;
			let mut cur_children = Some(&mut old_roots);

			while let Some(cur_index) = root_index.pop() {
				if let Some(children) = cur_children.take() {
					if root_index.is_empty() {
						root = Some(children.remove(cur_index));
					} else {
						cur_children = Some(&mut children[cur_index].children);
					}
				}
			}

			let mut root = root
				.expect("find_node_index_where will return array with at least one index; \
                         this results in at least one item in removed; qed");

			let mut removed = old_roots;

			// we found the deepest ancestor of the finalized block, so we prune
			// out any children that don't include the finalized block.
			let root_children = std::mem::replace(&mut root.children, Vec::new());
			let mut is_first = true;

			for child in root_children {
				if is_first &&
					(child.number == *number && child.hash == *hash ||
					 child.number < *number && is_descendent_of(&child.hash, hash).unwrap_or(false))
				{
					root.children.push(child);
					// assuming that the tree is well formed only one child should pass this requirement
					// due to ancestry restrictions (i.e. they must be different forks).
					is_first = false;
				} else {
					removed.push(child);
				}
			}

			self.roots = vec![root];

			removed
		} else {
			Vec::new()
		};

		self.rebalance();

		Ok(RemovedIterator { stack: removed })
	}
}

impl<H, N, V> ForkTree<H, N, V> where
	H: PartialEq,
	N: Ord,
{
	/// Create a new empty tree.
	pub fn new() -> ForkTree<H, N, V> {
		ForkTree {
			roots: Vec::new(),
			best_finalized_number: None,
		}
	}

	/// Rebalance the tree, i.e. sort child nodes by max branch depth
	/// (decreasing).
	///
	/// Most operations in the tree are performed with depth-first search
	/// starting from the leftmost node at every level, since this tree is meant
	/// to be used in a blockchain context, a good heuristic is that the node
	/// we'll be looking
	/// for at any point will likely be in one of the deepest chains (i.e. the
	/// longest ones).
	pub fn rebalance(&mut self) {
		self.roots.sort_by_key(|n| Reverse(n.max_depth()));
		for root in &mut self.roots {
			root.rebalance();
		}
	}

	/// Import a new node into the tree. The given function `is_descendent_of`
	/// should return `true` if the second hash (target) is a descendent of the
	/// first hash (base). This method assumes that nodes in the same branch are
	/// imported in order.
	///
	/// Returns `true` if the imported node is a root.
	pub fn import<F, E>(
		&mut self,
		mut hash: H,
		mut number: N,
		mut data: V,
		is_descendent_of: &F,
	) -> Result<bool, Error<E>>
		where E: std::error::Error,
			  F: Fn(&H, &H) -> Result<bool, E>,
	{
		if let Some(ref best_finalized_number) = self.best_finalized_number {
			if number <= *best_finalized_number {
				return Err(Error::Revert);
			}
		}

		for root in self.roots.iter_mut() {
			if root.hash == hash {
				return Err(Error::Duplicate);
			}

			match root.import(hash, number, data, is_descendent_of)? {
				Some((h, n, d)) => {
					hash = h;
					number = n;
					data = d;
				},
				None => return Ok(false),
			}
		}

		self.roots.push(Node {
			data,
			hash: hash,
			number: number,
			children: Vec::new(),
		});

		self.rebalance();

		Ok(true)
	}

	/// Iterates over the existing roots in the tree.
	pub fn roots(&self) -> impl Iterator<Item=(&H, &N, &V)> {
		self.roots.iter().map(|node| (&node.hash, &node.number, &node.data))
	}

	fn node_iter(&self) -> impl Iterator<Item=&Node<H, N, V>> {
		ForkTreeIterator { stack: self.roots.iter().collect() }
	}

	/// Iterates the nodes in the tree in pre-order.
	pub fn iter(&self) -> impl Iterator<Item=(&H, &N, &V)> {
		self.node_iter().map(|node| (&node.hash, &node.number, &node.data))
	}

	/// Find a node in the tree that is the deepest ancestor of the given
	/// block hash and which passes the given predicate. The given function
	/// `is_descendent_of` should return `true` if the second hash (target)
	/// is a descendent of the first hash (base).
	pub fn find_node_where<F, E, P>(
		&self,
		hash: &H,
		number: &N,
		is_descendent_of: &F,
		predicate: &P,
	) -> Result<Option<&Node<H, N, V>>, Error<E>> where
		E: std::error::Error,
		F: Fn(&H, &H) -> Result<bool, E>,
		P: Fn(&V) -> bool,
	{
		// search for node starting from all roots
		for root in self.roots.iter() {
			let node = root.find_node_where(hash, number, is_descendent_of, predicate)?;

			// found the node, early exit
			if let FindOutcome::Found(node) = node {
				return Ok(Some(node));
			}
		}

		Ok(None)
	}

	/// Map fork tree into values of new types.
	pub fn map<VT, F>(
		self,
		f: &mut F,
	) -> ForkTree<H, N, VT> where
		F: FnMut(&H, &N, V) -> VT,
	{
		let roots = self.roots
			.into_iter()
			.map(|root| {
				root.map(f)
			})
			.collect();

		ForkTree {
			roots,
			best_finalized_number: self.best_finalized_number,
		}
	}

	/// Same as [`find_node_where`](Self::find_node_where), but returns mutable reference.
	pub fn find_node_where_mut<F, E, P>(
		&mut self,
		hash: &H,
		number: &N,
		is_descendent_of: &F,
		predicate: &P,
	) -> Result<Option<&mut Node<H, N, V>>, Error<E>> where
		E: std::error::Error,
		F: Fn(&H, &H) -> Result<bool, E>,
		P: Fn(&V) -> bool,
	{
		// search for node starting from all roots
		for root in self.roots.iter_mut() {
			let node = root.find_node_where_mut(hash, number, is_descendent_of, predicate)?;

			// found the node, early exit
			if let FindOutcome::Found(node) = node {
				return Ok(Some(node));
			}
		}

		Ok(None)
	}

	/// Same as [`find_node_where`](Self::find_node_where), but returns indexes.
	pub fn find_node_index_where<F, E, P>(
		&self,
		hash: &H,
		number: &N,
		is_descendent_of: &F,
		predicate: &P,
	) -> Result<Option<Vec<usize>>, Error<E>> where
		E: std::error::Error,
		F: Fn(&H, &H) -> Result<bool, E>,
		P: Fn(&V) -> bool,
	{
		// search for node starting from all roots
		for (index, root) in self.roots.iter().enumerate() {
			let node = root.find_node_index_where(hash, number, is_descendent_of, predicate)?;

			// found the node, early exit
			if let FindOutcome::Found(mut node) = node {
				node.push(index);
				return Ok(Some(node));
			}
		}

		Ok(None)
	}

	/// Finalize a root in the tree and return it, return `None` in case no root
	/// with the given hash exists. All other roots are pruned, and the children
	/// of the finalized node become the new roots.
	pub fn finalize_root(&mut self, hash: &H) -> Option<V> {
		self.roots.iter().position(|node| node.hash == *hash)
			.map(|position| self.finalize_root_at(position))
	}

	/// Finalize root at given position. See `finalize_root` comment for details.
	fn finalize_root_at(&mut self, position: usize) -> V {
		let node = self.roots.swap_remove(position);
		self.roots = node.children;
		self.best_finalized_number = Some(node.number);
		return node.data;
	}

	/// Finalize a node in the tree. This method will make sure that the node
	/// being finalized is either an existing root (and return its data), or a
	/// node from a competing branch (not in the tree), tree pruning is done
	/// accordingly. The given function `is_descendent_of` should return `true`
	/// if the second hash (target) is a descendent of the first hash (base).
	pub fn finalize<F, E>(
		&mut self,
		hash: &H,
		number: N,
		is_descendent_of: &F,
	) -> Result<FinalizationResult<V>, Error<E>>
		where E: std::error::Error,
			  F: Fn(&H, &H) -> Result<bool, E>
	{
		if let Some(ref best_finalized_number) = self.best_finalized_number {
			if number <= *best_finalized_number {
				return Err(Error::Revert);
			}
		}

		// check if one of the current roots is being finalized
		if let Some(root) = self.finalize_root(hash) {
			return Ok(FinalizationResult::Changed(Some(root)));
		}

		// make sure we're not finalizing a descendent of any root
		for root in self.roots.iter() {
			if number > root.number && is_descendent_of(&root.hash, hash)? {
				return Err(Error::UnfinalizedAncestor);
			}
		}

		// we finalized a block earlier than any existing root (or possibly
		// another fork not part of the tree). make sure to only keep roots that
		// are part of the finalized branch
		let mut changed = false;
		self.roots.retain(|root| {
			let retain = root.number > number && is_descendent_of(hash, &root.hash).unwrap_or(false);

			if !retain {
				changed = true;
			}

			retain
		});

		self.best_finalized_number = Some(number);

		if changed {
			Ok(FinalizationResult::Changed(None))
		} else {
			Ok(FinalizationResult::Unchanged)
		}
	}

	/// Finalize a node in the tree and all its ancestors. The given function
	/// `is_descendent_of` should return `true` if the second hash (target) is
	// a descendent of the first hash (base).
	pub fn finalize_with_ancestors<F, E>(
		&mut self,
		hash: &H,
		number: N,
		is_descendent_of: &F,
	) -> Result<FinalizationResult<V>, Error<E>>
		where E: std::error::Error,
				F: Fn(&H, &H) -> Result<bool, E>
	{
		if let Some(ref best_finalized_number) = self.best_finalized_number {
			if number <= *best_finalized_number {
				return Err(Error::Revert);
			}
		}

		// check if one of the current roots is being finalized
		if let Some(root) = self.finalize_root(hash) {
			return Ok(FinalizationResult::Changed(Some(root)));
		}

		// we need to:
		// 1) remove all roots that are not ancestors AND not descendants of finalized block;
		// 2) if node is descendant - just leave it;
		// 3) if node is ancestor - 'open it'
		let mut changed = false;
		let mut idx = 0;
		while idx != self.roots.len() {
			let (is_finalized, is_descendant, is_ancestor) = {
				let root = &self.roots[idx];
				let is_finalized = root.hash == *hash;
				let is_descendant = !is_finalized
					&& root.number > number && is_descendent_of(hash, &root.hash).unwrap_or(false);
				let is_ancestor = !is_finalized && !is_descendant
					&& root.number < number && is_descendent_of(&root.hash, hash).unwrap_or(false);
				(is_finalized, is_descendant, is_ancestor)
			};

			// if we have met finalized root - open it and return
			if is_finalized {
				return Ok(FinalizationResult::Changed(Some(self.finalize_root_at(idx))));
			}

			// if node is descendant of finalized block - just leave it as is
			if is_descendant {
				idx += 1;
				continue;
			}

			// if node is ancestor of finalized block - remove it and continue with children
			if is_ancestor {
				let root = self.roots.swap_remove(idx);
				self.roots.extend(root.children);
				changed = true;
				continue;
			}

			// if node is neither ancestor, nor descendant of the finalized block - remove it
			self.roots.swap_remove(idx);
			changed = true;
		}

		self.best_finalized_number = Some(number);

		if changed {
			Ok(FinalizationResult::Changed(None))
		} else {
			Ok(FinalizationResult::Unchanged)
		}
	}

	/// Checks if any node in the tree is finalized by either finalizing the
	/// node itself or a child node that's not in the tree, guaranteeing that
	/// the node being finalized isn't a descendent of any of the node's
	/// children. Returns `Some(true)` if the node being finalized is a root,
	/// `Some(false)` if the node being finalized is not a root, and `None` if
	/// no node in the tree is finalized. The given `predicate` is checked on
	/// the prospective finalized root and must pass for finalization to occur.
	/// The given function `is_descendent_of` should return `true` if the second
	/// hash (target) is a descendent of the first hash (base).
	pub fn finalizes_any_with_descendent_if<F, P, E>(
		&self,
		hash: &H,
		number: N,
		is_descendent_of: &F,
		predicate: P,
	) -> Result<Option<bool>, Error<E>>
		where E: std::error::Error,
			  F: Fn(&H, &H) -> Result<bool, E>,
			  P: Fn(&V) -> bool,
	{
		if let Some(ref best_finalized_number) = self.best_finalized_number {
			if number <= *best_finalized_number {
				return Err(Error::Revert);
			}
		}

		// check if the given hash is equal or a descendent of any node in the
		// tree, if we find a valid node that passes the predicate then we must
		// ensure that we're not finalizing past any of its child nodes.
		for node in self.node_iter() {
			if predicate(&node.data) {
				if node.hash == *hash || is_descendent_of(&node.hash, hash)? {
					for node in node.children.iter() {
						if node.number <= number && is_descendent_of(&node.hash, &hash)? {
							return Err(Error::UnfinalizedAncestor);
						}
					}

					return Ok(Some(self.roots.iter().any(|root| root.hash == node.hash)));
				}
			}
		}

		Ok(None)
	}

	/// Finalize a root in the tree by either finalizing the node itself or a
	/// child node that's not in the tree, guaranteeing that the node being
	/// finalized isn't a descendent of any of the root's children. The given
	/// `predicate` is checked on the prospective finalized root and must pass for
	/// finalization to occur. The given function `is_descendent_of` should
	/// return `true` if the second hash (target) is a descendent of the first
	/// hash (base).
	pub fn finalize_with_descendent_if<F, P, E>(
		&mut self,
		hash: &H,
		number: N,
		is_descendent_of: &F,
		predicate: P,
	) -> Result<FinalizationResult<V>, Error<E>>
		where E: std::error::Error,
			  F: Fn(&H, &H) -> Result<bool, E>,
			  P: Fn(&V) -> bool,
	{
		if let Some(ref best_finalized_number) = self.best_finalized_number {
			if number <= *best_finalized_number {
				return Err(Error::Revert);
			}
		}

		// check if the given hash is equal or a a descendent of any root, if we
		// find a valid root that passes the predicate then we must ensure that
		// we're not finalizing past any children node.
		let mut position = None;
		for (i, root) in self.roots.iter().enumerate() {
			if predicate(&root.data) {
				if root.hash == *hash || is_descendent_of(&root.hash, hash)? {
					for node in root.children.iter() {
						if node.number <= number && is_descendent_of(&node.hash, &hash)? {
							return Err(Error::UnfinalizedAncestor);
						}
					}

					position = Some(i);
					break;
				}
			}
		}

		let node_data = position.map(|i| {
			let node = self.roots.swap_remove(i);
			self.roots = node.children;
			self.best_finalized_number = Some(node.number);
			node.data
		});

		// if the block being finalized is earlier than a given root, then it
		// must be its ancestor, otherwise we can prune the root. if there's a
		// root at the same height then the hashes must match. otherwise the
		// node being finalized is higher than the root so it must be its
		// descendent (in this case the node wasn't finalized earlier presumably
		// because the predicate didn't pass).
		let mut changed = false;
		self.roots.retain(|root| {
			let retain =
				root.number > number && is_descendent_of(hash, &root.hash).unwrap_or(false) ||
				root.number == number && root.hash == *hash ||
				is_descendent_of(&root.hash, hash).unwrap_or(false);

			if !retain {
				changed = true;
			}

			retain
		});

		self.best_finalized_number = Some(number);

		match (node_data, changed) {
			(Some(data), _) => Ok(FinalizationResult::Changed(Some(data))),
			(None, true) => Ok(FinalizationResult::Changed(None)),
			(None, false) => Ok(FinalizationResult::Unchanged),
		}
	}
}

// Workaround for: https://github.com/rust-lang/rust/issues/34537
mod node_implementation {
	use super::*;

	/// The outcome of a search within a node.
	pub enum FindOutcome<T> {
		// this is the node we were looking for.
		Found(T),
		// not the node we're looking for. contains a flag indicating
		// whether the node was a descendent. true implies the predicate failed.
		Failure(bool),
		// Abort search.
		Abort,
	}

	#[derive(Clone, Debug, Decode, Encode, PartialEq)]
	pub struct Node<H, N, V> {
		pub hash: H,
		pub number: N,
		pub data: V,
		pub children: Vec<Node<H, N, V>>,
	}

	impl<H: PartialEq, N: Ord, V> Node<H, N, V> {
		/// Rebalance the tree, i.e. sort child nodes by max branch depth (decreasing).
		pub fn rebalance(&mut self) {
			self.children.sort_by_key(|n| Reverse(n.max_depth()));
			for child in &mut self.children {
				child.rebalance();
			}
		}

		/// Finds the max depth among all branches descendent from this node.
		pub fn max_depth(&self) -> usize {
			let mut max = 0;

			for node in &self.children {
				max = node.max_depth().max(max)
			}

			max + 1
		}

		/// Map node data into values of new types.
		pub fn map<VT, F>(
			self,
			f: &mut F,
		) -> Node<H, N, VT> where
			F: FnMut(&H, &N, V) -> VT,
		{
			let children = self.children
				.into_iter()
				.map(|node| {
					node.map(f)
				})
				.collect();

			let vt = f(&self.hash, &self.number, self.data);
			Node {
				hash: self.hash,
				number: self.number,
				data: vt,
				children,
			}
		}

		pub fn import<F, E: std::error::Error>(
			&mut self,
			mut hash: H,
			mut number: N,
			mut data: V,
			is_descendent_of: &F,
		) -> Result<Option<(H, N, V)>, Error<E>>
			where E: fmt::Debug,
				  F: Fn(&H, &H) -> Result<bool, E>,
		{
			if self.hash == hash {
				return Err(Error::Duplicate);
			};

			if number <= self.number { return Ok(Some((hash, number, data))); }

			for node in self.children.iter_mut() {
				match node.import(hash, number, data, is_descendent_of)? {
					Some((h, n, d)) => {
						hash = h;
						number = n;
						data = d;
					},
					None => return Ok(None),
				}
			}

			if is_descendent_of(&self.hash, &hash)? {
				self.children.push(Node {
					data,
					hash: hash,
					number: number,
					children: Vec::new(),
				});

				Ok(None)
			} else {
				Ok(Some((hash, number, data)))
			}
		}

		/// Find a node in the tree that is the deepest ancestor of the given
		/// block hash which also passes the given predicate, backtracking
		/// when the predicate fails.
		/// The given function `is_descendent_of` should return `true` if the second hash (target)
		/// is a descendent of the first hash (base).
		///
		/// The returned indices are from last to first. The earliest index in the traverse path
		/// goes last, and the final index in the traverse path goes first. An empty list means
		/// that the current node is the result.
		pub fn find_node_index_where<F, P, E>(
			&self,
			hash: &H,
			number: &N,
			is_descendent_of: &F,
			predicate: &P,
		) -> Result<FindOutcome<Vec<usize>>, Error<E>>
			where E: std::error::Error,
				  F: Fn(&H, &H) -> Result<bool, E>,
				  P: Fn(&V) -> bool,
		{
			// stop searching this branch
			if *number < self.number {
				return Ok(FindOutcome::Failure(false));
			}

			let mut known_descendent_of = false;

			// continue depth-first search through all children
			for (i, node) in self.children.iter().enumerate() {
				// found node, early exit
				match node.find_node_index_where(hash, number, is_descendent_of, predicate)? {
					FindOutcome::Abort => return Ok(FindOutcome::Abort),
					FindOutcome::Found(mut x) => {
						x.push(i);
						return Ok(FindOutcome::Found(x))
					},
					FindOutcome::Failure(true) => {
						// if the block was a descendent of this child,
						// then it cannot be a descendent of any others,
						// so we don't search them.
						known_descendent_of = true;
						break;
					},
					FindOutcome::Failure(false) => {},
				}
			}

			// node not found in any of the descendents, if the node we're
			// searching for is a descendent of this node then we will stop the
			// search here, since there aren't any more children and we found
			// the correct node so we don't want to backtrack.
			let is_descendent_of = known_descendent_of || is_descendent_of(&self.hash, hash)?;
			if is_descendent_of {
				// if the predicate passes we return the node
				if predicate(&self.data) {
					return Ok(FindOutcome::Found(Vec::new()));
				}
			}

			// otherwise, tell our ancestor that we failed, and whether
			// the block was a descendent.
			Ok(FindOutcome::Failure(is_descendent_of))
		}

		/// Find a node in the tree that is the deepest ancestor of the given
		/// block hash which also passes the given predicate, backtracking
		/// when the predicate fails.
		/// The given function `is_descendent_of` should return `true` if the second hash (target)
		/// is a descendent of the first hash (base).
		pub fn find_node_where<F, P, E>(
			&self,
			hash: &H,
			number: &N,
			is_descendent_of: &F,
			predicate: &P,
		) -> Result<FindOutcome<&Node<H, N, V>>, Error<E>>
			where E: std::error::Error,
				  F: Fn(&H, &H) -> Result<bool, E>,
				  P: Fn(&V) -> bool,
		{
			let outcome = self.find_node_index_where(hash, number, is_descendent_of, predicate)?;

			match outcome {
				FindOutcome::Abort => Ok(FindOutcome::Abort),
				FindOutcome::Failure(f) => Ok(FindOutcome::Failure(f)),
				FindOutcome::Found(mut indexes) => {
					let mut cur = self;

					while let Some(i) = indexes.pop() {
						cur = &cur.children[i];
					}
					Ok(FindOutcome::Found(cur))
				},
			}
		}

		/// Find a node in the tree that is the deepest ancestor of the given
		/// block hash which also passes the given predicate, backtracking
		/// when the predicate fails.
		/// The given function `is_descendent_of` should return `true` if the second hash (target)
		/// is a descendent of the first hash (base).
		pub fn find_node_where_mut<F, P, E>(
			&mut self,
			hash: &H,
			number: &N,
			is_descendent_of: &F,
			predicate: &P,
		) -> Result<FindOutcome<&mut Node<H, N, V>>, Error<E>>
			where E: std::error::Error,
				  F: Fn(&H, &H) -> Result<bool, E>,
				  P: Fn(&V) -> bool,
		{
			let outcome = self.find_node_index_where(hash, number, is_descendent_of, predicate)?;

			match outcome {
				FindOutcome::Abort => Ok(FindOutcome::Abort),
				FindOutcome::Failure(f) => Ok(FindOutcome::Failure(f)),
				FindOutcome::Found(mut indexes) => {
					let mut cur = self;

					while let Some(i) = indexes.pop() {
						cur = &mut cur.children[i];
					}
					Ok(FindOutcome::Found(cur))
				},
			}
		}
	}
}

// Workaround for: https://github.com/rust-lang/rust/issues/34537
use node_implementation::{Node, FindOutcome};

struct ForkTreeIterator<'a, H, N, V> {
	stack: Vec<&'a Node<H, N, V>>,
}

impl<'a, H, N, V> Iterator for ForkTreeIterator<'a, H, N, V> {
	type Item = &'a Node<H, N, V>;

	fn next(&mut self) -> Option<Self::Item> {
		self.stack.pop().map(|node| {
			// child nodes are stored ordered by max branch height (decreasing),
			// we want to keep this ordering while iterating but since we're
			// using a stack for iterator state we need to reverse it.
			self.stack.extend(node.children.iter().rev());
			node
		})
	}
}

struct RemovedIterator<H, N, V> {
	stack: Vec<Node<H, N, V>>,
}

impl<H, N, V> Iterator for RemovedIterator<H, N, V> {
	type Item = (H, N, V);

	fn next(&mut self) -> Option<Self::Item> {
		self.stack.pop().map(|mut node| {
			// child nodes are stored ordered by max branch height (decreasing),
			// we want to keep this ordering while iterating but since we're
			// using a stack for iterator state we need to reverse it.
			let mut children = Vec::new();
			std::mem::swap(&mut children, &mut node.children);

			self.stack.extend(children.into_iter().rev());
			(node.hash, node.number, node.data)
		})
	}
}

#[cfg(test)]
mod test {
	use super::{FinalizationResult, ForkTree, Error};

	#[derive(Debug, PartialEq)]
	struct TestError;

	impl std::fmt::Display for TestError {
		fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
			write!(f, "TestError")
		}
	}

	impl std::error::Error for TestError {}

	fn test_fork_tree<'a>() -> (ForkTree<&'a str, u64, ()>, impl Fn(&&str, &&str) -> Result<bool, TestError>)  {
		let mut tree = ForkTree::new();

		//
		//     - B - C - D - E
		//    /
		//   /   - G
		//  /   /
		// A - F - H - I
		//          \
		//           - L - M
		//              \
		//               - O
		//  \
		//   — J - K
		//
		// (where N is not a part of fork tree)
		let is_descendent_of = |base: &&str, block: &&str| -> Result<bool, TestError> {
			let letters = vec!["B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "O"];
			match (*base, *block) {
				("A", b) => Ok(letters.into_iter().any(|n| n == b)),
				("B", b) => Ok(b == "C" || b == "D" || b == "E"),
				("C", b) => Ok(b == "D" || b == "E"),
				("D", b) => Ok(b == "E"),
				("E", _) => Ok(false),
				("F", b) => Ok(b == "G" || b == "H" || b == "I" || b == "L" || b == "M" || b == "O"),
				("G", _) => Ok(false),
				("H", b) => Ok(b == "I" || b == "L" || b == "M" || b == "O"),
				("I", _) => Ok(false),
				("J", b) => Ok(b == "K"),
				("K", _) => Ok(false),
				("L", b) => Ok(b == "M" || b == "O"),
				("M", _) => Ok(false),
				("O", _) => Ok(false),
				("0", _) => Ok(true),
				_ => Ok(false),
			}
		};

		tree.import("A", 1, (), &is_descendent_of).unwrap();

		tree.import("B", 2, (), &is_descendent_of).unwrap();
		tree.import("C", 3, (), &is_descendent_of).unwrap();
		tree.import("D", 4, (), &is_descendent_of).unwrap();
		tree.import("E", 5, (), &is_descendent_of).unwrap();

		tree.import("F", 2, (), &is_descendent_of).unwrap();
		tree.import("G", 3, (), &is_descendent_of).unwrap();

		tree.import("H", 3, (), &is_descendent_of).unwrap();
		tree.import("I", 4, (), &is_descendent_of).unwrap();
		tree.import("L", 4, (), &is_descendent_of).unwrap();
		tree.import("M", 5, (), &is_descendent_of).unwrap();
		tree.import("O", 5, (), &is_descendent_of).unwrap();

		tree.import("J", 2, (), &is_descendent_of).unwrap();
		tree.import("K", 3, (), &is_descendent_of).unwrap();

		(tree, is_descendent_of)
	}

	#[test]
	fn import_doesnt_revert() {
		let (mut tree, is_descendent_of) = test_fork_tree();

		tree.finalize_root(&"A");

		assert_eq!(
			tree.best_finalized_number,
			Some(1),
		);

		assert_eq!(
			tree.import("A", 1, (), &is_descendent_of),
			Err(Error::Revert),
		);
	}

	#[test]
	fn import_doesnt_add_duplicates() {
		let (mut tree, is_descendent_of) = test_fork_tree();

		assert_eq!(
			tree.import("A", 1, (), &is_descendent_of),
			Err(Error::Duplicate),
		);

		assert_eq!(
			tree.import("I", 4, (), &is_descendent_of),
			Err(Error::Duplicate),
		);

		assert_eq!(
			tree.import("G", 3, (), &is_descendent_of),
			Err(Error::Duplicate),
		);

		assert_eq!(
			tree.import("K", 3, (), &is_descendent_of),
			Err(Error::Duplicate),
		);
	}

	#[test]
	fn finalize_root_works() {
		let finalize_a = || {
			let (mut tree, ..) = test_fork_tree();

			assert_eq!(
				tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
				vec![("A", 1)],
			);

			// finalizing "A" opens up three possible forks
			tree.finalize_root(&"A");

			assert_eq!(
				tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
				vec![("B", 2), ("F", 2), ("J", 2)],
			);

			tree
		};

		{
			let mut tree = finalize_a();

			// finalizing "B" will progress on its fork and remove any other competing forks
			tree.finalize_root(&"B");

			assert_eq!(
				tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
				vec![("C", 3)],
			);

			// all the other forks have been pruned
			assert!(tree.roots.len() == 1);
		}

		{
			let mut tree = finalize_a();

			// finalizing "J" will progress on its fork and remove any other competing forks
			tree.finalize_root(&"J");

			assert_eq!(
				tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
				vec![("K", 3)],
			);

			// all the other forks have been pruned
			assert!(tree.roots.len() == 1);
		}
	}

	#[test]
	fn finalize_works() {
		let (mut tree, is_descendent_of) = test_fork_tree();

		let original_roots = tree.roots.clone();

		// finalizing a block prior to any in the node doesn't change the tree
		assert_eq!(
			tree.finalize(&"0", 0, &is_descendent_of),
			Ok(FinalizationResult::Unchanged),
		);

		assert_eq!(tree.roots, original_roots);

		// finalizing "A" opens up three possible forks
		assert_eq!(
			tree.finalize(&"A", 1, &is_descendent_of),
			Ok(FinalizationResult::Changed(Some(()))),
		);

		assert_eq!(
			tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
			vec![("B", 2), ("F", 2), ("J", 2)],
		);

		// finalizing anything lower than what we observed will fail
		assert_eq!(
			tree.best_finalized_number,
			Some(1),
		);

		assert_eq!(
			tree.finalize(&"Z", 1, &is_descendent_of),
			Err(Error::Revert),
		);

		// trying to finalize a node without finalizing its ancestors first will fail
		assert_eq!(
			tree.finalize(&"H", 3, &is_descendent_of),
			Err(Error::UnfinalizedAncestor),
		);

		// after finalizing "F" we can finalize "H"
		assert_eq!(
			tree.finalize(&"F", 2, &is_descendent_of),
			Ok(FinalizationResult::Changed(Some(()))),
		);

		assert_eq!(
			tree.finalize(&"H", 3, &is_descendent_of),
			Ok(FinalizationResult::Changed(Some(()))),
		);

		assert_eq!(
			tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
			vec![("I", 4), ("L", 4)],
		);

		// finalizing a node from another fork that isn't part of the tree clears the tree
		assert_eq!(
			tree.finalize(&"Z", 5, &is_descendent_of),
			Ok(FinalizationResult::Changed(None)),
		);

		assert!(tree.roots.is_empty());
	}

	#[test]
	fn finalize_with_ancestor_works() {
		let (mut tree, is_descendent_of) = test_fork_tree();

		let original_roots = tree.roots.clone();

		// finalizing a block prior to any in the node doesn't change the tree
		assert_eq!(
			tree.finalize_with_ancestors(&"0", 0, &is_descendent_of),
			Ok(FinalizationResult::Unchanged),
		);

		assert_eq!(tree.roots, original_roots);

		// finalizing "A" opens up three possible forks
		assert_eq!(
			tree.finalize_with_ancestors(&"A", 1, &is_descendent_of),
			Ok(FinalizationResult::Changed(Some(()))),
		);

		assert_eq!(
			tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
			vec![("B", 2), ("F", 2), ("J", 2)],
		);

		// finalizing H:
		// 1) removes roots that are not ancestors/descendants of H (B, J)
		// 2) opens root that is ancestor of H (F -> G+H)
		// 3) finalizes the just opened root H (H -> I + L)
		assert_eq!(
			tree.finalize_with_ancestors(&"H", 3, &is_descendent_of),
			Ok(FinalizationResult::Changed(Some(()))),
		);

		assert_eq!(
			tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
			vec![("I", 4), ("L", 4)],
		);

		assert_eq!(
			tree.best_finalized_number,
			Some(3),
		);

		// finalizing N (which is not a part of the tree):
		// 1) removes roots that are not ancestors/descendants of N (I)
		// 2) opens root that is ancestor of N (L -> M+O)
		// 3) removes roots that are not ancestors/descendants of N (O)
		// 4) opens root that is ancestor of N (M -> {})
		assert_eq!(
			tree.finalize_with_ancestors(&"N", 6, &is_descendent_of),
			Ok(FinalizationResult::Changed(None)),
		);

		assert_eq!(
			tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
			vec![],
		);

		assert_eq!(
			tree.best_finalized_number,
			Some(6),
		);
	}

	#[test]
	fn finalize_with_descendent_works() {
		#[derive(Debug, PartialEq)]
		struct Change { effective: u64 };

		let (mut tree, is_descendent_of) = {
			let mut tree = ForkTree::new();

			let is_descendent_of = |base: &&str, block: &&str| -> Result<bool, TestError> {

				//
				// A0 #1 - (B #2) - (C #5) - D #10 - E #15 - (F #100)
				//                            \
				//                             - (G #100)
				//
				// A1 #1
				//
				// Nodes B, C, F and G  are not part of the tree.
				match (*base, *block) {
					("A0", b) => Ok(b == "B" || b == "C" || b == "D" || b == "G"),
					("A1", _) => Ok(false),
					("C", b) => Ok(b == "D"),
					("D", b) => Ok(b == "E" || b == "F" || b == "G"),
					("E", b) => Ok(b == "F"),
					_ => Ok(false),
				}
			};

			tree.import("A0", 1, Change { effective: 5 }, &is_descendent_of).unwrap();
			tree.import("A1", 1, Change { effective: 5 }, &is_descendent_of).unwrap();
			tree.import("D", 10, Change { effective: 10 }, &is_descendent_of).unwrap();
			tree.import("E", 15, Change { effective: 50 }, &is_descendent_of).unwrap();

			(tree, is_descendent_of)
		};

		assert_eq!(
			tree.finalizes_any_with_descendent_if(
				&"B",
				2,
				&is_descendent_of,
				|c| c.effective <= 2,
			),
			Ok(None),
		);

		// finalizing "D" will finalize a block from the tree, but it can't be applied yet
		// since it is not a root change
		assert_eq!(
			tree.finalizes_any_with_descendent_if(
				&"D",
				10,
				&is_descendent_of,
				|c| c.effective == 10,
			),
			Ok(Some(false)),
		);

		// finalizing "B" doesn't finalize "A0" since the predicate doesn't pass,
		// although it will clear out "A1" from the tree
		assert_eq!(
			tree.finalize_with_descendent_if(
				&"B",
				2,
				&is_descendent_of,
				|c| c.effective <= 2,
			),
			Ok(FinalizationResult::Changed(None)),
		);

		assert_eq!(
			tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
			vec![("A0", 1)],
		);

		// finalizing "C" will finalize the node "A0" and prune it out of the tree
		assert_eq!(
			tree.finalizes_any_with_descendent_if(
				&"C",
				5,
				&is_descendent_of,
				|c| c.effective <= 5,
			),
			Ok(Some(true)),
		);

		assert_eq!(
			tree.finalize_with_descendent_if(
				&"C",
				5,
				&is_descendent_of,
				|c| c.effective <= 5,
			),
			Ok(FinalizationResult::Changed(Some(Change { effective: 5 }))),
		);

		assert_eq!(
			tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
			vec![("D", 10)],
		);

		// finalizing "F" will fail since it would finalize past "E" without finalizing "D" first
		assert_eq!(
			tree.finalizes_any_with_descendent_if(
				&"F",
				100,
				&is_descendent_of,
				|c| c.effective <= 100,
			),
			Err(Error::UnfinalizedAncestor),
		);

		// it will work with "G" though since it is not in the same branch as "E"
		assert_eq!(
			tree.finalizes_any_with_descendent_if(
				&"G",
				100,
				&is_descendent_of,
				|c| c.effective <= 100,
			),
			Ok(Some(true)),
		);

		assert_eq!(
			tree.finalize_with_descendent_if(
				&"G",
				100,
				&is_descendent_of,
				|c| c.effective <= 100,
			),
			Ok(FinalizationResult::Changed(Some(Change { effective: 10 }))),
		);

		// "E" will be pruned out
		assert_eq!(tree.roots().count(), 0);
	}

	#[test]
	fn iter_iterates_in_preorder() {
		let (tree, ..) = test_fork_tree();
		assert_eq!(
			tree.iter().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
			vec![
				("A", 1),
				("B", 2), ("C", 3), ("D", 4), ("E", 5),
				("F", 2),
				("G", 3),
				("H", 3), ("I", 4),
				("L", 4), ("M", 5), ("O", 5),
				("J", 2), ("K", 3)
			],
		);
	}

	#[test]
	fn minimizes_calls_to_is_descendent_of() {
		use std::sync::atomic::{AtomicUsize, Ordering};

		let n_is_descendent_of_calls = AtomicUsize::new(0);

		let is_descendent_of = |_: &&str, _: &&str| -> Result<bool, TestError> {
			n_is_descendent_of_calls.fetch_add(1, Ordering::SeqCst);
			Ok(true)
		};

		{
			// Deep tree where we want to call `finalizes_any_with_descendent_if`. The
			// search for the node should first check the predicate (which is cheaper) and
			// only then call `is_descendent_of`
			let mut tree = ForkTree::new();
			let letters = vec!["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K"];

			for (i, letter) in letters.iter().enumerate() {
				tree.import::<_, TestError>(*letter, i, i, &|_, _| Ok(true)).unwrap();
			}

			// "L" is a descendent of "K", but the predicate will only pass for "K",
			// therefore only one call to `is_descendent_of` should be made
			assert_eq!(
				tree.finalizes_any_with_descendent_if(
					&"L",
					11,
					&is_descendent_of,
					|i| *i == 10,
				),
				Ok(Some(false)),
			);

			assert_eq!(
				n_is_descendent_of_calls.load(Ordering::SeqCst),
				1,
			);
		}

		n_is_descendent_of_calls.store(0, Ordering::SeqCst);

		{
			// Multiple roots in the tree where we want to call `finalize_with_descendent_if`.
			// The search for the root node should first check the predicate (which is cheaper)
			// and only then call `is_descendent_of`
			let mut tree = ForkTree::new();
			let letters = vec!["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K"];

			for (i, letter) in letters.iter().enumerate() {
				tree.import::<_, TestError>(*letter, i, i, &|_, _| Ok(false)).unwrap();
			}

			// "L" is a descendent of "K", but the predicate will only pass for "K",
			// therefore only one call to `is_descendent_of` should be made
			assert_eq!(
				tree.finalize_with_descendent_if(
					&"L",
					11,
					&is_descendent_of,
					|i| *i == 10,
				),
				Ok(FinalizationResult::Changed(Some(10))),
			);

			assert_eq!(
				n_is_descendent_of_calls.load(Ordering::SeqCst),
				1,
			);
		}
	}

	#[test]
	fn find_node_works() {
		let (tree, is_descendent_of) = test_fork_tree();

		let node = tree.find_node_where(
			&"D",
			&4,
			&is_descendent_of,
			&|_| true,
		).unwrap().unwrap();

		assert_eq!(node.hash, "C");
		assert_eq!(node.number, 3);
	}

	#[test]
	fn map_works() {
		let (tree, _is_descendent_of) = test_fork_tree();

		let _tree = tree.map(&mut |_, _, _| ());
	}

	#[test]
	fn prune_works() {
		let (mut tree, is_descendent_of) = test_fork_tree();

		let removed = tree.prune(
			&"C",
			&3,
			&is_descendent_of,
			&|_| true,
		).unwrap();

		assert_eq!(
			tree.roots.iter().map(|node| node.hash).collect::<Vec<_>>(),
			vec!["B"],
		);

		assert_eq!(
			tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(),
			vec!["B", "C", "D", "E"],
		);

		assert_eq!(
			removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(),
			vec!["A", "F", "G", "H", "I", "L", "M", "O", "J", "K"]
		);

		let removed = tree.prune(
			&"E",
			&5,
			&is_descendent_of,
			&|_| true,
		).unwrap();

		assert_eq!(
			tree.roots.iter().map(|node| node.hash).collect::<Vec<_>>(),
			vec!["D"],
		);

		assert_eq!(
			tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(),
			vec!["D", "E"],
		);

		assert_eq!(
			removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(),
			vec!["B", "C"]
		);
	}

	#[test]
	fn find_node_backtracks_after_finding_highest_descending_node() {
		let mut tree = ForkTree::new();

		//
		// A - B
		//  \
		//   — C
		//
		let is_descendent_of = |base: &&str, block: &&str| -> Result<bool, TestError> {
			match (*base, *block) {
				("A", b) => Ok(b == "B" || b == "C" || b == "D"),
				("B", b) | ("C", b) => Ok(b == "D"),
				("0", _) => Ok(true),
				_ => Ok(false),
			}
		};

		tree.import("A", 1, 1, &is_descendent_of).unwrap();
		tree.import("B", 2, 2, &is_descendent_of).unwrap();
		tree.import("C", 2, 4, &is_descendent_of).unwrap();

		// when searching the tree we reach node `C`, but the
		// predicate doesn't pass. we should backtrack to `B`, but not to `A`,
		// since "B" fulfills the predicate.
		let node = tree.find_node_where(
			&"D",
			&3,
			&is_descendent_of,
			&|data| *data < 3,
		).unwrap();

		assert_eq!(node.unwrap().hash, "B");
	}

	#[test]
	fn tree_rebalance() {
		let (mut tree, _) = test_fork_tree();

		assert_eq!(
			tree.iter().map(|(h, _, _)| *h).collect::<Vec<_>>(),
			vec!["A", "B", "C", "D", "E", "F", "G", "H", "I", "L", "M", "O", "J", "K"],
		);

		// after rebalancing the tree we should iterate in preorder exploring
		// the longest forks first. check the ascii art above to understand the
		// expected output below.
		tree.rebalance();

		assert_eq!(
			tree.iter().map(|(h, _, _)| *h).collect::<Vec<_>>(),
			["A", "B", "C", "D", "E", "F", "H", "L", "M", "O", "I", "G", "J", "K"]
		);
	}
}