// This file is part of Substrate. // Copyright (C) 2019-2022 Parity Technologies (UK) Ltd. // SPDX-License-Identifier: Apache-2.0 // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. //! Utility library for managing tree-like ordered data with logic for pruning //! the tree while finalizing nodes. #![warn(missing_docs)] use codec::{Decode, Encode}; use std::{cmp::Reverse, fmt}; /// Error occurred when iterating with the tree. #[derive(Clone, Debug, PartialEq)] pub enum Error { /// 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 fmt::Display for Error { 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 std::error::Error for Error { fn cause(&self) -> Option<&dyn std::error::Error> { None } } impl From for Error { fn from(err: E) -> Error { Error::Client(err) } } /// Result of finalizing a node (that could be a part of the tree or not). #[derive(Debug, PartialEq)] pub enum FinalizationResult { /// The tree has changed, optionally return the value associated with the finalized node. Changed(Option), /// The tree has not changed. Unchanged, } /// Filtering action. #[derive(Debug, PartialEq)] pub enum FilterAction { /// Remove the node and its subtree. Remove, /// Maintain the node. KeepNode, /// Maintain the node and its subtree. KeepTree, } /// 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 { roots: Vec>, best_finalized_number: Option, } impl ForkTree where H: PartialEq, N: Ord, { /// Create a new empty tree. pub fn new() -> ForkTree { 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())); let mut stack: Vec<_> = self.roots.iter_mut().collect(); while let Some(node) = stack.pop() { node.children.sort_by_key(|n| Reverse(n.max_depth())); stack.extend(node.children.iter_mut()); } } /// 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. // WARNING: some users of this method (i.e. consensus epoch changes tree) currently silently // rely on a **post-order DFS** traversal. If we are using instead a top-down traversal method // then the `is_descendent_of` closure, when used after a warp-sync, may end up querying the // backend for a block (the one corresponding to the root) that is not present and thus will // return a wrong result. pub fn import( &mut self, hash: H, number: N, data: V, is_descendent_of: &F, ) -> Result> where E: std::error::Error, F: Fn(&H, &H) -> Result, { if let Some(ref best_finalized_number) = self.best_finalized_number { if number <= *best_finalized_number { return Err(Error::Revert) } } let (children, is_root) = match self.find_node_where_mut(&hash, &number, is_descendent_of, &|_| true)? { Some(parent) => (&mut parent.children, false), None => (&mut self.roots, true), }; if children.iter().any(|elem| elem.hash == hash) { return Err(Error::Duplicate) } children.push(Node { data, hash, number, children: Default::default() }); if children.len() == 1 { // Rebalance may be required only if we've extended the branch depth. self.rebalance(); } Ok(is_root) } /// Iterates over the existing roots in the tree. pub fn roots(&self) -> impl Iterator { self.roots.iter().map(|node| (&node.hash, &node.number, &node.data)) } fn node_iter(&self) -> impl Iterator> { // we need to reverse the order of roots to maintain the expected // ordering since the iterator uses a stack to track state. ForkTreeIterator { stack: self.roots.iter().rev().collect() } } /// Iterates the nodes in the tree in pre-order. pub fn iter(&self) -> impl Iterator { self.node_iter().map(|node| (&node.hash, &node.number, &node.data)) } /// Map fork tree into values of new types. /// /// Tree traversal technique (e.g. BFS vs DFS) is left as not specified and /// may be subject to change in the future. In other words, your predicates /// should not rely on the observed traversal technique currently in use. pub fn map(self, f: &mut F) -> ForkTree where F: FnMut(&H, &N, V) -> VT, { let mut queue: Vec<_> = self.roots.into_iter().rev().map(|node| (usize::MAX, node)).collect(); let mut next_queue = Vec::new(); let mut output = Vec::new(); while !queue.is_empty() { for (parent_index, node) in queue.drain(..) { let new_data = f(&node.hash, &node.number, node.data); let new_node = Node { hash: node.hash, number: node.number, data: new_data, children: Vec::with_capacity(node.children.len()), }; let node_id = output.len(); output.push((parent_index, new_node)); for child in node.children.into_iter().rev() { next_queue.push((node_id, child)); } } std::mem::swap(&mut queue, &mut next_queue); } let mut roots = Vec::new(); while let Some((parent_index, new_node)) = output.pop() { if parent_index == usize::MAX { roots.push(new_node); } else { output[parent_index].1.children.push(new_node); } } ForkTree { roots, best_finalized_number: self.best_finalized_number } } /// 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( &self, hash: &H, number: &N, is_descendent_of: &F, predicate: &P, ) -> Result>, Error> where E: std::error::Error, F: Fn(&H, &H) -> Result, P: Fn(&V) -> bool, { let maybe_path = self.find_node_index_where(hash, number, is_descendent_of, predicate)?; Ok(maybe_path.map(|path| { let children = path.iter().take(path.len() - 1).fold(&self.roots, |curr, &i| &curr[i].children); &children[path[path.len() - 1]] })) } /// Same as [`find_node_where`](ForkTree::find_node_where), but returns mutable reference. pub fn find_node_where_mut( &mut self, hash: &H, number: &N, is_descendent_of: &F, predicate: &P, ) -> Result>, Error> where E: std::error::Error, F: Fn(&H, &H) -> Result, P: Fn(&V) -> bool, { let maybe_path = self.find_node_index_where(hash, number, is_descendent_of, predicate)?; Ok(maybe_path.map(|path| { let children = path .iter() .take(path.len() - 1) .fold(&mut self.roots, |curr, &i| &mut curr[i].children); &mut children[path[path.len() - 1]] })) } /// Same as [`find_node_where`](ForkTree::find_node_where), but returns indices. /// /// The returned indices represent the full path to reach the matching node starting /// from first to last, i.e. the earliest index in the traverse path goes first, and the final /// index in the traverse path goes last. If a node is found that matches the predicate /// the returned path should always contain at least one index, otherwise `None` is /// returned. // WARNING: some users of this method (i.e. consensus epoch changes tree) currently silently // rely on a **post-order DFS** traversal. If we are using instead a top-down traversal method // then the `is_descendent_of` closure, when used after a warp-sync, will end up querying the // backend for a block (the one corresponding to the root) that is not present and thus will // return a wrong result. pub fn find_node_index_where( &self, hash: &H, number: &N, is_descendent_of: &F, predicate: &P, ) -> Result>, Error> where E: std::error::Error, F: Fn(&H, &H) -> Result, P: Fn(&V) -> bool, { let mut stack = vec![]; let mut root_idx = 0; let mut found = false; let mut is_descendent = false; while root_idx < self.roots.len() { if *number <= self.roots[root_idx].number { root_idx += 1; continue } // The second element in the stack tuple tracks what is the **next** children // index to search into. If we find an ancestor then we stop searching into // alternative branches and we focus on the current path up to the root. stack.push((&self.roots[root_idx], 0)); while let Some((node, i)) = stack.pop() { if i < node.children.len() && !is_descendent { stack.push((node, i + 1)); if node.children[i].number < *number { stack.push((&node.children[i], 0)); } } else if is_descendent || is_descendent_of(&node.hash, hash)? { is_descendent = true; if predicate(&node.data) { found = true; break } } } // If the element we are looking for is a descendent of the current root // then we can stop the search. if is_descendent { break } root_idx += 1; } Ok(if found { // The path is the root index followed by the indices of all the children // we were processing when we found the element (remember the stack // contains the index of the **next** children to process). let path: Vec<_> = std::iter::once(root_idx).chain(stack.iter().map(|(_, i)| *i - 1)).collect(); Some(path) } else { None }) } /// 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( &mut self, hash: &H, number: &N, is_descendent_of: &F, predicate: &P, ) -> Result, Error> where E: std::error::Error, F: Fn(&H, &H) -> Result, P: Fn(&V) -> bool, { let root_index = match self.find_node_index_where(hash, number, is_descendent_of, predicate)? { Some(idx) => idx, None => return Ok(RemovedIterator { stack: Vec::new() }), }; let mut old_roots = std::mem::take(&mut self.roots); let curr_children = root_index .iter() .take(root_index.len() - 1) .fold(&mut old_roots, |curr, idx| &mut curr[*idx].children); let mut root = curr_children.remove(root_index[root_index.len() - 1]); 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::take(&mut root.children); 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)?) { 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]; self.rebalance(); Ok(RemovedIterator { stack: removed }) } /// 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 { 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); 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( &mut self, hash: &H, number: N, is_descendent_of: &F, ) -> Result, Error> where E: std::error::Error, F: Fn(&H, &H) -> Result, { 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; let roots = std::mem::take(&mut self.roots); for root in roots { if root.number > number && is_descendent_of(hash, &root.hash)? { self.roots.push(root); } else { changed = true; } } 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( &mut self, hash: &H, number: N, is_descendent_of: &F, ) -> Result, Error> where E: std::error::Error, F: Fn(&H, &H) -> Result, { 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)?; let is_ancestor = !is_finalized && !is_descendant && root.number < number && is_descendent_of(&root.hash, hash)?; (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 node's descendent that's not in the tree, guaranteeing /// that the node being finalized isn't a descendent of (or equal to) 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( &self, hash: &H, number: N, is_descendent_of: &F, predicate: P, ) -> Result, Error> where E: std::error::Error, F: Fn(&H, &H) -> Result, 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) && (node.hash == *hash || is_descendent_of(&node.hash, hash)?) { for child in node.children.iter() { if child.number <= number && (child.hash == *hash || is_descendent_of(&child.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 /// node's descendent that's not in the tree, guaranteeing that the node /// being finalized isn't a descendent of (or equal to) 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( &mut self, hash: &H, number: N, is_descendent_of: &F, predicate: P, ) -> Result, Error> where E: std::error::Error, F: Fn(&H, &H) -> Result, 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) && (root.hash == *hash || is_descendent_of(&root.hash, hash)?) { for child in root.children.iter() { if child.number <= number && (child.hash == *hash || is_descendent_of(&child.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 }); // Retain only roots that are descendents of the finalized block (this // happens if the node has been properly finalized) or that are // ancestors (or equal) to the finalized block (in this case the node // wasn't finalized earlier presumably because the predicate didn't // pass). let mut changed = false; let roots = std::mem::take(&mut self.roots); for root in roots { let retain = root.number > number && is_descendent_of(hash, &root.hash)? || root.number == number && root.hash == *hash || is_descendent_of(&root.hash, hash)?; if retain { self.roots.push(root); } else { changed = true; } } 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), } } /// Remove from the tree some nodes (and their subtrees) using a `filter` predicate. /// /// The `filter` is called over tree nodes and returns a filter action: /// - `Remove` if the node and its subtree should be removed; /// - `KeepNode` if we should maintain the node and keep processing the tree. /// - `KeepTree` if we should maintain the node and its entire subtree. /// /// An iterator over all the pruned nodes is returned. pub fn drain_filter(&mut self, filter: F) -> impl Iterator where F: Fn(&H, &N, &V) -> FilterAction, { let mut removed = vec![]; let mut retained = Vec::new(); let mut queue: Vec<_> = std::mem::take(&mut self.roots) .into_iter() .rev() .map(|node| (usize::MAX, node)) .collect(); let mut next_queue = Vec::new(); while !queue.is_empty() { for (parent_idx, mut node) in queue.drain(..) { match filter(&node.hash, &node.number, &node.data) { FilterAction::KeepNode => { let node_idx = retained.len(); let children = std::mem::take(&mut node.children); retained.push((parent_idx, node)); for child in children.into_iter().rev() { next_queue.push((node_idx, child)); } }, FilterAction::KeepTree => { retained.push((parent_idx, node)); }, FilterAction::Remove => { removed.push(node); }, } } std::mem::swap(&mut queue, &mut next_queue); } while let Some((parent_idx, node)) = retained.pop() { if parent_idx == usize::MAX { self.roots.push(node); } else { retained[parent_idx].1.children.push(node); } } if !removed.is_empty() { self.rebalance(); } RemovedIterator { stack: removed } } } // Workaround for: https://github.com/rust-lang/rust/issues/34537 use node_implementation::Node; mod node_implementation { use super::*; #[derive(Clone, Debug, Decode, Encode, PartialEq)] pub struct Node { pub hash: H, pub number: N, pub data: V, pub children: Vec>, } impl Node { /// Finds the max depth among all branches descendent from this node. pub fn max_depth(&self) -> usize { let mut max: usize = 0; let mut stack = vec![(self, 0)]; while let Some((node, height)) = stack.pop() { if height > max { max = height; } node.children.iter().for_each(|n| stack.push((n, height + 1))); } max } } } struct ForkTreeIterator<'a, H, N, V> { stack: Vec<&'a Node>, } impl<'a, H, N, V> Iterator for ForkTreeIterator<'a, H, N, V> { type Item = &'a Node; fn next(&mut self) -> Option { 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 { stack: Vec>, } impl Iterator for RemovedIterator { type Item = (H, N, V); fn next(&mut self) -> Option { 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 children = std::mem::take(&mut node.children); self.stack.extend(children.into_iter().rev()); (node.hash, node.number, node.data) }) } } #[cfg(test)] mod test { use crate::FilterAction; use super::{Error, FinalizationResult, ForkTree}; #[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) { let mut tree = ForkTree::new(); #[rustfmt::skip] // // - B - C - D - E // / // / - G // / / // A - F - H - I // \ \ // \ - L - M - N // \ \ // \ - O // - J - K // // (where N is not a part of fork tree) // // NOTE: the tree will get automatically rebalance on import and won't be laid out like the // diagram above. the children will be ordered by subtree depth and the longest branches // will be on the leftmost side of the tree. let is_descendent_of = |base: &&str, block: &&str| -> Result { let letters = vec!["B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "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 == "N" || b == "O"), ("G", _) => Ok(false), ("H", b) => Ok(b == "I" || b == "L" || b == "M" || b == "N" || b == "O"), ("I", _) => Ok(false), ("J", b) => Ok(b == "K"), ("K", _) => Ok(false), ("L", b) => Ok(b == "M" || b == "O" || b == "N"), ("M", b) => Ok(b == "N"), ("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, *n)).collect::>(), vec![("A", 1)]); // finalizing "A" opens up three possible forks tree.finalize_root(&"A"); assert_eq!( tree.roots().map(|(h, n, _)| (*h, *n)).collect::>(), 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, *n)).collect::>(), 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, *n)).collect::>(), 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, *n)).collect::>(), 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, *n)).collect::>(), vec![("L", 4), ("I", 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, *n)).collect::>(), 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, *n)).collect::>(), vec![("L", 4), ("I", 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, *n)).collect::>(), 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 { // 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 == "E" || 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), } }; let is_root = tree.import("A0", 1, Change { effective: 5 }, &is_descendent_of).unwrap(); assert!(is_root); let is_root = tree.import("A1", 1, Change { effective: 5 }, &is_descendent_of).unwrap(); assert!(is_root); let is_root = tree.import("D", 10, Change { effective: 10 }, &is_descendent_of).unwrap(); assert!(!is_root); let is_root = tree.import("E", 15, Change { effective: 50 }, &is_descendent_of).unwrap(); assert!(!is_root); (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" is not allowed since it is not a root. assert_eq!( tree.finalize_with_descendent_if(&"D", 10, &is_descendent_of, |c| c.effective <= 10), Err(Error::UnfinalizedAncestor) ); // 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 "E" is not allowed since there are not finalized anchestors. assert_eq!( tree.finalizes_any_with_descendent_if(&"E", 15, &is_descendent_of, |c| c.effective == 10), Err(Error::UnfinalizedAncestor) ); // 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, *n)).collect::>(), 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, *n)).collect::>(), 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, *n)).collect::>(), vec![ ("A", 1), ("B", 2), ("C", 3), ("D", 4), ("E", 5), ("F", 2), ("H", 3), ("L", 4), ("M", 5), ("O", 5), ("I", 4), ("G", 3), ("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 { 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 map_works() { let (mut tree, _) = test_fork_tree(); // Extend the single root fork-tree to also excercise the roots order during map. let is_descendent_of = |_: &&str, _: &&str| -> Result { Ok(false) }; let is_root = tree.import("A1", 1, (), &is_descendent_of).unwrap(); assert!(is_root); let is_root = tree.import("A2", 1, (), &is_descendent_of).unwrap(); assert!(is_root); let old_tree = tree.clone(); let new_tree = tree.map(&mut |hash, _, _| hash.to_owned()); // Check content and order assert!(new_tree.iter().all(|(hash, _, data)| hash == data)); assert_eq!( old_tree.iter().map(|(hash, _, _)| *hash).collect::>(), new_tree.iter().map(|(hash, _, _)| *hash).collect::>(), ); } #[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!["B"]); assert_eq!( tree.iter().map(|(hash, _, _)| *hash).collect::>(), vec!["B", "C", "D", "E"], ); assert_eq!( removed.map(|(hash, _, _)| hash).collect::>(), vec!["A", "F", "H", "L", "M", "O", "I", "G", "J", "K"] ); let removed = tree.prune(&"E", &5, &is_descendent_of, &|_| true).unwrap(); assert_eq!(tree.roots.iter().map(|node| node.hash).collect::>(), vec!["D"]); assert_eq!(tree.iter().map(|(hash, _, _)| *hash).collect::>(), vec!["D", "E"]); assert_eq!(removed.map(|(hash, _, _)| hash).collect::>(), 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 { 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 rebalance_works() { let (mut tree, _) = test_fork_tree(); // the tree is automatically rebalanced on import, therefore we should iterate in preorder // exploring the longest forks first. check the ascii art above to understand the expected // output below. assert_eq!( tree.iter().map(|(h, _, _)| *h).collect::>(), vec!["A", "B", "C", "D", "E", "F", "H", "L", "M", "O", "I", "G", "J", "K"], ); // let's add a block "P" which is a descendent of block "O" let is_descendent_of = |base: &&str, block: &&str| -> Result { match (*base, *block) { (b, "P") => Ok(vec!["A", "F", "H", "L", "O"].into_iter().any(|n| n == b)), _ => Ok(false), } }; tree.import("P", 6, (), &is_descendent_of).unwrap(); // this should re-order the tree, since the branch "A -> B -> C -> D -> E" is no longer tied // with 5 blocks depth. additionally "O" should be visited before "M" now, since it has one // descendent "P" which makes that branch 6 blocks long. assert_eq!( tree.iter().map(|(h, _, _)| *h).collect::>(), ["A", "F", "H", "L", "O", "P", "M", "I", "G", "B", "C", "D", "E", "J", "K"] ); } #[test] fn drain_filter_works() { let (mut tree, _) = test_fork_tree(); let filter = |h: &&str, _: &u64, _: &()| match *h { "A" | "B" | "F" | "G" => FilterAction::KeepNode, "C" => FilterAction::KeepTree, "H" | "J" => FilterAction::Remove, _ => panic!("Unexpected filtering for node: {}", *h), }; let removed = tree.drain_filter(filter); assert_eq!( tree.iter().map(|(h, _, _)| *h).collect::>(), ["A", "B", "C", "D", "E", "F", "G"] ); assert_eq!( removed.map(|(h, _, _)| h).collect::>(), ["H", "L", "M", "O", "I", "J", "K"] ); } #[test] fn find_node_index_works() { let (tree, is_descendent_of) = test_fork_tree(); let path = tree .find_node_index_where(&"D", &4, &is_descendent_of, &|_| true) .unwrap() .unwrap(); assert_eq!(path, [0, 0, 0]); let path = tree .find_node_index_where(&"O", &5, &is_descendent_of, &|_| true) .unwrap() .unwrap(); assert_eq!(path, [0, 1, 0, 0]); let path = tree .find_node_index_where(&"N", &6, &is_descendent_of, &|_| true) .unwrap() .unwrap(); assert_eq!(path, [0, 1, 0, 0, 0]); } #[test] fn find_node_index_with_predicate_works() { let is_descendent_of = |parent: &char, child: &char| match *parent { 'A' => Ok(['B', 'C', 'D', 'E', 'F'].contains(child)), 'B' => Ok(['C', 'D'].contains(child)), 'C' => Ok(['D'].contains(child)), 'E' => Ok(['F'].contains(child)), 'D' | 'F' => Ok(false), _ => Err(TestError), }; // A(t) --- B(f) --- C(t) --- D(f) // \-- E(t) --- F(f) let mut tree: ForkTree = ForkTree::new(); tree.import('A', 1, true, &is_descendent_of).unwrap(); tree.import('B', 2, false, &is_descendent_of).unwrap(); tree.import('C', 3, true, &is_descendent_of).unwrap(); tree.import('D', 4, false, &is_descendent_of).unwrap(); tree.import('E', 2, true, &is_descendent_of).unwrap(); tree.import('F', 3, false, &is_descendent_of).unwrap(); let path = tree .find_node_index_where(&'D', &4, &is_descendent_of, &|&value| !value) .unwrap() .unwrap(); assert_eq!(path, [0, 0]); let path = tree .find_node_index_where(&'D', &4, &is_descendent_of, &|&value| value) .unwrap() .unwrap(); assert_eq!(path, [0, 0, 0]); let path = tree .find_node_index_where(&'F', &3, &is_descendent_of, &|&value| !value) .unwrap(); assert_eq!(path, None); let path = tree .find_node_index_where(&'F', &3, &is_descendent_of, &|&value| value) .unwrap() .unwrap(); assert_eq!(path, [0, 1]); } #[test] fn find_node_works() { let (tree, is_descendent_of) = test_fork_tree(); let node = tree.find_node_where(&"B", &2, &is_descendent_of, &|_| true).unwrap().unwrap(); assert_eq!((node.hash, node.number), ("A", 1)); let node = tree.find_node_where(&"D", &4, &is_descendent_of, &|_| true).unwrap().unwrap(); assert_eq!((node.hash, node.number), ("C", 3)); let node = tree.find_node_where(&"O", &5, &is_descendent_of, &|_| true).unwrap().unwrap(); assert_eq!((node.hash, node.number), ("L", 4)); let node = tree.find_node_where(&"N", &6, &is_descendent_of, &|_| true).unwrap().unwrap(); assert_eq!((node.hash, node.number), ("M", 5)); } #[test] fn post_order_traversal_requirement() { let (mut tree, is_descendent_of) = test_fork_tree(); // Test for the post-order DFS traversal requirement as specified by the // `find_node_index_where` and `import` comments. let is_descendent_of_for_post_order = |parent: &&str, child: &&str| match *parent { "A" => Err(TestError), "K" if *child == "Z" => Ok(true), _ => is_descendent_of(parent, child), }; // Post order traversal requirement for `find_node_index_where` let path = tree .find_node_index_where(&"N", &6, &is_descendent_of_for_post_order, &|_| true) .unwrap() .unwrap(); assert_eq!(path, [0, 1, 0, 0, 0]); // Post order traversal requirement for `import` let res = tree.import(&"Z", 100, (), &is_descendent_of_for_post_order); assert_eq!(res, Ok(false)); assert_eq!( tree.iter().map(|node| *node.0).collect::>(), vec!["A", "B", "C", "D", "E", "F", "H", "L", "M", "O", "I", "G", "J", "K", "Z"], ); } }