Epoch-Changes tree pruning was lagging by one epoch (paritytech#12567)

* Remove all not required nodes from the epoch-changes tree

Some outdated nodes were left there because of the predicate

* Test to excercise the fix

* Add a fork on genesis to the test

* Fix typo in comments

client/consensus/babe/src/tests.rs

@@ -869,12 +869,12 @@ fn importing_epoch_change_block_prunes_tree() {
 
 	// Create and import the canon chain and keep track of fork blocks (A, C, D)
 	// from the diagram above.
-	let canon_hashes = propose_and_import_blocks_wrap(BlockId::Number(0), 30);
+	let canon = propose_and_import_blocks_wrap(BlockId::Number(0), 30);
 
 	// Create the forks
-	let fork_1 = propose_and_import_blocks_wrap(BlockId::Hash(canon_hashes[0]), 10);
-	let fork_2 = propose_and_import_blocks_wrap(BlockId::Hash(canon_hashes[12]), 15);
-	let fork_3 = propose_and_import_blocks_wrap(BlockId::Hash(canon_hashes[18]), 10);
+	let fork_1 = propose_and_import_blocks_wrap(BlockId::Hash(canon[0]), 10);
+	let fork_2 = propose_and_import_blocks_wrap(BlockId::Hash(canon[12]), 15);
+	let fork_3 = propose_and_import_blocks_wrap(BlockId::Hash(canon[18]), 10);
 
 	// We should be tracking a total of 9 epochs in the fork tree
 	assert_eq!(epoch_changes.shared_data().tree().iter().count(), 9);
@@ -884,51 +884,31 @@ fn importing_epoch_change_block_prunes_tree() {
 
 	// We finalize block #13 from the canon chain, so on the next epoch
 	// change the tree should be pruned, to not contain F (#7).
-	client.finalize_block(canon_hashes[12], None, false).unwrap();
+	client.finalize_block(canon[12], None, false).unwrap();
 	propose_and_import_blocks_wrap(BlockId::Hash(client.chain_info().best_hash), 7);
 
-	// at this point no hashes from the first fork must exist on the tree
-	assert!(!epoch_changes
-		.shared_data()
-		.tree()
-		.iter()
-		.map(|(h, _, _)| h)
-		.any(|h| fork_1.contains(h)),);
+	let nodes: Vec<_> = epoch_changes.shared_data().tree().iter().map(|(h, _, _)| *h).collect();
 
-	// but the epoch changes from the other forks must still exist
-	assert!(epoch_changes
-		.shared_data()
-		.tree()
-		.iter()
-		.map(|(h, _, _)| h)
-		.any(|h| fork_2.contains(h)));
+	// no hashes from the first fork must exist on the tree
+	assert!(!nodes.iter().any(|h| fork_1.contains(h)));
 
-	assert!(epoch_changes
-		.shared_data()
-		.tree()
-		.iter()
-		.map(|(h, _, _)| h)
-		.any(|h| fork_3.contains(h)),);
+	// but the epoch changes from the other forks must still exist
+	assert!(nodes.iter().any(|h| fork_2.contains(h)));
+	assert!(nodes.iter().any(|h| fork_3.contains(h)));
 
 	// finalizing block #25 from the canon chain should prune out the second fork
-	client.finalize_block(canon_hashes[24], None, false).unwrap();
+	client.finalize_block(canon[24], None, false).unwrap();
 	propose_and_import_blocks_wrap(BlockId::Hash(client.chain_info().best_hash), 8);
 
-	// at this point no hashes from the second fork must exist on the tree
-	assert!(!epoch_changes
-		.shared_data()
-		.tree()
-		.iter()
-		.map(|(h, _, _)| h)
-		.any(|h| fork_2.contains(h)),);
+	let nodes: Vec<_> = epoch_changes.shared_data().tree().iter().map(|(h, _, _)| *h).collect();
 
-	// while epoch changes from the last fork should still exist
-	assert!(epoch_changes
-		.shared_data()
-		.tree()
-		.iter()
-		.map(|(h, _, _)| h)
-		.any(|h| fork_3.contains(h)),);
+	// no hashes from the other forks must exist on the tree
+	assert!(!nodes.iter().any(|h| fork_2.contains(h)));
+	assert!(!nodes.iter().any(|h| fork_3.contains(h)));
+
+	// Check that we contain the nodes that we care about
+	assert!(nodes.iter().any(|h| *h == canon[18]));
+	assert!(nodes.iter().any(|h| *h == canon[24]));
 }
 
 #[test]

client/consensus/epochs/src/lib.rs

@@ -518,8 +518,8 @@ where
 		let is_descendent_of = descendent_of_builder.build_is_descendent_of(None);
 
 		let predicate = |epoch: &PersistedEpochHeader<E>| match *epoch {
-			PersistedEpochHeader::Genesis(_, ref epoch_1) => slot >= epoch_1.end_slot,
-			PersistedEpochHeader::Regular(ref epoch_n) => slot >= epoch_n.end_slot,
+			PersistedEpochHeader::Genesis(ref epoch_0, _) => epoch_0.start_slot <= slot,
+			PersistedEpochHeader::Regular(ref epoch_n) => epoch_n.start_slot <= slot,
 		};
 
 		// prune any epochs which could not be _live_ as of the children of the
@@ -777,11 +777,6 @@ where
 		}
 	}
 
