diff --git a/utils/fork-tree/src/lib.rs b/utils/fork-tree/src/lib.rs index 45957f4390532..6f987fa798461 100644 --- a/utils/fork-tree/src/lib.rs +++ b/utils/fork-tree/src/lib.rs @@ -289,6 +289,11 @@ where /// 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. Here we are implementing a post-order DFS. pub fn find_node_index_where<F, E, P>( &self, hash: &H, @@ -301,30 +306,52 @@ where F: Fn(&H, &H) -> Result<bool, E>, P: Fn(&V) -> bool, { - let mut path = vec![]; - let mut children = &self.roots; - let mut i = 0; - let mut best_depth = 0; - - while i < children.len() { - let node = &children[i]; - if node.number < *number && is_descendent_of(&node.hash, hash)? { - path.push(i); - if predicate(&node.data) { - best_depth = path.len(); + 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 + } } - i = 0; - children = &node.children; - } else { - i += 1; } + + // 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 best_depth == 0 { - None - } else { - path.truncate(best_depth); + 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 }) } @@ -1468,6 +1495,9 @@ mod 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));