Prevent epochs pruning while on epoch 0

client/consensus/epochs/src/lib.rs

@@ -517,14 +517,16 @@ where
 	) -> Result<(), fork_tree::Error<D::Error>> {
 		let is_descendent_of = descendent_of_builder.build_is_descendent_of(None);
 
+		// Prevents removal of epoch0 data while we are in epoch1
+		// (these two epochs are kept in the same node)
 		let predicate = |epoch: &PersistedEpochHeader<E>| match *epoch {
-			PersistedEpochHeader::Genesis(ref epoch_0, _) => epoch_0.start_slot <= slot,
+			PersistedEpochHeader::Genesis(_, ref epoch_1) => epoch_1.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
+		// Prune any epochs which could not be _live_ as of the children of the
 		// finalized block, i.e. re-root the fork tree to the oldest ancestor of
-		// (hash, number) where epoch.end_slot() >= finalized_slot
+		// (hash, number) where `epoch.start_slot() <= finalized_slot`.
 		let removed = self.inner.prune(hash, &number, &is_descendent_of, &predicate)?;
 
 		for (hash, number, _) in removed {
@@ -1197,6 +1199,88 @@ mod tests {
 		assert_eq!(nodes, vec![b"B", b"C", b"F", b"G"]);
 	}
 
+	#[test]
+	fn near_genesis_prune_works() {
+		// [X]: announces next epoch change (i.e. adds a node in the epoch changes tree)
+		//
+		// 0--[A]--B--C--D--E--[G]--H--I--J--K--[L]
+		//                  +
+		//                  \--[F]
+
+		let is_descendent_of = |base: &Hash, block: &Hash| -> Result<bool, TestError> {
+			match (block, base) {
+				| (b"A", b"0") |
+				(b"B", b"0" | b"A") |
+				(b"C", b"0" | b"A" | b"B") |
+				(b"D", b"0" | b"A" | b"B" | b"C") |
+				(b"E", b"0" | b"A" | b"B" | b"C" | b"D") |
+				(b"F", b"0" | b"A" | b"B" | b"C" | b"D" | b"E") |
+				(b"G", b"0" | b"A" | b"B" | b"C" | b"D" | b"E") |
+				(b"H", b"0" | b"A" | b"B" | b"C" | b"D" | b"E" | b"G") |
+				(b"I", b"0" | b"A" | b"B" | b"C" | b"D" | b"E" | b"G" | b"H") |
+				(b"J", b"0" | b"A" | b"B" | b"C" | b"D" | b"E" | b"G" | b"H" | b"I") |
+				(b"K", b"0" | b"A" | b"B" | b"C" | b"D" | b"E" | b"G" | b"H" | b"I" | b"J") |
+				(
+					b"L",
+					b"0" | b"A" | b"B" | b"C" | b"D" | b"E" | b"G" | b"H" | b"I" | b"J" | b"K",
+				) => Ok(true),
+				_ => Ok(false),
+			}
+		};
+
+		let mut epoch_changes = EpochChanges::new();
+
+		let epoch = Epoch { start_slot: 278183811, duration: 5 };
+		let epoch = PersistedEpoch::Genesis(epoch.clone(), epoch.increment(()));
+
+		epoch_changes
+			.import(&is_descendent_of, *b"A", 1, Default::default(), IncrementedEpoch(epoch))
+			.unwrap();
+
+		let import_at = |epoch_changes: &mut EpochChanges<_, _, Epoch>,
+		                 slot,
+		                 hash: &Hash,
+		                 number,
+		                 parent_hash,
+		                 parent_number| {
+			let make_genesis = |slot| Epoch { start_slot: slot, duration: 5 };
+			// 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();
+		};
+
+		// Should not prune anything
+		epoch_changes.prune_finalized(&is_descendent_of, b"C", 3, 278183813).unwrap();
+
+		import_at(&mut epoch_changes, 278183816, b"G", 6, b"E", 5);
+		import_at(&mut epoch_changes, 278183816, b"F", 6, b"E", 5);
+
+		// Should not prune anything since we are on epoch0
+		epoch_changes.prune_finalized(&is_descendent_of, b"C", 3, 278183813).unwrap();
+		let mut list: Vec<_> = epoch_changes.inner.iter().map(|e| e.0).collect();
+		list.sort();
+		assert_eq!(list, vec![b"A", b"F", b"G"]);
+
+		import_at(&mut epoch_changes, 278183821, b"L", 11, b"K", 10);
+
+		// Should prune any fork of our ancestor not in the canonical chain (i.e. "F")
+		epoch_changes.prune_finalized(&is_descendent_of, b"J", 9, 278183819).unwrap();
+		let mut list: Vec<_> = epoch_changes.inner.iter().map(|e| e.0).collect();
+		list.sort();
+		assert_eq!(list, vec![b"A", b"G", b"L"]);
+	}
+
 	/// 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() {