Hotfix/clear stake delta on drain (#934)

* add comment also

* chore: fmt

* add back delta handling in helpers

* track removals in stake delta

* handle stake delta in ck swap func

* add tests back

* add tests for staking back

* add back test for ck swap

pallets/subtensor/src/coinbase/run_coinbase.rs

@@ -303,8 +303,8 @@ impl<T: Config> Pallet<T> {
         // --- 8 Iterate over each nominator.
         if total_viable_nominator_stake != 0 {
             for (nominator, nominator_stake) in Stake::<T>::iter_prefix(hotkey) {
-                // --- 9 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.
+                // --- 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));
 
@@ -331,7 +331,10 @@ impl<T: Config> Pallet<T> {
         let hotkey_new_tao: u64 = hotkey_take.saturating_add(remainder);
         Self::increase_stake_on_hotkey_account(hotkey, hotkey_new_tao);
 
-        // --- 14 Record new tao creation event and return the amount created.
+        // --- 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/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/swap/swap_coldkey.rs

@@ -169,7 +169,20 @@ impl<T: Config> Pallet<T> {
             weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
         }
 
-        // 4. Swap total coldkey stake.
+        // 4. Swap StakeDeltaSinceLastEmissionDrain
+        for hotkey in StakingHotkeys::<T>::get(old_coldkey) {
+            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.
         // TotalColdkeyStake: MAP ( coldkey ) --> u64 | Total stake of the coldkey.
         let old_coldkey_stake: u64 = TotalColdkeyStake::<T>::get(old_coldkey);
         // Get the stake of the new coldkey.
@@ -183,7 +196,7 @@ impl<T: Config> Pallet<T> {
         );
         weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
 
-        // 5. Swap StakingHotkeys.
+        // 6. Swap StakingHotkeys.
         // StakingHotkeys: MAP ( coldkey ) --> Vec<hotkeys> | Hotkeys staking for the coldkey.
         let old_staking_hotkeys: Vec<T::AccountId> = StakingHotkeys::<T>::get(old_coldkey);
         let mut new_staking_hotkeys: Vec<T::AccountId> = StakingHotkeys::<T>::get(new_coldkey);
@@ -197,7 +210,7 @@ impl<T: Config> Pallet<T> {
         StakingHotkeys::<T>::insert(new_coldkey, new_staking_hotkeys);
         weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
 
-        // 6. Swap hotkey owners.
+        // 7. Swap hotkey owners.
         // Owner: MAP ( hotkey ) --> coldkey | Owner of the hotkey.
         // OwnedHotkeys: MAP ( coldkey ) --> Vec<hotkeys> | Hotkeys owned by the coldkey.
         let old_owned_hotkeys: Vec<T::AccountId> = OwnedHotkeys::<T>::get(old_coldkey);
@@ -216,7 +229,7 @@ impl<T: Config> Pallet<T> {
         OwnedHotkeys::<T>::insert(new_coldkey, new_owned_hotkeys);
         weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
 
-        // 7. Transfer remaining balance.
+        // 8. Transfer remaining balance.
         // Balance: MAP ( coldkey ) --> u64 | Balance of the coldkey.
         // Transfer any remaining balance from old_coldkey to new_coldkey
         let remaining_balance = Self::get_coldkey_balance(old_coldkey);