Merge pull request #2008 from opentensor/auto_stake

Auto stake

Author: sam0x17

pallets/subtensor/src/coinbase/run_coinbase.rs

@@ -510,11 +510,12 @@ impl<T: Config> Pallet<T> {
                 );
                 continue;
             }
-
-            // Increase stake for miner.
+            let destination: T::AccountId;
+            let owner: T::AccountId = Owner::<T>::get(&hotkey);
+            destination = AutoStakeDestination::<T>::get(&owner).unwrap_or(hotkey.clone());
             Self::increase_stake_for_hotkey_and_coldkey_on_subnet(
-                &hotkey.clone(),
-                &Owner::<T>::get(hotkey.clone()),
+                &destination,
+                &owner,
                 netuid,
                 incentive,
             );

pallets/subtensor/src/lib.rs

@@ -1083,6 +1083,9 @@ 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 ) --> (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

@@ -854,7 +854,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(14))
+        .saturating_add(T::DbWeight::get().reads(15_u64))
         .saturating_add(T::DbWeight::get().writes(9)), DispatchClass::Operational, Pays::No))]
         pub fn swap_coldkey(
             origin: OriginFor<T>,
@@ -2049,5 +2049,30 @@ mod dispatches {
                 commit_reveal_version,
             )
         }
+
+        /// Set the autostake destination hotkey for a coldkey.
+        ///
+        /// The caller selects a hotkey where all future rewards
+        /// will be automatically staked.  
+        ///
+        /// # Args:
+        /// * `origin` - (<T as frame_system::Config>::Origin):
+        ///     - The signature of the caller's coldkey.
+        ///
+        /// * `hotkey` (T::AccountId):
+        ///     - The hotkey account to designate as the autostake destination.
+        #[pallet::call_index(114)]
+        #[pallet::weight((Weight::from_parts(64_530_000, 0)
+		.saturating_add(T::DbWeight::get().reads(7_u64))
+		.saturating_add(T::DbWeight::get().writes(2)), DispatchClass::Normal, Pays::Yes))]
+        pub fn set_coldkey_auto_stake_hotkey(
+            origin: T::RuntimeOrigin,
+            hotkey: T::AccountId,
+        ) -> DispatchResult {
+            let coldkey = ensure_signed(origin)?;
+            log::debug!("set_coldkey_auto_stake_hotkey( origin:{coldkey:?} hotkey:{hotkey:?} )");
+            AutoStakeDestination::<T>::insert(coldkey, hotkey.clone());
+            Ok(())
+        }
     }
 }

pallets/subtensor/src/swap/swap_coldkey.rs

@@ -178,6 +178,11 @@ 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);

pallets/subtensor/src/tests/coinbase.rs

@@ -2749,3 +2749,101 @@ fn test_coinbase_v3_liquidity_update() {
         assert!(liquidity_before < liquidity_after);
     });
 }
+
+#[test]
+fn test_incentive_is_autostaked_to_owner_destination() {
+    new_test_ext(1).execute_with(|| {
+        let subnet_owner_ck = U256::from(0);
+        let subnet_owner_hk = U256::from(1);
+
+        let miner_ck = U256::from(10);
+        let miner_hk = U256::from(11);
+        let dest_hk = U256::from(12);
+
+        Owner::<Test>::insert(miner_hk, miner_ck);
+        Owner::<Test>::insert(dest_hk, miner_ck);
+        OwnedHotkeys::<Test>::insert(miner_ck, vec![miner_hk, dest_hk]);
+
+        let netuid = add_dynamic_network(&subnet_owner_hk, &subnet_owner_ck);
+
+        Uids::<Test>::insert(netuid, miner_hk, 1);
+        Uids::<Test>::insert(netuid, dest_hk, 2);
+
+        // Set autostake destination for the miner's coldkey
+        AutoStakeDestination::<Test>::insert(miner_ck, dest_hk);
+
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_on_subnet(&miner_hk, netuid),
+            0.into()
+        );
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_on_subnet(&dest_hk, netuid),
+            0.into()
+        );
+
+        // Distribute an incentive to the miner hotkey
+        let mut incentives: BTreeMap<U256, AlphaCurrency> = BTreeMap::new();
+        let incentive: AlphaCurrency = 10_000_000u64.into();
+        incentives.insert(miner_hk, incentive);
+
+        SubtensorModule::distribute_dividends_and_incentives(
+            netuid,
+            AlphaCurrency::ZERO, // owner_cut
+            incentives,
+            BTreeMap::new(), // alpha_dividends
+            BTreeMap::new(), // tao_dividends
+        );
+
+        // Expect the stake to land on the destination hotkey (not the original miner hotkey)
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_on_subnet(&miner_hk, netuid),
+            0.into()
+        );
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_on_subnet(&dest_hk, netuid),
+            incentive
+        );
+    });
+}
+
+#[test]
+fn test_incentive_goes_to_hotkey_when_no_autostake_destination() {
+    new_test_ext(1).execute_with(|| {
+        let subnet_owner_ck = U256::from(0);
+        let subnet_owner_hk = U256::from(1);
+
+        let miner_ck = U256::from(20);
+        let miner_hk = U256::from(21);
+
+        Owner::<Test>::insert(miner_hk, miner_ck);
+        OwnedHotkeys::<Test>::insert(miner_ck, vec![miner_hk]);
+
+        let netuid = add_dynamic_network(&subnet_owner_hk, &subnet_owner_ck);
+
+        Uids::<Test>::insert(netuid, miner_hk, 1);
+
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_on_subnet(&miner_hk, netuid),
+            0.into()
+        );
+
+        // Distribute an incentive to the miner hotkey
+        let mut incentives: BTreeMap<U256, AlphaCurrency> = BTreeMap::new();
+        let incentive: AlphaCurrency = 5_000_000u64.into();
+        incentives.insert(miner_hk, incentive);
+
+        SubtensorModule::distribute_dividends_and_incentives(
+            netuid,
+            AlphaCurrency::ZERO, // owner_cut
+            incentives,
+            BTreeMap::new(), // alpha_dividends
+            BTreeMap::new(), // tao_dividends
+        );
+
+        // With no autostake destination, the incentive should be staked to the original hotkey
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_on_subnet(&miner_hk, netuid),
+            incentive
+        );
+    });
+}