Reimplement StakeDelta

pallets/subtensor/src/coinbase/block_step.rs

@@ -18,7 +18,10 @@ impl<T: Config + pallet_drand::Config> Pallet<T> {
         Self::run_coinbase(block_emission);
         // --- 4. Set pending children on the epoch; but only after the coinbase has been run.
         Self::try_set_pending_children(block_number);
-        // Return ok.
+
+        // Clear stake delta for netuids with running epoch this block
+        Self::clear_stake_delta(block_number);
+
         Ok(())
     }
 

pallets/subtensor/src/coinbase/run_coinbase.rs

@@ -508,10 +508,20 @@ impl<T: Config> Pallet<T> {
         for hotkey in hotkeys {
             // Get hotkey ALPHA on subnet.
             let alpha_stake: u64 = Self::get_stake_for_hotkey_on_subnet(hotkey, netuid);
+
             // Get hotkey TAO on root.
             let root_stake: u64 =
                 Self::get_stake_for_hotkey_on_subnet(hotkey, Self::get_root_netuid());
-            stake_map.insert(hotkey.clone(), (alpha_stake, root_stake));
+
+            let alpha_stake_delta = Self::get_stake_delta_for_hotkey_on_subnet(hotkey, netuid);
+            let root_stake_delta =
+                Self::get_stake_delta_for_hotkey_on_subnet(hotkey, Self::get_root_netuid());
+
+            let final_alpha_stake = alpha_stake.saturating_sub(alpha_stake_delta);
+
+            let final_root_stake = root_stake.saturating_sub(root_stake_delta);
+
+            stake_map.insert(hotkey.clone(), (final_alpha_stake, final_root_stake));
         }
         stake_map
     }

pallets/subtensor/src/lib.rs

@@ -909,6 +909,18 @@ pub mod pallet {
         50400
     }
 
+    #[pallet::type_value]
+    /// Default stake delta.
+    pub fn DefaultStakeDeltaMap() -> i128 {
+        0
+    }
+
+    #[pallet::type_value]
+    /// Default value for enable stake delta calculation.
+    pub fn DefaultEnableStakeDeltaCalculation() -> bool {
+        true
+    }
+
     #[pallet::storage]
     pub type MinActivityCutoff<T: Config> =
         StorageValue<_, u16, ValueQuery, DefaultMinActivityCutoff<T>>;
