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,24 +1106,32 @@ impl<T: Config> Pallet<T> {
                     DefaultStakingFee::<T>::get()
                 } else {
                     // Otherwise, calculate the fee based on the alpha estimate
-                    let mut fee = alpha_estimate
+                    // Here we are using TotalHotkeyAlphaLastEpoch, which is exactly the value that
+                    // was used to calculate AlphaDividendsPerSubnet
+                    let tao_estimate = U96F32::saturating_from_num(
+                        Self::sim_swap_alpha_for_tao(
+                            origin_netuid,
+                            alpha_estimate.saturating_to_num::<u64>(),
+                        )
+                        .unwrap_or(0),
+                    );
+                    let mut fee = tao_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
                         .saturating_to_num::<u64>();
 
                     // 0.005% per epoch matches to 44% annual in compound interest. Do not allow the fee
                     // to be lower than that. (1.00005^(365*20) ~= 1.44)
                     let apr_20_percent = U96F32::saturating_from_num(0.00005);
                     fee = fee.max(
-                        alpha_estimate
+                        tao_estimate
                             .saturating_mul(apr_20_percent)
                             .saturating_to_num::<u64>(),
                     );

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/staking.rs

@@ -771,6 +771,170 @@ fn test_remove_stake_total_issuance_no_change() {
     });
 }
 
+// cargo test --package pallet-subtensor --lib -- tests::staking::test_remove_prev_epoch_stake --exact --show-output --nocapture
+#[test]
+fn test_remove_prev_epoch_stake() {
+    new_test_ext(1).execute_with(|| {
+        let def_fee = DefaultStakingFee::<Test>::get();
+
+        // Test case: (amount_to_stake, AlphaDividendsPerSubnet, TotalHotkeyAlphaLastEpoch, expected_fee)
+        [
+            // No previous epoch stake and low hotkey stake
+            (
+                DefaultMinStake::<Test>::get() * 10,
+                0_u64,
+                1000_u64,
+                def_fee * 2,
+            ),
+            // Same, but larger amount to stake - we get 0.005% for unstake
+            (
+                1_000_000_000,
+                0_u64,
+                1000_u64,
+                (1_000_000_000_f64 * 0.00005) as u64 + def_fee,
+            ),
+            (
+                100_000_000_000,
+                0_u64,
+                1000_u64,
+                (100_000_000_000_f64 * 0.00005) as u64 + def_fee,
+            ),
+            // Lower previous epoch stake than current stake
+            // Staking/unstaking 100 TAO, divs / total = 0.1 => fee is 1 TAO
+            (
+                100_000_000_000,
+                1_000_000_000_u64,
+                10_000_000_000_u64,
+                (100_000_000_000_f64 * 0.1) as u64 + def_fee,
+            ),
+            // Staking/unstaking 100 TAO, divs / total = 0.001 => fee is 0.01 TAO
+            (
+                100_000_000_000,
+                10_000_000_u64,
+                10_000_000_000_u64,
+                (100_000_000_000_f64 * 0.001) as u64 + def_fee,
+            ),
+            // Higher previous epoch stake than current stake
+            (
+                1_000_000_000,
+                100_000_000_000_u64,
+                100_000_000_000_000_u64,
+                (1_000_000_000_f64 * 0.001) as u64 + def_fee,
+            ),
+        ]
+        .iter()
+        .for_each(
+            |(amount_to_stake, alpha_divs, hotkey_alpha, expected_fee)| {
+                let subnet_owner_coldkey = U256::from(1);
+                let subnet_owner_hotkey = U256::from(2);
+                let hotkey_account_id = U256::from(581337);
+                let coldkey_account_id = U256::from(81337);
+                let amount = *amount_to_stake;
+                let netuid: u16 = add_dynamic_network(&subnet_owner_hotkey, &subnet_owner_coldkey);
+                register_ok_neuron(netuid, hotkey_account_id, coldkey_account_id, 192213123);
+
+                // Give it some $$$ in his coldkey balance
+                SubtensorModule::add_balance_to_coldkey_account(&coldkey_account_id, amount);
+                AlphaDividendsPerSubnet::<Test>::insert(netuid, hotkey_account_id, *alpha_divs);
+                TotalHotkeyAlphaLastEpoch::<Test>::insert(hotkey_account_id, netuid, *hotkey_alpha);
+                let balance_before = SubtensorModule::get_coldkey_balance(&coldkey_account_id);
+
+                // Stake to hotkey account, and check if the result is ok
+                assert_ok!(SubtensorModule::add_stake(
+                    RuntimeOrigin::signed(coldkey_account_id),
+                    hotkey_account_id,
+                    netuid,
+                    amount
+                ));
+
+                // Remove all stake
+                let stake = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+                    &hotkey_account_id,
+                    &coldkey_account_id,
+                    netuid,
+                );
+
+                assert_ok!(SubtensorModule::remove_stake(
+                    RuntimeOrigin::signed(coldkey_account_id),
+                    hotkey_account_id,
+                    netuid,
+                    stake
+                ));
+
+                // Measure actual fee
+                let balance_after = SubtensorModule::get_coldkey_balance(&coldkey_account_id);
+                let actual_fee = balance_before - balance_after;
+
+                assert_abs_diff_eq!(actual_fee, *expected_fee, epsilon = *expected_fee / 100,);
+            },
+        );
+    });
+}
+
+// cargo test --package pallet-subtensor --lib -- tests::staking::test_staking_sets_div_variables --exact --show-output --nocapture
+#[test]
+fn test_staking_sets_div_variables() {
+    new_test_ext(1).execute_with(|| {
+        let subnet_owner_coldkey = U256::from(1);
+        let subnet_owner_hotkey = U256::from(2);
+        let hotkey_account_id = U256::from(581337);
+        let coldkey_account_id = U256::from(81337);
+        let amount = 100_000_000_000;
+        let netuid: u16 = add_dynamic_network(&subnet_owner_hotkey, &subnet_owner_coldkey);
+        let tempo = 10;
+        Tempo::<Test>::insert(netuid, tempo);
+        register_ok_neuron(netuid, hotkey_account_id, coldkey_account_id, 192213123);
+
+        // Give it some $$$ in his coldkey balance
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey_account_id, amount);
+
+        // Verify that divident variables are clear in the beginning
+        assert_eq!(
+            AlphaDividendsPerSubnet::<Test>::get(netuid, hotkey_account_id),
+            0
+        );
+        assert_eq!(
+            TotalHotkeyAlphaLastEpoch::<Test>::get(hotkey_account_id, netuid),
+            0
+        );
+
+        // Stake to hotkey account, and check if the result is ok
+        assert_ok!(SubtensorModule::add_stake(
+            RuntimeOrigin::signed(coldkey_account_id),
+            hotkey_account_id,
+            netuid,
+            amount
+        ));
+
+        // Verify that divident variables are still clear in the beginning
+        assert_eq!(
+            AlphaDividendsPerSubnet::<Test>::get(netuid, hotkey_account_id),
+            0
+        );
+        assert_eq!(
+            TotalHotkeyAlphaLastEpoch::<Test>::get(hotkey_account_id, netuid),
+            0
+        );
+
+        // Wait for 1 epoch
+        step_block(tempo + 1);
+
+        // Verify that divident variables have been set
+        let stake = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey_account_id,
+            &coldkey_account_id,
+            netuid,
+        );
+
+        assert!(AlphaDividendsPerSubnet::<Test>::get(netuid, hotkey_account_id) > 0);
+        assert_abs_diff_eq!(
+            TotalHotkeyAlphaLastEpoch::<Test>::get(hotkey_account_id, netuid),
+            stake,
+            epsilon = stake / 100_000
+        );
+    });
+}
+
 /***********************************************************
     staking::get_coldkey_balance() tests
 ************************************************************/
