Used CountedStorageMap in pallet-staking (paritytech#10233)

* Removed counters and used CountedStorageMap instead.

* Little refactoring

* Update frame/staking/src/migrations.rs

Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>

* Removed redundant code to update counter for validator & nominator.

* Removed redundant code to update counter for validator & nominator.

* Removed unreachable code to inject the hashed prefix for nominator & validator.

* Removed redundant check for nominator & validator count.

* Generated `fn prefix_hash` for `CountedStorageMap`.

* Applied changes from `cargo fmt`

* Possible correct implementation of migration code

* Implemented fn module_prefix, storage_prefix and prefix_hash.

* Removed counted_map.rs

* Renamed `fn storage_prefix` to `storage_counter_prefix`.

* Update frame/support/src/storage/types/counted_map.rs

* Update frame/bags-list/remote-tests/src/snapshot.rs

* Update frame/support/src/storage/types/counted_map.rs

* Fixed errors.

Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>
Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>

frame/bags-list/remote-tests/src/snapshot.rs

@@ -27,6 +27,7 @@ pub async fn execute<Runtime: crate::RuntimeT, Block: BlockT + DeserializeOwned>
 	ws_url: String,
 ) {
 	use frame_support::storage::generator::StorageMap;
+
 	let mut ext = Builder::<Block>::new()
 		.mode(Mode::Online(OnlineConfig {
 			transport: ws_url.to_string().into(),
@@ -38,10 +39,10 @@ pub async fn execute<Runtime: crate::RuntimeT, Block: BlockT + DeserializeOwned>
 		}))
 		.inject_hashed_prefix(&<pallet_staking::Bonded<Runtime>>::prefix_hash())
 		.inject_hashed_prefix(&<pallet_staking::Ledger<Runtime>>::prefix_hash())
-		.inject_hashed_prefix(&<pallet_staking::Validators<Runtime>>::prefix_hash())
-		.inject_hashed_prefix(&<pallet_staking::Nominators<Runtime>>::prefix_hash())
-		.inject_hashed_key(&<pallet_staking::CounterForNominators<Runtime>>::hashed_key())
-		.inject_hashed_key(&<pallet_staking::CounterForValidators<Runtime>>::hashed_key())
+		.inject_hashed_prefix(&<pallet_staking::Validators<Runtime>>::map_storage_final_prefix())
+		.inject_hashed_prefix(&<pallet_staking::Nominators<Runtime>>::map_storage_final_prefix())
+		.inject_hashed_key(&<pallet_staking::Validators<Runtime>>::counter_storage_final_key())
+		.inject_hashed_key(&<pallet_staking::Nominators<Runtime>>::counter_storage_final_key())
 		.build()
 		.await
 		.unwrap();

frame/staking/src/benchmarking.rs

@@ -112,8 +112,8 @@ pub fn create_validator_with_nominators<T: Config>(
 
 	assert_eq!(new_validators.len(), 1);
 	assert_eq!(new_validators[0], v_stash, "Our validator was not selected!");
-	assert_ne!(CounterForValidators::<T>::get(), 0);
-	assert_ne!(CounterForNominators::<T>::get(), 0);
+	assert_ne!(Validators::<T>::count(), 0);
+	assert_ne!(Nominators::<T>::count(), 0);
 
 	// Give Era Points
 	let reward = EraRewardPoints::<T::AccountId> {
@@ -922,8 +922,8 @@ mod tests {
 			let count_validators = Validators::<Test>::iter().count();
 			let count_nominators = Nominators::<Test>::iter().count();
 
-			assert_eq!(count_validators, CounterForValidators::<Test>::get() as usize);
-			assert_eq!(count_nominators, CounterForNominators::<Test>::get() as usize);
+			assert_eq!(count_validators, Validators::<Test>::count() as usize);
+			assert_eq!(count_nominators, Nominators::<Test>::count() as usize);
 
 			assert_eq!(count_validators, v as usize);
 			assert_eq!(count_nominators, n as usize);

frame/staking/src/migrations.rs

@@ -70,10 +70,22 @@ pub mod v8 {
 
 pub mod v7 {
 	use super::*;
+	use frame_support::generate_storage_alias;
+
+	generate_storage_alias!(Staking, CounterForValidators => Value<u32>);
+	generate_storage_alias!(Staking, CounterForNominators => Value<u32>);
 
 	pub fn pre_migrate<T: Config>() -> Result<(), &'static str> {
-		assert!(CounterForValidators::<T>::get().is_zero(), "CounterForValidators already set.");
-		assert!(CounterForNominators::<T>::get().is_zero(), "CounterForNominators already set.");
+		assert!(
+			CounterForValidators::get().unwrap().is_zero(),
+			"CounterForValidators already set."
+		);
+		assert!(
+			CounterForNominators::get().unwrap().is_zero(),
+			"CounterForNominators already set."
+		);
+		assert!(Validators::<T>::count().is_zero(), "Validators already set.");
+		assert!(Nominators::<T>::count().is_zero(), "Nominators already set.");
 		assert!(StorageVersion::<T>::get() == Releases::V6_0_0);
 		Ok(())
 	}
@@ -83,8 +95,8 @@ pub mod v7 {
 		let validator_count = Validators::<T>::iter().count() as u32;
 		let nominator_count = Nominators::<T>::iter().count() as u32;
 
-		CounterForValidators::<T>::put(validator_count);
-		CounterForNominators::<T>::put(nominator_count);
+		CounterForValidators::put(validator_count);
+		CounterForNominators::put(nominator_count);
 
 		StorageVersion::<T>::put(Releases::V7_0_0);
 		log!(info, "Completed staking migration to Releases::V7_0_0");

frame/staking/src/mock.rs

@@ -536,8 +536,8 @@ fn post_conditions() {
 fn check_count() {
 	let nominator_count = Nominators::<Test>::iter().count() as u32;
 	let validator_count = Validators::<Test>::iter().count() as u32;
-	assert_eq!(nominator_count, CounterForNominators::<Test>::get());
-	assert_eq!(validator_count, CounterForValidators::<Test>::get());
+	assert_eq!(nominator_count, Nominators::<Test>::count());
+	assert_eq!(validator_count, Validators::<Test>::count());
 
 	// the voters that the `SortedListProvider` list is storing for us.
 	let external_voters = <Test as Config>::SortedListProvider::count();

frame/staking/src/pallet/impls.rs

@@ -665,8 +665,8 @@ impl<T: Config> Pallet<T> {
 		maybe_max_len: Option<usize>,
 	) -> Vec<(T::AccountId, VoteWeight, Vec<T::AccountId>)> {
 		let max_allowed_len = {
-			let nominator_count = CounterForNominators::<T>::get() as usize;
-			let validator_count = CounterForValidators::<T>::get() as usize;
+			let nominator_count = Nominators::<T>::count() as usize;
+			let validator_count = Validators::<T>::count() as usize;
 			let all_voter_count = validator_count.saturating_add(nominator_count);
 			maybe_max_len.unwrap_or(all_voter_count).min(all_voter_count)
 		};
@@ -765,18 +765,15 @@ impl<T: Config> Pallet<T> {
 	}
 
 	/// This function will add a nominator to the `Nominators` storage map,
-	/// [`SortedListProvider`] and keep track of the `CounterForNominators`.
+	/// and [`SortedListProvider`].
 	///
 	/// If the nominator already exists, their nominations will be updated.
 	///
 	/// NOTE: you must ALWAYS use this function to add nominator or update their targets. Any access
-	/// to `Nominators`, its counter, or `VoterList` outside of this function is almost certainly
+	/// to `Nominators` or `VoterList` outside of this function is almost certainly
 	/// wrong.
 	pub fn do_add_nominator(who: &T::AccountId, nominations: Nominations<T::AccountId>) {
 		if !Nominators::<T>::contains_key(who) {
-			// maybe update the counter.
-			CounterForNominators::<T>::mutate(|x| x.saturating_inc());
-
 			// maybe update sorted list. Error checking is defensive-only - this should never fail.
 			if T::SortedListProvider::on_insert(who.clone(), Self::weight_of(who)).is_err() {
 				log!(warn, "attempt to insert duplicate nominator ({:#?})", who);
@@ -790,53 +787,46 @@ impl<T: Config> Pallet<T> {
 	}
 
 	/// This function will remove a nominator from the `Nominators` storage map,
-	/// [`SortedListProvider`] and keep track of the `CounterForNominators`.
+	/// and [`SortedListProvider`].
 	///
 	/// Returns true if `who` was removed from `Nominators`, otherwise false.
 	///
 	/// NOTE: you must ALWAYS use this function to remove a nominator from the system. Any access to
-	/// `Nominators`, its counter, or `VoterList` outside of this function is almost certainly
+	/// `Nominators` or `VoterList` outside of this function is almost certainly
 	/// wrong.
 	pub fn do_remove_nominator(who: &T::AccountId) -> bool {
 		if Nominators::<T>::contains_key(who) {
 			Nominators::<T>::remove(who);
-			CounterForNominators::<T>::mutate(|x| x.saturating_dec());
 			T::SortedListProvider::on_remove(who);
 			debug_assert_eq!(T::SortedListProvider::sanity_check(), Ok(()));
-			debug_assert_eq!(CounterForNominators::<T>::get(), T::SortedListProvider::count());
+			debug_assert_eq!(Nominators::<T>::count(), T::SortedListProvider::count());
 			true
 		} else {
 			false
 		}
 	}
 
-	/// This function will add a validator to the `Validators` storage map, and keep track of the
-	/// `CounterForValidators`.
+	/// This function will add a validator to the `Validators` storage map.
 	///
 	/// If the validator already exists, their preferences will be updated.
 	///
 	/// NOTE: you must ALWAYS use this function to add a validator to the system. Any access to
-	/// `Validators`, its counter, or `VoterList` outside of this function is almost certainly
+	/// `Validators` or `VoterList` outside of this function is almost certainly
 	/// wrong.
 	pub fn do_add_validator(who: &T::AccountId, prefs: ValidatorPrefs) {
-		if !Validators::<T>::contains_key(who) {
-			CounterForValidators::<T>::mutate(|x| x.saturating_inc())
-		}
 		Validators::<T>::insert(who, prefs);
 	}
 
-	/// This function will remove a validator from the `Validators` storage map,
-	/// and keep track of the `CounterForValidators`.
+	/// This function will remove a validator from the `Validators` storage map.
 	///
 	/// Returns true if `who` was removed from `Validators`, otherwise false.
 	///
 	/// NOTE: you must ALWAYS use this function to remove a validator from the system. Any access to
-	/// `Validators`, its counter, or `VoterList` outside of this function is almost certainly
+	/// `Validators` or `VoterList` outside of this function is almost certainly
 	/// wrong.
 	pub fn do_remove_validator(who: &T::AccountId) -> bool {
 		if Validators::<T>::contains_key(who) {
 			Validators::<T>::remove(who);
-			CounterForValidators::<T>::mutate(|x| x.saturating_dec());
 			true
 		} else {
 			false
@@ -865,14 +855,6 @@ impl<T: Config> ElectionDataProvider<T::AccountId, BlockNumberFor<T>> for Pallet
 	fn voters(
 		maybe_max_len: Option<usize>,
 	) -> data_provider::Result<Vec<(T::AccountId, VoteWeight, Vec<T::AccountId>)>> {
-		debug_assert!(<Nominators<T>>::iter().count() as u32 == CounterForNominators::<T>::get());
-		debug_assert!(<Validators<T>>::iter().count() as u32 == CounterForValidators::<T>::get());
-		debug_assert_eq!(
-			CounterForNominators::<T>::get(),
-			T::SortedListProvider::count(),
-			"voter_count must be accurate",
-		);
-
 		// This can never fail -- if `maybe_max_len` is `Some(_)` we handle it.
 		let voters = Self::get_npos_voters(maybe_max_len);
 		debug_assert!(maybe_max_len.map_or(true, |max| voters.len() <= max));
@@ -881,7 +863,7 @@ impl<T: Config> ElectionDataProvider<T::AccountId, BlockNumberFor<T>> for Pallet
 	}
 
 	fn targets(maybe_max_len: Option<usize>) -> data_provider::Result<Vec<T::AccountId>> {
-		let target_count = CounterForValidators::<T>::get();
+		let target_count = Validators::<T>::count();
 
 		// We can't handle this case yet -- return an error.
 		if maybe_max_len.map_or(false, |max_len| target_count > max_len as u32) {
@@ -967,10 +949,9 @@ impl<T: Config> ElectionDataProvider<T::AccountId, BlockNumberFor<T>> for Pallet
 	fn clear() {
 		<Bonded<T>>::remove_all(None);
 		<Ledger<T>>::remove_all(None);
-		<Validators<T>>::remove_all(None);
-		<Nominators<T>>::remove_all(None);
-		<CounterForNominators<T>>::kill();
-		<CounterForValidators<T>>::kill();
+		<Validators<T>>::remove_all();
+		<Nominators<T>>::remove_all();
+
 		T::SortedListProvider::unsafe_clear();
 	}
 
@@ -1284,7 +1265,7 @@ impl<T: Config> SortedListProvider<T::AccountId> for UseNominatorsMap<T> {
 		Box::new(Nominators::<T>::iter().map(|(n, _)| n))
 	}
 	fn count() -> u32 {
-		CounterForNominators::<T>::get()
+		Nominators::<T>::count()
 	}
 	fn contains(id: &T::AccountId) -> bool {
 		Nominators::<T>::contains_key(id)
@@ -1309,10 +1290,10 @@ impl<T: Config> SortedListProvider<T::AccountId> for UseNominatorsMap<T> {
 	fn sanity_check() -> Result<(), &'static str> {
 		Ok(())
 	}
+
 	fn unsafe_clear() {
 		// NOTE: Caller must ensure this doesn't lead to too many storage accesses. This is a
 		// condition of SortedListProvider::unsafe_clear.
-		Nominators::<T>::remove_all(None);
-		CounterForNominators::<T>::take();
+		Nominators::<T>::remove_all();
 	}
 }

frame/staking/src/pallet/mod.rs

@@ -234,16 +234,10 @@ pub mod pallet {
 		StorageMap<_, Twox64Concat, T::AccountId, RewardDestination<T::AccountId>, ValueQuery>;
 
 	/// The map from (wannabe) validator stash key to the preferences of that validator.
-	///
-	/// When updating this storage item, you must also update the `CounterForValidators`.
 	#[pallet::storage]
 	#[pallet::getter(fn validators)]
 	pub type Validators<T: Config> =
-		StorageMap<_, Twox64Concat, T::AccountId, ValidatorPrefs, ValueQuery>;
-
-	/// A tracker to keep count of the number of items in the `Validators` map.
-	#[pallet::storage]
-	pub type CounterForValidators<T> = StorageValue<_, u32, ValueQuery>;
+		CountedStorageMap<_, Twox64Concat, T::AccountId, ValidatorPrefs, ValueQuery>;
 
 	/// The maximum validator count before we stop allowing new validators to join.
 	///
@@ -252,16 +246,10 @@ pub mod pallet {
 	pub type MaxValidatorsCount<T> = StorageValue<_, u32, OptionQuery>;
 
 	/// The map from nominator stash key to the set of stash keys of all validators to nominate.
-	///
-	/// When updating this storage item, you must also update the `CounterForNominators`.
 	#[pallet::storage]
 	#[pallet::getter(fn nominators)]
 	pub type Nominators<T: Config> =
-		StorageMap<_, Twox64Concat, T::AccountId, Nominations<T::AccountId>>;
-
-	/// A tracker to keep count of the number of items in the `Nominators` map.
-	#[pallet::storage]
-	pub type CounterForNominators<T> = StorageValue<_, u32, ValueQuery>;
+		CountedStorageMap<_, Twox64Concat, T::AccountId, Nominations<T::AccountId>>;
 
 	/// The maximum nominator count before we stop allowing new validators to join.
 	///
@@ -570,7 +558,7 @@ pub mod pallet {
 			// all voters are reported to the `SortedListProvider`.
 			assert_eq!(
 				T::SortedListProvider::count(),
-				CounterForNominators::<T>::get(),
+				Nominators::<T>::count(),
 				"not all genesis stakers were inserted into sorted list provider, something is wrong."
 			);
 		}
@@ -987,7 +975,7 @@ pub mod pallet {
 				// the runtime.
 				if let Some(max_validators) = MaxValidatorsCount::<T>::get() {
 					ensure!(
-						CounterForValidators::<T>::get() < max_validators,
+						Validators::<T>::count() < max_validators,
 						Error::<T>::TooManyValidators
 					);
 				}
@@ -1027,7 +1015,7 @@ pub mod pallet {
 				// the runtime.
 				if let Some(max_nominators) = MaxNominatorsCount::<T>::get() {
 					ensure!(
-						CounterForNominators::<T>::get() < max_nominators,
+						Nominators::<T>::count() < max_nominators,
 						Error::<T>::TooManyNominators
 					);
 				}
@@ -1610,7 +1598,7 @@ pub mod pallet {
 				let min_active_bond = if Nominators::<T>::contains_key(&stash) {
 					let max_nominator_count =
 						MaxNominatorsCount::<T>::get().ok_or(Error::<T>::CannotChillOther)?;
-					let current_nominator_count = CounterForNominators::<T>::get();
+					let current_nominator_count = Nominators::<T>::count();
 					ensure!(
 						threshold * max_nominator_count < current_nominator_count,
 						Error::<T>::CannotChillOther
@@ -1619,7 +1607,7 @@ pub mod pallet {
 				} else if Validators::<T>::contains_key(&stash) {
 					let max_validator_count =
 						MaxValidatorsCount::<T>::get().ok_or(Error::<T>::CannotChillOther)?;
-					let current_validator_count = CounterForValidators::<T>::get();
+					let current_validator_count = Validators::<T>::count();
 					ensure!(
 						threshold * max_validator_count < current_validator_count,
 						Error::<T>::CannotChillOther

frame/staking/src/testing_utils.rs

@@ -36,12 +36,11 @@ const SEED: u32 = 0;
 
 /// This function removes all validators and nominators from storage.
 pub fn clear_validators_and_nominators<T: Config>() {
-	Validators::<T>::remove_all(None);
-	CounterForValidators::<T>::kill();
+	Validators::<T>::remove_all();
 
 	// whenever we touch nominators counter we should update `T::SortedListProvider` as well.
-	Nominators::<T>::remove_all(None);
-	CounterForNominators::<T>::kill();
+	Nominators::<T>::remove_all();
+
 	// NOTE: safe to call outside block production
 	T::SortedListProvider::unsafe_clear();
 }

frame/staking/src/tests.rs

@@ -4300,8 +4300,8 @@ fn chill_other_works() {
 		.min_nominator_bond(1_000)
 		.min_validator_bond(1_500)
 		.build_and_execute(|| {
-			let initial_validators = CounterForValidators::<Test>::get();
-			let initial_nominators = CounterForNominators::<Test>::get();
+			let initial_validators = Validators::<Test>::count();
+			let initial_nominators = Nominators::<Test>::count();
 			for i in 0..15 {
 				let a = 4 * i;
 				let b = 4 * i + 1;
@@ -4424,8 +4424,8 @@ fn chill_other_works() {
 			));
 
 			// 16 people total because tests start with 2 active one
-			assert_eq!(CounterForNominators::<Test>::get(), 15 + initial_nominators);
-			assert_eq!(CounterForValidators::<Test>::get(), 15 + initial_validators);
+			assert_eq!(Nominators::<Test>::count(), 15 + initial_nominators);
+			assert_eq!(Validators::<Test>::count(), 15 + initial_validators);
 
 			// Users can now be chilled down to 7 people, so we try to remove 9 of them (starting
 			// with 16)
@@ -4437,23 +4437,23 @@ fn chill_other_works() {
 			}
 
 			// chill a nominator. Limit is not reached, not chill-able
-			assert_eq!(CounterForNominators::<Test>::get(), 7);
+			assert_eq!(Nominators::<Test>::count(), 7);
 			assert_noop!(
 				Staking::chill_other(Origin::signed(1337), 1),
 				Error::<Test>::CannotChillOther
 			);
 			// chill a validator. Limit is reached, chill-able.
-			assert_eq!(CounterForValidators::<Test>::get(), 9);
+			assert_eq!(Validators::<Test>::count(), 9);
 			assert_ok!(Staking::chill_other(Origin::signed(1337), 3));
 		})
 }
 
 #[test]
 fn capped_stakers_works() {
 	ExtBuilder::default().build_and_execute(|| {
-		let validator_count = CounterForValidators::<Test>::get();
+		let validator_count = Validators::<Test>::count();
 		assert_eq!(validator_count, 3);
-		let nominator_count = CounterForNominators::<Test>::get();
+		let nominator_count = Nominators::<Test>::count();
 		assert_eq!(nominator_count, 1);
 
 		// Change the maximums