@@ -1176,6 +1188,24 @@ pub mod pallet {
     pub type TokenSymbol<T: Config> =
         StorageMap<_, Identity, u16, Vec<u8>, ValueQuery, DefaultUnicodeVecU8<T>>;
 
+    #[pallet::storage]
+    /// Map ( netuid, hot ) --> StakeDeltaMap| Stake added/removed since last emission drain for coldkey.
+    pub type StakeDeltaSinceLastEmissionDrain<T: Config> = StorageDoubleMap<
+        _,
+        Blake2_128Concat,
+        u16,
+        Identity,
+        T::AccountId,
+        i128,
+        ValueQuery,
+        DefaultStakeDeltaMap,
+    >;
+
+    #[pallet::storage]
+    /// --- ITEM( enables effect of the saved stake delta )
+    pub type EnableStakeDeltaCalculation<T> =
+        StorageValue<_, bool, ValueQuery, DefaultEnableStakeDeltaCalculation>;
+
     /// ============================
     /// ==== Global Parameters =====
     /// ============================

pallets/subtensor/src/staking/stake_utils.rs

@@ -2,7 +2,7 @@ use super::*;
 //use frame_system::pallet_prelude::BlockNumberFor;
 use safe_math::*;
 use share_pool::{SharePool, SharePoolDataOperations};
-//use sp_runtime::Saturating;
+use sp_runtime::SaturatedConversion;
 use sp_std::ops::Neg;
 use substrate_fixed::types::{I64F64, I96F32, U64F64, U96F32, U110F18};
 
@@ -159,7 +159,7 @@ impl<T: Config> Pallet<T> {
 
         // Step 2: Get the global tao stake for the hotkey
         let tao_stake = I64F64::saturating_from_num(Self::get_tao_inherited_for_hotkey_on_subnet(
-            hotkey, netuid,
+            hotkey, netuid, false,
         ));
         log::debug!("tao_stake: {:?}", tao_stake);
 
@@ -185,9 +185,13 @@ impl<T: Config> Pallet<T> {
             .map(|uid| {
                 if Keys::<T>::contains_key(netuid, uid) {
                     let hotkey: T::AccountId = Keys::<T>::get(netuid, uid);
-                    I64F64::saturating_from_num(Self::get_inherited_for_hotkey_on_subnet(
-                        &hotkey, netuid,
-                    ))
+                    let initial_stake = Self::get_inherited_for_hotkey_on_subnet(&hotkey, netuid);
+
+                    let stake_delta = Self::get_stake_delta_for_hotkey_on_subnet(&hotkey, netuid);
+
+                    let final_stake = initial_stake.saturating_sub(stake_delta);
+
+                    I64F64::saturating_from_num(final_stake)
                 } else {
                     I64F64::saturating_from_num(0)
                 }
@@ -202,7 +206,7 @@ impl<T: Config> Pallet<T> {
                 if Keys::<T>::contains_key(netuid, uid) {
                     let hotkey: T::AccountId = Keys::<T>::get(netuid, uid);
                     I64F64::saturating_from_num(Self::get_tao_inherited_for_hotkey_on_subnet(
-                        &hotkey, netuid,
+                        &hotkey, netuid, true,
                     ))
                 } else {
                     I64F64::saturating_from_num(0)
@@ -243,17 +247,32 @@ impl<T: Config> Pallet<T> {
     /// # Arguments
     /// * `hotkey` - AccountId of the hotkey whose total inherited stake is to be calculated.
     /// * `netuid` - Network unique identifier specifying the subnet context.
+    /// * `use_stake_delta` - Whether to use stake delta for root stake for hotkey.
     ///
     /// # Returns
     /// * `u64`: The total inherited alpha for the hotkey on the subnet after considering the
     ///   stakes allocated to children and inherited from parents.
     ///
     /// # Note
     /// This function uses saturating arithmetic to prevent overflows.
-    pub fn get_tao_inherited_for_hotkey_on_subnet(hotkey: &T::AccountId, netuid: u16) -> u64 {
-        let initial_tao: U96F32 = U96F32::saturating_from_num(
-            Self::get_stake_for_hotkey_on_subnet(hotkey, Self::get_root_netuid()),
-        );
+    pub fn get_tao_inherited_for_hotkey_on_subnet(
+        hotkey: &T::AccountId,
+        netuid: u16,
+        use_stake_delta: bool,
+    ) -> u64 {
+        let initial_tao: U96F32 = U96F32::saturating_from_num({
+            let initial_tao = Self::get_stake_for_hotkey_on_subnet(hotkey, Self::get_root_netuid());
+
+            if use_stake_delta {
+                let stake_delta =
+                    Self::get_stake_delta_for_hotkey_on_subnet(hotkey, Self::get_root_netuid());
+
+                // modified tao
+                initial_tao.saturating_sub(stake_delta)
+            } else {
+                initial_tao
+            }
+        });
 
         // Initialize variables to track alpha allocated to children and inherited from parents.
         let mut tao_to_children: U96F32 = U96F32::saturating_from_num(0);
@@ -298,9 +317,23 @@ impl<T: Config> Pallet<T> {
         // Step 4: Calculate the total tao inherited from parents.
         for (proportion, parent) in parents {
             // Retrieve the parent's total stake on this subnet.
-            let parent_tao: U96F32 = U96F32::saturating_from_num(
-                Self::get_stake_for_hotkey_on_subnet(&parent, Self::get_root_netuid()),
-            );
+            let initial_tao =
+                Self::get_stake_for_hotkey_on_subnet(&parent, Self::get_root_netuid());
+
+            let parent_tao = U96F32::saturating_from_num({
+                if use_stake_delta {
+                    let stake_delta = Self::get_stake_delta_for_hotkey_on_subnet(
+                        &parent,
+                        Self::get_root_netuid(),
+                    );
+
+                    // modified tao
+                    initial_tao.saturating_sub(stake_delta)
+                } else {
+                    initial_tao
+                }
+            });
+
             log::trace!(
                 "Parent tao for parent {:?} on subnet {}: {:?}",
                 parent,
@@ -344,6 +377,31 @@ impl<T: Config> Pallet<T> {
         finalized_tao.saturating_to_num::<u64>()
     }
 
+    pub fn get_stake_delta_for_hotkey_on_subnet(hotkey: &T::AccountId, netuid: u16) -> u64 {
+        let enable_stake_delta = EnableStakeDeltaCalculation::<T>::get();
+
+        if !enable_stake_delta {
+            return 0;
+        }
+
+        let total_stake = StakeDeltaSinceLastEmissionDrain::<T>::get(netuid, hotkey);
+
+        // We only care about positive stake delta
+        total_stake.max(0).saturated_into()
+    }
+
+    // Remove all stake delta after the epoch calculation
+    pub fn clear_stake_delta(block_number: u64) {
+        let subnets: Vec<u16> = Self::get_all_subnet_netuids();
+
+        for netuid in subnets.into_iter() {
+            if Self::should_run_epoch(netuid, block_number) {
+                // We don't use the operation result
+                let _ = StakeDeltaSinceLastEmissionDrain::<T>::clear_prefix(netuid, u32::MAX, None);
+            }
+        }
+    }
+
     pub fn get_inherited_for_hotkey_on_subnet(hotkey: &T::AccountId, netuid: u16) -> u64 {
         // Step 1: Retrieve the initial total stake (alpha) for the hotkey on the specified subnet.
         let initial_alpha: U96F32 =
@@ -787,6 +845,11 @@ impl<T: Config> Pallet<T> {
         let actual_alpha_decrease =
             Self::decrease_stake_for_hotkey_and_coldkey_on_subnet(hotkey, coldkey, netuid, alpha);
 
+        // Track this substraction in the stake delta.
+        StakeDeltaSinceLastEmissionDrain::<T>::mutate(netuid, hotkey, |stake_delta| {
+            *stake_delta = stake_delta.saturating_sub(actual_alpha_decrease.into());
+        });
+
         // Step 2: Swap the alpha for TAO.
         let tao: u64 = Self::swap_alpha_for_tao(netuid, actual_alpha_decrease);
 
@@ -856,6 +919,11 @@ impl<T: Config> Pallet<T> {
             actual_alpha = Self::increase_stake_for_hotkey_and_coldkey_on_subnet(
                 hotkey, coldkey, netuid, alpha,
             );
+
+            // Track this addition in the stake delta.
+            StakeDeltaSinceLastEmissionDrain::<T>::mutate(netuid, hotkey, |stake_delta| {
+                *stake_delta = stake_delta.saturating_add(actual_alpha.into());
+            });
         }
 
         // Step 4. Increase Tao reserves by the fee amount.

pallets/subtensor/src/tests/children.rs

@@ -2812,6 +2812,7 @@ fn test_childkey_take_drain() {
             let proportion: u64 = u64::MAX / 2;
 
             // Add network, register hotkeys, and setup network parameters
+            EnableStakeDeltaCalculation::<Test>::put(false);
             add_network(netuid, subnet_tempo, 0);
             register_ok_neuron(netuid, child_hotkey, child_coldkey, 0);
             register_ok_neuron(netuid, parent_hotkey, parent_coldkey, 1);

pallets/subtensor/src/tests/coinbase.rs

@@ -947,11 +947,11 @@ fn test_get_root_children() {
 
         // Assert Alice and Bob TAO inherited stakes
         assert_eq!(
-            SubtensorModule::get_tao_inherited_for_hotkey_on_subnet(&alice, alpha),
+            SubtensorModule::get_tao_inherited_for_hotkey_on_subnet(&alice, alpha, false),
             0
         );
         assert_eq!(
-            SubtensorModule::get_tao_inherited_for_hotkey_on_subnet(&bob, alpha),
+            SubtensorModule::get_tao_inherited_for_hotkey_on_subnet(&bob, alpha, false),
             bob_root_stake + alice_root_stake
         );
 

pallets/subtensor/src/tests/staking.rs

@@ -1925,6 +1925,13 @@ fn test_staking_sets_div_variables() {
 
         // Wait for 1 epoch
         step_block(tempo + 1);
+        // Stake delta will mitigate stake movement
+        assert_eq!(
+            AlphaDividendsPerSubnet::<Test>::get(netuid, hotkey_account_id),
+            0
+        );
+        // Wait for 2 epoch
+        step_block(tempo + 1);
 
         // Verify that divident variables have been set
         let stake = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(