epoch-changes: separate epoch header with epoch data (#4881)

* fork-tree: prune returns all pruned node data

* epoch-changes: split EpochHeader vs epoch data

* EpochChanges::viable_epoch and add missing comments

* Incoperate the new epoch_changes interface for BABE

* Fix BABE tests

* Fix fork-tree pruning issue

* Fix tests

* Fix pruning algorithm

* fork-tree: implement map function for mapping one value type to another

* Add migration script for new epoch changes scheme

* Update utils/fork-tree/src/lib.rs

Co-Authored-By: André Silva <andre.beat@gmail.com>

* Update client/consensus/slots/src/lib.rs

Co-Authored-By: André Silva <andre.beat@gmail.com>

* Remove authorities_len.is_none check, which is duplicate of unwrap_or(false)

* Update client/consensus/epochs/src/lib.rs

Co-Authored-By: André Silva <andre.beat@gmail.com>

* Update client/consensus/epochs/src/lib.rs

Co-Authored-By: André Silva <andre.beat@gmail.com>

* No trailing ; for return statement

* Use VERSION_KEY for migration

* Fix issues that removed nodes are not properly added into removed list

* Add comments indicating end_slot is non-inclusive

* fork-tree: use &mut F for map type declaration

* Add tests for v0 epoch_changes migration

* Fix babe RPC tests

Co-authored-by: André Silva <andre.beat@gmail.com>

substrate/bin/node/cli/src/service.rs

@@ -509,12 +509,11 @@ mod tests {
 					service.transaction_pool()
 				);
 
-				let epoch = babe_link.epoch_changes().lock().epoch_for_child_of(
+				let epoch_descriptor = babe_link.epoch_changes().lock().epoch_descriptor_for_child_of(
 					descendent_query(&*service.client()),
 					&parent_hash,
 					parent_number,
 					slot_num,
-					|slot| babe_link.config().genesis_epoch(slot)
 				).unwrap().unwrap();
 
 				let mut digest = Digest::<H256>::default();
@@ -564,7 +563,7 @@ mod tests {
 				params.body = Some(new_body);
 				params.intermediates.insert(
 					Cow::from(INTERMEDIATE_KEY),
-					Box::new(BabeIntermediate { epoch }) as Box<dyn Any>,
+					Box::new(BabeIntermediate::<Block> { epoch_descriptor }) as Box<dyn Any>,
 				);
 				params.fork_choice = Some(ForkChoiceStrategy::LongestChain);
 

substrate/client/consensus/aura/src/lib.rs

@@ -230,8 +230,8 @@ impl<B, C, E, I, P, Error, SO> sc_consensus_slots::SimpleSlotWorker<B> for AuraW
 		authorities(self.client.as_ref(), &BlockId::Hash(header.hash()))
 	}
 
-	fn authorities_len(&self, epoch_data: &Self::EpochData) -> usize {
-		epoch_data.len()
+	fn authorities_len(&self, epoch_data: &Self::EpochData) -> Option<usize> {
+		Some(epoch_data.len())
 	}
 
 	fn claim_slot(

substrate/client/consensus/babe/rpc/src/lib.rs

@@ -179,7 +179,7 @@ fn epoch_data<B, C, SC>(
 		SC: SelectChain<B>,
 {
 	let parent = select_chain.best_chain()?;
-	epoch_changes.lock().epoch_for_child_of(
+	epoch_changes.lock().epoch_data_for_child_of(
 		descendent_query(&**client),
 		&parent.hash(),
 		parent.number().clone(),
@@ -187,7 +187,6 @@ fn epoch_data<B, C, SC>(
 		|slot| babe_config.genesis_epoch(slot),
 	)
 		.map_err(|e| Error::Consensus(ConsensusError::ChainLookup(format!("{:?}", e))))?
-		.map(|e| e.into_inner())
 		.ok_or(Error::Consensus(ConsensusError::InvalidAuthoritiesSet))
 }
 

substrate/client/consensus/babe/src/aux_schema.rs

@@ -25,10 +25,12 @@ use sc_client_api::backend::AuxStore;
 use sp_blockchain::{Result as ClientResult, Error as ClientError};
 use sp_runtime::traits::Block as BlockT;
 use sp_consensus_babe::BabeBlockWeight;
-use sc_consensus_epochs::{EpochChangesFor, SharedEpochChanges};
+use sc_consensus_epochs::{EpochChangesFor, SharedEpochChanges, migration::EpochChangesForV0};
 use crate::Epoch;
 
-const BABE_EPOCH_CHANGES: &[u8] = b"babe_epoch_changes";
+const BABE_EPOCH_CHANGES_VERSION: &[u8] = b"babe_epoch_changes_version";
+const BABE_EPOCH_CHANGES_KEY: &[u8] = b"babe_epoch_changes";
+const BABE_EPOCH_CHANGES_CURRENT_VERSION: u32 = 1;
 
 fn block_weight_key<H: Encode>(block_hash: H) -> Vec<u8> {
 	(b"block_weight", block_hash).encode()
@@ -52,14 +54,30 @@ fn load_decode<B, T>(backend: &B, key: &[u8]) -> ClientResult<Option<T>>
 pub(crate) fn load_epoch_changes<Block: BlockT, B: AuxStore>(
 	backend: &B,
 ) -> ClientResult<SharedEpochChanges<Block, Epoch>> {
-	let epoch_changes = load_decode::<_, EpochChangesFor<Block, Epoch>>(backend, BABE_EPOCH_CHANGES)?
-		.map(|v| Arc::new(Mutex::new(v)))
-		.unwrap_or_else(|| {
-			info!(target: "babe",
-				"Creating empty BABE epoch changes on what appears to be first startup."
-			);
-			SharedEpochChanges::<Block, Epoch>::default()
-		});
+	let version = load_decode::<_, u32>(backend, BABE_EPOCH_CHANGES_VERSION)?;
+
+	let maybe_epoch_changes = match version {
+		None => load_decode::<_, EpochChangesForV0<Block, Epoch>>(
+			backend,
+			BABE_EPOCH_CHANGES_KEY,
+		)?.map(|v0| v0.migrate()),
+		Some(BABE_EPOCH_CHANGES_CURRENT_VERSION) => load_decode::<_, EpochChangesFor<Block, Epoch>>(
+			backend,
+			BABE_EPOCH_CHANGES_KEY,
+		)?,
+		Some(other) => {
+			return Err(ClientError::Backend(
+				format!("Unsupported BABE DB version: {:?}", other)
+			))
+		},
+	};
+
+	let epoch_changes = Arc::new(Mutex::new(maybe_epoch_changes.unwrap_or_else(|| {
+		info!(target: "babe",
+			  "Creating empty BABE epoch changes on what appears to be first startup."
+		);
+		EpochChangesFor::<Block, Epoch>::default()
+	})));
 
 	// rebalance the tree after deserialization. this isn't strictly necessary
 	// since the tree is now rebalanced on every update operation. but since the
@@ -77,10 +95,13 @@ pub(crate) fn write_epoch_changes<Block: BlockT, F, R>(
 ) -> R where
 	F: FnOnce(&[(&'static [u8], &[u8])]) -> R,
 {
-	let encoded_epoch_changes = epoch_changes.encode();
-	write_aux(
-		&[(BABE_EPOCH_CHANGES, encoded_epoch_changes.as_slice())],
-	)
+	BABE_EPOCH_CHANGES_CURRENT_VERSION.using_encoded(|version| {
+		let encoded_epoch_changes = epoch_changes.encode();
+		write_aux(
+			&[(BABE_EPOCH_CHANGES_KEY, encoded_epoch_changes.as_slice()),
+			  (BABE_EPOCH_CHANGES_VERSION, version)],
+		)
+	})
 }
 
 /// Write the cumulative chain-weight of a block ot aux storage.
@@ -91,7 +112,6 @@ pub(crate) fn write_block_weight<H: Encode, F, R>(
 ) -> R where
 	F: FnOnce(&[(Vec<u8>, &[u8])]) -> R,
 {
-
 	let key = block_weight_key(block_hash);
 	block_weight.using_encoded(|s|
 		write_aux(
@@ -107,3 +127,72 @@ pub(crate) fn load_block_weight<H: Encode, B: AuxStore>(
 ) -> ClientResult<Option<BabeBlockWeight>> {
 	load_decode(backend, block_weight_key(block_hash).as_slice())
 }
+
+#[cfg(test)]
+mod test {
+	use super::*;
+	use crate::Epoch;
+	use fork_tree::ForkTree;
+	use substrate_test_runtime_client;
+	use sp_core::H256;
+	use sp_runtime::traits::NumberFor;
+	use sc_consensus_epochs::{PersistedEpoch, PersistedEpochHeader, EpochHeader};
+	use sp_consensus::Error as ConsensusError;
+	use sc_network_test::Block as TestBlock;
+
+	#[test]
+	fn load_decode_from_v0_epoch_changes() {
+		let epoch = Epoch {
+			start_slot: 0,
+			authorities: vec![],
+			randomness: [0; 32],
+			epoch_index: 1,
+			duration: 100,
+		};
+		let client = substrate_test_runtime_client::new();
+		let mut v0_tree = ForkTree::<H256, NumberFor<TestBlock>, _>::new();
+		v0_tree.import::<_, ConsensusError>(
+			Default::default(),
+			Default::default(),
+			PersistedEpoch::Regular(epoch),
+			&|_, _| Ok(false), // Test is single item only so this can be set to false.
+		).unwrap();
+
+		client.insert_aux(
+			&[(BABE_EPOCH_CHANGES_KEY,
+			   &EpochChangesForV0::<TestBlock, Epoch>::from_raw(v0_tree).encode()[..])],
+			&[],
+		).unwrap();
+
+		assert_eq!(
+			load_decode::<_, u32>(&client, BABE_EPOCH_CHANGES_VERSION).unwrap(),
+			None,
+		);
+
+		let epoch_changes = load_epoch_changes::<TestBlock, _>(&client).unwrap();
+
+		assert!(
+			epoch_changes.lock()
+				.tree()
+				.iter()
+				.map(|(_, _, epoch)| epoch.clone())
+				.collect::<Vec<_>>() ==
+				vec![PersistedEpochHeader::Regular(EpochHeader {
+					start_slot: 0,
+					end_slot: 100,
+				})],
+		); // PersistedEpochHeader does not implement Debug, so we use assert! directly.
+
+		write_epoch_changes::<TestBlock, _, _>(
+			&epoch_changes.lock(),
+			|values| {
+				client.insert_aux(values, &[]).unwrap();
+			},
+		);
+
+		assert_eq!(
+			load_decode::<_, u32>(&client, BABE_EPOCH_CHANGES_VERSION).unwrap(),
+			Some(1),
+		);
+	}
+}

substrate/client/consensus/babe/src/lib.rs

@@ -78,7 +78,7 @@ use sp_runtime::{
 	generic::{BlockId, OpaqueDigestItemId}, Justification,
 	traits::{Block as BlockT, Header, DigestItemFor, Zero},
 };
-use sp_api::ProvideRuntimeApi;
+use sp_api::{ProvideRuntimeApi, NumberFor};
 use sc_keystore::KeyStorePtr;
 use parking_lot::Mutex;
 use sp_core::Pair;
@@ -104,7 +104,7 @@ use sc_consensus_slots::{
 	SlotWorker, SlotInfo, SlotCompatible, StorageChanges, CheckedHeader, check_equivocation,
 };
 use sc_consensus_epochs::{
-	descendent_query, ViableEpoch, SharedEpochChanges, EpochChangesFor, Epoch as EpochT
+	descendent_query, SharedEpochChanges, EpochChangesFor, Epoch as EpochT, ViableEpochDescriptor,
 };
 use sp_blockchain::{
 	Result as ClientResult, Error as ClientError,
@@ -231,9 +231,9 @@ macro_rules! babe_info {
 
 
 /// Intermediate value passed to block importer.
-pub struct BabeIntermediate {
-	/// The epoch data, if available.
-	pub epoch: ViableEpoch<Epoch>,
+pub struct BabeIntermediate<B: BlockT> {
+	/// The epoch descriptor.
+	pub epoch_descriptor: ViableEpochDescriptor<B::Hash, NumberFor<B>, Epoch>,
 }
 
 /// Intermediate key for Babe engine.
@@ -402,7 +402,7 @@ impl<B, C, E, I, Error, SO> sc_consensus_slots::SimpleSlotWorker<B> for BabeWork
 	SO: SyncOracle + Send + Clone,
 	Error: std::error::Error + Send + From<ConsensusError> + From<I::Error> + 'static,
 {
-	type EpochData = ViableEpoch<Epoch>;
+	type EpochData = ViableEpochDescriptor<B::Hash, NumberFor<B>, Epoch>;
 	type Claim = (PreDigest, AuthorityPair);
 	type SyncOracle = SO;
 	type CreateProposer = Pin<Box<
@@ -424,31 +424,35 @@ impl<B, C, E, I, Error, SO> sc_consensus_slots::SimpleSlotWorker<B> for BabeWork
 		parent: &B::Header,
 		slot_number: u64,
 	) -> Result<Self::EpochData, ConsensusError> {
-		self.epoch_changes.lock().epoch_for_child_of(
+		self.epoch_changes.lock().epoch_descriptor_for_child_of(
 			descendent_query(&*self.client),
 			&parent.hash(),
 			parent.number().clone(),
 			slot_number,
-			|slot| self.config.genesis_epoch(slot)
 		)
 			.map_err(|e| ConsensusError::ChainLookup(format!("{:?}", e)))?
 			.ok_or(sp_consensus::Error::InvalidAuthoritiesSet)
 	}
 
-	fn authorities_len(&self, epoch_data: &Self::EpochData) -> usize {
-		epoch_data.as_ref().authorities.len()
+	fn authorities_len(&self, epoch_descriptor: &Self::EpochData) -> Option<usize> {
+		self.epoch_changes.lock()
+			.viable_epoch(&epoch_descriptor, |slot| self.config.genesis_epoch(slot))
+			.map(|epoch| epoch.as_ref().authorities.len())
 	}
 
 	fn claim_slot(
 		&self,
 		_parent_header: &B::Header,
 		slot_number: SlotNumber,
-		epoch_data: &ViableEpoch<Epoch>,
+		epoch_descriptor: &ViableEpochDescriptor<B::Hash, NumberFor<B>, Epoch>,
 	) -> Option<Self::Claim> {
 		debug!(target: "babe", "Attempting to claim slot {}", slot_number);
 		let s = authorship::claim_slot(
 			slot_number,
-			epoch_data.as_ref(),
+			self.epoch_changes.lock().viable_epoch(
+				&epoch_descriptor,
+				|slot| self.config.genesis_epoch(slot)
+			)?.as_ref(),
 			&*self.config,
 			&self.keystore,
 		);
@@ -478,7 +482,7 @@ impl<B, C, E, I, Error, SO> sc_consensus_slots::SimpleSlotWorker<B> for BabeWork
 		Self::Claim,
 		Self::EpochData,
 	) -> sp_consensus::BlockImportParams<B, I::Transaction> + Send> {
-		Box::new(|header, header_hash, body, storage_changes, (_, pair), epoch| {
+		Box::new(|header, header_hash, body, storage_changes, (_, pair), epoch_descriptor| {
 			// sign the pre-sealed hash of the block and then
 			// add it to a digest item.
 			let signature = pair.sign(header_hash.as_ref());
@@ -490,7 +494,7 @@ impl<B, C, E, I, Error, SO> sc_consensus_slots::SimpleSlotWorker<B> for BabeWork
 			import_block.storage_changes = Some(storage_changes);
 			import_block.intermediates.insert(
 				Cow::from(INTERMEDIATE_KEY),
-				Box::new(BabeIntermediate { epoch }) as Box<dyn Any>,
+				Box::new(BabeIntermediate::<B> { epoch_descriptor }) as Box<dyn Any>,
 			);
 
 			import_block
@@ -729,18 +733,19 @@ impl<Block, Client> Verifier<Block> for BabeVerifier<Block, Client> where
 			.map_err(Error::<Block>::FetchParentHeader)?;
 
 		let pre_digest = find_pre_digest::<Block>(&header)?;
-		let epoch = {
-			let epoch_changes = self.epoch_changes.lock();
-			epoch_changes.epoch_for_child_of(
-				descendent_query(&*self.client),
-				&parent_hash,
-				parent_header_metadata.number,
-				pre_digest.slot_number(),
-				|slot| self.config.genesis_epoch(slot),
-			)
-				.map_err(|e| Error::<Block>::ForkTree(Box::new(e)))?
-				.ok_or_else(|| Error::<Block>::FetchEpoch(parent_hash))?
-		};
+		let epoch_changes = self.epoch_changes.lock();
+		let epoch_descriptor = epoch_changes.epoch_descriptor_for_child_of(
+			descendent_query(&*self.client),
+			&parent_hash,
+			parent_header_metadata.number,
+			pre_digest.slot_number(),
+		)
+			.map_err(|e| Error::<Block>::ForkTree(Box::new(e)))?
+			.ok_or_else(|| Error::<Block>::FetchEpoch(parent_hash))?;
+		let viable_epoch = epoch_changes.viable_epoch(
+			&epoch_descriptor,
+			|slot| self.config.genesis_epoch(slot)
+		).ok_or_else(|| Error::<Block>::FetchEpoch(parent_hash))?;
 
 		// We add one to the current slot to allow for some small drift.
 		// FIXME #1019 in the future, alter this queue to allow deferring of headers
@@ -748,7 +753,7 @@ impl<Block, Client> Verifier<Block> for BabeVerifier<Block, Client> where
 			header: header.clone(),
 			pre_digest: Some(pre_digest.clone()),
 			slot_now: slot_now + 1,
-			epoch: epoch.as_ref(),
+			epoch: viable_epoch.as_ref(),
 			config: &self.config,
 		};
 
@@ -808,7 +813,7 @@ impl<Block, Client> Verifier<Block> for BabeVerifier<Block, Client> where
 				import_block.justification = justification;
 				import_block.intermediates.insert(
 					Cow::from(INTERMEDIATE_KEY),
-					Box::new(BabeIntermediate { epoch }) as Box<dyn Any>,
+					Box::new(BabeIntermediate::<Block> { epoch_descriptor }) as Box<dyn Any>,
 				);
 				import_block.post_hash = Some(hash);
 
@@ -946,7 +951,7 @@ impl<Block, Client, Inner> BlockImport<Block> for BabeBlockImport<Block, Client,
 		// if this is the first block in its chain for that epoch.
 		//
 		// also provides the total weight of the chain, including the imported block.
-		let (epoch, first_in_epoch, parent_weight) = {
+		let (epoch_descriptor, first_in_epoch, parent_weight) = {
 			let parent_weight = if *parent_header.number() == Zero::zero() {
 				0
 			} else {
@@ -957,13 +962,13 @@ impl<Block, Client, Inner> BlockImport<Block> for BabeBlockImport<Block, Client,
 					))?
 			};
 
-			let intermediate = block.take_intermediate::<BabeIntermediate>(
+			let intermediate = block.take_intermediate::<BabeIntermediate<Block>>(
 				INTERMEDIATE_KEY
 			)?;
 
-			let epoch = intermediate.epoch;
-			let first_in_epoch = parent_slot < epoch.as_ref().start_slot;
-			(epoch, first_in_epoch, parent_weight)
+			let epoch_descriptor = intermediate.epoch_descriptor;
+			let first_in_epoch = parent_slot < epoch_descriptor.start_slot();
+			(epoch_descriptor, first_in_epoch, parent_weight)
 		};
 
 		let total_weight = parent_weight + pre_digest.added_weight();
@@ -994,12 +999,23 @@ impl<Block, Client, Inner> BlockImport<Block> for BabeBlockImport<Block, Client,
 		let info = self.client.info();
 
 		if let Some(next_epoch_descriptor) = next_epoch_digest {
-			let next_epoch = epoch.increment(next_epoch_descriptor);
-
 			old_epoch_changes = Some(epoch_changes.clone());
 
+			let viable_epoch = epoch_changes.viable_epoch(
+				&epoch_descriptor,
+				|slot| self.config.genesis_epoch(slot),
+			).ok_or_else(|| {
+				ConsensusError::ClientImport(Error::<Block>::FetchEpoch(parent_hash).into())
+			})?;
+
 			babe_info!("New epoch {} launching at block {} (block slot {} >= start slot {}).",
-				epoch.as_ref().epoch_index, hash, slot_number, epoch.as_ref().start_slot);
+					   viable_epoch.as_ref().epoch_index,
+					   hash,
+					   slot_number,
+					   viable_epoch.as_ref().start_slot);
+
+			let next_epoch = viable_epoch.increment(next_epoch_descriptor);
+
 			babe_info!("Next epoch starts at slot {}", next_epoch.as_ref().start_slot);
 
 			// prune the tree of epochs not part of the finalized chain or
@@ -1227,7 +1243,8 @@ pub mod test_helpers {
 			HeaderMetadata<B, Error = ClientError>,
 		C::Api: BabeApi<B>,
 	{
-		let epoch = link.epoch_changes.lock().epoch_for_child_of(
+		let epoch_changes = link.epoch_changes.lock();
+		let epoch = epoch_changes.epoch_data_for_child_of(
 			descendent_query(client),
 			&parent.hash(),
 			parent.number().clone(),
@@ -1237,7 +1254,7 @@ pub mod test_helpers {
 
 		authorship::claim_slot(
 			slot_number,
-			epoch.as_ref(),
+			&epoch,
 			&link.config,
 			keystore,
 		).map(|(digest, _)| digest)

substrate/client/consensus/babe/src/tests.rs

@@ -122,7 +122,7 @@ impl DummyProposer {
 		// figure out if we should add a consensus digest, since the test runtime
 		// doesn't.
 		let epoch_changes = self.factory.epoch_changes.lock();
-		let epoch = epoch_changes.epoch_for_child_of(
+		let epoch = epoch_changes.epoch_data_for_child_of(
 			descendent_query(&*self.factory.client),
 			&self.parent_hash,
 			self.parent_number,
@@ -130,8 +130,7 @@ impl DummyProposer {
 			|slot| self.factory.config.genesis_epoch(slot),
 		)
 			.expect("client has data to find epoch")
-			.expect("can compute epoch for baked block")
-			.into_inner();
+			.expect("can compute epoch for baked block");
 
 		let first_in_epoch = self.parent_slot < epoch.start_slot;
 		if first_in_epoch {
@@ -421,7 +420,7 @@ fn run_one_test(
 					panic!("Verification failed for {:?}: {}", h, e);
 				}
 			}
-	
+
 			Poll::<()>::Pending
 		}),
 		future::select(future::join_all(import_notifications), future::join_all(babe_futures))
@@ -566,12 +565,11 @@ fn propose_and_import_block<Transaction>(
 
 	let mut block = futures::executor::block_on(proposer.propose_with(pre_digest)).unwrap().block;
 
-	let epoch = proposer_factory.epoch_changes.lock().epoch_for_child_of(
+	let epoch_descriptor = proposer_factory.epoch_changes.lock().epoch_descriptor_for_child_of(
 		descendent_query(&*proposer_factory.client),
 		&parent_hash,
 		*parent.number(),
 		slot_number,
-		|slot| proposer_factory.config.genesis_epoch(slot)
 	).unwrap().unwrap();
 
 	let seal = {
@@ -595,7 +593,7 @@ fn propose_and_import_block<Transaction>(
 	import.body = Some(block.extrinsics);
 	import.intermediates.insert(
 		Cow::from(INTERMEDIATE_KEY),
-		Box::new(BabeIntermediate { epoch }) as Box<dyn Any>,
+		Box::new(BabeIntermediate::<TestBlock> { epoch_descriptor }) as Box<dyn Any>,
 	);
 	import.fork_choice = Some(ForkChoiceStrategy::LongestChain);
 	let import_result = block_import.import_block(import, Default::default()).unwrap();
@@ -637,13 +635,13 @@ fn importing_block_one_sets_genesis_epoch() {
 	let genesis_epoch = data.link.config.genesis_epoch(999);
 
 	let epoch_changes = data.link.epoch_changes.lock();
-	let epoch_for_second_block = epoch_changes.epoch_for_child_of(
+	let epoch_for_second_block = epoch_changes.epoch_data_for_child_of(
 		descendent_query(&*client),
 		&block_hash,
 		1,
 		1000,
 		|slot| data.link.config.genesis_epoch(slot),
-	).unwrap().unwrap().into_inner();
+	).unwrap().unwrap();
 
 	assert_eq!(epoch_for_second_block, genesis_epoch);
 }

substrate/client/consensus/epochs/src/lib.rs

@@ -16,7 +16,9 @@
 
 //! Generic utilities for epoch-based consensus engines.
 
-use std::{sync::Arc, ops::Add};
+pub mod migration;
+
+use std::{sync::Arc, ops::Add, collections::BTreeMap, borrow::{Borrow, BorrowMut}};
 use parking_lot::Mutex;
 use codec::{Encode, Decode};
 use fork_tree::ForkTree;
@@ -67,60 +69,126 @@ impl<'a, H, Block> IsDescendentOfBuilder<Block::Hash>
 }
 
 /// Epoch data, distinguish whether it is genesis or not.
+///
+/// Once an epoch is created, it must have a known `start_slot` and `end_slot`, which cannot be
+/// changed. Consensus engine may modify any other data in the epoch, if needed.
 pub trait Epoch {
 	/// Descriptor for the next epoch.
 	type NextEpochDescriptor;
 	/// Type of the slot number.
-	type SlotNumber: Ord;
-
-	/// Increment the epoch data, using the next epoch descriptor.
-	fn increment(&self, descriptor: Self::NextEpochDescriptor) -> Self;
+	type SlotNumber: Ord + Copy;
 
+	/// The starting slot of the epoch.
+	fn start_slot(&self) -> Self::SlotNumber;
 	/// Produce the "end slot" of the epoch. This is NOT inclusive to the epoch,
 	/// i.e. the slots covered by the epoch are `self.start_slot() .. self.end_slot()`.
 	fn end_slot(&self) -> Self::SlotNumber;
-	/// Produce the "start slot" of the epoch.
-	fn start_slot(&self) -> Self::SlotNumber;
+	/// Increment the epoch data, using the next epoch descriptor.
+	fn increment(&self, descriptor: Self::NextEpochDescriptor) -> Self;
 }
 
-/// An unimported genesis epoch.
-pub struct UnimportedGenesisEpoch<Epoch>(Epoch);
+impl<'a, E: Epoch> From<&'a E> for EpochHeader<E> {
+	fn from(epoch: &'a E) -> EpochHeader<E> {
+		Self {
+			start_slot: epoch.start_slot(),
+			end_slot: epoch.end_slot(),
+		}
+	}
+}
+
+/// Header of epoch data, consisting of start and end slot.
+#[derive(Eq, PartialEq, Encode, Decode, Debug)]
+pub struct EpochHeader<E: Epoch> {
+	/// The starting slot of the epoch.
+	pub start_slot: E::SlotNumber,
+	/// The end slot of the epoch. This is NOT inclusive to the epoch,
+	/// i.e. the slots covered by the epoch are `self.start_slot() .. self.end_slot()`.
+	pub end_slot: E::SlotNumber,
+}
+
+impl<E: Epoch> Clone for EpochHeader<E> {
+	fn clone(&self) -> Self {
+		Self {
+			start_slot: self.start_slot,
+			end_slot: self.end_slot,
+		}
+	}
+}
+
+/// Position of the epoch identifier.
+#[derive(PartialEq, Eq, PartialOrd, Ord, Copy, Clone, Debug)]
+pub enum EpochIdentifierPosition {
+	/// The identifier points to a genesis epoch `epoch_0`.
+	Genesis0,
+	/// The identifier points to a genesis epoch `epoch_1`.
+	Genesis1,
+	/// The identifier points to a regular epoch.
+	Regular,
+}
+
+/// Epoch identifier.
+#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Debug)]
+pub struct EpochIdentifier<Hash, Number> {
+	/// Location of the epoch.
+	pub position: EpochIdentifierPosition,
+	/// Hash of the block when the epoch is signaled.
+	pub hash: Hash,
+	/// Number of the block when the epoch is signaled.
+	pub number: Number,
+}
 
 /// The viable epoch under which a block can be verified.
 ///
 /// If this is the first non-genesis block in the chain, then it will
 /// hold an `UnimportedGenesis` epoch.
-pub enum ViableEpoch<Epoch> {
-	/// Genesis viable epoch data.
-	Genesis(UnimportedGenesisEpoch<Epoch>),
+pub enum ViableEpoch<E, ERef = E> {
+	/// Unimported genesis viable epoch data.
+	UnimportedGenesis(E),
 	/// Regular viable epoch data.
-	Regular(Epoch),
+	Signaled(ERef),
 }
 
-impl<Epoch> From<Epoch> for ViableEpoch<Epoch> {
-	fn from(epoch: Epoch) -> ViableEpoch<Epoch> {
-		ViableEpoch::Regular(epoch)
-	}
-}
-
-impl<Epoch> AsRef<Epoch> for ViableEpoch<Epoch> {
-	fn as_ref(&self) -> &Epoch {
+impl<E, ERef> AsRef<E> for ViableEpoch<E, ERef> where
+	ERef: Borrow<E>,
+{
+	fn as_ref(&self) -> &E {
 		match *self {
-			ViableEpoch::Genesis(UnimportedGenesisEpoch(ref e)) => e,
-			ViableEpoch::Regular(ref e) => e,
+			ViableEpoch::UnimportedGenesis(ref e) => e,
+			ViableEpoch::Signaled(ref e) => e.borrow(),
 		}
 	}
 }
 
-impl<Epoch> ViableEpoch<Epoch> where
-	Epoch: crate::Epoch + Clone,
+impl<E, ERef> AsMut<E> for ViableEpoch<E, ERef> where
+	ERef: BorrowMut<E>,
+{
+	fn as_mut(&mut self) -> &mut E {
+		match *self {
+			ViableEpoch::UnimportedGenesis(ref mut e) => e,
+			ViableEpoch::Signaled(ref mut e) => e.borrow_mut(),
+		}
+	}
+}
+
+impl<E, ERef> ViableEpoch<E, ERef> where
+	E: Epoch + Clone,
+	ERef: Borrow<E>,
 {
 	/// Extract the underlying epoch, disregarding the fact that a genesis
 	/// epoch may be unimported.
-	pub fn into_inner(self) -> Epoch {
+	pub fn into_cloned_inner(self) -> E {
 		match self {
-			ViableEpoch::Genesis(UnimportedGenesisEpoch(e)) => e,
-			ViableEpoch::Regular(e) => e,
+			ViableEpoch::UnimportedGenesis(e) => e,
+			ViableEpoch::Signaled(e) => e.borrow().clone(),
+		}
+	}
+
+	/// Get cloned value for the viable epoch.
+	pub fn into_cloned(self) -> ViableEpoch<E, E> {
+		match self {
+			ViableEpoch::UnimportedGenesis(e) =>
+				ViableEpoch::UnimportedGenesis(e),
+			ViableEpoch::Signaled(e) => ViableEpoch::Signaled(e.borrow().clone()),
 		}
 	}
 
@@ -128,36 +196,84 @@ impl<Epoch> ViableEpoch<Epoch> where
 	/// into the fork-tree.
 	pub fn increment(
 		&self,
-		next_descriptor: Epoch::NextEpochDescriptor
-	) -> IncrementedEpoch<Epoch> {
+		next_descriptor: E::NextEpochDescriptor
+	) -> IncrementedEpoch<E> {
 		let next = self.as_ref().increment(next_descriptor);
 		let to_persist = match *self {
-			ViableEpoch::Genesis(UnimportedGenesisEpoch(ref epoch_0)) =>
+			ViableEpoch::UnimportedGenesis(ref epoch_0) =>
 				PersistedEpoch::Genesis(epoch_0.clone(), next),
-			ViableEpoch::Regular(_) => PersistedEpoch::Regular(next),
+			ViableEpoch::Signaled(_) => PersistedEpoch::Regular(next),
 		};
 
 		IncrementedEpoch(to_persist)
 	}
 }
 
-/// The data type encoded on disk.
-#[derive(Clone, Encode, Decode)]
-pub enum PersistedEpoch<Epoch> {
+/// Descriptor for a viable epoch.
+#[derive(PartialEq, Eq, Clone, Debug)]
+pub enum ViableEpochDescriptor<Hash, Number, E: Epoch> {
+	/// The epoch is an unimported genesis, with given start slot number.
+	UnimportedGenesis(E::SlotNumber),
+	/// The epoch is signaled and has been imported, with given identifier and header.
+	Signaled(EpochIdentifier<Hash, Number>, EpochHeader<E>)
+}
+
+impl<Hash, Number, E: Epoch> ViableEpochDescriptor<Hash, Number, E> {
+	/// Start slot of the descriptor.
+	pub fn start_slot(&self) -> E::SlotNumber {
+		match self {
+			Self::UnimportedGenesis(start_slot) => *start_slot,
+			Self::Signaled(_, header) => header.start_slot,
+		}
+	}
+}
+
+/// Persisted epoch stored in EpochChanges.
+#[derive(Clone, Encode, Decode, Debug)]
+pub enum PersistedEpoch<E: Epoch> {
 	/// Genesis persisted epoch data. epoch_0, epoch_1.
-	Genesis(Epoch, Epoch),
+	Genesis(E, E),
 	/// Regular persisted epoch data. epoch_n.
-	Regular(Epoch),
+	Regular(E),
+}
+
+impl<'a, E: Epoch> From<&'a PersistedEpoch<E>> for PersistedEpochHeader<E> {
+	fn from(epoch: &'a PersistedEpoch<E>) -> Self {
+		match epoch {
+			PersistedEpoch::Genesis(ref epoch_0, ref epoch_1) =>
+				PersistedEpochHeader::Genesis(epoch_0.into(), epoch_1.into()),
+			PersistedEpoch::Regular(ref epoch_n) =>
+				PersistedEpochHeader::Regular(epoch_n.into()),
+		}
+	}
+}
+
+/// Persisted epoch header stored in ForkTree.
+#[derive(Encode, Decode, PartialEq, Eq)]
+pub enum PersistedEpochHeader<E: Epoch> {
+	/// Genesis persisted epoch header. epoch_0, epoch_1.
+	Genesis(EpochHeader<E>, EpochHeader<E>),
+	/// Regular persisted epoch header. epoch_n.
+	Regular(EpochHeader<E>),
+}
+
+impl<E: Epoch> Clone for PersistedEpochHeader<E> {
+	fn clone(&self) -> Self {
+		match self {
+			Self::Genesis(epoch_0, epoch_1) => Self::Genesis(epoch_0.clone(), epoch_1.clone()),
+			Self::Regular(epoch_n) => Self::Regular(epoch_n.clone()),
+		}
+	}
 }
 
 /// A fresh, incremented epoch to import into the underlying fork-tree.
 ///
 /// Create this with `ViableEpoch::increment`.
 #[must_use = "Freshly-incremented epoch must be imported with `EpochChanges::import`"]
-pub struct IncrementedEpoch<Epoch>(PersistedEpoch<Epoch>);
+pub struct IncrementedEpoch<E: Epoch>(PersistedEpoch<E>);
 
-impl<Epoch> AsRef<Epoch> for IncrementedEpoch<Epoch> {
-	fn as_ref(&self) -> &Epoch {
+impl<E: Epoch> AsRef<E> for IncrementedEpoch<E> {
+	fn as_ref(&self) -> &E {
 		match self.0 {
 			PersistedEpoch::Genesis(_, ref epoch_1) => epoch_1,
 			PersistedEpoch::Regular(ref epoch_n) => epoch_n,
@@ -181,8 +297,9 @@ impl<Epoch> AsRef<Epoch> for IncrementedEpoch<Epoch> {
 ///
 /// Further epochs (epoch_2, ..., epoch_n) each get their own entry.
 #[derive(Clone, Encode, Decode)]
-pub struct EpochChanges<Hash, Number, Epoch> {
-	inner: ForkTree<Hash, Number, PersistedEpoch<Epoch>>,
+pub struct EpochChanges<Hash, Number, E: Epoch> {
+	inner: ForkTree<Hash, Number, PersistedEpochHeader<E>>,
+	epochs: BTreeMap<(Hash, Number), PersistedEpoch<E>>,
 }
 
 // create a fake header hash which hasn't been included in the chain.
@@ -194,19 +311,18 @@ fn fake_head_hash<H: AsRef<[u8]> + AsMut<[u8]> + Clone>(parent_hash: &H) -> H {
 	h
 }
 
-impl<Hash, Number, Epoch> Default for EpochChanges<Hash, Number, Epoch> where
-	Hash: PartialEq,
+impl<Hash, Number, E: Epoch> Default for EpochChanges<Hash, Number, E> where
+	Hash: PartialEq + Ord,
 	Number: Ord,
 {
 	fn default() -> Self {
-		EpochChanges { inner: ForkTree::new() }
+		EpochChanges { inner: ForkTree::new(), epochs: BTreeMap::new() }
 	}
 }
 
-impl<Hash, Number, Epoch> EpochChanges<Hash, Number, Epoch> where
-	Hash: PartialEq + AsRef<[u8]> + AsMut<[u8]> + Copy,
+impl<Hash, Number, E: Epoch> EpochChanges<Hash, Number, E> where
+	Hash: PartialEq + Ord + AsRef<[u8]> + AsMut<[u8]> + Copy,
 	Number: Ord + One + Zero + Add<Output=Number> + Copy,
-	Epoch: crate::Epoch + Clone,
 {
 	/// Create a new epoch change.
 	pub fn new() -> Self {
@@ -227,45 +343,162 @@ impl<Hash, Number, Epoch> EpochChanges<Hash, Number, Epoch> where
 		descendent_of_builder: D,
 		hash: &Hash,
 		number: Number,
-		slot: Epoch::SlotNumber,
+		slot: E::SlotNumber,
 	) -> Result<(), fork_tree::Error<D::Error>> {
 		let is_descendent_of = descendent_of_builder
 			.build_is_descendent_of(None);
 
-		let predicate = |epoch: &PersistedEpoch<Epoch>| match *epoch {
-			PersistedEpoch::Genesis(_, ref epoch_1) =>
-				slot >= epoch_1.end_slot(),
-			PersistedEpoch::Regular(ref epoch_n) =>
-				slot >= epoch_n.end_slot(),
+		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,
 		};
 
 		// 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
-		self.inner.prune(
+		let removed = self.inner.prune(
 			hash,
 			&number,
 			&is_descendent_of,
 			&predicate,
 		)?;
 
+		for (hash, number, _) in removed {
+			self.epochs.remove(&(hash, number));
+		}
+
 		Ok(())
 	}
 
+	/// Get a reference to an epoch with given identifier.
+	pub fn epoch(&self, id: &EpochIdentifier<Hash, Number>) -> Option<&E> {
+		self.epochs.get(&(id.hash, id.number))
+			.and_then(|v| {
+				match v {
+					PersistedEpoch::Genesis(ref epoch_0, _)
+						if id.position == EpochIdentifierPosition::Genesis0 => Some(epoch_0),
+					PersistedEpoch::Genesis(_, ref epoch_1)
+						if id.position == EpochIdentifierPosition::Genesis1 => Some(epoch_1),
+					PersistedEpoch::Regular(ref epoch_n)
+						if id.position == EpochIdentifierPosition::Regular => Some(epoch_n),
+					_ => None,
+				}
+			})
+	}
+
+	/// Get a reference to a viable epoch with given descriptor.
+	pub fn viable_epoch<G>(
+		&self,
+		descriptor: &ViableEpochDescriptor<Hash, Number, E>,
+		make_genesis: G,
+	) -> Option<ViableEpoch<E, &E>> where
+		G: FnOnce(E::SlotNumber) -> E
+	{
+		match descriptor {
+			ViableEpochDescriptor::UnimportedGenesis(slot_number) => {
+				Some(ViableEpoch::UnimportedGenesis(make_genesis(*slot_number)))
+			},
+			ViableEpochDescriptor::Signaled(identifier, _) => {
+				self.epoch(&identifier).map(ViableEpoch::Signaled)
+			},
+		}
+	}
+
+	/// Get a mutable reference to an epoch with given identifier.
+	pub fn epoch_mut(&mut self, id: &EpochIdentifier<Hash, Number>) -> Option<&mut E> {
+		self.epochs.get_mut(&(id.hash, id.number))
+			.and_then(|v| {
+				match v {
+					PersistedEpoch::Genesis(ref mut epoch_0, _)
+						if id.position == EpochIdentifierPosition::Genesis0 => Some(epoch_0),
+					PersistedEpoch::Genesis(_, ref mut epoch_1)
+						if id.position == EpochIdentifierPosition::Genesis1 => Some(epoch_1),
+					PersistedEpoch::Regular(ref mut epoch_n)
+						if id.position == EpochIdentifierPosition::Regular => Some(epoch_n),
+					_ => None,
+				}
+			})
+	}
+
+	/// Get a mutable reference to a viable epoch with given descriptor.
+	pub fn viable_epoch_mut<G>(
+		&mut self,
+		descriptor: &ViableEpochDescriptor<Hash, Number, E>,
+		make_genesis: G,
+	) -> Option<ViableEpoch<E, &mut E>> where
+		G: FnOnce(E::SlotNumber) -> E
+	{
+		match descriptor {
+			ViableEpochDescriptor::UnimportedGenesis(slot_number) => {
+				Some(ViableEpoch::UnimportedGenesis(make_genesis(*slot_number)))
+			},
+			ViableEpochDescriptor::Signaled(identifier, _) => {
+				self.epoch_mut(&identifier).map(ViableEpoch::Signaled)
+			},
+		}
+	}
+
+	/// Get the epoch data from an epoch descriptor.
+	///
+	/// Note that this function ignores the fact that an genesis epoch might need to be imported.
+	/// Mostly useful for testing.
+	pub fn epoch_data<G>(
+		&self,
+		descriptor: &ViableEpochDescriptor<Hash, Number, E>,
+		make_genesis: G
+	) -> Option<E> where
+		G: FnOnce(E::SlotNumber) -> E,
+		E: Clone,
+	{
+		match descriptor {
+			ViableEpochDescriptor::UnimportedGenesis(slot_number) => {
+				Some(make_genesis(*slot_number))
+			},
+			ViableEpochDescriptor::Signaled(identifier, _) => {
+				self.epoch(&identifier).cloned()
+			},
+		}
+	}
+
+	/// Finds the epoch data for a child of the given block. Similar to
+	/// `epoch_descriptor_for_child_of` but returns the full data.
+	///
+	/// Note that this function ignores the fact that an genesis epoch might need to be imported.
+	/// Mostly useful for testing.
+	pub fn epoch_data_for_child_of<D: IsDescendentOfBuilder<Hash>, G>(
+		&self,
+		descendent_of_builder: D,
+		parent_hash: &Hash,
+		parent_number: Number,
+		slot_number: E::SlotNumber,
+		make_genesis: G,
+	) -> Result<Option<E>, fork_tree::Error<D::Error>> where
+		G: FnOnce(E::SlotNumber) -> E,
+		E: Clone,
+	{
+		let descriptor = self.epoch_descriptor_for_child_of(
+			descendent_of_builder,
+			parent_hash,
+			parent_number,
+			slot_number
+		)?;
+
+		Ok(descriptor.and_then(|des| self.epoch_data(&des, make_genesis)))
+	}
+
 	/// Finds the epoch for a child of the given block, assuming the given slot number.
 	///
 	/// If the returned epoch is an `UnimportedGenesis` epoch, it should be imported into the
 	/// tree.
-	pub fn epoch_for_child_of<D: IsDescendentOfBuilder<Hash>, G>(
+	pub fn epoch_descriptor_for_child_of<D: IsDescendentOfBuilder<Hash>>(
 		&self,
 		descendent_of_builder: D,
 		parent_hash: &Hash,
 		parent_number: Number,
-		slot_number: Epoch::SlotNumber,
-		make_genesis: G,
-	) -> Result<Option<ViableEpoch<Epoch>>, fork_tree::Error<D::Error>>
-		where G: FnOnce(Epoch::SlotNumber) -> Epoch
-	{
+		slot_number: E::SlotNumber,
+	) -> Result<Option<ViableEpochDescriptor<Hash, Number, E>>, fork_tree::Error<D::Error>> {
 		// find_node_where will give you the node in the fork-tree which is an ancestor
 		// of the `parent_hash` by default. if the last epoch was signalled at the parent_hash,
 		// then it won't be returned. we need to create a new fake chain head hash which
@@ -277,8 +510,7 @@ impl<Hash, Number, Epoch> EpochChanges<Hash, Number, Epoch> where
 
 		if parent_number == Zero::zero() {
 			// need to insert the genesis epoch.
-			let genesis_epoch = make_genesis(slot_number);
-			return Ok(Some(ViableEpoch::Genesis(UnimportedGenesisEpoch(genesis_epoch))));
+			return Ok(Some(ViableEpochDescriptor::UnimportedGenesis(slot_number)))
 		}
 
 		// We want to find the deepest node in the tree which is an ancestor
@@ -286,11 +518,11 @@ impl<Hash, Number, Epoch> EpochChanges<Hash, Number, Epoch> where
 		// slot of our block. The genesis special-case doesn't need to look
 		// at epoch_1 -- all we're doing here is figuring out which node
 		// we need.
-		let predicate = |epoch: &PersistedEpoch<Epoch>| match *epoch {
-			PersistedEpoch::Genesis(ref epoch_0, _) =>
-				epoch_0.start_slot() <= slot_number,
-			PersistedEpoch::Regular(ref epoch_n) =>
-				epoch_n.start_slot() <= slot_number,
+		let predicate = |epoch: &PersistedEpochHeader<E>| match *epoch {
+			PersistedEpochHeader::Genesis(ref epoch_0, _) =>
+				epoch_0.start_slot <= slot_number,
+			PersistedEpochHeader::Regular(ref epoch_n) =>
+				epoch_n.start_slot <= slot_number,
 		};
 
 		self.inner.find_node_where(
@@ -299,18 +531,27 @@ impl<Hash, Number, Epoch> EpochChanges<Hash, Number, Epoch> where
 			&is_descendent_of,
 			&predicate,
 		)
-			.map(|n| n.map(|node| ViableEpoch::Regular(match node.data {
-				// Ok, we found our node.
-				// and here we figure out which of the internal epochs
-				// of a genesis node to use based on their start slot.
-				PersistedEpoch::Genesis(ref epoch_0, ref epoch_1) =>
-					if epoch_1.start_slot() <= slot_number {
-						epoch_1.clone()
-					} else {
-						epoch_0.clone()
-					},
-				PersistedEpoch::Regular(ref epoch_n) => epoch_n.clone(),
-			})))
+			.map(|n| {
+				n.map(|node| (match node.data {
+					// Ok, we found our node.
+					// and here we figure out which of the internal epochs
+					// of a genesis node to use based on their start slot.
+					PersistedEpochHeader::Genesis(ref epoch_0, ref epoch_1) =>
+						if epoch_1.start_slot <= slot_number {
+							(EpochIdentifierPosition::Genesis1, epoch_1.clone())
+						} else {
+							(EpochIdentifierPosition::Genesis0, epoch_0.clone())
+						},
+					PersistedEpochHeader::Regular(ref epoch_n) =>
+						(EpochIdentifierPosition::Regular, epoch_n.clone()),
+				}, node)).map(|((position, header), node)| {
+					ViableEpochDescriptor::Signaled(EpochIdentifier {
+						position,
+						hash: node.hash,
+						number: node.number
+					}, header)
+				})
+			})
 	}
 
 	/// Import a new epoch-change, signalled at the given block.
@@ -324,26 +565,30 @@ impl<Hash, Number, Epoch> EpochChanges<Hash, Number, Epoch> where
 		hash: Hash,
 		number: Number,
 		parent_hash: Hash,
-		epoch: IncrementedEpoch<Epoch>,
+		epoch: IncrementedEpoch<E>,
 	) -> Result<(), fork_tree::Error<D::Error>> {
 		let is_descendent_of = descendent_of_builder
 			.build_is_descendent_of(Some((hash, parent_hash)));
+		let header = PersistedEpochHeader::<E>::from(&epoch.0);
 
 		let res = self.inner.import(
 			hash,
 			number,
-			epoch.0,
+			header,
 			&is_descendent_of,
 		);
 
 		match res {
-			Ok(_) | Err(fork_tree::Error::Duplicate) => Ok(()),
+			Ok(_) | Err(fork_tree::Error::Duplicate) => {
+				self.epochs.insert((hash, number), epoch.0);
+				Ok(())
+			},
 			Err(e) => Err(e),
 		}
 	}
 
 	/// Return the inner fork tree.
-	pub fn tree(&self) -> &ForkTree<Hash, Number, PersistedEpoch<Epoch>> {
+	pub fn tree(&self) -> &ForkTree<Hash, Number, PersistedEpochHeader<E>> {
 		&self.inner
 	}
 }
@@ -443,39 +688,34 @@ mod tests {
 			}
 		};
 
-		let make_genesis = |slot| Epoch {
-			start_slot: slot,
-			duration: 100,
-		};
-
-		let epoch_changes = EpochChanges::new();
-		let genesis_epoch = epoch_changes.epoch_for_child_of(
+		let epoch_changes = EpochChanges::<_, _, Epoch>::new();
+		let genesis_epoch = epoch_changes.epoch_descriptor_for_child_of(
 			&is_descendent_of,
 			b"0",
 			0,
 			10101,
-			&make_genesis,
 		).unwrap().unwrap();
 
 		match genesis_epoch {
-			ViableEpoch::Genesis(_) => {},
+			ViableEpochDescriptor::UnimportedGenesis(slot_number) => {
+				assert_eq!(slot_number, 10101u64);
+			},
 			_ => panic!("should be unimported genesis"),
 		};
-		assert_eq!(genesis_epoch.as_ref(), &make_genesis(10101));
 
-		let genesis_epoch_2 = epoch_changes.epoch_for_child_of(
+		let genesis_epoch_2 = epoch_changes.epoch_descriptor_for_child_of(
 			&is_descendent_of,
 			b"0",
 			0,
 			10102,
-			&make_genesis,
 		).unwrap().unwrap();
 
 		match genesis_epoch_2 {
-			ViableEpoch::Genesis(_) => {},
+			ViableEpochDescriptor::UnimportedGenesis(slot_number) => {
+				assert_eq!(slot_number, 10102u64);
+			},
 			_ => panic!("should be unimported genesis"),
 		};
-		assert_eq!(genesis_epoch_2.as_ref(), &make_genesis(10102));
 	}
 
 	#[test]
@@ -499,18 +739,20 @@ mod tests {
 			duration: 100,
 		};
 
-		let mut epoch_changes = EpochChanges::new();
-		let genesis_epoch = epoch_changes.epoch_for_child_of(
+		let mut epoch_changes = EpochChanges::<_, _, Epoch>::new();
+		let genesis_epoch = epoch_changes.epoch_descriptor_for_child_of(
 			&is_descendent_of,
 			b"0",
 			0,
 			100,
-			&make_genesis,
 		).unwrap().unwrap();
 
-		assert_eq!(genesis_epoch.as_ref(), &make_genesis(100));
+		assert_eq!(genesis_epoch, ViableEpochDescriptor::UnimportedGenesis(100));
 
-		let import_epoch_1 = genesis_epoch.increment(());
+		let import_epoch_1 = epoch_changes
+			.viable_epoch(&genesis_epoch, &make_genesis)
+			.unwrap()
+			.increment(());
 		let epoch_1 = import_epoch_1.as_ref().clone();
 
 		epoch_changes.import(
@@ -520,7 +762,7 @@ mod tests {
 			*b"0",
 			import_epoch_1,
 		).unwrap();
-		let genesis_epoch = genesis_epoch.into_inner();
+		let genesis_epoch = epoch_changes.epoch_data(&genesis_epoch, &make_genesis).unwrap();
 
 		assert!(is_descendent_of(b"0", b"A").unwrap());
 
@@ -529,13 +771,13 @@ mod tests {
 
 		{
 			// x is still within the genesis epoch.
-			let x = epoch_changes.epoch_for_child_of(
+			let x = epoch_changes.epoch_data_for_child_of(
 				&is_descendent_of,
 				b"A",
 				1,
 				end_slot - 1,
 				&make_genesis,
-			).unwrap().unwrap().into_inner();
+			).unwrap().unwrap();
 
 			assert_eq!(x, genesis_epoch);
 		}
@@ -543,13 +785,13 @@ mod tests {
 		{
 			// x is now at the next epoch, because the block is now at the
 			// start slot of epoch 1.
-			let x = epoch_changes.epoch_for_child_of(
+			let x = epoch_changes.epoch_data_for_child_of(
 				&is_descendent_of,
 				b"A",
 				1,
 				end_slot,
 				&make_genesis,
-			).unwrap().unwrap().into_inner();
+			).unwrap().unwrap();
 
 			assert_eq!(x, epoch_1);
 		}
@@ -557,13 +799,13 @@ mod tests {
 		{
 			// x is now at the next epoch, because the block is now after
 			// start slot of epoch 1.
-			let x = epoch_changes.epoch_for_child_of(
+			let x = epoch_changes.epoch_data_for_child_of(
 				&is_descendent_of,
 				b"A",
 				1,
 				epoch_1.end_slot() - 1,
 				&make_genesis,
-			).unwrap().unwrap().into_inner();
+			).unwrap().unwrap();
 
 			assert_eq!(x, epoch_1);
 		}
@@ -596,47 +838,54 @@ mod tests {
 
 		// insert genesis epoch for A
 		{
-			let genesis_epoch_a = epoch_changes.epoch_for_child_of(
+			let genesis_epoch_a_descriptor = epoch_changes.epoch_descriptor_for_child_of(
 				&is_descendent_of,
 				b"0",
 				0,
 				100,
-				&make_genesis,
 			).unwrap().unwrap();
 
+			let incremented_epoch = epoch_changes
+				.viable_epoch(&genesis_epoch_a_descriptor, &make_genesis)
+				.unwrap()
+				.increment(next_descriptor.clone());
+
 			epoch_changes.import(
 				&is_descendent_of,
 				*b"A",
 				1,
 				*b"0",
-				genesis_epoch_a.increment(next_descriptor.clone()),
+				incremented_epoch,
 			).unwrap();
-
 		}
 
 		// insert genesis epoch for X
 		{
-			let genesis_epoch_x = epoch_changes.epoch_for_child_of(
+			let genesis_epoch_x_descriptor = epoch_changes.epoch_descriptor_for_child_of(
 				&is_descendent_of,
 				b"0",
 				0,
 				1000,
-				&make_genesis,
 			).unwrap().unwrap();
 
+			let incremented_epoch = epoch_changes
+				.viable_epoch(&genesis_epoch_x_descriptor, &make_genesis)
+				.unwrap()
+				.increment(next_descriptor.clone());
+
 			epoch_changes.import(
 				&is_descendent_of,
 				*b"X",
 				1,
 				*b"0",
-				genesis_epoch_x.increment(next_descriptor.clone()),
+				incremented_epoch,
 			).unwrap();
 		}
 
 		// now check that the genesis epochs for our respective block 1s
 		// respect the chain structure.
 		{
-			let epoch_for_a_child = epoch_changes.epoch_for_child_of(
+			let epoch_for_a_child = epoch_changes.epoch_data_for_child_of(
 				&is_descendent_of,
 				b"A",
 				1,
@@ -644,9 +893,9 @@ mod tests {
 				&make_genesis,
 			).unwrap().unwrap();
 
-			assert_eq!(epoch_for_a_child.into_inner(), make_genesis(100));
+			assert_eq!(epoch_for_a_child, make_genesis(100));
 
-			let epoch_for_x_child = epoch_changes.epoch_for_child_of(
+			let epoch_for_x_child = epoch_changes.epoch_data_for_child_of(
 				&is_descendent_of,
 				b"X",
 				1,
@@ -654,9 +903,9 @@ mod tests {
 				&make_genesis,
 			).unwrap().unwrap();
 
-			assert_eq!(epoch_for_x_child.into_inner(), make_genesis(1000));
+			assert_eq!(epoch_for_x_child, make_genesis(1000));
 
-			let epoch_for_x_child_before_genesis = epoch_changes.epoch_for_child_of(
+			let epoch_for_x_child_before_genesis = epoch_changes.epoch_data_for_child_of(
 				&is_descendent_of,
 				b"X",
 				1,

substrate/client/consensus/epochs/src/migration.rs

@@ -0,0 +1,55 @@
+// Copyright 2019-2020 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Migration types for epoch changes.
+
+use std::collections::BTreeMap;
+use codec::{Encode, Decode};
+use fork_tree::ForkTree;
+use sp_runtime::traits::{Block as BlockT, NumberFor};
+use crate::{Epoch, EpochChanges, PersistedEpoch, PersistedEpochHeader};
+
+/// Legacy definition of epoch changes.
+#[derive(Clone, Encode, Decode)]
+pub struct EpochChangesV0<Hash, Number, E: Epoch> {
+	inner: ForkTree<Hash, Number, PersistedEpoch<E>>,
+}
+
+/// Type alias for legacy definition of epoch changes.
+pub type EpochChangesForV0<Block, Epoch> = EpochChangesV0<<Block as BlockT>::Hash, NumberFor<Block>, Epoch>;
+
+impl<Hash, Number, E: Epoch> EpochChangesV0<Hash, Number, E> where
+	Hash: PartialEq + Ord + Copy,
+	Number: Ord + Copy,
+{
+	/// Create a new value of this type from raw.
+	pub fn from_raw(inner: ForkTree<Hash, Number, PersistedEpoch<E>>) -> Self {
+		Self { inner }
+	}
+
+	/// Migrate the type into current epoch changes definition.
+	pub fn migrate(self) -> EpochChanges<Hash, Number, E> {
+		let mut epochs = BTreeMap::new();
+
+		let inner = self.inner.map(&mut |hash, number, data| {
+			let header = PersistedEpochHeader::from(&data);
+			epochs.insert((*hash, *number), data);
+			header
+		});
+
+		EpochChanges { inner, epochs }
+	}
+}

substrate/client/consensus/slots/src/lib.rs

@@ -94,7 +94,8 @@ pub trait SimpleSlotWorker<B: BlockT> {
 	fn epoch_data(&self, header: &B::Header, slot_number: u64) -> Result<Self::EpochData, sp_consensus::Error>;
 
 	/// Returns the number of authorities given the epoch data.
-	fn authorities_len(&self, epoch_data: &Self::EpochData) -> usize;
+	/// None indicate that the authorities information is incomplete.
+	fn authorities_len(&self, epoch_data: &Self::EpochData) -> Option<usize>;
 
 	/// Tries to claim the given slot, returning an object with claim data if successful.
 	fn claim_slot(
@@ -194,7 +195,10 @@ pub trait SimpleSlotWorker<B: BlockT> {
 
 		let authorities_len = self.authorities_len(&epoch_data);
 
-		if !self.force_authoring() && self.sync_oracle().is_offline() && authorities_len > 1 {
+		if !self.force_authoring() &&
+			self.sync_oracle().is_offline() &&
+			authorities_len.map(|a| a > 1).unwrap_or(false)
+		{
 			debug!(target: self.logging_target(), "Skipping proposal slot. Waiting for the network.");
 			telemetry!(
 				CONSENSUS_DEBUG;

substrate/utils/fork-tree/src/lib.rs

@@ -93,41 +93,77 @@ impl<H, N, V> ForkTree<H, N, V> where
 	/// node. Otherwise the tree remains unchanged. The given function
 	/// `is_descendent_of` should return `true` if the second hash (target) is a
 	/// descendent of the first hash (base).
+	///
+	/// Returns all pruned node data.
 	pub fn prune<F, E, P>(
 		&mut self,
 		hash: &H,
 		number: &N,
 		is_descendent_of: &F,
 		predicate: &P,
-	) -> Result<(), Error<E>>
+	) -> Result<impl Iterator<Item=(H, N, V)>, Error<E>>
 		where E: std::error::Error,
 			  F: Fn(&H, &H) -> Result<bool, E>,
 			  P: Fn(&V) -> bool,
 	{
-		let new_root = self.find_node_where(
+		let new_root_index = self.find_node_index_where(
 			hash,
 			number,
 			is_descendent_of,
 			predicate,
 		)?;
 
-		if let Some(root) = new_root {
-			let mut root = root.clone();
+		let removed = if let Some(mut root_index) = new_root_index {
+			let mut old_roots = std::mem::replace(&mut self.roots, Vec::new());
+
+			let mut root = None;
+			let mut cur_children = Some(&mut old_roots);
+
+			while let Some(cur_index) = root_index.pop() {
+				if let Some(children) = cur_children.take() {
+					if root_index.is_empty() {
+						root = Some(children.remove(cur_index));
+					} else {
+						cur_children = Some(&mut children[cur_index].children);
+					}
+				}
+			}
+
+			let mut root = root
+				.expect("find_node_index_where will return array with at least one index; \
+                         this results in at least one item in removed; qed");
+
+			let mut removed = old_roots;
 
 			// we found the deepest ancestor of the finalized block, so we prune
 			// out any children that don't include the finalized block.
-			let children = std::mem::replace(&mut root.children, Vec::new());
-			root.children = children.into_iter().filter(|node| {
-				node.number == *number && node.hash == *hash ||
-					node.number < *number && is_descendent_of(&node.hash, hash).unwrap_or(false)
-			}).take(1).collect();
+			let root_children = std::mem::replace(&mut root.children, Vec::new());
+			let mut is_first = true;
+
+			for child in root_children {
+				if is_first &&
+					(child.number == *number && child.hash == *hash ||
+					 child.number < *number && is_descendent_of(&child.hash, hash).unwrap_or(false))
+				{
+					root.children.push(child);
+					// assuming that the tree is well formed only one child should pass this requirement
+					// due to ancestry restrictions (i.e. they must be different forks).
+					is_first = false;
+				} else {
+					removed.push(child);
+				}
+			}
 
 			self.roots = vec![root];
-		}
+
+			removed
+		} else {
+			Vec::new()
+		};
 
 		self.rebalance();
 
-		Ok(())
+		Ok(RemovedIterator { stack: removed })
 	}
 }
 
@@ -250,6 +286,26 @@ impl<H, N, V> ForkTree<H, N, V> where
 		Ok(None)
 	}
 
+	/// Map fork tree into values of new types.
+	pub fn map<VT, F>(
+		self,
+		f: &mut F,
+	) -> ForkTree<H, N, VT> where
+		F: FnMut(&H, &N, V) -> VT,
+	{
+		let roots = self.roots
+			.into_iter()
+			.map(|root| {
+				root.map(f)
+			})
+			.collect();
+
+		ForkTree {
+			roots,
+			best_finalized_number: self.best_finalized_number,
+		}
+	}
+
 	/// Same as [`find_node_where`](Self::find_node_where), but returns mutable reference.
 	pub fn find_node_where_mut<F, E, P>(
 		&mut self,
@@ -275,6 +331,32 @@ impl<H, N, V> ForkTree<H, N, V> where
 		Ok(None)
 	}
 
+	/// Same as [`find_node_where`](Self::find_node_where), but returns indexes.
+	pub fn find_node_index_where<F, E, P>(
+		&self,
+		hash: &H,
+		number: &N,
+		is_descendent_of: &F,
+		predicate: &P,
+	) -> Result<Option<Vec<usize>>, Error<E>> where
+		E: std::error::Error,
+		F: Fn(&H, &H) -> Result<bool, E>,
+		P: Fn(&V) -> bool,
+	{
+		// search for node starting from all roots
+		for (index, root) in self.roots.iter().enumerate() {
+			let node = root.find_node_index_where(hash, number, is_descendent_of, predicate)?;
+
+			// found the node, early exit
+			if let FindOutcome::Found(mut node) = node {
+				node.push(index);
+				return Ok(Some(node));
+			}
+		}
+
+		Ok(None)
+	}
+
 	/// Finalize a root in the tree and return it, return `None` in case no root
 	/// with the given hash exists. All other roots are pruned, and the children
 	/// of the finalized node become the new roots.
@@ -588,6 +670,29 @@ mod node_implementation {
 			max + 1
 		}
 
+		/// Map node data into values of new types.
+		pub fn map<VT, F>(
+			self,
+			f: &mut F,
+		) -> Node<H, N, VT> where
+			F: FnMut(&H, &N, V) -> VT,
+		{
+			let children = self.children
+				.into_iter()
+				.map(|node| {
+					node.map(f)
+				})
+				.collect();
+
+			let vt = f(&self.hash, &self.number, self.data);
+			Node {
+				hash: self.hash,
+				number: self.number,
+				data: vt,
+				children,
+			}
+		}
+
 		pub fn import<F, E: std::error::Error>(
 			&mut self,
 			mut hash: H,
@@ -780,6 +885,27 @@ impl<'a, H, N, V> Iterator for ForkTreeIterator<'a, H, N, V> {
 	}
 }
 
+struct RemovedIterator<H, N, V> {
+	stack: Vec<Node<H, N, V>>,
+}
+
+impl<H, N, V> Iterator for RemovedIterator<H, N, V> {
+	type Item = (H, N, V);
+
+	fn next(&mut self) -> Option<Self::Item> {
+		self.stack.pop().map(|mut node| {
+			// child nodes are stored ordered by max branch height (decreasing),
+			// we want to keep this ordering while iterating but since we're
+			// using a stack for iterator state we need to reverse it.
+			let mut children = Vec::new();
+			std::mem::swap(&mut children, &mut node.children);
+
+			self.stack.extend(children.into_iter().rev());
+			(node.hash, node.number, node.data)
+		})
+	}
+}
+
 #[cfg(test)]
 mod test {
 	use super::{FinalizationResult, ForkTree, Error};
@@ -805,7 +931,7 @@ mod test {
 		//  /   /
 		// A - F - H - I
 		//          \
-		//           - L - M - N
+		//           - L - M
 		//              \
 		//               - O
 		//  \
@@ -813,22 +939,21 @@ mod test {
 		//
 		// (where N is not a part of fork tree)
 		let is_descendent_of = |base: &&str, block: &&str| -> Result<bool, TestError> {
-			let letters = vec!["B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L"];
+			let letters = vec!["B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "O"];
 			match (*base, *block) {
 				("A", b) => Ok(letters.into_iter().any(|n| n == b)),
 				("B", b) => Ok(b == "C" || b == "D" || b == "E"),
 				("C", b) => Ok(b == "D" || b == "E"),
 				("D", b) => Ok(b == "E"),
 				("E", _) => Ok(false),
-				("F", b) => Ok(b == "G" || b == "H" || b == "I" || b == "L" || b == "M" || b == "N" || b == "O"),
+				("F", b) => Ok(b == "G" || b == "H" || b == "I" || b == "L" || b == "M" || b == "O"),
 				("G", _) => Ok(false),
 				("H", b) => Ok(b == "I" || b == "L" || b == "M" || b == "O"),
 				("I", _) => Ok(false),
 				("J", b) => Ok(b == "K"),
 				("K", _) => Ok(false),
-				("L", b) => Ok(b == "M" || b == "O" || b == "N"),
-				("M", b) => Ok(b == "N"),
-				("N", _) => Ok(false),
+				("L", b) => Ok(b == "M" || b == "O"),
+				("M", _) => Ok(false),
 				("O", _) => Ok(false),
 				("0", _) => Ok(true),
 				_ => Ok(false),
@@ -1324,11 +1449,18 @@ mod test {
 		assert_eq!(node.number, 3);
 	}
 
+	#[test]
+	fn map_works() {
+		let (tree, _is_descendent_of) = test_fork_tree();
+
+		let _tree = tree.map(&mut |_, _, _| ());
+	}
+
 	#[test]
 	fn prune_works() {
 		let (mut tree, is_descendent_of) = test_fork_tree();
 
-		tree.prune(
+		let removed = tree.prune(
 			&"C",
 			&3,
 			&is_descendent_of,
@@ -1345,7 +1477,12 @@ mod test {
 			vec!["B", "C", "D", "E"],
 		);
 
-		tree.prune(
+		assert_eq!(
+			removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(),
+			vec!["A", "F", "G", "H", "I", "L", "M", "O", "J", "K"]
+		);
+
+		let removed = tree.prune(
 			&"E",
 			&5,
 			&is_descendent_of,
@@ -1361,6 +1498,11 @@ mod test {
 			tree.iter().map(|(hash, _, _)| *hash).collect::<Vec<_>>(),
 			vec!["D", "E"],
 		);
+
+		assert_eq!(
+			removed.map(|(hash, _, _)| hash).collect::<Vec<_>>(),
+			vec!["B", "C"]
+		);
 	}
 
 	#[test]