From 9ceb416940ae356dff9da16d8202a788ea952e2f Mon Sep 17 00:00:00 2001
From: Lldenaurois <Ljdenaurois@gmail.com>
Date: Mon, 2 Aug 2021 15:11:17 -0400
Subject: [PATCH 1/4] runtime/parachains: Keep track of included candidates in 
 module

---
 runtime/parachains/src/disputes.rs       | 115 +++++++++++-----------
 runtime/parachains/src/paras_inherent.rs |  18 +++-
 runtime/parachains/src/shared.rs         | 116 ++++++++++++++++++++++-
 3 files changed, 192 insertions(+), 57 deletions(-)

diff --git a/runtime/parachains/src/disputes.rs b/runtime/parachains/src/disputes.rs
index 43388c64b01c..3bad15b90aa9 100644
--- a/runtime/parachains/src/disputes.rs
+++ b/runtime/parachains/src/disputes.rs
@@ -17,6 +17,7 @@
 //! Runtime component for handling disputes of parachain candidates.
 
 use crate::{
+	shared,
 	configuration::{self, HostConfiguration},
 	initializer::SessionChangeNotification,
 	session_info,
@@ -31,7 +32,7 @@ use primitives::v1::{
 	ValidDisputeStatementKind, ValidatorId, ValidatorIndex, ValidatorSignature,
 };
 use sp_runtime::{
-	traits::{AppVerify, One, Saturating, Zero},
+	traits::{AppVerify, Saturating},
 	DispatchError, RuntimeDebug, SaturatedConversion,
 };
 use sp_std::{collections::btree_set::BTreeSet, prelude::*};
@@ -116,10 +117,10 @@ pub trait DisputesHandler<BlockNumber> {
 	) -> Result<Vec<(SessionIndex, CandidateHash)>, DispatchError>;
 
 	/// Note that the given candidate has been included.
-	fn note_included(
+	fn process_included(
 		session: SessionIndex,
 		candidate_hash: CandidateHash,
-		included_in: BlockNumber,
+		revert_to: BlockNumber,
 	);
 
 	/// Whether the given candidate could be invalid, i.e. there is an ongoing
@@ -151,10 +152,10 @@ impl<BlockNumber> DisputesHandler<BlockNumber> for () {
 		Ok(Vec::new())
 	}
 
-	fn note_included(
+	fn process_included(
 		_session: SessionIndex,
 		_candidate_hash: CandidateHash,
-		_included_in: BlockNumber,
+		_revert_to: BlockNumber,
 	) {
 	}
 
@@ -186,12 +187,12 @@ impl<T: Config> DisputesHandler<T::BlockNumber> for pallet::Pallet<T> {
 		pallet::Pallet::<T>::provide_multi_dispute_data(statement_sets)
 	}
 
-	fn note_included(
+	fn process_included(
 		session: SessionIndex,
 		candidate_hash: CandidateHash,
-		included_in: T::BlockNumber,
+		revert_to: T::BlockNumber,
 	) {
-		pallet::Pallet::<T>::note_included(session, candidate_hash, included_in)
+		pallet::Pallet::<T>::process_included(session, candidate_hash, revert_to);
 	}
 
 	fn could_be_invalid(session: SessionIndex, candidate_hash: CandidateHash) -> bool {
@@ -243,18 +244,6 @@ pub mod pallet {
 		DisputeState<T::BlockNumber>,
 	>;
 
-	/// All included blocks on the chain, as well as the block number in this chain that
-	/// should be reverted back to if the candidate is disputed and determined to be invalid.
-	#[pallet::storage]
-	pub(super) type Included<T: Config> = StorageDoubleMap<
-		_,
-		Twox64Concat,
-		SessionIndex,
-		Blake2_128Concat,
-		CandidateHash,
-		T::BlockNumber,
-	>;
-
 	/// Maps session indices to a vector indicating the number of potentially-spam disputes
 	/// each validator is participating in. Potentially-spam disputes are remote disputes which have
 	/// fewer than `byzantine_threshold + 1` validators.
@@ -565,7 +554,11 @@ impl<T: Config> Pallet<T> {
 				dispute.concluded_at = Some(now);
 				<Disputes<T>>::insert(session_index, candidate_hash, &dispute);
 
-				if <Included<T>>::contains_key(&session_index, &candidate_hash) {
+				let is_local = <shared::Pallet<T>>::is_included_candidate(
+					&session_index,
+					&candidate_hash,
+				);
+				if is_local {
 					// Local disputes don't count towards spam.
 
 					weight += T::DbWeight::get().reads_writes(1, 1);
@@ -632,7 +625,7 @@ impl<T: Config> Pallet<T> {
 
 				// This is larger, and will be extracted to the `shared` module for more proper pruning.
 				// TODO: https://github.com/paritytech/polkadot/issues/3469
-				<Included<T>>::remove_prefix(to_prune, None);
+				shared::Pallet::<T>::prune_included_candidates(to_prune);
 				SpamSlots::<T>::remove(to_prune);
 			}
 
@@ -787,7 +780,7 @@ impl<T: Config> Pallet<T> {
 		};
 
 		// Apply spam slot changes. Bail early if too many occupied.
-		let is_local = <Included<T>>::contains_key(&set.session, &set.candidate_hash);
+		let is_local = <shared::Pallet<T>>::is_included_candidate(&set.session, &set.candidate_hash);
 		if !is_local {
 			let mut spam_slots: Vec<u32> =
 				SpamSlots::<T>::get(&set.session).unwrap_or_else(|| vec![0; n_validators]);
@@ -919,7 +912,10 @@ impl<T: Config> Pallet<T> {
 		};
 
 		// Apply spam slot changes. Bail early if too many occupied.
-		let is_local = <Included<T>>::contains_key(&set.session, &set.candidate_hash);
+		let is_local = <shared::Pallet<T>>::is_included_candidate(
+			&set.session,
+			&set.candidate_hash,
+		);
 		if !is_local {
 			let mut spam_slots: Vec<u32> =
 				SpamSlots::<T>::get(&set.session).unwrap_or_else(|| vec![0; n_validators]);
@@ -993,7 +989,10 @@ impl<T: Config> Pallet<T> {
 
 		// Freeze if just concluded against some local candidate
 		if summary.new_flags.contains(DisputeStateFlags::AGAINST_SUPERMAJORITY) {
-			if let Some(revert_to) = <Included<T>>::get(&set.session, &set.candidate_hash) {
+			if let Some((revert_to, _)) = <shared::Pallet<T>>::get_included_candidate(
+				&set.session,
+				&set.candidate_hash,
+			) {
 				Self::revert_and_freeze(revert_to);
 			}
 		}
@@ -1006,19 +1005,11 @@ impl<T: Config> Pallet<T> {
 		<Disputes<T>>::iter().collect()
 	}
 
-	pub(crate) fn note_included(
+	pub(crate) fn process_included(
 		session: SessionIndex,
 		candidate_hash: CandidateHash,
-		included_in: T::BlockNumber,
+		revert_to: T::BlockNumber,
 	) {
-		if included_in.is_zero() {
-			return
-		}
-
-		let revert_to = included_in - One::one();
-
-		<Included<T>>::insert(&session, &candidate_hash, revert_to);
-
 		// If we just included a block locally which has a live dispute, decrement spam slots
 		// for any involved validators, if the dispute is not already confirmed by f + 1.
 		if let Some(state) = <Disputes<T>>::get(&session, candidate_hash) {
@@ -1120,6 +1111,7 @@ fn check_signature(
 #[cfg(test)]
 mod tests {
 	use super::*;
+	use primitives::v1::CoreIndex;
 	use crate::mock::{
 		new_test_ext, AccountId, AllPallets, Initializer, MockGenesisConfig, System, Test,
 		PUNISH_VALIDATORS_AGAINST, PUNISH_VALIDATORS_FOR, PUNISH_VALIDATORS_INCONCLUSIVE,
@@ -1171,6 +1163,21 @@ mod tests {
 		}
 	}
 
+	fn include_candidate(session: SessionIndex, candidate_hash: &CandidateHash, block_number: BlockNumber) {
+		if let Some(revert_to) = shared::Pallet::<Test>::note_included_candidate(
+			session,
+			candidate_hash.clone(),
+			block_number,
+			CoreIndex(0),
+		) {
+			Pallet::<Test>::process_included(
+				session,
+				candidate_hash.clone(),
+				revert_to,
+			);
+		}
+	}
+
 	#[test]
 	fn test_dispute_state_flag_from_state() {
 		assert_eq!(
@@ -1469,13 +1476,13 @@ mod tests {
 			let v0 = <ValidatorId as CryptoType>::Pair::generate().0;
 
 			let candidate_hash = CandidateHash(sp_core::H256::repeat_byte(1));
-			Pallet::<Test>::note_included(0, candidate_hash.clone(), 0);
-			Pallet::<Test>::note_included(1, candidate_hash.clone(), 1);
-			Pallet::<Test>::note_included(2, candidate_hash.clone(), 2);
-			Pallet::<Test>::note_included(3, candidate_hash.clone(), 3);
-			Pallet::<Test>::note_included(4, candidate_hash.clone(), 4);
-			Pallet::<Test>::note_included(5, candidate_hash.clone(), 5);
-			Pallet::<Test>::note_included(6, candidate_hash.clone(), 5);
+			include_candidate(0, &candidate_hash, 0);
+			include_candidate(1, &candidate_hash, 1);
+			include_candidate(2, &candidate_hash, 2);
+			include_candidate(3, &candidate_hash, 3);
+			include_candidate(4, &candidate_hash, 4);
+			include_candidate(5, &candidate_hash, 5);
+			include_candidate(6, &candidate_hash, 5);
 
 			run_to_block(7, |b| {
 				// a new session at each block
@@ -1485,13 +1492,13 @@ mod tests {
 			// current session is 7,
 			// we keep for dispute_period + 1 session and we remove in on_finalize
 			// thus we keep info for session 3, 4, 5, 6, 7.
-			assert_eq!(Included::<Test>::iter_prefix(0).count(), 0);
-			assert_eq!(Included::<Test>::iter_prefix(1).count(), 0);
-			assert_eq!(Included::<Test>::iter_prefix(2).count(), 0);
-			assert_eq!(Included::<Test>::iter_prefix(3).count(), 1);
-			assert_eq!(Included::<Test>::iter_prefix(4).count(), 1);
-			assert_eq!(Included::<Test>::iter_prefix(5).count(), 1);
-			assert_eq!(Included::<Test>::iter_prefix(6).count(), 1);
+			assert_eq!(shared::Pallet::<Test>::included_candidates_iter_prefix(0).count(), 0);
+			assert_eq!(shared::Pallet::<Test>::included_candidates_iter_prefix(1).count(), 0);
+			assert_eq!(shared::Pallet::<Test>::included_candidates_iter_prefix(2).count(), 0);
+			assert_eq!(shared::Pallet::<Test>::included_candidates_iter_prefix(3).count(), 1);
+			assert_eq!(shared::Pallet::<Test>::included_candidates_iter_prefix(4).count(), 1);
+			assert_eq!(shared::Pallet::<Test>::included_candidates_iter_prefix(5).count(), 1);
+			assert_eq!(shared::Pallet::<Test>::included_candidates_iter_prefix(6).count(), 1);
 		});
 	}
 
@@ -1593,7 +1600,7 @@ mod tests {
 	}
 
 	#[test]
-	fn test_freeze_on_note_included() {
+	fn test_freeze_on_process_included() {
 		new_test_ext(Default::default()).execute_with(|| {
 			let v0 = <ValidatorId as CryptoType>::Pair::generate().0;
 
@@ -1623,7 +1630,7 @@ mod tests {
 			}];
 			assert!(Pallet::<Test>::provide_multi_dispute_data(stmts).is_ok());
 
-			Pallet::<Test>::note_included(3, candidate_hash.clone(), 3);
+			include_candidate(3, &candidate_hash, 3);
 			assert_eq!(Frozen::<Test>::get(), Some(2));
 		});
 	}
@@ -1640,7 +1647,7 @@ mod tests {
 
 			let candidate_hash = CandidateHash(sp_core::H256::repeat_byte(1));
 
-			Pallet::<Test>::note_included(3, candidate_hash.clone(), 3);
+			include_candidate(3, &candidate_hash, 3);
 
 			// v0 votes for 3
 			let stmts = vec![DisputeStatementSet {
@@ -1669,7 +1676,7 @@ mod tests {
 	// * provide_multi_dispute: with success scenario
 	// * disputes: correctness of datas
 	// * could_be_invalid: correctness of datas
-	// * note_included: decrement spam correctly
+	// * process_included: decrement spam correctly
 	// * spam slots: correctly incremented and decremented
 	// * ensure rewards and punishment are correctly called.
 	#[test]
@@ -1945,7 +1952,7 @@ mod tests {
 
 			// Ensure inclusion removes spam slots
 			assert_eq!(SpamSlots::<Test>::get(4), Some(vec![0, 0, 0, 1]));
-			Pallet::<Test>::note_included(4, candidate_hash.clone(), 4);
+			include_candidate(4, &candidate_hash, 4);
 			assert_eq!(SpamSlots::<Test>::get(4), Some(vec![0, 0, 0, 0]));
 
 			// Ensure the reward_validator function was correctly called
diff --git a/runtime/parachains/src/paras_inherent.rs b/runtime/parachains/src/paras_inherent.rs
index 82e21cbe1805..67daa096a4ec 100644
--- a/runtime/parachains/src/paras_inherent.rs
+++ b/runtime/parachains/src/paras_inherent.rs
@@ -211,8 +211,22 @@ pub mod pallet {
 
 			// Inform the disputes module of all included candidates.
 			let now = <frame_system::Pallet<T>>::block_number();
-			for (_, candidate_hash) in &freed_concluded {
-				T::DisputesHandler::note_included(current_session, *candidate_hash, now);
+			for (core_index, candidate_hash) in &freed_concluded {
+				let revert_to = match shared::Pallet::<T>::note_included_candidate(
+					current_session,
+					*candidate_hash,
+					now,
+					*core_index,
+				) {
+					Some(revert_to) => revert_to,
+					None => continue,
+				};
+
+				T::DisputesHandler::process_included(
+					current_session,
+					*candidate_hash,
+					revert_to,
+				);
 			}
 
 			// Handle timeouts for any availability core work.
diff --git a/runtime/parachains/src/shared.rs b/runtime/parachains/src/shared.rs
index 9b5aa5c59629..65cf7d0ae87f 100644
--- a/runtime/parachains/src/shared.rs
+++ b/runtime/parachains/src/shared.rs
@@ -20,8 +20,9 @@
 //! dependent on any of the other pallets.
 
 use frame_support::pallet_prelude::*;
-use primitives::v1::{SessionIndex, ValidatorId, ValidatorIndex};
+use primitives::v1::{CandidateHash, CoreIndex, SessionIndex, ValidatorId, ValidatorIndex};
 use sp_std::vec::Vec;
+use sp_runtime::traits::{One, Zero};
 
 use rand::{seq::SliceRandom, SeedableRng};
 use rand_chacha::ChaCha20Rng;
@@ -64,6 +65,18 @@ pub mod pallet {
 	#[pallet::getter(fn active_validator_keys)]
 	pub(super) type ActiveValidatorKeys<T: Config> = StorageValue<_, Vec<ValidatorId>, ValueQuery>;
 
+	/// All included blocks on the chain, mapped to the core index the candidate was assigned to
+	/// and the block number in this chain that should be reverted back to if the candidate is
+	/// disputed and determined to be invalid.
+	#[pallet::storage]
+	#[pallet::getter(fn included_candidates)]
+	pub(super) type IncludedCandidates<T: Config> = StorageDoubleMap<
+		_,
+		Twox64Concat, SessionIndex,
+		Blake2_128Concat, CandidateHash,
+		(T::BlockNumber, CoreIndex),
+	>;
+
 	#[pallet::call]
 	impl<T: Config> Pallet<T> {}
 }
@@ -114,6 +127,51 @@ impl<T: Config> Pallet<T> {
 		Self::session_index().saturating_add(SESSION_DELAY)
 	}
 
+	/// Records an included candidate, returning the block height that should be reverted to if the
+	/// block is found to be invalid. This method will return `None` if and only if `included_in`
+	/// is zero.
+	pub(crate) fn note_included_candidate(
+		session: SessionIndex,
+		candidate_hash: CandidateHash,
+		included_in: T::BlockNumber,
+		core_index: CoreIndex,
+	) -> Option<T::BlockNumber> {
+		if included_in.is_zero() { return None }
+
+		let revert_to = included_in - One::one();
+		<IncludedCandidates<T>>::insert(&session, &candidate_hash, (revert_to, core_index));
+
+		Some(revert_to)
+	}
+
+	/// Returns true if a candidate hash exists for the given session and candidate hash.
+	pub(crate) fn is_included_candidate(
+		session: &SessionIndex,
+		candidate_hash: &CandidateHash,
+	) -> bool {
+		<IncludedCandidates<T>>::contains_key(session, candidate_hash)
+	}
+
+	/// Returns the revert-to height and core index for a given candidate, if one exists.
+	pub(crate) fn get_included_candidate(
+		session: &SessionIndex,
+		candidate_hash: &CandidateHash,
+	) -> Option<(T::BlockNumber, CoreIndex)> {
+		<IncludedCandidates<T>>::get(session, candidate_hash)
+	}
+
+	/// Prunes all candidates that were included in the `to_prune` session.
+	pub(crate) fn prune_included_candidates(to_prune: SessionIndex) {
+		<IncludedCandidates<T>>::remove_prefix(to_prune, None);
+	}
+
+	#[cfg(test)]
+	pub(crate) fn included_candidates_iter_prefix(
+		session: SessionIndex,
+	) -> impl Iterator<Item=(CandidateHash, (T::BlockNumber, CoreIndex))> {
+		<IncludedCandidates<T>>::iter_prefix(session)
+	}
+
 	/// Test function for setting the current session index.
 	#[cfg(any(feature = "std", feature = "runtime-benchmarks", test))]
 	pub fn set_session_index(index: SessionIndex) {
@@ -227,4 +285,60 @@ mod tests {
 			);
 		});
 	}
+
+	#[test]
+	fn note_included_candidate_ignores_block_zero() {
+		new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+			let session = 1;
+			let candidate_hash = CandidateHash(sp_core::H256::repeat_byte(1));
+			assert!(ParasShared::note_included_candidate(session, candidate_hash, 0, CoreIndex(0)).is_none());
+			assert!(!ParasShared::is_included_candidate(&session, &candidate_hash));
+			assert!(ParasShared::get_included_candidate(&session, &candidate_hash).is_none());
+		});
+	}
+
+	#[test]
+	fn note_included_candidate_stores_block_non_zero_with_revert_to_minus_one() {
+		new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+			let session = 1;
+			let candidate_hash = CandidateHash(sp_core::H256::repeat_byte(1));
+			let block_number = 1;
+			let core_index = CoreIndex(2);
+			assert_eq!(
+				ParasShared::note_included_candidate(session, candidate_hash, block_number, core_index),
+				Some(block_number - 1),
+			);
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash));
+			assert_eq!(
+				ParasShared::get_included_candidate(&session, &candidate_hash),
+				Some((block_number - 1, core_index)),
+			);
+		});
+	}
+
+	#[test]
+	fn prune_included_candidate_removes_all_candidates_with_same_session() {
+		new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+			let candidate_hash1 = CandidateHash(sp_core::H256::repeat_byte(1));
+			let candidate_hash2 = CandidateHash(sp_core::H256::repeat_byte(2));
+			let candidate_hash3 = CandidateHash(sp_core::H256::repeat_byte(3));
+
+			assert!(ParasShared::note_included_candidate(1, candidate_hash1, 1, CoreIndex(0)).is_some());
+			assert!(ParasShared::note_included_candidate(1, candidate_hash2, 1, CoreIndex(0)).is_some());
+			assert!(ParasShared::note_included_candidate(2, candidate_hash3, 2, CoreIndex(0)).is_some());
+
+			assert_eq!(ParasShared::included_candidates_iter_prefix(1).count(), 2);
+			assert_eq!(ParasShared::included_candidates_iter_prefix(2).count(), 1);
+
+			ParasShared::prune_included_candidates(1);
+
+			assert_eq!(ParasShared::included_candidates_iter_prefix(1).count(), 0);
+			assert_eq!(ParasShared::included_candidates_iter_prefix(2).count(), 1);
+
+			ParasShared::prune_included_candidates(2);
+
+			assert_eq!(ParasShared::included_candidates_iter_prefix(1).count(), 0);
+			assert_eq!(ParasShared::included_candidates_iter_prefix(2).count(), 0);
+		});
+	}
 }

From 58f8b5f6b8f440e48029237807fb6c78394d1dc3 Mon Sep 17 00:00:00 2001
From: Lldenaurois <Ljdenaurois@gmail.com>
Date: Tue, 3 Aug 2021 13:58:41 -0400
Subject: [PATCH 2/4] Rustfmt changes

---
 runtime/parachains/src/disputes.rs | 40 +++++++++++++-----------------
 runtime/parachains/src/shared.rs   | 36 +++++++++++++++++++--------
 2 files changed, 43 insertions(+), 33 deletions(-)

diff --git a/runtime/parachains/src/disputes.rs b/runtime/parachains/src/disputes.rs
index 3bad15b90aa9..631bc74f63e8 100644
--- a/runtime/parachains/src/disputes.rs
+++ b/runtime/parachains/src/disputes.rs
@@ -17,10 +17,9 @@
 //! Runtime component for handling disputes of parachain candidates.
 
 use crate::{
-	shared,
 	configuration::{self, HostConfiguration},
 	initializer::SessionChangeNotification,
-	session_info,
+	session_info, shared,
 };
 use bitvec::{bitvec, order::Lsb0 as BitOrderLsb0};
 use frame_support::{ensure, traits::Get, weights::Weight};
@@ -554,10 +553,8 @@ impl<T: Config> Pallet<T> {
 				dispute.concluded_at = Some(now);
 				<Disputes<T>>::insert(session_index, candidate_hash, &dispute);
 
-				let is_local = <shared::Pallet<T>>::is_included_candidate(
-					&session_index,
-					&candidate_hash,
-				);
+				let is_local =
+					<shared::Pallet<T>>::is_included_candidate(&session_index, &candidate_hash);
 				if is_local {
 					// Local disputes don't count towards spam.
 
@@ -780,7 +777,8 @@ impl<T: Config> Pallet<T> {
 		};
 
 		// Apply spam slot changes. Bail early if too many occupied.
-		let is_local = <shared::Pallet<T>>::is_included_candidate(&set.session, &set.candidate_hash);
+		let is_local =
+			<shared::Pallet<T>>::is_included_candidate(&set.session, &set.candidate_hash);
 		if !is_local {
 			let mut spam_slots: Vec<u32> =
 				SpamSlots::<T>::get(&set.session).unwrap_or_else(|| vec![0; n_validators]);
@@ -912,10 +910,8 @@ impl<T: Config> Pallet<T> {
 		};
 
 		// Apply spam slot changes. Bail early if too many occupied.
-		let is_local = <shared::Pallet<T>>::is_included_candidate(
-			&set.session,
-			&set.candidate_hash,
-		);
+		let is_local =
+			<shared::Pallet<T>>::is_included_candidate(&set.session, &set.candidate_hash);
 		if !is_local {
 			let mut spam_slots: Vec<u32> =
 				SpamSlots::<T>::get(&set.session).unwrap_or_else(|| vec![0; n_validators]);
@@ -989,10 +985,9 @@ impl<T: Config> Pallet<T> {
 
 		// Freeze if just concluded against some local candidate
 		if summary.new_flags.contains(DisputeStateFlags::AGAINST_SUPERMAJORITY) {
-			if let Some((revert_to, _)) = <shared::Pallet<T>>::get_included_candidate(
-				&set.session,
-				&set.candidate_hash,
-			) {
+			if let Some((revert_to, _)) =
+				<shared::Pallet<T>>::get_included_candidate(&set.session, &set.candidate_hash)
+			{
 				Self::revert_and_freeze(revert_to);
 			}
 		}
@@ -1111,7 +1106,6 @@ fn check_signature(
 #[cfg(test)]
 mod tests {
 	use super::*;
-	use primitives::v1::CoreIndex;
 	use crate::mock::{
 		new_test_ext, AccountId, AllPallets, Initializer, MockGenesisConfig, System, Test,
 		PUNISH_VALIDATORS_AGAINST, PUNISH_VALIDATORS_FOR, PUNISH_VALIDATORS_INCONCLUSIVE,
@@ -1122,7 +1116,7 @@ mod tests {
 		traits::{OnFinalize, OnInitialize},
 	};
 	use frame_system::InitKind;
-	use primitives::v1::BlockNumber;
+	use primitives::v1::{BlockNumber, CoreIndex};
 	use sp_core::{crypto::CryptoType, Pair};
 
 	// All arguments for `initializer::on_new_session`
@@ -1163,18 +1157,18 @@ mod tests {
 		}
 	}
 
-	fn include_candidate(session: SessionIndex, candidate_hash: &CandidateHash, block_number: BlockNumber) {
+	fn include_candidate(
+		session: SessionIndex,
+		candidate_hash: &CandidateHash,
+		block_number: BlockNumber,
+	) {
 		if let Some(revert_to) = shared::Pallet::<Test>::note_included_candidate(
 			session,
 			candidate_hash.clone(),
 			block_number,
 			CoreIndex(0),
 		) {
-			Pallet::<Test>::process_included(
-				session,
-				candidate_hash.clone(),
-				revert_to,
-			);
+			Pallet::<Test>::process_included(session, candidate_hash.clone(), revert_to);
 		}
 	}
 
diff --git a/runtime/parachains/src/shared.rs b/runtime/parachains/src/shared.rs
index 65cf7d0ae87f..bb2ce8d2a3ea 100644
--- a/runtime/parachains/src/shared.rs
+++ b/runtime/parachains/src/shared.rs
@@ -21,8 +21,8 @@
 
 use frame_support::pallet_prelude::*;
 use primitives::v1::{CandidateHash, CoreIndex, SessionIndex, ValidatorId, ValidatorIndex};
-use sp_std::vec::Vec;
 use sp_runtime::traits::{One, Zero};
+use sp_std::vec::Vec;
 
 use rand::{seq::SliceRandom, SeedableRng};
 use rand_chacha::ChaCha20Rng;
@@ -72,8 +72,10 @@ pub mod pallet {
 	#[pallet::getter(fn included_candidates)]
 	pub(super) type IncludedCandidates<T: Config> = StorageDoubleMap<
 		_,
-		Twox64Concat, SessionIndex,
-		Blake2_128Concat, CandidateHash,
+		Twox64Concat,
+		SessionIndex,
+		Blake2_128Concat,
+		CandidateHash,
 		(T::BlockNumber, CoreIndex),
 	>;
 
@@ -136,7 +138,9 @@ impl<T: Config> Pallet<T> {
 		included_in: T::BlockNumber,
 		core_index: CoreIndex,
 	) -> Option<T::BlockNumber> {
-		if included_in.is_zero() { return None }
+		if included_in.is_zero() {
+			return None
+		}
 
 		let revert_to = included_in - One::one();
 		<IncludedCandidates<T>>::insert(&session, &candidate_hash, (revert_to, core_index));
@@ -168,7 +172,7 @@ impl<T: Config> Pallet<T> {
 	#[cfg(test)]
 	pub(crate) fn included_candidates_iter_prefix(
 		session: SessionIndex,
-	) -> impl Iterator<Item=(CandidateHash, (T::BlockNumber, CoreIndex))> {
+	) -> impl Iterator<Item = (CandidateHash, (T::BlockNumber, CoreIndex))> {
 		<IncludedCandidates<T>>::iter_prefix(session)
 	}
 
@@ -291,7 +295,8 @@ mod tests {
 		new_test_ext(MockGenesisConfig::default()).execute_with(|| {
 			let session = 1;
 			let candidate_hash = CandidateHash(sp_core::H256::repeat_byte(1));
-			assert!(ParasShared::note_included_candidate(session, candidate_hash, 0, CoreIndex(0)).is_none());
+			assert!(ParasShared::note_included_candidate(session, candidate_hash, 0, CoreIndex(0))
+				.is_none());
 			assert!(!ParasShared::is_included_candidate(&session, &candidate_hash));
 			assert!(ParasShared::get_included_candidate(&session, &candidate_hash).is_none());
 		});
@@ -305,7 +310,12 @@ mod tests {
 			let block_number = 1;
 			let core_index = CoreIndex(2);
 			assert_eq!(
-				ParasShared::note_included_candidate(session, candidate_hash, block_number, core_index),
+				ParasShared::note_included_candidate(
+					session,
+					candidate_hash,
+					block_number,
+					core_index
+				),
 				Some(block_number - 1),
 			);
 			assert!(ParasShared::is_included_candidate(&session, &candidate_hash));
@@ -323,9 +333,15 @@ mod tests {
 			let candidate_hash2 = CandidateHash(sp_core::H256::repeat_byte(2));
 			let candidate_hash3 = CandidateHash(sp_core::H256::repeat_byte(3));
 
-			assert!(ParasShared::note_included_candidate(1, candidate_hash1, 1, CoreIndex(0)).is_some());
-			assert!(ParasShared::note_included_candidate(1, candidate_hash2, 1, CoreIndex(0)).is_some());
-			assert!(ParasShared::note_included_candidate(2, candidate_hash3, 2, CoreIndex(0)).is_some());
+			assert!(
+				ParasShared::note_included_candidate(1, candidate_hash1, 1, CoreIndex(0)).is_some()
+			);
+			assert!(
+				ParasShared::note_included_candidate(1, candidate_hash2, 1, CoreIndex(0)).is_some()
+			);
+			assert!(
+				ParasShared::note_included_candidate(2, candidate_hash3, 2, CoreIndex(0)).is_some()
+			);
 
 			assert_eq!(ParasShared::included_candidates_iter_prefix(1).count(), 2);
 			assert_eq!(ParasShared::included_candidates_iter_prefix(2).count(), 1);

From 9bf8442954eecc7cfb51b0158876991ba843cc2f Mon Sep 17 00:00:00 2001
From: Lldenaurois <Ljdenaurois@gmail.com>
Date: Mon, 9 Aug 2021 14:06:32 -0400
Subject: [PATCH 3/4] runtime/parachains: Add pruning logic

---
 runtime/parachains/src/disputes.rs       |  10 +-
 runtime/parachains/src/paras_inherent.rs |   9 +-
 runtime/parachains/src/shared.rs         | 243 +++++++++++++++++++++--
 3 files changed, 233 insertions(+), 29 deletions(-)

diff --git a/runtime/parachains/src/disputes.rs b/runtime/parachains/src/disputes.rs
index 631bc74f63e8..175307394b25 100644
--- a/runtime/parachains/src/disputes.rs
+++ b/runtime/parachains/src/disputes.rs
@@ -619,10 +619,12 @@ impl<T: Config> Pallet<T> {
 			for to_prune in to_prune {
 				// This should be small, as disputes are rare, so `None` is fine.
 				<Disputes<T>>::remove_prefix(to_prune, None);
-
-				// This is larger, and will be extracted to the `shared` module for more proper pruning.
-				// TODO: https://github.com/paritytech/polkadot/issues/3469
-				shared::Pallet::<T>::prune_included_candidates(to_prune);
+				// Mark the session as pruneable so that its candidates can be incrementally
+				// removed over the course of many block inclusions.
+				shared::Pallet::<T>::mark_session_pruneable(to_prune);
+				// TODO(ladi): remove this call, currently allows unit tests to pass. Need to
+				// figure out how to invoke paras_inherent::enter in run_to_block.
+				shared::Pallet::<T>::prune_ancient_sessions(shared::MAX_CANDIDATES_TO_PRUNE);
 				SpamSlots::<T>::remove(to_prune);
 			}
 
diff --git a/runtime/parachains/src/paras_inherent.rs b/runtime/parachains/src/paras_inherent.rs
index 67daa096a4ec..e9d1116280e5 100644
--- a/runtime/parachains/src/paras_inherent.rs
+++ b/runtime/parachains/src/paras_inherent.rs
@@ -222,11 +222,7 @@ pub mod pallet {
 					None => continue,
 				};
 
-				T::DisputesHandler::process_included(
-					current_session,
-					*candidate_hash,
-					revert_to,
-				);
+				T::DisputesHandler::process_included(current_session, *candidate_hash, revert_to);
 			}
 
 			// Handle timeouts for any availability core work.
@@ -281,6 +277,9 @@ pub mod pallet {
 			// And track that we've finished processing the inherent for this block.
 			Included::<T>::set(Some(()));
 
+			// Prune candidates incrementally with each block inclusion.
+			shared::Pallet::<T>::prune_ancient_sessions(shared::MAX_CANDIDATES_TO_PRUNE);
+
 			Ok(Some(
 				MINIMAL_INCLUSION_INHERENT_WEIGHT +
 					(backed_candidates_len * BACKED_CANDIDATE_WEIGHT),
diff --git a/runtime/parachains/src/shared.rs b/runtime/parachains/src/shared.rs
index bb2ce8d2a3ea..30a21f39ad67 100644
--- a/runtime/parachains/src/shared.rs
+++ b/runtime/parachains/src/shared.rs
@@ -36,6 +36,11 @@ pub use pallet::*;
 // which guarantees that at least one full session has passed before any changes are applied.
 pub(crate) const SESSION_DELAY: SessionIndex = 2;
 
+// `MAX_CANDIDATES_TO_PRUNE` is used to upper bound the number of ancient candidates that can be
+// pruned when a new block is included. This is used to distribute the cost of pruning over
+// multiple blocks rather than pruning all historical candidates upon starting a new session.
+pub(crate) const MAX_CANDIDATES_TO_PRUNE: usize = 200;
+
 #[frame_support::pallet]
 pub mod pallet {
 	use super::*;
@@ -79,6 +84,15 @@ pub mod pallet {
 		(T::BlockNumber, CoreIndex),
 	>;
 
+	/// The set of all past sessions that have yet to fully pruned. Sessions are added to this set
+	/// upon moving to new current session, and removed after all of their included candidates
+	/// have been removed. This set is tracked independently so that the cost of removing included
+	/// candidates can be amortized over multiple blocks rather than removing all candidates upon
+	/// transitioning to a new session.
+	#[pallet::storage]
+	#[pallet::getter(fn pruneable_sessions)]
+	pub(super) type PruneableSessions<T: Config> = StorageMap<_, Twox64Concat, SessionIndex, ()>;
+
 	#[pallet::call]
 	impl<T: Config> Pallet<T> {}
 }
@@ -129,6 +143,55 @@ impl<T: Config> Pallet<T> {
 		Self::session_index().saturating_add(SESSION_DELAY)
 	}
 
+	/// Adds a session that is no longer in the dispute window as pruneable. Subsequent block
+	/// inclusions will incrementally remove old candidates to avoid taking the performance hit of
+	/// removing all candidates at once.
+	pub(crate) fn mark_session_pruneable(session: SessionIndex) {
+		PruneableSessions::<T>::insert(session, ());
+	}
+
+	/// Prunes up to `max_candidates_to_prune` candidates from `IncludedCandidates` that belong to
+	/// non-active sessions.
+	pub(crate) fn prune_ancient_sessions(max_candidates_to_prune: usize) {
+		let mut to_prune = Vec::new();
+		let mut n_candidates = 0;
+		let mut incomplete_session = None;
+		for session in PruneableSessions::<T>::iter_keys() {
+			let mut hashes = Vec::new();
+			for candidate_hash in IncludedCandidates::<T>::iter_key_prefix(session) {
+				// Exit condition when this session still has more candidates to prune; mark the
+				// session as incomplete so the remaining candidates can be pruned in a subsequent
+				// invocation.
+				if n_candidates >= max_candidates_to_prune {
+					incomplete_session = Some(session);
+					break
+				}
+				hashes.push(candidate_hash);
+				n_candidates += 1;
+			}
+
+			to_prune.push((session, hashes));
+			// Exit condition when all candidates from this session were selected for pruning.
+			if n_candidates >= max_candidates_to_prune {
+				break
+			}
+		}
+
+		for (session, candidate_hashes) in to_prune {
+			for candidate_hash in candidate_hashes {
+				IncludedCandidates::<T>::remove(session, candidate_hash);
+			}
+
+			// Prune the session only if it was not marked as incomplete.
+			match incomplete_session {
+				Some(incomplete_session) if incomplete_session != session =>
+					PruneableSessions::<T>::remove(session),
+				None => PruneableSessions::<T>::remove(session),
+				_ => {},
+			}
+		}
+	}
+
 	/// Records an included candidate, returning the block height that should be reverted to if the
 	/// block is found to be invalid. This method will return `None` if and only if `included_in`
 	/// is zero.
@@ -164,9 +227,9 @@ impl<T: Config> Pallet<T> {
 		<IncludedCandidates<T>>::get(session, candidate_hash)
 	}
 
-	/// Prunes all candidates that were included in the `to_prune` session.
-	pub(crate) fn prune_included_candidates(to_prune: SessionIndex) {
-		<IncludedCandidates<T>>::remove_prefix(to_prune, None);
+	#[cfg(test)]
+	pub(crate) fn is_pruneable_session(session: &SessionIndex) -> bool {
+		<PruneableSessions<T>>::contains_key(session)
 	}
 
 	#[cfg(test)]
@@ -327,34 +390,174 @@ mod tests {
 	}
 
 	#[test]
-	fn prune_included_candidate_removes_all_candidates_with_same_session() {
+	fn prune_ancient_sessions_no_incomplete_session() {
 		new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+			let session = 1;
 			let candidate_hash1 = CandidateHash(sp_core::H256::repeat_byte(1));
 			let candidate_hash2 = CandidateHash(sp_core::H256::repeat_byte(2));
-			let candidate_hash3 = CandidateHash(sp_core::H256::repeat_byte(3));
+			let block_number = 1;
+			let core_index = CoreIndex(0);
 
-			assert!(
-				ParasShared::note_included_candidate(1, candidate_hash1, 1, CoreIndex(0)).is_some()
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session,
+					candidate_hash1,
+					block_number,
+					core_index,
+				),
+				Some(block_number - 1),
 			);
-			assert!(
-				ParasShared::note_included_candidate(1, candidate_hash2, 1, CoreIndex(0)).is_some()
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session,
+					candidate_hash2,
+					block_number,
+					core_index,
+				),
+				Some(block_number - 1),
+			);
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash2));
+
+			// Prune before any sessions are marked pruneable.
+			ParasShared::prune_ancient_sessions(2);
+
+			// Both candidates should still exist.
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash2));
+
+			// Mark the candidates' session as pruneable.
+			ParasShared::mark_session_pruneable(session);
+			assert!(ParasShared::is_pruneable_session(&session));
+
+			// Prune the sessions, which should remove both candidates. The session should not be
+			// marked as incomplete since there are exactly two candidates.
+			ParasShared::prune_ancient_sessions(2);
+
+			assert!(!ParasShared::is_pruneable_session(&session));
+			assert!(!ParasShared::is_included_candidate(&session, &candidate_hash1));
+			assert!(!ParasShared::is_included_candidate(&session, &candidate_hash2));
+		})
+	}
+
+	#[test]
+	fn prune_ancient_sessions_incomplete_session() {
+		new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+			let session = 1;
+			let candidate_hash1 = CandidateHash(sp_core::H256::repeat_byte(1));
+			let candidate_hash2 = CandidateHash(sp_core::H256::repeat_byte(2));
+			let block_number = 1;
+			let core_index = CoreIndex(0);
+
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session,
+					candidate_hash1,
+					block_number,
+					core_index,
+				),
+				Some(block_number - 1),
 			);
-			assert!(
-				ParasShared::note_included_candidate(2, candidate_hash3, 2, CoreIndex(0)).is_some()
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session,
+					candidate_hash2,
+					block_number,
+					core_index,
+				),
+				Some(block_number - 1),
 			);
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash2));
 
-			assert_eq!(ParasShared::included_candidates_iter_prefix(1).count(), 2);
-			assert_eq!(ParasShared::included_candidates_iter_prefix(2).count(), 1);
+			// Prune before any sessions are marked pruneable.
+			ParasShared::prune_ancient_sessions(1);
 
-			ParasShared::prune_included_candidates(1);
+			// Both candidates should still exist.
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash2));
 
-			assert_eq!(ParasShared::included_candidates_iter_prefix(1).count(), 0);
-			assert_eq!(ParasShared::included_candidates_iter_prefix(2).count(), 1);
+			// Mark the candidates' session as pruneable.
+			ParasShared::mark_session_pruneable(session);
+			assert!(ParasShared::is_pruneable_session(&session));
 
-			ParasShared::prune_included_candidates(2);
+			// Prune the sessions, which should remove one of the candidates. The session will be
+			// marked as incomplete so the session should remain unpruned.
+			ParasShared::prune_ancient_sessions(1);
 
-			assert_eq!(ParasShared::included_candidates_iter_prefix(1).count(), 0);
-			assert_eq!(ParasShared::included_candidates_iter_prefix(2).count(), 0);
-		});
+			assert!(ParasShared::is_pruneable_session(&session));
+			assert!(!ParasShared::is_included_candidate(&session, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash2));
+		})
+	}
+
+	#[test]
+	fn prune_ancient_sessions_complete_and_incomplete_sessions() {
+		new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+			let session1 = 1;
+			let session2 = 2;
+			let candidate_hash1 = CandidateHash(sp_core::H256::repeat_byte(1));
+			let candidate_hash2 = CandidateHash(sp_core::H256::repeat_byte(2));
+			let candidate_hash3 = CandidateHash(sp_core::H256::repeat_byte(3));
+			let block_number1 = 1;
+			let block_number2 = 2;
+			let core_index = CoreIndex(0);
+
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session1,
+					candidate_hash1,
+					block_number1,
+					core_index,
+				),
+				Some(block_number1 - 1),
+			);
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session2,
+					candidate_hash2,
+					block_number2,
+					core_index,
+				),
+				Some(block_number2 - 1),
+			);
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session2,
+					candidate_hash3,
+					block_number2,
+					core_index,
+				),
+				Some(block_number2 - 1),
+			);
+			assert!(ParasShared::is_included_candidate(&session1, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session2, &candidate_hash2));
+			assert!(ParasShared::is_included_candidate(&session2, &candidate_hash3));
+
+			// Prune before any sessions are marked pruneable.
+			ParasShared::prune_ancient_sessions(2);
+
+			// Both candidates should still exist.
+			assert!(ParasShared::is_included_candidate(&session1, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session2, &candidate_hash2));
+			assert!(ParasShared::is_included_candidate(&session2, &candidate_hash3));
+
+			// Mark the candidates' session as pruneable.
+			ParasShared::mark_session_pruneable(session1);
+			ParasShared::mark_session_pruneable(session2);
+			assert!(ParasShared::is_pruneable_session(&session1));
+			assert!(ParasShared::is_pruneable_session(&session2));
+
+			// Prune the sessions, which should remove one candidate from each session. The first
+			// session should be pruned while the second session will be marked as incomplete, and
+			// so should remain in the set of pruneable sessions.
+			ParasShared::prune_ancient_sessions(2);
+
+			assert!(!ParasShared::is_pruneable_session(&session1));
+			assert!(ParasShared::is_pruneable_session(&session2));
+			assert!(!ParasShared::is_included_candidate(&session1, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session2, &candidate_hash2));
+			assert!(!ParasShared::is_included_candidate(&session2, &candidate_hash3));
+		})
 	}
 }

From f8e65cdc19bb25b633a417d13437458081b68a8b Mon Sep 17 00:00:00 2001
From: Lldenaurois <Ljdenaurois@gmail.com>
Date: Tue, 10 Aug 2021 19:08:42 +0000
Subject: [PATCH 4/4] runtime/parachains - Add HistoricalBabeVrfs to shared
 module

---
 Cargo.lock                               |   2 +
 runtime/parachains/Cargo.toml            |   2 +
 runtime/parachains/src/paras_inherent.rs | 119 +++++++++++++++++++++++
 runtime/parachains/src/shared.rs         |  72 ++++++++++++--
 4 files changed, 188 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 985413bae76f..4cf753a9fdff 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -6971,6 +6971,8 @@ dependencies = [
  "serde_json",
  "sp-api",
  "sp-application-crypto",
+ "sp-consensus-babe",
+ "sp-consensus-vrf",
  "sp-core",
  "sp-inherents",
  "sp-io",
diff --git a/runtime/parachains/Cargo.toml b/runtime/parachains/Cargo.toml
index 49016163cadd..fb7f6968113d 100644
--- a/runtime/parachains/Cargo.toml
+++ b/runtime/parachains/Cargo.toml
@@ -21,6 +21,8 @@ sp-runtime = { git = "https://github.com/paritytech/substrate", branch = "master
 sp-session = { git = "https://github.com/paritytech/substrate", branch = "master", default-features = false }
 sp-staking = { git = "https://github.com/paritytech/substrate", branch = "master", default-features = false }
 sp-core = { git = "https://github.com/paritytech/substrate", branch = "master", default-features = false }
+sp-consensus-babe = { git = "https://github.com/paritytech/substrate", branch = "master", default-features = false }
+sp-consensus-vrf = { git = "https://github.com/paritytech/substrate", branch = "master", default-features = false }
 sp-keystore = { git = "https://github.com/paritytech/substrate", branch = "master", optional = true }
 
 pallet-authority-discovery = { git = "https://github.com/paritytech/substrate", branch = "master", default-features = false }
diff --git a/runtime/parachains/src/paras_inherent.rs b/runtime/parachains/src/paras_inherent.rs
index e9d1116280e5..c4c042eeb7cc 100644
--- a/runtime/parachains/src/paras_inherent.rs
+++ b/runtime/parachains/src/paras_inherent.rs
@@ -268,6 +268,8 @@ pub mod pallet {
 				<scheduler::Pallet<T>>::group_validators,
 			)?;
 
+			let vrfs = <shared::Pallet<T>>::extract_vrfs();
+			let _ = <shared::Pallet<T>>::note_vrfs(current_session, now, vrfs);
 			// Note which of the scheduled cores were actually occupied by a backed candidate.
 			<scheduler::Pallet<T>>::occupied(&occupied);
 
@@ -337,6 +339,11 @@ fn limit_backed_candidates<T: Config>(
 mod tests {
 	use super::*;
 
+	use primitives::v1::{CandidateHash, CoreIndex, Hash, SessionIndex, Slot};
+
+	use keyring::Sr25519Keyring;
+	use sp_consensus_vrf::schnorrkel::{VRFOutput, VRFProof};
+
 	use crate::mock::{new_test_ext, MockGenesisConfig, System, Test};
 
 	mod limit_backed_candidates {
@@ -505,5 +512,117 @@ mod tests {
 				assert_eq!(post_info.actual_weight.unwrap(), expected_weight);
 			});
 		}
+
+		// used for generating assignments where the validity of the VRF doesn't matter.
+		fn garbage_vrf(keyring: Sr25519Keyring) -> (VRFOutput, VRFProof) {
+			use sp_consensus_babe::Transcript;
+			let key = keyring.pair();
+			let key = key.as_ref();
+
+			use sp_runtime::ConsensusEngineId;
+			const BABE_ENGINE_ID: ConsensusEngineId = *b"BABE";
+			let (o, p, _) = key.vrf_sign(Transcript::new(&BABE_ENGINE_ID));
+			(VRFOutput(o.to_output()), VRFProof(p))
+		}
+
+		fn setup_vrf(
+			incomplete_session: bool,
+		) -> (
+			SessionIndex,
+			<Test as frame_system::Config>::BlockNumber,
+			<Test as frame_system::Config>::BlockNumber,
+		) {
+			use sp_consensus_babe::digests::{
+				CompatibleDigestItem, PreDigest, SecondaryVRFPreDigest,
+			};
+			use sp_runtime::{Digest, DigestItem};
+			let slot = Slot::from(1);
+			let mut digest = Digest::<<Test as frame_system::Config>::Hash>::default();
+
+			let (vrf_output, vrf_proof) = garbage_vrf(Sr25519Keyring::Alice);
+			digest.push(DigestItem::babe_pre_digest(PreDigest::SecondaryVRF(
+				SecondaryVRFPreDigest { authority_index: 0, slot, vrf_output, vrf_proof },
+			)));
+
+			let (vrf_output, vrf_proof) = garbage_vrf(Sr25519Keyring::Bob);
+			digest.push(DigestItem::babe_pre_digest(PreDigest::SecondaryVRF(
+				SecondaryVRFPreDigest { authority_index: 0, slot, vrf_output, vrf_proof },
+			)));
+
+			let (vrf_output, vrf_proof) = garbage_vrf(Sr25519Keyring::Charlie);
+			digest.push(DigestItem::babe_pre_digest(PreDigest::SecondaryVRF(
+				SecondaryVRFPreDigest { authority_index: 0, slot, vrf_output, vrf_proof },
+			)));
+
+			let block: <Test as frame_system::Config>::BlockNumber = 1;
+			let next_block: <Test as frame_system::Config>::BlockNumber = 2;
+
+			System::initialize(&block, &Default::default(), &digest, frame_system::InitKind::Full);
+
+			let vrfs = shared::Pallet::<Test>::extract_vrfs();
+			assert_eq!(vrfs.len(), 3);
+
+			let session: SessionIndex = 1;
+			assert!(shared::Pallet::<Test>::note_vrfs(session, block, vrfs));
+
+			// Assert that all VRFs have been stored
+			assert_eq!(shared::Pallet::<Test>::historical_vrfs(session, block).len(), 3);
+
+			System::initialize(
+				&next_block,
+				&Default::default(),
+				&digest,
+				frame_system::InitKind::Full,
+			);
+
+			let vrfs = shared::Pallet::<Test>::extract_vrfs();
+			assert_eq!(vrfs.len(), 3);
+
+			assert!(shared::Pallet::<Test>::note_vrfs(session, next_block, vrfs));
+
+			// Assert that all VRFs have been stored
+			assert_eq!(shared::Pallet::<Test>::historical_vrfs(session, block).len(), 3);
+			assert_eq!(shared::Pallet::<Test>::historical_vrfs(session, next_block).len(), 3);
+
+			if incomplete_session {
+				for i in 0..crate::shared::MAX_CANDIDATES_TO_PRUNE + 1 {
+					let candidate_hash = CandidateHash(Hash::repeat_byte(i as u8));
+					shared::Pallet::<Test>::note_included_candidate(
+						session,
+						candidate_hash,
+						block,
+						CoreIndex(0),
+					);
+				}
+			}
+
+			// Mark session as pruneable, prune that session
+			shared::Pallet::<Test>::mark_session_pruneable(session);
+			shared::Pallet::<Test>::prune_ancient_sessions(1);
+
+			(session, block, next_block)
+		}
+
+		#[test]
+		fn test_vrf_prune_on_complete_session() {
+			new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+				let (session, block, next_block) = setup_vrf(false);
+				// Since the session is complete, ensure that all VRFs have been removed.
+				// Note that if the session was incomplete, we would only prune the first block
+				assert_eq!(shared::Pallet::<Test>::historical_vrfs(session, block).len(), 0);
+				assert_eq!(shared::Pallet::<Test>::historical_vrfs(session, next_block).len(), 0);
+			});
+		}
+
+		#[test]
+		fn test_vrf_prune_on_incomplete_session() {
+			new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+				let (session, block, next_block) = setup_vrf(true);
+				// Since the session is incomplete, ensure that only one block has been removed.
+				// Note that if the session was complete, we would prune all blocks
+				assert_eq!(shared::Pallet::<Test>::historical_vrfs(session, block).len(), 0);
+				assert_eq!(shared::Pallet::<Test>::historical_vrfs(session, next_block).len(), 3);
+			});
+		}
 	}
 }
diff --git a/runtime/parachains/src/shared.rs b/runtime/parachains/src/shared.rs
index 30a21f39ad67..d056ee8d564a 100644
--- a/runtime/parachains/src/shared.rs
+++ b/runtime/parachains/src/shared.rs
@@ -21,6 +21,7 @@
 
 use frame_support::pallet_prelude::*;
 use primitives::v1::{CandidateHash, CoreIndex, SessionIndex, ValidatorId, ValidatorIndex};
+use sp_consensus_babe::digests::{CompatibleDigestItem, PreDigest};
 use sp_runtime::traits::{One, Zero};
 use sp_std::vec::Vec;
 
@@ -93,6 +94,17 @@ pub mod pallet {
 	#[pallet::getter(fn pruneable_sessions)]
 	pub(super) type PruneableSessions<T: Config> = StorageMap<_, Twox64Concat, SessionIndex, ()>;
 
+	#[pallet::storage]
+	#[pallet::getter(fn historical_babe_vrfs)]
+	pub(super) type HistoricalBabeVrfs<T: Config> = StorageDoubleMap<
+		_,
+		Blake2_128Concat,
+		SessionIndex,
+		Twox64Concat,
+		T::BlockNumber,
+		Vec<Vec<u8>>,
+	>;
+
 	#[pallet::call]
 	impl<T: Config> Pallet<T> {}
 }
@@ -153,8 +165,11 @@ impl<T: Config> Pallet<T> {
 	/// Prunes up to `max_candidates_to_prune` candidates from `IncludedCandidates` that belong to
 	/// non-active sessions.
 	pub(crate) fn prune_ancient_sessions(max_candidates_to_prune: usize) {
-		let mut to_prune = Vec::new();
+		let mut to_prune_candidates = Vec::new();
+		let mut to_prune_blocks = Vec::new();
+
 		let mut n_candidates = 0;
+
 		let mut incomplete_session = None;
 		for session in PruneableSessions::<T>::iter_keys() {
 			let mut hashes = Vec::new();
@@ -170,24 +185,33 @@ impl<T: Config> Pallet<T> {
 				n_candidates += 1;
 			}
 
-			to_prune.push((session, hashes));
+			for block_number in HistoricalBabeVrfs::<T>::iter_key_prefix(session) {
+				to_prune_blocks.push((session, block_number));
+				if incomplete_session.is_some() {
+					break
+				}
+			}
+
+			to_prune_candidates.push((session, hashes));
 			// Exit condition when all candidates from this session were selected for pruning.
 			if n_candidates >= max_candidates_to_prune {
 				break
 			}
 		}
 
-		for (session, candidate_hashes) in to_prune {
+		for (session, block_number) in to_prune_blocks {
+			HistoricalBabeVrfs::<T>::remove(session, block_number);
+		}
+
+		for (session, candidate_hashes) in to_prune_candidates {
 			for candidate_hash in candidate_hashes {
 				IncludedCandidates::<T>::remove(session, candidate_hash);
 			}
 
 			// Prune the session only if it was not marked as incomplete.
 			match incomplete_session {
-				Some(incomplete_session) if incomplete_session != session =>
-					PruneableSessions::<T>::remove(session),
-				None => PruneableSessions::<T>::remove(session),
-				_ => {},
+				Some(incomplete_session) if incomplete_session == session => {},
+				_ => PruneableSessions::<T>::remove(session),
 			}
 		}
 	}
@@ -239,6 +263,40 @@ impl<T: Config> Pallet<T> {
 		<IncludedCandidates<T>>::iter_prefix(session)
 	}
 
+	/// Records an included historical BABE VRF
+	pub fn extract_vrfs() -> Vec<Vec<u8>> {
+		<frame_system::Pallet<T>>::digest()
+			.logs
+			.into_iter()
+			.filter_map(|d| {
+				d.as_babe_pre_digest()
+					.map(|p| match p {
+						PreDigest::Primary(primary) => Some(primary.vrf_output.encode()),
+						PreDigest::SecondaryVRF(secondary) => Some(secondary.vrf_output.encode()),
+						PreDigest::SecondaryPlain(_) => None,
+					})
+					.flatten()
+			})
+			.collect()
+	}
+
+	pub fn note_vrfs(
+		session: SessionIndex,
+		block_number: T::BlockNumber,
+		vrfs: Vec<Vec<u8>>,
+	) -> bool {
+		if HistoricalBabeVrfs::<T>::get(session, block_number).is_none() {
+			HistoricalBabeVrfs::<T>::insert(session, block_number, vrfs);
+			true
+		} else {
+			false
+		}
+	}
+
+	pub fn historical_vrfs(session: SessionIndex, block_number: T::BlockNumber) -> Vec<Vec<u8>> {
+		HistoricalBabeVrfs::<T>::get(session, block_number).unwrap_or_default()
+	}
+
 	/// Test function for setting the current session index.
 	#[cfg(any(feature = "std", feature = "runtime-benchmarks", test))]
 	pub fn set_session_index(index: SessionIndex) {