Fix node lookup on fork-tree after a warp-sync

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));