calculate staking fee based on everything

Author: camfairchild

pallets/subtensor/src/staking/move_stake.rs

@@ -339,8 +339,10 @@ impl<T: Config> Pallet<T> {
 
         // Unstake from the origin subnet, returning TAO (or a 1:1 equivalent).
         let fee = Self::calculate_staking_fee(
-            origin_netuid,
-            origin_hotkey,
+            Some((origin_hotkey, origin_netuid)),
+            origin_coldkey,
+            Some((destination_hotkey, destination_netuid)),
+            destination_coldkey,
             I96F32::saturating_from_num(alpha_amount),
         )
         .safe_div(2);

pallets/subtensor/src/staking/remove_stake.rs

@@ -59,8 +59,10 @@ impl<T: Config> Pallet<T> {
 
         // 3. Swap the alpba to tao and update counters for this subnet.
         let fee = Self::calculate_staking_fee(
-            netuid,
-            &hotkey,
+            Some((&hotkey, netuid)),
+            &coldkey,
+            None,
+            &coldkey,
             I96F32::saturating_from_num(alpha_unstaked),
         );
         let tao_unstaked: u64 =
@@ -133,8 +135,10 @@ impl<T: Config> Pallet<T> {
             let alpha_unstaked =
                 Self::get_stake_for_hotkey_and_coldkey_on_subnet(&hotkey, &coldkey, netuid);
             let fee = Self::calculate_staking_fee(
-                netuid,
-                &hotkey,
+                Some((&hotkey, netuid)),
+                &coldkey,
+                None,
+                &coldkey,
                 I96F32::saturating_from_num(alpha_unstaked),
             );
 
@@ -208,8 +212,10 @@ impl<T: Config> Pallet<T> {
                 let alpha_unstaked =
                     Self::get_stake_for_hotkey_and_coldkey_on_subnet(&hotkey, &coldkey, netuid);
                 let fee = Self::calculate_staking_fee(
-                    netuid,
-                    &hotkey,
+                    Some((&hotkey, netuid)),
+                    &coldkey,
+                    None,
+                    &coldkey,
                     I96F32::saturating_from_num(alpha_unstaked),
                 );
 
@@ -315,8 +321,10 @@ impl<T: Config> Pallet<T> {
 
         // 4. Swap the alpha to tao and update counters for this subnet.
         let fee = Self::calculate_staking_fee(
-            netuid,
-            &hotkey,
+            Some((&hotkey, netuid)),
+            &coldkey,
+            None,
+            &coldkey,
             I96F32::saturating_from_num(alpha_unstaked),
         );
         let tao_unstaked =

pallets/subtensor/src/staking/stake_utils.rs

@@ -1073,24 +1073,48 @@ impl<T: Config> Pallet<T> {
     }
 
     pub(crate) fn calculate_staking_fee(
-        netuid: u16,
-        hotkey: &T::AccountId,
+        origin: Option<(&T::AccountId, u16)>,
+        _origin_coldkey: &T::AccountId,
+        destination: Option<(&T::AccountId, u16)>,
+        _destination_coldkey: &T::AccountId,
         alpha_estimate: I96F32,
     ) -> u64 {
-        if (netuid == Self::get_root_netuid()) || (SubnetMechanism::<T>::get(netuid)) == 0 {
-            DefaultStakingFee::<T>::get()
-        } else {
-            let fee = alpha_estimate
-                .saturating_mul(
-                    I96F32::saturating_from_num(AlphaDividendsPerSubnet::<T>::get(netuid, &hotkey))
-                        .safe_div(I96F32::saturating_from_num(TotalHotkeyAlpha::<T>::get(
-                            &hotkey, netuid,
-                        ))),
-                )
-                .saturating_mul(Self::get_alpha_price(netuid)) // fee needs to be in TAO
-                .saturating_to_num::<u64>();
+        match origin {
+            // If origin is defined, we are removing/moving stake
+            Some((origin_hotkey, origin_netuid)) => {
+                if let Some((_destination_hotkey, destination_netuid)) = destination {
+                    // This is a stake move/swap/transfer
+                    if destination_netuid == origin_netuid {
+                        // If destination is on the same subnet, use the default fee
+                        return DefaultStakingFee::<T>::get();
+                    }
+                }
 
-            fee.max(DefaultStakingFee::<T>::get())
+                if origin_netuid == Self::get_root_netuid()
+                    || SubnetMechanism::<T>::get(origin_netuid) == 0
+                {
+                    // If the origin netuid is root, or the subnet mechanism is 0, use the default fee
+                    DefaultStakingFee::<T>::get()
+                } else {
+                    // Otherwise, calculate the fee based on the alpha estimate
+                    let fee = alpha_estimate
+                        .saturating_mul(
+                            I96F32::saturating_from_num(AlphaDividendsPerSubnet::<T>::get(
+                                origin_netuid,
+                                &origin_hotkey,
+                            ))
+                            .safe_div(I96F32::saturating_from_num(
+                                TotalHotkeyAlpha::<T>::get(&origin_hotkey, origin_netuid),
+                            )),
+                        )
+                        .saturating_mul(Self::get_alpha_price(origin_netuid)) // fee needs to be in TAO
+                        .saturating_to_num::<u64>();
+
+                    fee.max(DefaultStakingFee::<T>::get())
+                }
+            }
+            // If origin is not defined, we are adding stake; use default fee
+            None => DefaultStakingFee::<T>::get(),
         }
     }
 }

pallets/subtensor/src/tests/swap_coldkey.rs

@@ -540,7 +540,13 @@ fn test_swap_concurrent_modifications() {
         let additional_stake = 500_000_000_000;
         let initial_stake_alpha =
             I96F32::from(initial_stake).saturating_mul(SubtensorModule::get_alpha_price(netuid));
-        let fee = SubtensorModule::calculate_staking_fee(netuid, &hotkey, initial_stake_alpha);
+        let fee = SubtensorModule::calculate_staking_fee(
+            None,
+            &new_coldkey,
+            Some((&hotkey, netuid)),
+            &new_coldkey,
+            initial_stake_alpha,
+        );
 
         // Setup initial state
         add_network(netuid, 1, 1);