Merge branch 'devnet-ready' into devnet-ready-with-nodes

Author: Roman

pallets/subtensor/runtime-api/src/lib.rs

@@ -46,6 +46,7 @@ sp_api::decl_runtime_apis! {
         fn get_stake_info_for_coldkey( coldkey_account: AccountId32 ) -> Vec<StakeInfo<AccountId32>>;
         fn get_stake_info_for_coldkeys( coldkey_accounts: Vec<AccountId32> ) -> Vec<(AccountId32, Vec<StakeInfo<AccountId32>>)>;
         fn get_stake_info_for_hotkey_coldkey_netuid( hotkey_account: AccountId32, coldkey_account: AccountId32, netuid: u16 ) -> Option<StakeInfo<AccountId32>>;
+        fn get_stake_fee( origin: Option<(AccountId32, u16)>, origin_coldkey_account: AccountId32, destination: Option<(AccountId32, u16)>, destination_coldkey_account: AccountId32, amount: u64 ) -> u64;
     }
 
     pub trait SubnetRegistrationRuntimeApi {

pallets/subtensor/src/rpc_info/stake_info.rs

@@ -2,6 +2,7 @@ use super::*;
 use frame_support::pallet_prelude::{Decode, Encode};
 extern crate alloc;
 use codec::Compact;
+use substrate_fixed::types::I96F32;
 
 #[freeze_struct("5cfb3c84c3af3116")]
 #[derive(Decode, Encode, PartialEq, Eq, Clone, Debug, TypeInfo)]
@@ -112,4 +113,24 @@ impl<T: Config> Pallet<T> {
             is_registered,
         })
     }
+
+    pub fn get_stake_fee(
+        origin: Option<(T::AccountId, u16)>,
+        _origin_coldkey_account: T::AccountId,
+        _destination: Option<(T::AccountId, u16)>,
+        _destination_coldkey_account: T::AccountId,
+        amount: u64,
+    ) -> u64 {
+        match origin {
+            Some((origin_hotkey_account, origin_netuid)) => Self::calculate_staking_fee(
+                origin_netuid,
+                &origin_hotkey_account,
+                I96F32::saturating_from_num(amount),
+            ),
+            None => {
+                // Adding stake (comes from no netuid)
+                DefaultStakingFee::<T>::get()
+            }
+        }
+    }
 }

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

@@ -623,3 +623,152 @@ fn test_try_associate_hotkey() {
         assert_eq!(SubtensorModule::get_owned_hotkeys(&coldkey2).len(), 0);
     });
 }
