Use CountedStorageMap instead of StorageMap for Jurors in court pallet

runtime/src/lib.rs

@@ -77,7 +77,7 @@ type Executive = frame_executive::Executive<
     frame_system::ChainContext<Runtime>,
     Runtime,
     AllPalletsWithSystem,
-    zrml_prediction_markets::migrations::RemoveDisputesOfResolvedMarkets<Runtime>,
+    zrml_court::migrations::JurorsCountedStorageMapMigration<Runtime>,
 >;
 
 type Header = generic::Header<BlockNumber, BlakeTwo256>;

zrml/court/src/lib.rs

@@ -12,6 +12,7 @@ mod benchmarks;
 mod court_pallet_api;
 mod juror;
 mod juror_status;
+pub mod migrations;
 mod mock;
 mod tests;
 pub mod weights;
@@ -32,7 +33,7 @@ mod pallet {
     use core::marker::PhantomData;
     use frame_support::{
         dispatch::DispatchResult,
-        pallet_prelude::{StorageDoubleMap, StorageMap, StorageValue, ValueQuery},
+        pallet_prelude::{CountedStorageMap, StorageDoubleMap, StorageValue, ValueQuery},
         traits::{
             BalanceStatus, Currency, Get, Hooks, IsType, NamedReservableCurrency, Randomness,
             StorageVersion,
@@ -58,7 +59,7 @@ mod pallet {
     const INITIAL_JURORS_NUM: usize = 3;
     const MAX_RANDOM_JURORS: usize = 13;
     /// The current storage version.
-    const STORAGE_VERSION: StorageVersion = StorageVersion::new(0);
+    const STORAGE_VERSION: StorageVersion = StorageVersion::new(1);
     // Weight used to increase the number of jurors for subsequent disputes
     // of the same market
     const SUBSEQUENT_JURORS_FACTOR: usize = 2;
@@ -93,7 +94,7 @@ mod pallet {
             if Jurors::<T>::get(&who).is_some() {
                 return Err(Error::<T>::JurorAlreadyExists.into());
             }
-            let jurors_num = Jurors::<T>::iter().count();
+            let jurors_num = Jurors::<T>::count() as usize;
             let jurors_num_plus_one = jurors_num.checked_add(1).ok_or(ArithmeticError::Overflow)?;
             let stake = Self::current_required_stake(jurors_num_plus_one);
             CurrencyOf::<T>::reserve_named(&RESERVE_ID, &who, stake)?;
@@ -525,7 +526,7 @@ mod pallet {
 
     /// Accounts that stake funds to decide outcomes.
     #[pallet::storage]
-    pub type Jurors<T: Config> = StorageMap<_, Blake2_128Concat, T::AccountId, Juror>;
+    pub type Jurors<T: Config> = CountedStorageMap<_, Blake2_128Concat, T::AccountId, Juror>;
 
     /// An extra layer of pseudo randomness.
     #[pallet::storage]

zrml/court/src/migrations.rs

@@ -0,0 +1,101 @@
+use crate::{Config, Jurors, Pallet};
+use frame_support::{
+    dispatch::Weight,
+    log,
+    pallet_prelude::PhantomData,
+    traits::{Get, OnRuntimeUpgrade, StorageVersion},
+};
+
+const COURT_REQUIRED_STORAGE_VERSION: u16 = 0;
+const COURT_NEXT_STORAGE_VERSION: u16 = 1;
+
+pub struct JurorsCountedStorageMapMigration<T>(PhantomData<T>);
+
+impl<T: Config> OnRuntimeUpgrade for JurorsCountedStorageMapMigration<T> {
+    fn on_runtime_upgrade() -> Weight
+    where
+        T: Config,
+    {
+        let mut total_weight = T::DbWeight::get().reads(1);
+
+        if StorageVersion::get::<Pallet<T>>() != COURT_REQUIRED_STORAGE_VERSION {
+            log::info!("Skipping storage migration of jurors of court already up to date");
+            return total_weight;
+        }
+        log::info!("Starting storage migration of jurors of court");
+        Jurors::<T>::initialize_counter();
+        total_weight = total_weight.saturating_add(
+            T::DbWeight::get().writes(Jurors::<T>::count().saturating_add(1) as u64),
+        );
+
+        StorageVersion::new(COURT_NEXT_STORAGE_VERSION).put::<Pallet<T>>();
+        total_weight = total_weight.saturating_add(T::DbWeight::get().writes(1));
+
+        log::info!("Completed storage migration of jurors of court");
+        total_weight
+    }
+
+    #[cfg(feature = "try-runtime")]
+    fn pre_upgrade() -> Result<(), &'static str> {
+        Ok(())
+    }
+
+    #[cfg(feature = "try-runtime")]
+    fn post_upgrade() -> Result<(), &'static str> {
+        let court_storage_version = StorageVersion::get::<Pallet<T>>();
+        assert_eq!(
+            court_storage_version, COURT_NEXT_STORAGE_VERSION,
+            "found unexpected court pallet storage version. Found: {:?}. Expected: {:?}",
+            court_storage_version, COURT_NEXT_STORAGE_VERSION,
+        );
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{mock::*, Juror};
+    use frame_support::Hashable;
+
+    #[test]
+    fn test_on_runtime_upgrade_on_untouched_chain() {
+        ExtBuilder::default().build().execute_with(|| {
+            setup_chain();
+            JurorsCountedStorageMapMigration::<Runtime>::on_runtime_upgrade();
+        });
+    }
+
+    #[test]
+    fn on_runtime_upgrade_updates_storage_versions() {
+        ExtBuilder::default().build().execute_with(|| {
+            setup_chain();
+            JurorsCountedStorageMapMigration::<Runtime>::on_runtime_upgrade();
+            assert_eq!(StorageVersion::get::<Pallet<Runtime>>(), COURT_NEXT_STORAGE_VERSION,);
+        });
+    }
+
+    #[test]
+    fn test_on_runtime_upgrade() {
+        ExtBuilder::default().build().execute_with(|| {
+            setup_chain();
+            let alice_hash = ALICE.blake2_128_concat();
+            let alice_juror = Juror { status: crate::JurorStatus::Ok };
+            frame_support::migration::put_storage_value(
+                b"Court",
+                b"Jurors",
+                &alice_hash,
+                alice_juror.clone(),
+            );
+            assert_eq!(Jurors::<Runtime>::count(), 0_u32);
+            JurorsCountedStorageMapMigration::<Runtime>::on_runtime_upgrade();
+            assert_eq!(Jurors::<Runtime>::count(), 1_u32);
+            assert_eq!(Jurors::<Runtime>::get(&ALICE), Some(alice_juror));
+        });
+    }
+
+    fn setup_chain() {
+        StorageVersion::new(COURT_REQUIRED_STORAGE_VERSION).put::<Pallet<Runtime>>();
+    }
+}

zrml/prediction-markets/src/lib.rs

@@ -67,7 +67,6 @@
 extern crate alloc;
 
 mod benchmarks;
-pub mod migrations;
 pub mod mock;
 mod tests;
 pub mod weights;