From c96cdb58bc6a14c99ccf2b344212fde69767fc78 Mon Sep 17 00:00:00 2001
From: Davide Galassi <davxy@datawok.net>
Date: Fri, 27 May 2022 13:29:31 +0200
Subject: [PATCH] Fork-Tree import requires post-order DFS traversal (#11531)

* Fork-tree insert requires post-order dfs traversal

* Add dedicated test for methods requireing post-order traversal
---
 utils/fork-tree/src/lib.rs | 64 ++++++++++++++++++++++++++------------
 1 file changed, 44 insertions(+), 20 deletions(-)

diff --git a/utils/fork-tree/src/lib.rs b/utils/fork-tree/src/lib.rs
index 6f987fa798461..34c9b85f56e40 100644
--- a/utils/fork-tree/src/lib.rs
+++ b/utils/fork-tree/src/lib.rs
@@ -126,6 +126,11 @@ where
 	/// 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<F, E>(
 		&mut self,
 		hash: H,
@@ -143,29 +148,20 @@ where
 			}
 		}
 
-		let mut children = &mut self.roots;
-		let mut i = 0;
-		while i < children.len() {
-			let child = &children[i];
-			if child.hash == hash {
-				return Err(Error::Duplicate)
-			}
-			if child.number < number && is_descendent_of(&child.hash, &hash)? {
-				children = &mut children[i].children;
-				i = 0;
-			} else {
-				i += 1;
-			}
+		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)
 		}
 
-		let is_first = children.is_empty();
 		children.push(Node { data, hash, number, children: Default::default() });
 
-		// Quick way to check if the pushed node is a root
-		let is_root = children.as_ptr() == self.roots.as_ptr();
-
-		if is_first {
-			// Rebalance is required only if we've extended the branch depth.
+		if children.len() == 1 {
+			// Rebalance may be required only if we've extended the branch depth.
 			self.rebalance();
 		}
 
@@ -293,7 +289,7 @@ where
 	// 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.
+	// return a wrong result.
 	pub fn find_node_index_where<F, E, P>(
 		&self,
 		hash: &H,
@@ -1507,4 +1503,32 @@ mod test {
 		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<_>>(),
+			vec!["A", "B", "C", "D", "E", "F", "H", "L", "M", "O", "I", "G", "J", "K", "Z"],
+		);
+	}
 }