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

utils/fork-tree/src/lib.rs

@@ -289,6 +289,10 @@ 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. Babe epoch changes tree) currently silently
+	// rely on a **post-order DFS** traversal of the tree. In particular via the
+	// `is_descendent_of` that when used after a warp-sync may query the backend for a
+	// root that is not present.
 	pub fn find_node_index_where<F, E, P>(
 		&self,
 		hash: &H,
@@ -301,30 +305,49 @@ 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() {
+			// The second element in the tuple tracks what is the next children index
+			// to search into. Once all the children are processed we check the node.
+			// We stop searching into alternative children as soon as we have found
+			// ancestor of the node we're looking for
+			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));
+					stack.push((&node.children[i], 0));
+				} else if node.number < *number &&
+					(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 node we are looking for is found to be a descendent of the current
+			// root then we can stop searching under the other roots.
+			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 +1491,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));