Try-state for Referenda pallet (paritytech#13949)

* Try-state for Referenda pallet

* fix & more tests

* checking more stuff

* remove log

* two more tests

* Update frame/referenda/src/lib.rs

Co-authored-by: Muharem Ismailov <[email protected]>

* fixes

* new check

* merge fixes

* fix warning

* remove check

* Update frame/referenda/src/lib.rs

Co-authored-by: Gonçalo Pestana <[email protected]>

* separate into multiple functions

* clean up

* unnecessary return value specified

* fix

* Update frame/referenda/src/lib.rs

* fmt

* remove import

* fmt

* fix CI

* Update frame/referenda/src/lib.rs

Co-authored-by: Liam Aharon <[email protected]>

* last fixes

* :P

* fmt

---------

Co-authored-by: Muharem Ismailov <[email protected]>
Co-authored-by: Gonçalo Pestana <[email protected]>
Co-authored-by: Liam Aharon <[email protected]>

frame/referenda/src/benchmarking.rs

@@ -659,7 +659,7 @@ benchmarks_instance_pallet! {
 
 	impl_benchmark_test_suite!(
 		Referenda,
-		crate::mock::new_test_ext(),
+		crate::mock::ExtBuilder::default().build(),
 		crate::mock::Test
 	);
 }

frame/referenda/src/lib.rs

@@ -66,6 +66,7 @@
 
 use codec::{Codec, Encode};
 use frame_support::{
+	dispatch::DispatchResult,
 	ensure,
 	traits::{
 		schedule::{
@@ -416,6 +417,15 @@ pub mod pallet {
 		PreimageNotExist,
 	}
 
+	#[pallet::hooks]
+	impl<T: Config<I>, I: 'static> Hooks<BlockNumberFor<T>> for Pallet<T, I> {
+		#[cfg(feature = "try-runtime")]
+		fn try_state(_n: BlockNumberFor<T>) -> Result<(), sp_runtime::TryRuntimeError> {
+			Self::do_try_state()?;
+			Ok(())
+		}
+	}
+
 	#[pallet::call]
 	impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		/// Propose a referendum on a privileged action.
@@ -1032,7 +1042,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 	/// - If it's ready to be decided, start deciding;
 	/// - If it's not ready to be decided and non-deciding timeout has passed, fail;
 	/// - If it's ongoing and passing, ensure confirming; if at end of confirmation period, pass.
-	/// - If it's ongoing and not passing, stop confirning; if it has reached end time, fail.
+	/// - If it's ongoing and not passing, stop confirming; if it has reached end time, fail.
 	///
 	/// Weight will be a bit different depending on what it does, but it's designed so as not to
 	/// differ dramatically, especially if `MaxQueue` is kept small. In particular _there are no
@@ -1284,4 +1294,87 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 			Self::deposit_event(Event::<T, I>::MetadataCleared { index, hash });
 		}
 	}
+
+	/// Ensure the correctness of the state of this pallet.
+	///
+	/// The following assertions must always apply.
+	///
+	/// General assertions:
+	///
+	/// * [`ReferendumCount`] must always be equal to the number of referenda in
+	///   [`ReferendumInfoFor`].
+	/// * Referendum indices in [`MetadataOf`] must also be stored in [`ReferendumInfoFor`].
+	#[cfg(any(feature = "try-runtime", test))]
+	fn do_try_state() -> Result<(), sp_runtime::TryRuntimeError> {
+		ensure!(
+			ReferendumCount::<T, I>::get() as usize ==
+				ReferendumInfoFor::<T, I>::iter_keys().count(),
+			"Number of referenda in `ReferendumInfoFor` is different than `ReferendumCount`"
+		);
+
+		MetadataOf::<T, I>::iter_keys().try_for_each(|referendum_index| -> DispatchResult {
+			ensure!(
+				ReferendumInfoFor::<T, I>::contains_key(referendum_index),
+				"Referendum indices in `MetadataOf` must also be stored in `ReferendumInfoOf`"
+			);
+			Ok(())
+		})?;
+
+		Self::try_state_referenda_info()?;
+		Self::try_state_tracks()?;
+
+		Ok(())
+	}
+
+	/// Looking at referenda info:
+	///
+	/// - Data regarding ongoing phase:
+	///
+	/// * There must exist track info for the track of the referendum.
+	/// * The deciding stage has to begin before confirmation period.
+	/// * If alarm is set the nudge call has to be at most [`UndecidingTimeout`] blocks away
+	///  from the submission block.
+	#[cfg(any(feature = "try-runtime", test))]
+	fn try_state_referenda_info() -> Result<(), sp_runtime::TryRuntimeError> {
+		ReferendumInfoFor::<T, I>::iter().try_for_each(|(_, referendum)| {
+			match referendum {
+				ReferendumInfo::Ongoing(status) => {
+					ensure!(
+						Self::track(status.track).is_some(),
+						"No track info for the track of the referendum."
+					);
+
+					if let Some(deciding) = status.deciding {
+						ensure!(
+							deciding.since <
+								deciding.confirming.unwrap_or(BlockNumberFor::<T>::max_value()),
+							"Deciding status cannot begin before confirming stage."
+						)
+					}
+				},
+				_ => {},
+			}
+			Ok(())
+		})
+	}
+
+	/// Looking at tracks:
+	///
+	/// * The referendum indices stored in [`TrackQueue`] must exist as keys in the
+	///  [`ReferendumInfoFor`] storage map.
+	#[cfg(any(feature = "try-runtime", test))]
+	fn try_state_tracks() -> Result<(), sp_runtime::TryRuntimeError> {
+		T::Tracks::tracks().iter().try_for_each(|track| {
+			TrackQueue::<T, I>::get(track.0).iter().try_for_each(
+				|(referendum_index, _)| -> Result<(), sp_runtime::TryRuntimeError> {
+					ensure!(
+					ReferendumInfoFor::<T, I>::contains_key(referendum_index),
+					"`ReferendumIndex` inside the `TrackQueue` should be a key in `ReferendumInfoFor`"
+				);
+					Ok(())
+				},
+			)?;
+			Ok(())
+		})
+	}
 }

