Split PrimaryPreDigest and SecondaryPreDigest (#5373)

* Split PrimaryPreDigest and SecondaryPreDigest

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

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

* Update primitives/consensus/babe/src/digests.rs

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

* Update primitives/consensus/babe/src/digests.rs

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

* Update primitives/consensus/babe/src/digests.rs

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

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

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

@@ -21,7 +21,7 @@ use sp_consensus_babe::{
 	AuthorityId, BabeAuthorityWeight, BABE_ENGINE_ID, BABE_VRF_PREFIX,
 	SlotNumber, AuthorityPair, BabeConfiguration
 };
-use sp_consensus_babe::digests::PreDigest;
+use sp_consensus_babe::digests::{PreDigest, PrimaryPreDigest, SecondaryPreDigest};
 use sp_consensus_vrf::schnorrkel::{VRFOutput, VRFProof};
 use sp_core::{U256, blake2_256};
 use codec::Encode;
@@ -128,10 +128,10 @@ fn claim_secondary_slot(
 		})
 	{
 		if pair.public() == *expected_author {
-			let pre_digest = PreDigest::Secondary {
+			let pre_digest = PreDigest::Secondary(SecondaryPreDigest {
 				slot_number,
 				authority_index: authority_index as u32,
-			};
+			});
 
 			return Some((pre_digest, pair));
 		}
@@ -200,12 +200,12 @@ fn claim_primary_slot(
 		let pre_digest = get_keypair(&pair)
 			.vrf_sign_after_check(transcript, |inout| super::authorship::check_primary_threshold(inout, threshold))
 			.map(|s| {
-				PreDigest::Primary {
+				PreDigest::Primary(PrimaryPreDigest {
 					slot_number,
 					vrf_output: VRFOutput(s.0.to_output()),
 					vrf_proof: VRFProof(s.1),
 					authority_index: authority_index as u32,
-				}
+				})
 			});
 
 		// early exit on first successful claim

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

@@ -62,7 +62,9 @@ pub use sp_consensus_babe::{
 	BabeApi, ConsensusLog, BABE_ENGINE_ID, SlotNumber, BabeConfiguration,
 	AuthorityId, AuthorityPair, AuthoritySignature,
 	BabeAuthorityWeight, VRF_OUTPUT_LENGTH,
-	digests::{PreDigest, CompatibleDigestItem, NextEpochDescriptor},
+	digests::{
+		CompatibleDigestItem, NextEpochDescriptor, PreDigest, PrimaryPreDigest, SecondaryPreDigest,
+	},
 };
 pub use sp_consensus::SyncOracle;
 use std::{
@@ -580,10 +582,10 @@ fn find_pre_digest<B: BlockT>(header: &B::Header) -> Result<PreDigest, Error<B>>
 	// genesis block doesn't contain a pre digest so let's generate a
 	// dummy one to not break any invariants in the rest of the code
 	if header.number().is_zero() {
-		return Ok(PreDigest::Secondary {
+		return Ok(PreDigest::Secondary(SecondaryPreDigest {
 			slot_number: 0,
 			authority_index: 0,
-		});
+		}));
 	}
 
 	let mut pre_digest: Option<_> = None;

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

@@ -553,10 +553,10 @@ fn propose_and_import_block<Transaction>(
 	let pre_digest = sp_runtime::generic::Digest {
 		logs: vec![
 			Item::babe_pre_digest(
-				PreDigest::Secondary {
+				PreDigest::Secondary(SecondaryPreDigest {
 					authority_index: 0,
 					slot_number,
-				},
+				}),
 			),
 		],
 	};

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

@@ -17,9 +17,10 @@
 //! Verification for BABE headers.
 use sp_runtime::{traits::Header, traits::DigestItemFor};
 use sp_core::{Pair, Public};
-use sp_consensus_babe::{AuthoritySignature, SlotNumber, AuthorityIndex, AuthorityPair, AuthorityId};
-use sp_consensus_babe::digests::{PreDigest, CompatibleDigestItem};
-use sp_consensus_vrf::schnorrkel::{VRFOutput, VRFProof};
+use sp_consensus_babe::{AuthoritySignature, SlotNumber, AuthorityPair, AuthorityId};
+use sp_consensus_babe::digests::{
+	PreDigest, PrimaryPreDigest, SecondaryPreDigest, CompatibleDigestItem
+};
 use sc_consensus_slots::CheckedHeader;
 use log::{debug, trace};
 use super::{find_pre_digest, babe_err, Epoch, BlockT, Error};
@@ -93,27 +94,23 @@ pub(super) fn check_header<B: BlockT + Sized>(
 	};
 
 	match &pre_digest {
-		PreDigest::Primary { vrf_output, vrf_proof, authority_index, slot_number } => {
+		PreDigest::Primary(primary) => {
 			debug!(target: "babe", "Verifying Primary block");
 
-			let digest = (vrf_output, vrf_proof, *authority_index, *slot_number);
-
 			check_primary_header::<B>(
 				pre_hash,
-				digest,
+				primary,
 				sig,
 				&epoch,
 				config.c,
 			)?;
 		},
-		PreDigest::Secondary { authority_index, slot_number } if config.secondary_slots => {
+		PreDigest::Secondary(secondary) if config.secondary_slots => {
 			debug!(target: "babe", "Verifying Secondary block");
 
-			let digest = (*authority_index, *slot_number);
-
 			check_secondary_header::<B>(
 				pre_hash,
-				digest,
+				secondary,
 				sig,
 				&epoch,
 			)?;
@@ -143,25 +140,23 @@ pub(super) struct VerifiedHeaderInfo<B: BlockT> {
 /// its parent since it is a primary block.
 fn check_primary_header<B: BlockT + Sized>(
 	pre_hash: B::Hash,
-	pre_digest: (&VRFOutput, &VRFProof, AuthorityIndex, SlotNumber),
+	pre_digest: &PrimaryPreDigest,
 	signature: AuthoritySignature,
 	epoch: &Epoch,
 	c: (u64, u64),
 ) -> Result<(), Error<B>> {
-	let (vrf_output, vrf_proof, authority_index, slot_number) = pre_digest;
-
-	let author = &epoch.authorities[authority_index as usize].0;
+	let author = &epoch.authorities[pre_digest.authority_index as usize].0;
 
 	if AuthorityPair::verify(&signature, pre_hash, &author) {
 		let (inout, _) = {
 			let transcript = make_transcript(
 				&epoch.randomness,
-				slot_number,
+				pre_digest.slot_number,
 				epoch.epoch_index,
 			);
 
 			schnorrkel::PublicKey::from_bytes(author.as_slice()).and_then(|p| {
-				p.vrf_verify(transcript, vrf_output, vrf_proof)
+				p.vrf_verify(transcript, &pre_digest.vrf_output, &pre_digest.vrf_proof)
 			}).map_err(|s| {
 				babe_err(Error::VRFVerificationFailed(s))
 			})?
@@ -170,7 +165,7 @@ fn check_primary_header<B: BlockT + Sized>(
 		let threshold = calculate_primary_threshold(
 			c,
 			&epoch.authorities,
-			authority_index as usize,
+			pre_digest.authority_index as usize,
 		);
 
 		if !check_primary_threshold(&inout, threshold) {
@@ -189,21 +184,19 @@ fn check_primary_header<B: BlockT + Sized>(
 /// compared to its parent since it is a secondary block.
 fn check_secondary_header<B: BlockT>(
 	pre_hash: B::Hash,
-	pre_digest: (AuthorityIndex, SlotNumber),
+	pre_digest: &SecondaryPreDigest,
 	signature: AuthoritySignature,
 	epoch: &Epoch,
 ) -> Result<(), Error<B>> {
-	let (authority_index, slot_number) = pre_digest;
-
 	// check the signature is valid under the expected authority and
 	// chain state.
 	let expected_author = secondary_slot_author(
-		slot_number,
+		pre_digest.slot_number,
 		&epoch.authorities,
 		epoch.randomness,
 	).ok_or_else(|| Error::NoSecondaryAuthorExpected)?;
 
-	let author = &epoch.authorities[authority_index as usize].0;
+	let author = &epoch.authorities[pre_digest.authority_index as usize].0;
 
 	if expected_author != author {
 		return Err(Error::InvalidAuthor(expected_author.clone(), author.clone()));

substrate/frame/babe/src/lib.rs

@@ -209,12 +209,7 @@ impl<T: Trait> FindAuthor<u32> for Module<T> {
 		for (id, mut data) in digests.into_iter() {
 			if id == BABE_ENGINE_ID {
 				let pre_digest: RawPreDigest = RawPreDigest::decode(&mut data).ok()?;
-				return Some(match pre_digest {
-					RawPreDigest::Primary { authority_index, .. } =>
-						authority_index,
-					RawPreDigest::Secondary { authority_index, .. } =>
-						authority_index,
-				});
+				return Some(pre_digest.authority_index())
 			}
 		}
 
@@ -426,11 +421,11 @@ impl<T: Trait> Module<T> {
 
 			CurrentSlot::put(digest.slot_number());
 
-			if let RawPreDigest::Primary { vrf_output, .. } = digest {
+			if let RawPreDigest::Primary(primary) = digest {
 				// place the VRF output into the `Initialized` storage item
 				// and it'll be put onto the under-construction randomness
 				// later, once we've decided which epoch this block is in.
-				Some(vrf_output)
+				Some(primary.vrf_output)
 			} else {
 				None
 			}

substrate/frame/babe/src/tests.rs

@@ -35,12 +35,14 @@ fn make_pre_digest(
 	vrf_output: RawVRFOutput,
 	vrf_proof: RawVRFProof,
 ) -> Digest {
-	let digest_data = sp_consensus_babe::digests::RawPreDigest::Primary {
-		authority_index,
-		slot_number,
-		vrf_output,
-		vrf_proof,
-	};
+	let digest_data = sp_consensus_babe::digests::RawPreDigest::Primary(
+		sp_consensus_babe::digests::RawPrimaryPreDigest {
+			authority_index,
+			slot_number,
+			vrf_output,
+			vrf_proof,
+		}
+	);
 	let log = DigestItem::PreRuntime(sp_consensus_babe::BABE_ENGINE_ID, digest_data.encode());
 	Digest { logs: vec![log] }
 }

substrate/primitives/consensus/babe/src/digests.rs

@@ -27,40 +27,67 @@ use codec::{Decode, Encode};
 #[cfg(feature = "std")]
 use codec::Codec;
 use sp_std::vec::Vec;
+use sp_runtime::RuntimeDebug;
 use sp_consensus_vrf::schnorrkel::{self, Randomness};
 #[cfg(feature = "std")]
 use sp_consensus_vrf::schnorrkel::SignatureError;
 
+/// Raw BABE primary slot assignment pre-digest.
+#[derive(Clone, RuntimeDebug, Encode, Decode)]
+pub struct RawPrimaryPreDigest<VRFOutput=schnorrkel::RawVRFOutput, VRFProof=schnorrkel::RawVRFProof> {
+	/// Authority index
+	pub authority_index: super::AuthorityIndex,
+	/// Slot number
+	pub slot_number: SlotNumber,
+	/// VRF output
+	pub vrf_output: VRFOutput,
+	/// VRF proof
+	pub vrf_proof: VRFProof,
+}
+
+#[cfg(feature = "std")]
+/// BABE primary slot assignment pre-digest for std environment.
+pub type PrimaryPreDigest = RawPrimaryPreDigest<schnorrkel::VRFOutput, schnorrkel::VRFProof>;
+
+#[cfg(feature = "std")]
+impl TryFrom<RawPrimaryPreDigest> for PrimaryPreDigest {
+	type Error = SignatureError;
+
+	fn try_from(raw: RawPrimaryPreDigest) -> Result<PrimaryPreDigest, SignatureError> {
+		Ok(PrimaryPreDigest {
+			authority_index: raw.authority_index,
+			slot_number: raw.slot_number,
+			vrf_output: raw.vrf_output.try_into()?,
+			vrf_proof: raw.vrf_proof.try_into()?,
+		})
+	}
+}
+
+/// BABE secondary slot assignment pre-digest.
+#[derive(Clone, RuntimeDebug, Encode, Decode)]
+pub struct SecondaryPreDigest {
+	/// Authority index
+	///
+	/// This is not strictly-speaking necessary, since the secondary slots
+	/// are assigned based on slot number and epoch randomness. But including
+	/// it makes things easier for higher-level users of the chain data to
+	/// be aware of the author of a secondary-slot block.
+	pub authority_index: super::AuthorityIndex,
+	/// Slot number
+	pub slot_number: SlotNumber,
+}
+
 /// A BABE pre-runtime digest. This contains all data required to validate a
 /// block and for the BABE runtime module. Slots can be assigned to a primary
 /// (VRF based) and to a secondary (slot number based).
-#[derive(Clone, Debug, Encode, Decode)]
+#[derive(Clone, RuntimeDebug, Encode, Decode)]
 pub enum RawPreDigest<VRFOutput=schnorrkel::RawVRFOutput, VRFProof=schnorrkel::RawVRFProof> {
 	/// A primary VRF-based slot assignment.
 	#[codec(index = "1")]
-	Primary {
-		/// Authority index
-		authority_index: super::AuthorityIndex,
-		/// Slot number
-		slot_number: SlotNumber,
-		/// VRF output
-		vrf_output: VRFOutput,
-		/// VRF proof
-		vrf_proof: VRFProof,
-	},
+	Primary(RawPrimaryPreDigest<VRFOutput, VRFProof>),
 	/// A secondary deterministic slot assignment.
 	#[codec(index = "2")]
-	Secondary {
-		/// Authority index
-		///
-		/// This is not strictly-speaking necessary, since the secondary slots
-		/// are assigned based on slot number and epoch randomness. But including
-		/// it makes things easier for higher-level users of the chain data to
-		/// be aware of the author of a secondary-slot block.
-		authority_index: super::AuthorityIndex,
-		/// Slot number
-		slot_number: SlotNumber,
-	},
+	Secondary(SecondaryPreDigest),
 }
 
 #[cfg(feature = "std")]
@@ -71,16 +98,16 @@ impl<VRFOutput, VRFProof> RawPreDigest<VRFOutput, VRFProof> {
 	/// Returns the slot number of the pre digest.
 	pub fn authority_index(&self) -> AuthorityIndex {
 		match self {
-			RawPreDigest::Primary { authority_index, .. } => *authority_index,
-			RawPreDigest::Secondary { authority_index, .. } => *authority_index,
+			RawPreDigest::Primary(primary) => primary.authority_index,
+			RawPreDigest::Secondary(secondary) => secondary.authority_index,
 		}
 	}
 
 	/// Returns the slot number of the pre digest.
 	pub fn slot_number(&self) -> SlotNumber {
 		match self {
-			RawPreDigest::Primary { slot_number, .. } => *slot_number,
-			RawPreDigest::Secondary { slot_number, .. } => *slot_number,
+			RawPreDigest::Primary(primary) => primary.slot_number,
+			RawPreDigest::Secondary(secondary) => secondary.slot_number,
 		}
 	}
 
@@ -88,8 +115,8 @@ impl<VRFOutput, VRFProof> RawPreDigest<VRFOutput, VRFProof> {
 	/// of the chain.
 	pub fn added_weight(&self) -> crate::BabeBlockWeight {
 		match self {
-			RawPreDigest::Primary { .. } => 1,
-			RawPreDigest::Secondary { .. } => 0,
+			RawPreDigest::Primary(_) => 1,
+			RawPreDigest::Secondary(_) => 0,
 		}
 	}
 }
@@ -100,25 +127,15 @@ impl TryFrom<RawPreDigest> for PreDigest {
 
 	fn try_from(raw: RawPreDigest) -> Result<PreDigest, SignatureError> {
 		Ok(match raw {
-			RawPreDigest::Primary { authority_index, slot_number, vrf_output, vrf_proof } =>
-				RawPreDigest::Primary {
-					authority_index,
-					slot_number,
-					vrf_output: vrf_output.try_into()?,
-					vrf_proof: vrf_proof.try_into()?,
-				},
-			RawPreDigest::Secondary { authority_index, slot_number } =>
-				RawPreDigest::Secondary {
-					authority_index,
-					slot_number,
-				}
+			RawPreDigest::Primary(primary) => PreDigest::Primary(primary.try_into()?),
+			RawPreDigest::Secondary(secondary) => PreDigest::Secondary(secondary),
 		})
 	}
 }
 
 /// Information about the next epoch. This is broadcast in the first block
 /// of the epoch.
-#[derive(Decode, Encode, Default, PartialEq, Eq, Clone, sp_runtime::RuntimeDebug)]
+#[derive(Decode, Encode, Default, PartialEq, Eq, Clone, RuntimeDebug)]
 pub struct NextEpochDescriptor {
 	/// The authorities.
 	pub authorities: Vec<(AuthorityId, BabeAuthorityWeight)>,

substrate/primitives/consensus/vrf/src/schnorrkel.rs

@@ -17,6 +17,7 @@
 //! Schnorrkel-based VRF.
 
 use codec::{Encode, Decode};
+#[cfg(feature = "std")]
 use sp_core::U512;
 use sp_runtime::RuntimeDebug;
 use sp_std::ops::{Deref, DerefMut};