babe: pass epoch data via intermediates

client/consensus/babe/src/lib.rs

@@ -63,7 +63,10 @@ pub use sp_consensus_babe::{
 	CompatibleDigestItem,
 };
 pub use sp_consensus::SyncOracle;
-use std::{collections::HashMap, sync::Arc, u64, pin::Pin, time::{Instant, Duration}};
+use std::{
+	collections::HashMap, sync::Arc, u64, pin::Pin, time::{Instant, Duration},
+	any::Any, borrow::Cow
+};
 use sp_consensus_babe;
 use sp_consensus::{ImportResult, CanAuthorWith};
 use sp_consensus::import_queue::{
@@ -101,7 +104,7 @@ use log::{warn, debug, info, trace};
 use sc_consensus_slots::{
 	SlotWorker, SlotInfo, SlotCompatible, StorageChanges, CheckedHeader, check_equivocation,
 };
-use epoch_changes::descendent_query;
+use epoch_changes::{descendent_query, ViableEpoch};
 use sp_blockchain::{
 	Result as ClientResult, Error as ClientError,
 	HeaderBackend, ProvideCache, HeaderMetadata
@@ -121,7 +124,6 @@ pub use sp_consensus_babe::{
 };
 pub use epoch_changes::{EpochChanges, EpochChangesFor, SharedEpochChanges};
 
-
 #[derive(derive_more::Display, Debug)]
 enum Error<B: BlockT> {
 	#[display(fmt = "Multiple BABE pre-runtime digests, rejecting!")]
@@ -197,6 +199,16 @@ macro_rules! babe_info {
 	};
 }
 
+
+/// Intermediate value passed to block importer.
+pub struct BabeIntermediate {
+	/// The epoch data, if available.
+	pub epoch: Option<ViableEpoch>,
+}
+
+/// Intermediate key for Babe engine.
+pub static INTERMEDIATE_KEY: &[u8] = b"babe1";
+
 /// A slot duration. Create with `get_or_compute`.
 // FIXME: Once Rust has higher-kinded types, the duplication between this
 // and `super::babe::Config` can be eliminated.
@@ -451,7 +463,14 @@ impl<B, C, E, I, Error, SO> sc_consensus_slots::SimpleSlotWorker<B> for BabeWork
 				storage_changes: Some(storage_changes),
 				finalized: false,
 				auxiliary: Vec::new(), // block-weight is written in block import.
-				intermediates: Default::default(),
+				intermediates: {
+					let mut intermediates = HashMap::new();
+					intermediates.insert(
+						Cow::from(INTERMEDIATE_KEY),
+						Box::new(BabeIntermediate { epoch: None }) as Box<dyn Any>,
+					);
+					intermediates
+				},
 				fork_choice: None,
 				allow_missing_state: false,
 				import_existing: false,
@@ -796,6 +815,14 @@ impl<B, E, Block, RA, PRA> Verifier<Block> for BabeVerifier<B, E, Block, RA, PRA
 					"babe.checked_and_importing";
 					"pre_header" => ?pre_header);
 
+				let mut intermediates = HashMap::new();
+				intermediates.insert(
+					Cow::from(INTERMEDIATE_KEY),
+					Box::new(BabeIntermediate {
+						epoch: Some(epoch),
+					}) as Box<dyn Any>,
+				);
+
 				let block_import_params = BlockImportParams {
 					origin,
 					header: pre_header,
@@ -805,7 +832,7 @@ impl<B, E, Block, RA, PRA> Verifier<Block> for BabeVerifier<B, E, Block, RA, PRA
 					finalized: false,
 					justification,
 					auxiliary: Vec::new(),
-					intermediates: Default::default(),
+					intermediates,
 					fork_choice: None,
 					allow_missing_state: false,
 					import_existing: false,
@@ -963,20 +990,29 @@ impl<B, E, Block, I, RA, PRA> BlockImport<Block> for BabeBlockImport<B, E, Block
 					))?
 			};
 
-			let epoch = epoch_changes.epoch_for_child_of(
-				descendent_query(&*self.client),
-				&parent_hash,
-				*parent_header.number(),
-				slot_number,
-				|slot| self.config.genesis_epoch(slot),
-			)
-				.map_err(|e: fork_tree::Error<sp_blockchain::Error>| ConsensusError::ChainLookup(
-					babe_err(Error::<Block>::CouldNotLookUpEpoch(Box::new(e))).into()
-				))?
-				.ok_or_else(|| ConsensusError::ClientImport(
-					babe_err(Error::<Block>::BlockNotValid(hash)).into()
-				))?;
-
+			let intermediate = block.take_intermediate::<BabeIntermediate>(
+				INTERMEDIATE_KEY
+			)?;
+
+			let epoch = match intermediate.epoch {
+				Some(epoch) => epoch,
+				None => {
+					epoch_changes.epoch_for_child_of(
+						descendent_query(&*self.client),
+						&parent_hash,
+						*parent_header.number(),
+						slot_number,
+						|slot| self.config.genesis_epoch(slot),
+					)
+						.map_err(|e: fork_tree::Error<sp_blockchain::Error>|
+								 ConsensusError::ChainLookup(
+									 babe_err(Error::<Block>::CouldNotLookUpEpoch(Box::new(e))).into()
+								 ))?
+						.ok_or_else(|| ConsensusError::ClientImport(
+							babe_err(Error::<Block>::BlockNotValid(hash)).into()
+						))?
+				},
+			};
 			let first_in_epoch = parent_slot < epoch.as_ref().start_slot;
 			(epoch, first_in_epoch, parent_weight)
 		};

client/consensus/babe/src/tests.rs

@@ -594,7 +594,14 @@ fn propose_and_import_block<Transaction>(
 			storage_changes: None,
 			finalized: false,
 			auxiliary: Vec::new(),
-			intermediates: Default::default(),
+			intermediates: {
+				let mut intermediates = HashMap::new();
+				intermediates.insert(
+					Cow::from(INTERMEDIATE_KEY),
+					Box::new(BabeIntermediate { epoch: None }) as Box<dyn Any>,
+				);
+				intermediates
+			},
 			fork_choice: Some(ForkChoiceStrategy::LongestChain),
 			allow_missing_state: false,
 			import_existing: false,