diff --git a/substrate/utils/fork-tree/src/lib.rs b/substrate/utils/fork-tree/src/lib.rs index 1aa085c3da..8ce0b729c6 100644 --- a/substrate/utils/fork-tree/src/lib.rs +++ b/substrate/utils/fork-tree/src/lib.rs @@ -200,7 +200,7 @@ impl ForkTree where data, hash: hash, number: number, - children: Vec::new(), + children: Vec::new(), }); self.rebalance(); @@ -232,10 +232,10 @@ impl ForkTree where number: &N, is_descendent_of: &F, predicate: &P, - ) -> Result>, Error> - where E: std::error::Error, - F: Fn(&H, &H) -> Result, - P: Fn(&V) -> bool, + ) -> Result>, Error> where + E: std::error::Error, + F: Fn(&H, &H) -> Result, + P: Fn(&V) -> bool, { // search for node starting from all roots for root in self.roots.iter() { @@ -250,6 +250,31 @@ impl ForkTree where Ok(None) } + /// Same as [`find_node_where`](Self::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, + { + // 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) + } + /// 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. @@ -609,16 +634,17 @@ mod node_implementation { /// 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). - // FIXME: it would be useful if this returned a mutable reference but - // rustc can't deal with lifetimes properly. an option would be to try - // an iterative definition instead. - pub fn find_node_where( + /// + /// 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( &self, hash: &H, number: &N, is_descendent_of: &F, predicate: &P, - ) -> Result>, Error> + ) -> Result>, Error> where E: std::error::Error, F: Fn(&H, &H) -> Result, P: Fn(&V) -> bool, @@ -631,11 +657,14 @@ mod node_implementation { let mut known_descendent_of = false; // continue depth-first search through all children - for node in self.children.iter() { + for (i, node) in self.children.iter().enumerate() { // found node, early exit - match node.find_node_where(hash, number, is_descendent_of, predicate)? { + match node.find_node_index_where(hash, number, is_descendent_of, predicate)? { FindOutcome::Abort => return Ok(FindOutcome::Abort), - FindOutcome::Found(x) => return Ok(FindOutcome::Found(x)), + 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, @@ -655,7 +684,7 @@ mod node_implementation { if is_descendent_of { // if the predicate passes we return the node if predicate(&self.data) { - return Ok(FindOutcome::Found(self)); + return Ok(FindOutcome::Found(Vec::new())); } } @@ -663,6 +692,70 @@ mod node_implementation { // 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( + &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 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( + &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 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)) + }, + } + } } }