frame/referenda/src/migration.rs

@@ -199,17 +199,19 @@ pub mod test {
 
 	#[test]
 	fn migration_v0_to_v1_works() {
-		new_test_ext().execute_with(|| {
+		ExtBuilder::default().build_and_execute(|| {
 			// create and insert into the storage an ongoing referendum v0.
 			let status_v0 = create_status_v0();
 			let ongoing_v0 = v0::ReferendumInfoOf::<T, ()>::Ongoing(status_v0.clone());
+			ReferendumCount::<T, ()>::mutate(|x| x.saturating_inc());
 			v0::ReferendumInfoFor::<T, ()>::insert(2, ongoing_v0);
 			// create and insert into the storage an approved referendum v0.
 			let approved_v0 = v0::ReferendumInfoOf::<T, ()>::Approved(
 				123,
 				Deposit { who: 1, amount: 10 },
 				Some(Deposit { who: 2, amount: 20 }),
 			);
+			ReferendumCount::<T, ()>::mutate(|x| x.saturating_inc());
 			v0::ReferendumInfoFor::<T, ()>::insert(5, approved_v0);
 			// run migration from v0 to v1.
 			v1::MigrateV0ToV1::<T, ()>::on_runtime_upgrade();

frame/referenda/src/mock.rs

@@ -225,23 +225,32 @@ impl Config for Test {
 	type Tracks = TestTracksInfo;
 	type Preimages = Preimage;
 }
+pub struct ExtBuilder {}
 
-pub fn new_test_ext() -> sp_io::TestExternalities {
-	let mut t = frame_system::GenesisConfig::<Test>::default().build_storage().unwrap();
-	let balances = vec![(1, 100), (2, 100), (3, 100), (4, 100), (5, 100), (6, 100)];
-	pallet_balances::GenesisConfig::<Test> { balances }
-		.assimilate_storage(&mut t)
-		.unwrap();
-	let mut ext = sp_io::TestExternalities::new(t);
-	ext.execute_with(|| System::set_block_number(1));
-	ext
+impl Default for ExtBuilder {
+	fn default() -> Self {
+		Self {}
+	}
 }
 
-/// Execute the function two times, with `true` and with `false`.
-#[allow(dead_code)]
-pub fn new_test_ext_execute_with_cond(execute: impl FnOnce(bool) -> () + Clone) {
-	new_test_ext().execute_with(|| (execute.clone())(false));
-	new_test_ext().execute_with(|| execute(true));
+impl ExtBuilder {
+	pub fn build(self) -> sp_io::TestExternalities {
+		let mut t = frame_system::GenesisConfig::<Test>::default().build_storage().unwrap();
+		let balances = vec![(1, 100), (2, 100), (3, 100), (4, 100), (5, 100), (6, 100)];
+		pallet_balances::GenesisConfig::<Test> { balances }
+			.assimilate_storage(&mut t)
+			.unwrap();
+		let mut ext = sp_io::TestExternalities::new(t);
+		ext.execute_with(|| System::set_block_number(1));
+		ext
+	}
+
+	pub fn build_and_execute(self, test: impl FnOnce() -> ()) {
+		self.build().execute_with(|| {
+			test();
+			Referenda::do_try_state().unwrap();
+		})
+	}
 }
 
 #[derive(Encode, Debug, Decode, TypeInfo, Eq, PartialEq, Clone, MaxEncodedLen)]