Hotfix/only exclude add stake from emission no packing (#807)

* add impl

* add tests

* add impl

* chore: clippy

* oops fix test

* bump spec

* add test from PR #805

* Revert "bump spec"

This reverts commit 00ad4bf.

* chore: fmt

* remove lastaddstakeincrease logic

* fix comment command

* coinbase test with no deltas

* add tests

* renumber and remove

* fix for childkey set weights

* fix test for last commit

* clippy: allow manual inspect

* chore: fmt

* add shebang

* chore: clippy

* move assert up

* move clippy allow to toml

* add fuzzy assert

* fix neg delta test

* Fix tests with correct take percent

* fix typo

Cargo.toml

@@ -41,6 +41,7 @@ indexing-slicing = "deny"
 arithmetic-side-effects = "deny"
 type_complexity = "allow"
 unwrap-used = "deny"
+manual_inspect = "allow"
 
 [workspace.dependencies]
 cargo-husky = { version = "1", default-features = false }

pallets/subtensor/src/coinbase/run_coinbase.rs

@@ -249,6 +249,20 @@ impl<T: Config> Pallet<T> {
         });
     }
 
+    /// Calculates the nonviable stake for a nominator.
+    /// The nonviable stake is the stake that was added by the nominator since the last emission drain.
+    /// This stake will not receive emission until the next emission drain.
+    /// Note: if the stake delta is below zero, we return zero. We don't allow more stake than the nominator has.
+    pub fn get_nonviable_stake(hotkey: &T::AccountId, nominator: &T::AccountId) -> u64 {
+        let stake_delta = StakeDeltaSinceLastEmissionDrain::<T>::get(hotkey, nominator);
+        if stake_delta.is_negative() {
+            0
+        } else {
+            // Should never fail the into, but we handle it anyway.
+            stake_delta.try_into().unwrap_or(u64::MAX)
+        }
+    }
+
     //. --- 4. Drains the accumulated hotkey emission through to the nominators. The hotkey takes a proportion of the emission.
     /// The remainder is drained through to the nominators keeping track of the last stake increase event to ensure that the hotkey does not
     /// gain more emission than it's stake since the last drain.
