Refactor TryState

Author: ales-otf

pallets/subtensor/src/macros/hooks.rs

@@ -4,6 +4,9 @@ use frame_support::pallet_macros::pallet_section;
 /// This can later be imported into the pallet using [`import_section`].
 #[pallet_section]
 mod hooks {
+	#[cfg(feature = "try-runtime")]
+	use crate::utils::try_state::TryState;
+
     // ================
     // ==== Hooks =====
     // ================
@@ -82,7 +85,9 @@ mod hooks {
 
         #[cfg(feature = "try-runtime")]
         fn try_state(_n: BlockNumberFor<T>) -> Result<(), sp_runtime::TryRuntimeError> {
-            Self::check_accounting_invariants()?;
+            TryState::<T>::check_total_issuance()?;
+			// Disabled: https://github.com/opentensor/subtensor/pull/1166
+            // TryState::<T>::check_total_stake()?;
             Ok(())
         }
     }

pallets/subtensor/src/migrations/migrate_init_total_issuance.rs

@@ -53,9 +53,10 @@ pub(crate) fn migrate_init_total_issuance<T: Config>() -> Weight {
 
 pub mod initialise_total_issuance {
     use frame_support::pallet_prelude::Weight;
-    use frame_support::traits::{fungible, OnRuntimeUpgrade};
-    use sp_core::Get;
+    use frame_support::traits::OnRuntimeUpgrade;
 
+    #[cfg(feature = "try-runtime")]
+    use crate::utils::try_state::TryState;
     use crate::*;
 
     pub struct Migration<T: Config>(PhantomData<T>);
@@ -79,7 +80,7 @@ pub mod initialise_total_issuance {
         #[cfg(feature = "try-runtime")]
         fn post_upgrade(_state: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
             // Verify that all accounting invariants are satisfied after the migration
-            crate::Pallet::<T>::check_accounting_invariants()?;
+            TryState::<T>::check_total_issuance()?;
             Ok(())
         }
     }

pallets/subtensor/src/utils/mod.rs

@@ -2,4 +2,5 @@ use super::*;
 pub mod identity;
 pub mod misc;
 pub mod rate_limiting;
+#[cfg(feature = "try-runtime")]
 pub mod try_state;

pallets/subtensor/src/utils/try_state.rs

@@ -1,42 +1,17 @@
-use super::*;
-
-impl<T: Config> Pallet<T> {
-    /// Checks if the accounting invariants for [`TotalStake`], [`TotalSubnetLocked`], and [`TotalIssuance`] are correct.
-    ///
-    /// This function verifies that:
-    /// 1. The sum of all stakes matches the [`TotalStake`].
-    /// 2. The [`TotalSubnetLocked`] is correctly calculated.
-    /// 3. The [`TotalIssuance`] equals the sum of currency issuance, total stake, and total subnet locked.
-    ///
-    /// # Returns
-    ///
-    /// Returns `Ok(())` if all invariants are correct, otherwise returns an error.
-    #[cfg(feature = "try-runtime")]
-    pub fn check_accounting_invariants() -> Result<(), sp_runtime::TryRuntimeError> {
+use core::marker::PhantomData;
         use frame_support::traits::fungible::Inspect;
 
-        // Disabled: https://github.com/opentensor/subtensor/pull/1166
-        //
-        // // Calculate the total staked amount
-        // let total_staked = SubnetTAO::<T>::iter().fold(0u64, |acc, (netuid, stake)| {
-        //     let acc = acc.saturating_add(stake);
-
-        //     if netuid == Self::get_root_netuid() {
-        //         // root network doesn't have initial pool TAO
-        //         acc
-        //     } else {
-        //         acc.saturating_sub(POOL_INITIAL_TAO)
-        //     }
-        // });
+use crate::subnets::subnet::POOL_INITIAL_TAO;
+use super::*;
 
-        // // Verify that the calculated total stake matches the stored TotalStake
-        // ensure!(
-        //     total_staked == TotalStake::<T>::get(),
-        //     "TotalStake does not match total staked",
-        // );
+pub(crate) struct TryState<T: Config>(PhantomData<T>);
 
+impl<T: Config> TryState<T> {
+	/// Checks [`TotalIssuance`] equals the sum of currency issuance, total stake, and total subnet
+	/// locked.
+    pub(crate) fn check_total_issuance() -> Result<(), sp_runtime::TryRuntimeError> {
         // Get the total subnet locked amount
-        let total_subnet_locked = Self::get_total_subnet_locked();
+        let total_subnet_locked = Pallet::<T>::get_total_subnet_locked();
 
         // Get the total currency issuance
         let currency_issuance = T::Currency::total_issuance();
@@ -59,11 +34,37 @@ impl<T: Config> Pallet<T> {
             expected_total_issuance.checked_sub(total_issuance)
         }
         .expect("LHS > RHS");
+
         ensure!(
             diff <= delta,
             "TotalIssuance diff greater than allowable delta",
         );
 
-        Ok(())
+		Ok(())
     }
+
+	/// Checks the sum of all stakes matches the [`TotalStake`].
+    pub(crate) fn check_total_stake() -> Result<(), sp_runtime::TryRuntimeError> {
+        // Calculate the total staked amount
+        let total_staked = SubnetTAO::<T>::iter().fold(0u64, |acc, (netuid, stake)| {
+            let acc = acc.saturating_add(stake);
+
+            if netuid == Pallet::<T>::get_root_netuid() {
+                // root network doesn't have initial pool TAO
+                acc
+            } else {
+                acc.saturating_sub(POOL_INITIAL_TAO)
+            }
+        });
+
+		log::warn!("total_staked: {}, TotalStake: {}", total_staked, TotalStake::<T>::get());
+
+        // Verify that the calculated total stake matches the stored TotalStake
+        ensure!(
+            total_staked == TotalStake::<T>::get(),
+            "TotalStake does not match total staked",
+        );
+
+		Ok(())
+	}
 }

scripts/try-runtime-upgrade.sh

@@ -49,7 +49,7 @@ do_try_runtime() {
   fi
 
   eval "try-runtime --runtime $runtime_wasm_path on-runtime-upgrade \
-    --no-weight-warnings --disable-spec-version-check --disable-idempotency-checks --checks=all \
+    --disable-spec-version-check --disable-idempotency-checks --checks=all \
     $chain_state"
 }