Fix stake deltas in remove stake

pallets/subtensor/src/coinbase/run_coinbase.rs

@@ -301,10 +301,10 @@ impl<T: Config> Pallet<T> {
                     nominator_stake.saturating_sub(Self::get_nonviable_stake(hotkey, &nominator));
 
                 // --- 10 Calculate this nominator's share of the emission.
-                let nominator_emission: I64F64 = I64F64::from_num(viable_nominator_stake)
+                let nominator_emission: I64F64 = I64F64::from_num(emission_minus_take)
+                    .saturating_mul(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));
+                    .unwrap_or(I64F64::from_num(0));
 
                 // --- 11 Increase the stake for the nominator.
                 Self::increase_stake_on_coldkey_hotkey_account(
@@ -323,7 +323,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/tests/staking.rs

@@ -2306,3 +2306,106 @@ fn test_get_total_delegated_stake_exclude_owner_stake() {
         );
     });
 }
+
+// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test staking -- test_stake_delta_tracks_adds_and_removes --exact --nocapture
+#[test]
+fn test_stake_delta_tracks_adds_and_removes() {
+    new_test_ext(1).execute_with(|| {
+        let netuid = 1u16;
+        let delegate_coldkey = U256::from(1);
+        let delegate_hotkey = U256::from(2);
+        let delegator = U256::from(3);
+
+        let owner_stake = 1000;
+        let owner_added_stake = 123;
+        let owner_removed_stake = 456;
+        // Add more than removed to test that the delta is updated correctly
+        let owner_adds_more_stake = owner_removed_stake + 1;
+
+        let delegator_added_stake = 999;
+
+        // Set stake rate limit very high
+        TargetStakesPerInterval::<Test>::put(1e9 as u64);
+
+        add_network(netuid, 0, 0);
+        register_ok_neuron(netuid, delegate_hotkey, delegate_coldkey, 0);
+        // Give extra stake to the owner
+        SubtensorModule::increase_stake_on_coldkey_hotkey_account(
+            &delegate_coldkey,
+            &delegate_hotkey,
+            owner_stake,
+        );
+
+        // Register as a delegate
+        assert_ok!(SubtensorModule::become_delegate(
+            RuntimeOrigin::signed(delegate_coldkey),
+            delegate_hotkey
+        ));
+
+        // Verify that the stake delta is empty
+        assert_eq!(
+            StakeDeltaSinceLastEmissionDrain::<Test>::get(delegate_hotkey, delegate_coldkey),
+            0
+        );
+
+        // Give the coldkey some balance; extra just in case
+        SubtensorModule::add_balance_to_coldkey_account(
+            &delegate_coldkey,
+            owner_added_stake + owner_adds_more_stake,
+        );
+
+        // Add some stake
+        assert_ok!(SubtensorModule::add_stake(
+            RuntimeOrigin::signed(delegate_coldkey),
+            delegate_hotkey,
+            owner_added_stake
+        ));
+
+        // Verify that the stake delta is correct
+        assert_eq!(
+            StakeDeltaSinceLastEmissionDrain::<Test>::get(delegate_hotkey, delegate_coldkey),
+            i128::from(owner_added_stake)
+        );
+
+        // Add some stake from a delegator
+        SubtensorModule::add_balance_to_coldkey_account(&delegator, delegator_added_stake);
+        assert_ok!(SubtensorModule::add_stake(
+            RuntimeOrigin::signed(delegator),
+            delegate_hotkey,
+            delegator_added_stake
+        ));
+
+        // Verify that the stake delta is unchanged for the owner
+        assert_eq!(
+            StakeDeltaSinceLastEmissionDrain::<Test>::get(delegate_hotkey, delegate_coldkey),
+            i128::from(owner_added_stake)
+        );
+
+        // Remove some stake
+        assert_ok!(SubtensorModule::remove_stake(
+            RuntimeOrigin::signed(delegate_coldkey),
+            delegate_hotkey,
+            owner_removed_stake
+        ));
+
+        // Verify that the stake delta is correct
+        assert_eq!(
+            StakeDeltaSinceLastEmissionDrain::<Test>::get(delegate_hotkey, delegate_coldkey),
+            i128::from(owner_added_stake).saturating_sub_unsigned(owner_removed_stake.into())
+        );
+
+        // Add more stake than was removed
+        assert_ok!(SubtensorModule::add_stake(
+            RuntimeOrigin::signed(delegate_coldkey),
+            delegate_hotkey,
+            owner_adds_more_stake
+        ));
+
+        // Verify that the stake delta is correct
+        assert_eq!(
+            StakeDeltaSinceLastEmissionDrain::<Test>::get(delegate_hotkey, delegate_coldkey),
+            i128::from(owner_added_stake)
+                .saturating_add_unsigned((owner_adds_more_stake - owner_removed_stake).into())
+        );
+    });
+}