Use last epoch hotkey alpha for fee calculation

pallets/subtensor/src/coinbase/run_coinbase.rs

@@ -454,17 +454,21 @@ impl<T: Config> Pallet<T> {
             log::debug!("hotkey: {:?} alpha_take: {:?}", hotkey, alpha_take);
             Self::increase_stake_for_hotkey_and_coldkey_on_subnet(
                 &hotkey,
-                &Owner::<T>::get(hotkey.clone()),
+                &Owner::<T>::get(&hotkey),
                 netuid,
                 tou64!(alpha_take),
             );
             // Give all other nominators.
             log::debug!("hotkey: {:?} alpha_divs: {:?}", hotkey, alpha_divs);
-            Self::increase_stake_for_hotkey_on_subnet(&hotkey.clone(), netuid, tou64!(alpha_divs));
+            Self::increase_stake_for_hotkey_on_subnet(&hotkey, netuid, tou64!(alpha_divs));
             // Record dividends for this hotkey.
-            AlphaDividendsPerSubnet::<T>::mutate(netuid, hotkey.clone(), |divs| {
+            AlphaDividendsPerSubnet::<T>::mutate(netuid, &hotkey, |divs| {
                 *divs = divs.saturating_add(tou64!(alpha_divs));
             });
+            // Record total hotkey alpha based on which this value of AlphaDividendsPerSubnet
+            // was calculated
+            let total_hotkey_alpha = TotalHotkeyAlpha::<T>::get(&hotkey, netuid);
+            TotalHotkeyAlphaLastEpoch::<T>::insert(hotkey, netuid, total_hotkey_alpha);
         }
 
         // Distribute root tao divs.

pallets/subtensor/src/lib.rs

@@ -1024,6 +1024,17 @@ pub mod pallet {
         ValueQuery,
         DefaultZeroU64<T>,
     >;
+    #[pallet::storage] // --- DMAP ( hot, netuid ) --> alpha | Returns the total amount of alpha a hotkey owned in the last epoch.
+    pub type TotalHotkeyAlphaLastEpoch<T: Config> = StorageDoubleMap<
+        _,
+        Blake2_128Concat,
+        T::AccountId,
+        Identity,
+        u16,
+        u64,
+        ValueQuery,
+        DefaultZeroU64<T>,
+    >;
     #[pallet::storage]
     /// DMAP ( hot, netuid ) --> total_alpha_shares | Returns the number of alpha shares for a hotkey on a subnet.
     pub type TotalHotkeyShares<T: Config> = StorageDoubleMap<

pallets/subtensor/src/staking/stake_utils.rs

@@ -1106,14 +1106,16 @@ impl<T: Config> Pallet<T> {
                     DefaultStakingFee::<T>::get()
                 } else {
                     // Otherwise, calculate the fee based on the alpha estimate
+                    // Here we are using TotalHotkeyAlphaLastEpoch, which is exactly the value that
+                    // was used to calculate AlphaDividendsPerSubnet
                     let mut fee = alpha_estimate
                         .saturating_mul(
                             U96F32::saturating_from_num(AlphaDividendsPerSubnet::<T>::get(
                                 origin_netuid,
                                 &origin_hotkey,
                             ))
                             .safe_div(U96F32::saturating_from_num(
-                                TotalHotkeyAlpha::<T>::get(&origin_hotkey, origin_netuid),
+                                TotalHotkeyAlphaLastEpoch::<T>::get(&origin_hotkey, origin_netuid),
                             )),
                         )
                         .saturating_mul(Self::get_alpha_price(origin_netuid)) // fee needs to be in TAO

pallets/subtensor/src/swap/swap_hotkey.rs

@@ -434,6 +434,45 @@ impl<T: Config> Pallet<T> {
             }
         }
 
+        // 16. Swap dividend records
+        TotalHotkeyAlphaLastEpoch::<T>::iter_prefix(old_hotkey)
+            .drain()
+            .for_each(|(netuid, old_alpha)| {
+                // 16.1 Swap TotalHotkeyAlphaLastEpoch
+                let new_total_hotkey_alpha =
+                    TotalHotkeyAlphaLastEpoch::<T>::get(new_hotkey, netuid);
+                TotalHotkeyAlphaLastEpoch::<T>::insert(
+                    new_hotkey,
+                    netuid,
+                    old_alpha.saturating_add(new_total_hotkey_alpha),
+                );
+                weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
+
+                // 16.2 Swap AlphaDividendsPerSubnet
+                let old_hotkey_alpha_dividends =
+                    AlphaDividendsPerSubnet::<T>::get(netuid, old_hotkey);
+                let new_hotkey_alpha_dividends =
+                    AlphaDividendsPerSubnet::<T>::get(netuid, new_hotkey);
+                AlphaDividendsPerSubnet::<T>::remove(netuid, old_hotkey);
+                AlphaDividendsPerSubnet::<T>::insert(
+                    netuid,
+                    new_hotkey,
+                    old_hotkey_alpha_dividends.saturating_add(new_hotkey_alpha_dividends),
+                );
+                weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
+
+                // 16.3 Swap TaoDividendsPerSubnet
+                let old_hotkey_tao_dividends = TaoDividendsPerSubnet::<T>::get(netuid, old_hotkey);
+                let new_hotkey_tao_dividends = TaoDividendsPerSubnet::<T>::get(netuid, new_hotkey);
+                TaoDividendsPerSubnet::<T>::remove(netuid, old_hotkey);
+                TaoDividendsPerSubnet::<T>::insert(
+                    netuid,
+                    new_hotkey,
+                    old_hotkey_tao_dividends.saturating_add(new_hotkey_tao_dividends),
+                );
+                weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
+            });
+
         // Return successful after swapping all the relevant terms.
         Ok(())
     }

pallets/subtensor/src/tests/swap_hotkey.rs

@@ -897,15 +897,26 @@ fn test_swap_stake_success() {
 
         // Initialize staking variables for old_hotkey
         TotalHotkeyAlpha::<Test>::insert(old_hotkey, netuid, amount);
+        TotalHotkeyAlphaLastEpoch::<Test>::insert(old_hotkey, netuid, amount * 2);
         TotalHotkeyShares::<Test>::insert(old_hotkey, netuid, U64F64::from_num(shares));
         Alpha::<Test>::insert((old_hotkey, coldkey, netuid), U64F64::from_num(amount));
+        AlphaDividendsPerSubnet::<Test>::insert(netuid, old_hotkey, amount);
+        TaoDividendsPerSubnet::<Test>::insert(netuid, old_hotkey, amount);
 
         // Perform the swap
         SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight);
 
         // Verify the swap
         assert_eq!(TotalHotkeyAlpha::<Test>::get(old_hotkey, netuid), 0);
         assert_eq!(TotalHotkeyAlpha::<Test>::get(new_hotkey, netuid), amount);
+        assert_eq!(
+            TotalHotkeyAlphaLastEpoch::<Test>::get(old_hotkey, netuid),
+            0
+        );
+        assert_eq!(
+            TotalHotkeyAlphaLastEpoch::<Test>::get(new_hotkey, netuid),
+            amount * 2
+        );
         assert_eq!(
             TotalHotkeyShares::<Test>::get(old_hotkey, netuid),
             U64F64::from_num(0)
@@ -922,6 +933,16 @@ fn test_swap_stake_success() {
             Alpha::<Test>::get((new_hotkey, coldkey, netuid)),
             U64F64::from_num(amount)
         );
+        assert_eq!(AlphaDividendsPerSubnet::<Test>::get(netuid, old_hotkey), 0);
+        assert_eq!(
+            AlphaDividendsPerSubnet::<Test>::get(netuid, new_hotkey),
+            amount
+        );
+        assert_eq!(TaoDividendsPerSubnet::<Test>::get(netuid, old_hotkey), 0);
+        assert_eq!(
+            TaoDividendsPerSubnet::<Test>::get(netuid, new_hotkey),
+            amount
+        );
     });
 }