-	/// Return the inner fork tree.
-	pub fn tree(&self) -> &ForkTree<Hash, Number, PersistedEpochHeader<E>> {
-		&self.inner
-	}
-
 	/// Reset to a specified pair of epochs, as if they were announced at blocks `parent_hash` and
 	/// `hash`.
 	pub fn reset(&mut self, parent_hash: Hash, hash: Hash, number: Number, current: E, next: E) {
@@ -832,6 +827,11 @@ where
 			self.epochs.remove(&(h, n));
 		});
 	}
+
+	/// Return the inner fork tree (mostly useful for testing)
+	pub fn tree(&self) -> &ForkTree<Hash, Number, PersistedEpochHeader<E>> {
+		&self.inner
+	}
 }
 
 /// Type alias to produce the epoch-changes tree from a block type.
@@ -1114,6 +1114,89 @@ mod tests {
 		}
 	}
 
+	#[test]
+	fn prune_removes_stale_nodes() {
+		//     +---D       +-------F
+		//     |           |
+		// 0---A---B--(x)--C--(y)--G
+		// |       |
+		// +---H   +-------E
+		//
+		//  Test parameters:
+		//  - epoch duration: 100
+		//
+		//  We are going to prune the tree at:
+		//  - 'x', a node between B and C
+		//  - 'y', a node between C and G
+
+		let is_descendent_of = |base: &Hash, block: &Hash| -> Result<bool, TestError> {
+			match (base, block) {
+				(b"0", _) => Ok(true),
+				(b"A", b) => Ok(b != b"0"),
+				(b"B", b) => Ok(b != b"0" && b != b"A" && b != b"D"),
+				(b"C", b) => Ok(b == b"F" || b == b"G" || b == b"y"),
+				(b"x", b) => Ok(b == b"C" || b == b"F" || b == b"G" || b == b"y"),
+				(b"y", b) => Ok(b == b"G"),
+				_ => Ok(false),
+			}
+		};
+
+		let mut epoch_changes = EpochChanges::new();
+
+		let mut import_at = |slot, hash: &Hash, number, parent_hash, parent_number| {
+			let make_genesis = |slot| Epoch { start_slot: slot, duration: 100 };
+			// Get epoch descriptor valid for 'slot'
+			let epoch_descriptor = epoch_changes
+				.epoch_descriptor_for_child_of(&is_descendent_of, parent_hash, parent_number, slot)
+				.unwrap()
+				.unwrap();
+			// Increment it
+			let next_epoch_desc = epoch_changes
+				.viable_epoch(&epoch_descriptor, &make_genesis)
+				.unwrap()
+				.increment(());
+			// Assign it to hash/number
+			epoch_changes
+				.import(&is_descendent_of, *hash, number, *parent_hash, next_epoch_desc)
+				.unwrap();
+		};
+
+		import_at(100, b"A", 10, b"0", 0);
+		import_at(200, b"B", 20, b"A", 10);
+		import_at(300, b"C", 30, b"B", 20);
+		import_at(200, b"D", 20, b"A", 10);
+		import_at(300, b"E", 30, b"B", 20);
+		import_at(400, b"F", 40, b"C", 30);
+		import_at(400, b"G", 40, b"C", 30);
+		import_at(100, b"H", 10, b"0", 0);
+
+		let mut nodes: Vec<_> = epoch_changes.tree().iter().map(|(h, _, _)| h).collect();
+		nodes.sort();
+		assert_eq!(nodes, vec![b"A", b"B", b"C", b"D", b"E", b"F", b"G", b"H"]);
+
+		// Finalize block 'x' @ number 25, slot 230
+		// This should prune all nodes imported by blocks with a number < 25 that are not
+		// ancestors of 'x' and all nodes before the one holding the epoch information
+		// to which 'x' belongs to (i.e. before A).
+
+		epoch_changes.prune_finalized(&is_descendent_of, b"x", 25, 230).unwrap();
+
+		let mut nodes: Vec<_> = epoch_changes.tree().iter().map(|(h, _, _)| h).collect();
+		nodes.sort();
+		assert_eq!(nodes, vec![b"A", b"B", b"C", b"E", b"F", b"G"]);
+
+		// Finalize block y @ number 35, slot 330
+		// This should prune all nodes imported by blocks with a number < 35 that are not
+		// ancestors of 'y' and all nodes before the one holding the epoch information
+		// to which 'y' belongs to (i.e. before B).
+
+		epoch_changes.prune_finalized(&is_descendent_of, b"y", 35, 330).unwrap();
+
+		let mut nodes: Vec<_> = epoch_changes.tree().iter().map(|(h, _, _)| h).collect();
+		nodes.sort();
+		assert_eq!(nodes, vec![b"B", b"C", b"F", b"G"]);
+	}
+
 	/// Test that ensures that the gap is not enabled when we import multiple genesis blocks.
 	#[test]
 	fn gap_is_not_enabled_when_multiple_genesis_epochs_are_imported() {