set auto hotkey per subnet

pallets/subtensor/rpc/src/lib.rs

@@ -92,6 +92,13 @@ pub trait SubtensorCustomApi<BlockHash> {
         metagraph_index: Vec<u16>,
         at: Option<BlockHash>,
     ) -> RpcResult<Vec<u8>>;
+    #[method(name = "subnetInfo_getColdkeyAutoStakeHotkey")]
+    fn get_coldkey_auto_stake_hotkey(
+        &self,
+        coldkey: AccountId32,
+        netuid: NetUid,
+        at: Option<BlockHash>,
+    ) -> RpcResult<Vec<u8>>;
     #[method(name = "subnetInfo_getSelectiveMechagraph")]
     fn get_selective_mechagraph(
         &self,
@@ -474,6 +481,24 @@ where
         }
     }
 
+    fn get_coldkey_auto_stake_hotkey(
+        &self,
+        coldkey: AccountId32,
+        netuid: NetUid,
+        at: Option<<Block as BlockT>::Hash>,
+    ) -> RpcResult<Vec<u8>> {
+        let api = self.client.runtime_api();
+        let at = at.unwrap_or_else(|| self.client.info().best_hash);
+
+        match api.get_coldkey_auto_stake_hotkey(at, coldkey, netuid) {
+            Ok(result) => Ok(result.encode()),
+            Err(e) => Err(Error::RuntimeError(format!(
+                "Unable to get coldkey auto stake hotkey: {e:?}"
+            ))
+            .into()),
+        }
+    }
+
     fn get_selective_mechagraph(
         &self,
         netuid: NetUid,

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

@@ -45,6 +45,7 @@ sp_api::decl_runtime_apis! {
         fn get_dynamic_info(netuid: NetUid) -> Option<DynamicInfo<AccountId32>>;
         fn get_subnet_state(netuid: NetUid) -> Option<SubnetState<AccountId32>>;
         fn get_selective_metagraph(netuid: NetUid, metagraph_indexes: Vec<u16>) -> Option<SelectiveMetagraph<AccountId32>>;
+        fn get_coldkey_auto_stake_hotkey(coldkey: AccountId32, netuid: NetUid) -> Option<AccountId32>;
         fn get_selective_mechagraph(netuid: NetUid, subid: MechId, metagraph_indexes: Vec<u16>) -> Option<SelectiveMetagraph<AccountId32>>;
         fn get_subnet_to_prune() -> Option<NetUid>;
     }

pallets/subtensor/src/benchmarks.rs

@@ -1560,9 +1560,21 @@ mod pallet_benchmarks {
     #[benchmark]
     fn set_coldkey_auto_stake_hotkey() {
         let coldkey: T::AccountId = whitelisted_caller();
-        let hot: T::AccountId = account("A", 0, 1);
+        let netuid = NetUid::from(1);
+        let hotkey: T::AccountId = account("A", 0, 1);
+        SubtokenEnabled::<T>::insert(netuid, true);
+        Subtensor::<T>::init_new_network(netuid, 1);
+        let amount = 900_000_000_000;
+
+        Subtensor::<T>::add_balance_to_coldkey_account(&coldkey.clone(), amount);
+
+        assert_ok!(Subtensor::<T>::burned_register(
+            RawOrigin::Signed(coldkey.clone()).into(),
+            netuid,
+            hotkey.clone()
+        ));
 
         #[extrinsic_call]
-        _(RawOrigin::Signed(coldkey.clone()), hot.clone());
+        _(RawOrigin::Signed(coldkey.clone()), netuid, hotkey.clone());
     }
 }

pallets/subtensor/src/coinbase/run_coinbase.rs

@@ -514,7 +514,7 @@ impl<T: Config> Pallet<T> {
             }
 
             let owner: T::AccountId = Owner::<T>::get(&hotkey);
-            let maybe_dest = AutoStakeDestination::<T>::get(&owner);
+            let maybe_dest = AutoStakeDestination::<T>::get(&owner, netuid);
 
             // Always stake but only emit event if autostake is set.
             let destination = maybe_dest.clone().unwrap_or(hotkey.clone());

pallets/subtensor/src/lib.rs

@@ -1140,9 +1140,16 @@ pub mod pallet {
     #[pallet::storage] // --- MAP ( cold ) --> Vec<hot> | Returns the vector of hotkeys controlled by this coldkey.
     pub type OwnedHotkeys<T: Config> =
         StorageMap<_, Blake2_128Concat, T::AccountId, Vec<T::AccountId>, ValueQuery>;
-    #[pallet::storage] // --- MAP ( cold ) --> hot | Returns the hotkey a coldkey will autostake to with mining rewards.
-    pub type AutoStakeDestination<T: Config> =
-        StorageMap<_, Blake2_128Concat, T::AccountId, T::AccountId, OptionQuery>;
+    #[pallet::storage] // --- DMAP ( cold, netuid )--> hot | Returns the hotkey a coldkey will autostake to with mining rewards.
+    pub type AutoStakeDestination<T: Config> = StorageDoubleMap<
+        _,
+        Blake2_128Concat,
+        T::AccountId,
+        Identity,
+        NetUid,
+        T::AccountId,
+        OptionQuery,
+    >;
 
     #[pallet::storage] // --- DMAP ( cold ) --> (block_expected, new_coldkey) | Maps coldkey to the block to swap at and new coldkey.
     pub type ColdkeySwapScheduled<T: Config> = StorageMap<

pallets/subtensor/src/macros/dispatches.rs

@@ -1086,7 +1086,7 @@ mod dispatches {
         /// Weight is calculated based on the number of database reads and writes.
         #[pallet::call_index(71)]
         #[pallet::weight((Weight::from_parts(161_700_000, 0)
-        .saturating_add(T::DbWeight::get().reads(15_u64))
+        .saturating_add(T::DbWeight::get().reads(16_u64))
         .saturating_add(T::DbWeight::get().writes(9)), DispatchClass::Operational, Pays::No))]
         pub fn swap_coldkey(
             origin: OriginFor<T>,
@@ -2292,16 +2292,36 @@ mod dispatches {
         /// * `hotkey` (T::AccountId):
         ///     - The hotkey account to designate as the autostake destination.
         #[pallet::call_index(114)]
-        #[pallet::weight((Weight::from_parts(5_170_000, 0)
-		.saturating_add(T::DbWeight::get().reads(0_u64))
+        #[pallet::weight((Weight::from_parts(29_930_000, 0)
+		.saturating_add(T::DbWeight::get().reads(3_u64))
 		.saturating_add(T::DbWeight::get().writes(1)), DispatchClass::Normal, Pays::No))]
         pub fn set_coldkey_auto_stake_hotkey(
             origin: T::RuntimeOrigin,
+            netuid: NetUid,
             hotkey: T::AccountId,
         ) -> DispatchResult {
             let coldkey = ensure_signed(origin)?;
+            ensure!(Self::if_subnet_exist(netuid), Error::<T>::SubnetNotExists);
+            ensure!(
+                Uids::<T>::contains_key(netuid, &hotkey),
+                Error::<T>::HotKeyNotRegisteredInSubNet
+            );
+
+            let current_hotkey = AutoStakeDestination::<T>::get(coldkey.clone(), netuid);
+            if let Some(current_hotkey) = current_hotkey {
+                ensure!(
+                    current_hotkey != hotkey,
+                    Error::<T>::SameAutoStakeHotkeyAlreadySet
+                );
+            }
+
+            AutoStakeDestination::<T>::insert(coldkey.clone(), netuid, hotkey.clone());
 
-            AutoStakeDestination::<T>::insert(coldkey, hotkey.clone());
+            Self::deposit_event(Event::AutoStakeDestinationSet {
+                coldkey,
+                netuid,
+                hotkey,
+            });
 
             Ok(())
         }