@@ -268,71 +282,67 @@ impl<T: Config> Pallet<T> {
         // --- 1.0 Drain the hotkey emission.
         PendingdHotkeyEmission::<T>::insert(hotkey, 0);
 
-        // --- 2 Retrieve the last time this hotkey's emissions were drained.
-        let last_emission_drain: u64 = LastHotkeyEmissionDrain::<T>::get(hotkey);
-
-        // --- 3 Update the block value to the current block number.
+        // --- 2 Update the block value to the current block number.
         LastHotkeyEmissionDrain::<T>::insert(hotkey, block_number);
 
-        // --- 4 Retrieve the total stake for the hotkey from all nominations.
+        // --- 3 Retrieve the total stake for the hotkey from all nominations.
         let total_hotkey_stake: u64 = Self::get_total_stake_for_hotkey(hotkey);
 
-        // --- 5 Calculate the emission take for the hotkey.
+        // --- 4 Calculate the emission take for the hotkey.
         let take_proportion: I64F64 = I64F64::from_num(Delegates::<T>::get(hotkey))
             .saturating_div(I64F64::from_num(u16::MAX));
         let hotkey_take: u64 =
             (take_proportion.saturating_mul(I64F64::from_num(emission))).to_num::<u64>();
 
-        // --- 6 Compute the remaining emission after deducting the hotkey's take.
+        // --- 5 Compute the remaining emission after deducting the hotkey's take.
         let emission_minus_take: u64 = emission.saturating_sub(hotkey_take);
 
-        // --- 7 Calculate the remaining emission after the hotkey's take.
+        // --- 6 Calculate the remaining emission after the hotkey's take.
         let mut remainder: u64 = emission_minus_take;
 
-        // --- 8 Iterate over each nominator and get all viable stake.
+        // --- 7 Iterate over each nominator and get all viable stake.
         let mut total_viable_nominator_stake: u64 = total_hotkey_stake;
-        for (nominator, nominator_stake) in Stake::<T>::iter_prefix(hotkey) {
-            if false && LastAddStakeIncrease::<T>::get(hotkey, nominator) > last_emission_drain {
-                total_viable_nominator_stake =
-                    total_viable_nominator_stake.saturating_sub(nominator_stake);
-            }
+        for (nominator, _) in Stake::<T>::iter_prefix(hotkey) {
+            let nonviable_nomintaor_stake = Self::get_nonviable_stake(hotkey, &nominator);
+
+            total_viable_nominator_stake =
+                total_viable_nominator_stake.saturating_sub(nonviable_nomintaor_stake);
         }
 
-        // --- 9 Iterate over each nominator.
+        // --- 8 Iterate over each nominator.
         if total_viable_nominator_stake != 0 {
             for (nominator, nominator_stake) in Stake::<T>::iter_prefix(hotkey) {
-                // --- 10 Check if the stake was manually increased by the user since the last emission drain for this hotkey.
-                // If it was, skip this nominator as they will not receive their proportion of the emission.
-                if false
-                    && LastAddStakeIncrease::<T>::get(hotkey, nominator.clone())
-                        > last_emission_drain
-                {
-                    continue;
-                }
+                // --- 9 Skip emission for any stake the was added by the nominator since the last emission drain.
+                // This means the nominator will get emission on existing stake, but not on new stake, until the next emission drain.
+                let viable_nominator_stake =
+                    nominator_stake.saturating_sub(Self::get_nonviable_stake(hotkey, &nominator));
 
-                // --- 11 Calculate this nominator's share of the emission.
-                let nominator_emission: I64F64 = I64F64::from_num(nominator_stake)
+                // --- 10 Calculate this nominator's share of the emission.
+                let nominator_emission: I64F64 = I64F64::from_num(viable_nominator_stake)
                     .checked_div(I64F64::from_num(total_viable_nominator_stake))
                     .unwrap_or(I64F64::from_num(0))
                     .saturating_mul(I64F64::from_num(emission_minus_take));
 
-                // --- 12 Increase the stake for the nominator.
+                // --- 11 Increase the stake for the nominator.
                 Self::increase_stake_on_coldkey_hotkey_account(
                     &nominator,
                     hotkey,
                     nominator_emission.to_num::<u64>(),
                 );
 
-                // --- 13* Record event and Subtract the nominator's emission from the remainder.
+                // --- 12* Record event and Subtract the nominator's emission from the remainder.
                 total_new_tao = total_new_tao.saturating_add(nominator_emission.to_num::<u64>());
                 remainder = remainder.saturating_sub(nominator_emission.to_num::<u64>());
             }
         }
 
-        // --- 14 Finally, add the stake to the hotkey itself, including its take and the remaining emission.
+        // --- 13 Finally, add the stake to the hotkey itself, including its take and the remaining emission.
         let hotkey_new_tao: u64 = hotkey_take.saturating_add(remainder);
         Self::increase_stake_on_hotkey_account(hotkey, hotkey_new_tao);
 
+        // --- 14 Reset the stake delta for the hotkey.
+        let _ = StakeDeltaSinceLastEmissionDrain::<T>::clear_prefix(hotkey, u32::MAX, None);
+
         // --- 15 Record new tao creation event and return the amount created.
         total_new_tao = total_new_tao.saturating_add(hotkey_new_tao);
         total_new_tao

pallets/subtensor/src/lib.rs

@@ -226,6 +226,11 @@ pub mod pallet {
         0
     }
     #[pallet::type_value]
+    /// Default stake delta.
+    pub fn DefaultStakeDelta<T: Config>() -> i128 {
+        0
+    }
+    #[pallet::type_value]
     /// Default stakes per interval.
     pub fn DefaultStakesPerInterval<T: Config>() -> (u64, u64) {
         (0, 0)
@@ -770,6 +775,18 @@ pub mod pallet {
         DefaultAccountTake<T>,
     >;
     #[pallet::storage]
+    /// Map ( hot, cold ) --> stake: i128 | Stake added/removed since last emission drain.
+    pub type StakeDeltaSinceLastEmissionDrain<T: Config> = StorageDoubleMap<
+        _,
+        Blake2_128Concat,
+        T::AccountId,
+        Identity,
+        T::AccountId,
+        i128,
+        ValueQuery,
+        DefaultStakeDelta<T>,
+    >;
+    #[pallet::storage]
     /// DMAP ( parent, netuid ) --> Vec<(proportion,child)>
     pub type ChildKeys<T: Config> = StorageDoubleMap<
         _,
@@ -1250,7 +1267,7 @@ pub mod pallet {
         /// Returns the transaction priority for setting weights.
         pub fn get_priority_set_weights(hotkey: &T::AccountId, netuid: u16) -> u64 {
             if let Ok(uid) = Self::get_uid_for_net_and_hotkey(netuid, hotkey) {
-                let _stake = Self::get_total_stake_for_hotkey(hotkey);
+                let _stake = Self::get_stake_for_hotkey_on_subnet(hotkey, netuid);
                 let current_block_number: u64 = Self::get_current_block_as_u64();
                 let default_priority: u64 =
                     current_block_number.saturating_sub(Self::get_last_update_for_uid(netuid, uid));
@@ -1260,9 +1277,9 @@ pub mod pallet {
         }
 
         /// Is the caller allowed to set weights
-        pub fn check_weights_min_stake(hotkey: &T::AccountId) -> bool {
+        pub fn check_weights_min_stake(hotkey: &T::AccountId, netuid: u16) -> bool {
             // Blacklist weights transactions for low stake peers.
-            Self::get_total_stake_for_hotkey(hotkey) >= Self::get_weights_min_stake()
+            Self::get_stake_for_hotkey_on_subnet(hotkey, netuid) >= Self::get_weights_min_stake()
         }
 
         /// Helper function to check if register is allowed
@@ -1355,8 +1372,8 @@ where
         Pallet::<T>::get_priority_set_weights(who, netuid)
     }
 
-    pub fn check_weights_min_stake(who: &T::AccountId) -> bool {
-        Pallet::<T>::check_weights_min_stake(who)
+    pub fn check_weights_min_stake(who: &T::AccountId, netuid: u16) -> bool {
+        Pallet::<T>::check_weights_min_stake(who, netuid)
     }
 }
 
@@ -1394,7 +1411,7 @@ where
     ) -> TransactionValidity {
         match call.is_sub_type() {
             Some(Call::commit_weights { netuid, .. }) => {
-                if Self::check_weights_min_stake(who) {
+                if Self::check_weights_min_stake(who, *netuid) {
                     let priority: u64 = Self::get_priority_set_weights(who, *netuid);
                     Ok(ValidTransaction {
                         priority,
@@ -1406,7 +1423,7 @@ where
                 }
             }
             Some(Call::reveal_weights { netuid, .. }) => {
-                if Self::check_weights_min_stake(who) {
+                if Self::check_weights_min_stake(who, *netuid) {
                     let priority: u64 = Self::get_priority_set_weights(who, *netuid);
                     Ok(ValidTransaction {
                         priority,
@@ -1418,7 +1435,7 @@ where
                 }
             }
             Some(Call::set_weights { netuid, .. }) => {
-                if Self::check_weights_min_stake(who) {
+                if Self::check_weights_min_stake(who, *netuid) {
                     let priority: u64 = Self::get_priority_set_weights(who, *netuid);
                     Ok(ValidTransaction {
                         priority,
@@ -1430,7 +1447,7 @@ where
                 }
             }
             Some(Call::set_root_weights { netuid, hotkey, .. }) => {
-                if Self::check_weights_min_stake(hotkey) {
+                if Self::check_weights_min_stake(hotkey, *netuid) {
                     let priority: u64 = Self::get_priority_set_weights(hotkey, *netuid);
                     Ok(ValidTransaction {
                         priority,

pallets/subtensor/src/staking/add_stake.rs

@@ -70,8 +70,10 @@ impl<T: Config> Pallet<T> {
             Error::<T>::StakeRateLimitExceeded
         );
 
-        // Set the last time the stake increased for nominator drain protection.
-        LastAddStakeIncrease::<T>::insert(&hotkey, &coldkey, Self::get_current_block_as_u64());
+        // Track this addition in the stake delta.
+        StakeDeltaSinceLastEmissionDrain::<T>::mutate(&hotkey, &coldkey, |stake_delta| {
+            *stake_delta = stake_delta.saturating_add_unsigned(stake_to_be_added as u128);
+        });
 
         // If coldkey is not owner of the hotkey, it's a nomination stake.
         if !Self::coldkey_owns_hotkey(&coldkey, &hotkey) {

pallets/subtensor/src/staking/helpers.rs

@@ -297,6 +297,9 @@ impl<T: Config> Pallet<T> {
         staking_hotkeys.retain(|h| h != hotkey);
         StakingHotkeys::<T>::insert(coldkey, staking_hotkeys);
 
+        // Update stake delta
+        StakeDeltaSinceLastEmissionDrain::<T>::remove(hotkey, coldkey);
+
         current_stake
     }
 
@@ -431,6 +434,9 @@ impl<T: Config> Pallet<T> {
 
             // Add the balance to the coldkey account.
             Self::add_balance_to_coldkey_account(&delegate_coldkey_i, stake_i);
+
+            // Remove stake delta
+            StakeDeltaSinceLastEmissionDrain::<T>::remove(hotkey, &delegate_coldkey_i);
         }
     }
 }

pallets/subtensor/src/staking/remove_stake.rs

@@ -76,6 +76,11 @@ impl<T: Config> Pallet<T> {
         // We remove the balance from the hotkey.
         Self::decrease_stake_on_coldkey_hotkey_account(&coldkey, &hotkey, stake_to_be_removed);
 
+        // Track this removal in the stake delta.
+        StakeDeltaSinceLastEmissionDrain::<T>::mutate(&hotkey, &coldkey, |stake_delta| {
+            *stake_delta = stake_delta.saturating_sub_unsigned(stake_to_be_removed as u128);
+        });
+
         // We add the balance to the coldkey.  If the above fails we will not credit this coldkey.
         Self::add_balance_to_coldkey_account(&coldkey, stake_to_be_removed);
 

pallets/subtensor/src/subnets/uids.rs

@@ -117,7 +117,7 @@ impl<T: Config> Pallet<T> {
     ///
     pub fn get_stake_for_uid_and_subnetwork(netuid: u16, neuron_uid: u16) -> u64 {
         if let Ok(hotkey) = Self::get_hotkey_for_net_and_uid(netuid, neuron_uid) {
-            Self::get_total_stake_for_hotkey(&hotkey)
+            Self::get_stake_for_hotkey_on_subnet(&hotkey, netuid)
         } else {
             0
         }

pallets/subtensor/src/swap/swap_coldkey.rs

@@ -169,12 +169,17 @@ impl<T: Config> Pallet<T> {
             weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
         }
 
-        // 4. Swap LastAddStakeIncrease.
+        // 4. Swap StakeDeltaSinceLastEmissionDrain
         for hotkey in StakingHotkeys::<T>::get(old_coldkey) {
-            let last_add_stake_increase = LastAddStakeIncrease::<T>::get(&hotkey, old_coldkey);
-            LastAddStakeIncrease::<T>::remove(&hotkey, old_coldkey);
-            LastAddStakeIncrease::<T>::insert(&hotkey, new_coldkey, last_add_stake_increase);
-            weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
+            let old_stake_delta = StakeDeltaSinceLastEmissionDrain::<T>::get(&hotkey, old_coldkey);
+            let new_stake_delta = StakeDeltaSinceLastEmissionDrain::<T>::get(&hotkey, new_coldkey);
+            StakeDeltaSinceLastEmissionDrain::<T>::insert(
+                &hotkey,
+                new_coldkey,
+                new_stake_delta.saturating_add(old_stake_delta),
+            );
+            StakeDeltaSinceLastEmissionDrain::<T>::remove(&hotkey, old_coldkey);
+            weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
         }
 
         // 5. Swap total coldkey stake.

pallets/subtensor/src/swap/swap_hotkey.rs

@@ -207,24 +207,13 @@ impl<T: Config> Pallet<T> {
             weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
         }
 
-        // 9.1. swap PendingdHotkeyEmission
+        // 9. swap PendingdHotkeyEmission
         if PendingdHotkeyEmission::<T>::contains_key(old_hotkey) {
             let old_pending_hotkey_emission = PendingdHotkeyEmission::<T>::get(old_hotkey);
             PendingdHotkeyEmission::<T>::remove(old_hotkey);
             PendingdHotkeyEmission::<T>::insert(new_hotkey, old_pending_hotkey_emission);
             weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
         }
-        // 9.2. swap LastAddStakeIncrease
-        // Swap double map prefix
-        let old_coldkey_entries: Vec<(T::AccountId, u64)> =
-            LastAddStakeIncrease::<T>::iter_prefix(old_hotkey).collect();
-        if !old_coldkey_entries.is_empty() {
-            for (coldkey, last_add_stake_increase) in old_coldkey_entries {
-                LastAddStakeIncrease::<T>::remove(old_hotkey, &coldkey);
-                LastAddStakeIncrease::<T>::insert(new_hotkey, &coldkey, last_add_stake_increase);
-                weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
-            }
-        }
 
         // 10. Swap all subnet specific info.
         let all_netuids: Vec<u16> = Self::get_all_subnet_netuids();
@@ -363,6 +352,19 @@ impl<T: Config> Pallet<T> {
             }
         }
 
+        // 13. Swap Stake Delta for all coldkeys.
+        for (coldkey, stake_delta) in StakeDeltaSinceLastEmissionDrain::<T>::iter_prefix(old_hotkey)
+        {
+            let new_stake_delta = StakeDeltaSinceLastEmissionDrain::<T>::get(new_hotkey, &coldkey);
+            StakeDeltaSinceLastEmissionDrain::<T>::insert(
+                new_hotkey,
+                &coldkey,
+                new_stake_delta.saturating_add(stake_delta),
+            );
+            StakeDeltaSinceLastEmissionDrain::<T>::remove(old_hotkey, &coldkey);
+            weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
+        }
+
         // Return successful after swapping all the relevant terms.
         Ok(())
     }