[WIP]store pending root alpha

pallets/subtensor/src/coinbase/run_coinbase.rs

@@ -631,6 +631,10 @@ impl<T: Config> Pallet<T> {
                 tou64!(root_alpha).into(),
             );
 
+            PendingRootAlpha::<T>::mutate(&hotkey, |alpha| {
+                *alpha = alpha.saturating_add(tou64!(root_alpha).into())
+            });
+
             // Record root alpha dividends for this validator on this subnet.
             RootAlphaDividendsPerSubnet::<T>::mutate(netuid, &hotkey, |divs| {
                 *divs = divs.saturating_add(tou64!(root_alpha).into());

pallets/subtensor/src/lib.rs

@@ -1644,6 +1644,11 @@ pub mod pallet {
     pub type PendingRootAlphaDivs<T> =
         StorageMap<_, Identity, NetUid, AlphaCurrency, ValueQuery, DefaultZeroAlpha<T>>;
 
+    /// --- MAP ( hotkey ) --> pending_root_alpha
+    #[pallet::storage]
+    pub type PendingRootAlpha<T: Config> =
+        StorageMap<_, Blake2_128Concat, T::AccountId, u128, ValueQuery, DefaultZeroU128<T>>;
+
     /// --- MAP ( netuid ) --> pending_owner_cut
     #[pallet::storage]
     pub type PendingOwnerCut<T> =

pallets/subtensor/src/migrations/migrate_init_pending_root_alpha.rs

@@ -0,0 +1,79 @@
+use super::*;
+use alloc::{collections::BTreeMap, string::String};
+use frame_support::{traits::Get, weights::Weight};
+use substrate_fixed::types::I96F32;
+
+/// Migration to initialize PendingRootAlpha storage based on RootClaimed storage.
+/// This aggregates all RootClaimed values across all netuids and coldkeys for each hotkey.
+pub fn migrate_init_pending_root_alpha<T: Config>() -> Weight {
+    let migration_name = b"migrate_init_pending_root_alpha".to_vec();
+    let mut weight = T::DbWeight::get().reads(1);
+
+    // Check if the migration has already run
+    if HasMigrationRun::<T>::get(&migration_name) {
+        log::info!(
+            "Migration '{:?}' has already run. Skipping.",
+            String::from_utf8_lossy(&migration_name)
+        );
+        return weight;
+    }
+
+    log::info!(
+        "Running migration '{}'",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    // Aggregate RootClaimable values by hotkey
+    let mut root_claimable_alpha_map: BTreeMap<T::AccountId, u128> = BTreeMap::new();
+    for (hotkey, root_claimable) in RootClaimable::<T>::iter() {
+        // Sum rates as I96F32 (not converting to u128 first, which would truncate)
+        let claimable_rate: I96F32 = root_claimable
+            .values()
+            .fold(I96F32::from(0), |acc, x| acc.saturating_add(*x));
+
+        let root_stake = Pallet::<T>::get_stake_for_hotkey_on_subnet(&hotkey, NetUid::ROOT);
+        let total = claimable_rate.saturating_mul(I96F32::saturating_from_num(root_stake));
+
+        root_claimable_alpha_map.insert(hotkey, total.saturating_to_num::<u128>());
+        weight = weight.saturating_add(T::DbWeight::get().reads(1));
+    }
+
+    // Aggregate RootClaimed values by hotkey
+    // Key: hotkey, Value: sum of all RootClaimed values for that hotkey
+    let mut root_claimed_alpha_map: BTreeMap<T::AccountId, u128> = BTreeMap::new();
+
+    // Iterate over all RootClaimed entries: (netuid, hotkey, coldkey) -> claimed_value
+    for ((_netuid, hotkey, _coldkey), claimed_value) in RootClaimed::<T>::iter() {
+        // Aggregate the claimed value for this hotkey
+        root_claimed_alpha_map
+            .entry(hotkey.clone())
+            .and_modify(|total| *total = total.saturating_add(claimed_value))
+            .or_insert(claimed_value);
+
+        // Account for read operation
+        weight = weight.saturating_add(T::DbWeight::get().reads(1));
+    }
+
+    // Set PendingRootAlpha for each hotkey
+    let mut migrated_count = 0u64;
+    for (hotkey, claimable) in root_claimable_alpha_map {
+        let claimed = root_claimed_alpha_map.get(&hotkey).unwrap_or(&0);
+        let pending = claimable.saturating_sub(*claimed);
+        PendingRootAlpha::<T>::insert(&hotkey, pending);
+        migrated_count = migrated_count.saturating_add(1);
+    }
+
+    weight = weight.saturating_add(T::DbWeight::get().writes(migrated_count));
+
+    log::info!(
+        "Migration '{}' completed successfully. Initialized PendingRootAlpha for {} hotkeys.",
+        String::from_utf8_lossy(&migration_name),
+        migrated_count
+    );
+
+    // Mark the migration as completed
+    HasMigrationRun::<T>::insert(&migration_name, true);
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    weight
+}

pallets/subtensor/src/migrations/mod.rs

@@ -20,6 +20,7 @@ pub mod migrate_fix_is_network_member;
 pub mod migrate_fix_root_subnet_tao;
 pub mod migrate_fix_root_tao_and_alpha_in;
 pub mod migrate_fix_staking_hot_keys;
+pub mod migrate_init_pending_root_alpha;
 pub mod migrate_init_tao_flow;
 pub mod migrate_init_total_issuance;
 pub mod migrate_kappa_map_to_default;

pallets/subtensor/src/staking/claim_root.rs

@@ -1,7 +1,7 @@
 use super::*;
 use frame_support::weights::Weight;
 use sp_core::Get;
-use sp_std::collections::btree_set::BTreeSet;
+use sp_std::collections::{btree_map::BTreeMap, btree_set::BTreeSet};
 use substrate_fixed::types::I96F32;
 use subtensor_swap_interface::SwapHandler;
 
@@ -207,6 +207,10 @@ impl<T: Config> Pallet<T> {
         RootClaimed::<T>::mutate((netuid, hotkey, coldkey), |root_claimed| {
             *root_claimed = root_claimed.saturating_add(owed_u64.into());
         });
+
+        PendingRootAlpha::<T>::mutate(hotkey, |value| {
+            *value = value.saturating_sub(owed_u64.into());
+        });
     }
 
     fn root_claim_on_subnet_weight(_root_claim_type: RootClaimTypeEnum) -> Weight {
@@ -257,11 +261,14 @@ impl<T: Config> Pallet<T> {
             let root_claimed: u128 = RootClaimed::<T>::get((netuid, hotkey, coldkey));
 
             // Increase root claimed based on the claimable rate.
-            let new_root_claimed = root_claimed.saturating_add(
-                claimable_rate
-                    .saturating_mul(I96F32::from(u64::from(amount)))
-                    .saturating_to_num(),
-            );
+            let added_amount = claimable_rate
+                .saturating_mul(I96F32::from(u64::from(amount)))
+                .saturating_to_num();
+            let new_root_claimed = root_claimed.saturating_add(added_amount);
+
+            PendingRootAlpha::<T>::mutate(hotkey, |value| {
+                *value = value.saturating_sub(added_amount);
+            });
 
             // Set the new root claimed value.
             RootClaimed::<T>::insert((netuid, hotkey, coldkey), new_root_claimed);
@@ -284,11 +291,14 @@ impl<T: Config> Pallet<T> {
             let root_claimed: u128 = RootClaimed::<T>::get((netuid, hotkey, coldkey));
 
             // Decrease root claimed based on the claimable rate.
-            let new_root_claimed = root_claimed.saturating_sub(
-                claimable_rate
-                    .saturating_mul(I96F32::from(u64::from(amount)))
-                    .saturating_to_num(),
-            );
+            let subed_amount = claimable_rate
+                .saturating_mul(I96F32::from(u64::from(amount)))
+                .saturating_to_num();
+            let new_root_claimed = root_claimed.saturating_sub(subed_amount);
+
+            PendingRootAlpha::<T>::mutate(hotkey, |value| {
+                *value = value.saturating_add(subed_amount);
+            });
 
             // Set the new root_claimed value.
             RootClaimed::<T>::insert((netuid, hotkey, coldkey), new_root_claimed);
@@ -388,16 +398,62 @@ impl<T: Config> Pallet<T> {
         RootClaimable::<T>::insert(new_hotkey, dst_root_claimable);
     }
 
+    pub fn transfer_pending_root_alpha_for_new_hotkey(
+        old_hotkey: &T::AccountId,
+        new_hotkey: &T::AccountId,
+    ) {
+        let src_pending_root_alpha = PendingRootAlpha::<T>::get(old_hotkey);
+        let dst_pending_root_alpha = PendingRootAlpha::<T>::get(new_hotkey);
+        PendingRootAlpha::<T>::remove(old_hotkey);
+        PendingRootAlpha::<T>::insert(
+            new_hotkey,
+            dst_pending_root_alpha.saturating_add(src_pending_root_alpha),
+        );
+    }
+
     /// Claim all root dividends for subnet and remove all associated data.
     pub fn finalize_all_subnet_root_dividends(netuid: NetUid) {
         let hotkeys = RootClaimable::<T>::iter_keys().collect::<Vec<_>>();
 
+        let mut root_claimable_alpha_map: BTreeMap<T::AccountId, u128> = BTreeMap::new();
+
         for hotkey in hotkeys.iter() {
+            let root_stake = Self::get_stake_for_hotkey_on_subnet(hotkey, NetUid::ROOT);
+            let claimable_rate = RootClaimable::<T>::get(hotkey)
+                .values()
+                .fold(I96F32::from(0), |acc, x| acc.saturating_add(*x));
+            let total = claimable_rate.saturating_mul(I96F32::saturating_from_num(root_stake));
+            // let pending_root_alpha = PendingRootAlpha::<T>::get(hotkey);
+
+            root_claimable_alpha_map.insert(hotkey.clone(), total.saturating_to_num::<u128>());
+
             RootClaimable::<T>::mutate(hotkey, |claimable| {
                 claimable.remove(&netuid);
             });
         }
 
+        let mut root_claimed_alpha_map: BTreeMap<T::AccountId, u128> = BTreeMap::new();
+
+        for ((_netuid, hotkey, _coldkey), root_claimed) in RootClaimed::<T>::iter() {
+            if !hotkeys.contains(&hotkey) {
+                continue;
+            }
+
+            root_claimed_alpha_map
+                .entry(hotkey.clone())
+                .and_modify(|total| *total = total.saturating_add(root_claimed))
+                .or_insert(root_claimed);
+        }
+
         let _ = RootClaimed::<T>::clear_prefix((netuid,), u32::MAX, None);
+
+        for (hotkey, claimable) in root_claimable_alpha_map {
+            let claimed = root_claimed_alpha_map.get(&hotkey).unwrap_or(&0);
+            // still some root alpha not claimed
+            let pending = claimable.saturating_sub(*claimed);
+            PendingRootAlpha::<T>::mutate(&hotkey, |value| {
+                *value = value.saturating_sub(pending);
+            });
+        }
     }
 }

pallets/subtensor/src/swap/swap_hotkey.rs

@@ -531,6 +531,9 @@ impl<T: Config> Pallet<T> {
             }
         }
 
+        // 9.3 update pending root alpha for the hotkeys.
+        Self::transfer_pending_root_alpha_for_new_hotkey(old_hotkey, new_hotkey);
+
         Ok(())
     }
 }