pallets/subtensor/src/macros/errors.rs

@@ -256,6 +256,8 @@ mod errors {
         CannotAffordLockCost,
         /// exceeded the rate limit for associating an EVM key.
         EvmKeyAssociateRateLimitExceeded,
+        /// Same auto stake hotkey already set
+        SameAutoStakeHotkeyAlreadySet,
         /// The UID map for the subnet could not be cleared
         UidMapCouldNotBeCleared,
         /// Trimming would exceed the max immune neurons percentage

pallets/subtensor/src/macros/events.rs

@@ -442,5 +442,19 @@ mod events {
 
         /// The minimum allowed UIDs for a subnet have been set.
         MinAllowedUidsSet(NetUid, u16),
+
+        /// The auto stake destination has been set.
+        ///
+        /// - **coldkey**: The account ID of the coldkey.
+        /// - **netuid**: The network identifier.
+        /// - **hotkey**: The account ID of the hotkey.
+        AutoStakeDestinationSet {
+            /// The account ID of the coldkey.
+            coldkey: T::AccountId,
+            /// The network identifier.
+            netuid: NetUid,
+            /// The account ID of the hotkey.
+            hotkey: T::AccountId,
+        },
     }
 }

pallets/subtensor/src/migrations/migrate_auto_stake_destination.rs