+
+#[test]
+fn test_stake_fee_api() {
+    new_test_ext(1).execute_with(|| {
+        let hotkey1 = U256::from(1);
+        let coldkey1 = U256::from(2);
+        let hotkey2 = U256::from(3);
+        let coldkey2 = U256::from(4);
+
+        let netuid0 = 1;
+        let netuid1 = 2;
+        let root_netuid = SubtensorModule::get_root_netuid();
+
+        let alpha_divs = 100_000_000_000;
+        let total_hotkey_alpha = 100_000_000_000;
+        let tao_in = 100_000_000_000; // 100 TAO
+        let reciprocal_price = 2; // 1 / price
+        let stake_amount = 100_000_000_000;
+
+        let default_fee = DefaultStakingFee::<Test>::get();
+
+        // Setup alpha out
+        SubnetAlphaOut::<Test>::insert(netuid0, 100_000_000_000);
+        SubnetAlphaOut::<Test>::insert(netuid1, 100_000_000_000);
+        // Set pools using price
+        SubnetAlphaIn::<Test>::insert(netuid0, tao_in * reciprocal_price);
+        SubnetTAO::<Test>::insert(netuid0, tao_in);
+        SubnetAlphaIn::<Test>::insert(netuid1, tao_in * reciprocal_price);
+        SubnetTAO::<Test>::insert(netuid1, tao_in);
+
+        // Setup alpha divs for hotkey1
+        AlphaDividendsPerSubnet::<Test>::insert(netuid0, hotkey1, alpha_divs);
+        AlphaDividendsPerSubnet::<Test>::insert(netuid1, hotkey1, alpha_divs);
+
+        // Setup total hotkey alpha for hotkey1
+        TotalHotkeyAlpha::<Test>::insert(hotkey1, netuid0, total_hotkey_alpha);
+        TotalHotkeyAlpha::<Test>::insert(hotkey1, netuid1, total_hotkey_alpha);
+
+        // Test stake fee for add_stake
+        let stake_fee_0 = SubtensorModule::get_stake_fee(
+            None,
+            coldkey1,
+            Some((hotkey1, netuid0)),
+            coldkey1,
+            stake_amount,
+        ); // Default for adding stake
+        assert_eq!(stake_fee_0, default_fee);
+
+        // Test stake fee for remove on root
+        let stake_fee_1 = SubtensorModule::get_stake_fee(
+            Some((hotkey1, root_netuid)),
+            coldkey1,
+            None,
+            coldkey1,
+            stake_amount,
+        ); // Default for removing stake from root
+        assert_eq!(stake_fee_1, default_fee);
+
+        // Test stake fee for move from root to non-root
+        let stake_fee_2 = SubtensorModule::get_stake_fee(
+            Some((hotkey1, root_netuid)),
+            coldkey1,
+            Some((hotkey1, netuid0)),
+            coldkey1,
+            stake_amount,
+        ); // Default for moving stake from root to non-root
+        assert_eq!(stake_fee_2, default_fee);
+
+        // Test stake fee for move between hotkeys on root
+        let stake_fee_3 = SubtensorModule::get_stake_fee(
+            Some((hotkey1, root_netuid)),
+            coldkey1,
+            Some((hotkey2, root_netuid)),
+            coldkey1,
+            stake_amount,
+        ); // Default for moving stake between hotkeys on root
+        assert_eq!(stake_fee_3, default_fee);
+
+        // Test stake fee for move between coldkeys on root
+        let stake_fee_4 = SubtensorModule::get_stake_fee(
+            Some((hotkey1, root_netuid)),
+            coldkey1,
+            Some((hotkey1, root_netuid)),
+            coldkey2,
+            stake_amount,
+        ); // Default for moving stake between coldkeys on root
+        assert_eq!(stake_fee_4, default_fee);
+
+        // Test stake fee for *swap* from non-root to root
+        let stake_fee_5 = SubtensorModule::get_stake_fee(
+            Some((hotkey1, netuid0)),
+            coldkey1,
+            Some((hotkey1, root_netuid)),
+            coldkey1,
+            stake_amount,
+        ); // Charged a dynamic fee
+        let dynamic_fee_5 = SubtensorModule::calculate_staking_fee(
+            netuid0,
+            &hotkey1,
+            I96F32::from_num(stake_amount),
+        );
+        assert_eq!(stake_fee_5, dynamic_fee_5);
+
+        // Test stake fee for move between hotkeys on non-root
+        let stake_fee_6 = SubtensorModule::get_stake_fee(
+            Some((hotkey1, netuid0)),
+            coldkey1,
+            Some((hotkey2, netuid0)),
+            coldkey1,
+            stake_amount,
+        ); // Charged a dynamic fee
+        let dynamic_fee_6 = SubtensorModule::calculate_staking_fee(
+            netuid0,
+            &hotkey1,
+            I96F32::from_num(stake_amount),
+        );
+        assert_eq!(stake_fee_6, dynamic_fee_6);
+
+        // Test stake fee for move between coldkeys on non-root
+        let stake_fee_7 = SubtensorModule::get_stake_fee(
+            Some((hotkey1, netuid0)),
+            coldkey1,
+            Some((hotkey1, netuid0)),
+            coldkey2,
+            stake_amount,
+        ); // Charged a dynamic fee
+        let dynamic_fee_7 = SubtensorModule::calculate_staking_fee(
+            netuid0,
+            &hotkey1,
+            I96F32::from_num(stake_amount),
+        );
+        assert_eq!(stake_fee_7, dynamic_fee_7);
+
+        // Test stake fee for *swap* from non-root to non-root
+        let stake_fee_8 = SubtensorModule::get_stake_fee(
+            Some((hotkey1, netuid0)),
+            coldkey1,
+            Some((hotkey1, netuid1)),
+            coldkey1,
+            stake_amount,
+        ); // Charged a dynamic fee
+        let dynamic_fee_8 = SubtensorModule::calculate_staking_fee(
+            netuid0,
+            &hotkey1,
+            I96F32::from_num(stake_amount),
+        );
+        assert_eq!(stake_fee_8, dynamic_fee_8);
+    });
+}

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

runtime/src/lib.rs

@@ -205,7 +205,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
     //   `spec_version`, and `authoring_version` are the same between Wasm and native.
     // This value is set to 100 to notify Polkadot-JS App (https://polkadot.js.org/apps) to use
     //   the compatible custom types.
-    spec_version: 248,
+    spec_version: 249,
     impl_version: 1,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 1,
@@ -2066,6 +2066,10 @@ impl_runtime_apis! {
         fn get_stake_info_for_hotkey_coldkey_netuid( hotkey_account: AccountId32, coldkey_account: AccountId32, netuid: u16 ) -> Option<StakeInfo<AccountId32>> {
             SubtensorModule::get_stake_info_for_hotkey_coldkey_netuid( hotkey_account, coldkey_account, netuid )
         }
+
+        fn get_stake_fee( origin: Option<(AccountId32, u16)>, origin_coldkey_account: AccountId32, destination: Option<(AccountId32, u16)>, destination_coldkey_account: AccountId32, amount: u64 ) -> u64 {
+            SubtensorModule::get_stake_fee( origin, origin_coldkey_account, destination, destination_coldkey_account, amount )
+        }
     }
 
     impl subtensor_custom_rpc_runtime_api::SubnetRegistrationRuntimeApi<Block> for Runtime {