@@ -2300,7 +2464,7 @@ fn test_remove_stake_fee_realistic_values() {
         SubnetTAO::<Test>::insert(netuid, tao_reserve.to_num::<u64>());
         SubnetAlphaIn::<Test>::insert(netuid, alpha_in.to_num::<u64>());
         AlphaDividendsPerSubnet::<Test>::insert(netuid, hotkey, alpha_divs);
-        let current_price = SubtensorModule::get_alpha_price(netuid).to_num::<f64>();
+        TotalHotkeyAlphaLastEpoch::<Test>::insert(hotkey, netuid, alpha_to_unstake);
 
         // Add stake first time to init TotalHotkeyAlpha
         SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
@@ -2310,17 +2474,18 @@ fn test_remove_stake_fee_realistic_values() {
             alpha_to_unstake,
         );
 
-        // Estimate fees
-        let mut expected_fee: f64 = current_price * alpha_divs as f64;
-        if expected_fee < alpha_to_unstake as f64 * 0.00005 {
-            expected_fee = alpha_to_unstake as f64 * 0.00005;
-        }
-
         // Remove stake to measure fee
         let balance_before = SubtensorModule::get_coldkey_balance(&coldkey);
         let expected_tao_no_fee =
             SubtensorModule::sim_swap_alpha_for_tao(netuid, alpha_to_unstake).unwrap();
 
+        // Estimate fees
+        let mut expected_fee =
+            expected_tao_no_fee as f64 * alpha_divs as f64 / alpha_to_unstake as f64;
+        if expected_fee < expected_tao_no_fee as f64 * 0.00005 {
+            expected_fee = expected_tao_no_fee as f64 * 0.00005;
+        }
+
         assert_ok!(SubtensorModule::remove_stake(
             RuntimeOrigin::signed(coldkey),
             hotkey,
@@ -3942,7 +4107,7 @@ fn test_remove_99_9991_per_cent_stake_removes_all() {
         assert_abs_diff_eq!(
             SubtensorModule::get_coldkey_balance(&coldkey_account_id),
             amount - fee,
-            epsilon = 10000,
+            epsilon = 100000,
         );
         assert_eq!(
             SubtensorModule::get_total_stake_for_hotkey(&hotkey_account_id),

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
+        );
     });
 }