@@ -0,0 +1,82 @@
+use super::*;
+use crate::AccountIdOf;
+use frame_support::{
+    IterableStorageMap,
+    pallet_prelude::{Blake2_128Concat, OptionQuery},
+    storage_alias,
+    traits::Get,
+    weights::Weight,
+};
+use scale_info::prelude::string::String;
+
+/// Module containing deprecated storage format for AutoStakeDestination
+pub mod deprecated_auto_stake_destination_format {
+    use super::*;
+
+    #[storage_alias]
+    pub(super) type AutoStakeDestination<T: Config> =
+        StorageMap<Pallet<T>, Blake2_128Concat, AccountIdOf<T>, AccountIdOf<T>, OptionQuery>;
+}
+
+/// Migrate the AutoStakeDestination map from single map to double map format
+pub fn migrate_auto_stake_destination<T: Config>() -> Weight {
+    use deprecated_auto_stake_destination_format as old;
+
+    let migration_name = b"migrate_auto_stake_destination".to_vec();
+    let mut weight = T::DbWeight::get().reads(1);
+
+    if HasMigrationRun::<T>::get(&migration_name) {
+        log::info!(
+            "Migration '{:?}' has already run. Skipping.",
+            String::from_utf8_lossy(&migration_name)
+        );
+        return weight;
+    }
+
+    log::info!(
+        "Running migration '{}'",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    // ------------------------------
+    // Step 1: Migrate AutoStakeDestination entries
+    // ------------------------------
+
+    let curr_keys: Vec<AccountIdOf<T>> = old::AutoStakeDestination::<T>::iter_keys().collect();
+    let root_netuid = NetUid::ROOT;
+    let netuids: Vec<NetUid> = <NetworksAdded<T> as IterableStorageMap<NetUid, bool>>::iter()
+        .map(|(netuid, _)| netuid)
+        .collect();
+
+    for coldkey in &curr_keys {
+        weight.saturating_accrue(T::DbWeight::get().reads(1));
+
+        if let Some(hotkey) = old::AutoStakeDestination::<T>::get(coldkey) {
+            for netuid in netuids.iter() {
+                if *netuid == root_netuid {
+                    continue;
+                }
+                AutoStakeDestination::<T>::insert(coldkey, netuid, hotkey.clone());
+            }
+
+            old::AutoStakeDestination::<T>::remove(coldkey);
+
+            weight.saturating_accrue(T::DbWeight::get().writes(netuids.len() as u64));
+        }
+    }
+
+    // ------------------------------
+    // Step 2: Mark Migration as Completed
+    // ------------------------------
+
+    HasMigrationRun::<T>::insert(&migration_name, true);
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    log::info!(
+        "Migration '{:?}' completed successfully. {} entries migrated.",
+        String::from_utf8_lossy(&migration_name),
+        curr_keys.len()
+    );
+
+    weight
+}

pallets/subtensor/src/migrations/mod.rs

@@ -4,6 +4,7 @@ use frame_support::pallet_prelude::Weight;
 use sp_io::KillStorageResult;
 use sp_io::hashing::twox_128;
 use sp_io::storage::clear_prefix;
+pub mod migrate_auto_stake_destination;
 pub mod migrate_chain_identity;
 pub mod migrate_coldkey_swap_scheduled;
 pub mod migrate_commit_reveal_settings;

pallets/subtensor/src/rpc_info/subnet_info.rs

@@ -403,4 +403,11 @@ impl<T: Config> Pallet<T> {
             user_liquidity_enabled,
         })
     }
+
+    pub fn get_coldkey_auto_stake_hotkey(
+        coldkey: T::AccountId,
+        netuid: NetUid,
+    ) -> Option<T::AccountId> {
+        AutoStakeDestination::<T>::get(coldkey, netuid)
+    }
 }

pallets/subtensor/src/swap/swap_coldkey.rs

@@ -155,6 +155,12 @@ impl<T: Config> Pallet<T> {
                 SubnetOwner::<T>::insert(netuid, new_coldkey.clone());
             }
             weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
+
+            if let Some(old_auto_stake_hotkey) = AutoStakeDestination::<T>::get(old_coldkey, netuid)
+            {
+                AutoStakeDestination::<T>::remove(old_coldkey, netuid);
+                AutoStakeDestination::<T>::insert(new_coldkey, netuid, old_auto_stake_hotkey);
+            }
         }
 
         // 3. Swap Stake.
@@ -178,11 +184,6 @@ impl<T: Config> Pallet<T> {
             weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
         }
 
-        if let Some(old_auto_stake_hotkey) = AutoStakeDestination::<T>::get(old_coldkey) {
-            AutoStakeDestination::<T>::remove(old_coldkey);
-            AutoStakeDestination::<T>::insert(new_coldkey, old_auto_stake_hotkey);
-        }
-
         // 4. Swap TotalColdkeyAlpha (DEPRECATED)
         // for netuid in Self::get_all_subnet_netuids() {
         //     let old_alpha_stake: u64 = TotalColdkeyAlpha::<T>::get(old_coldkey, netuid);