Add lock-aware stake transfer extrinsic

pallets/subtensor/src/macros/dispatches.rs

@@ -1558,6 +1558,32 @@ mod dispatches {
             )
         }
 
+        /// Transfers locked or unlocked stake from one coldkey to another, within one subnet,
+        /// while keeping the same hotkey.
+        ///
+        /// If `locked` is true, the call transfers at most the currently locked alpha and moves
+        /// the corresponding lock state. If `locked` is false, it transfers at most the currently
+        /// unlocked alpha.
+        #[pallet::call_index(139)]
+        #[pallet::weight(<T as crate::pallet::Config>::WeightInfo::transfer_stake())]
+        pub fn transfer_stake_lock_aware(
+            origin: OriginFor<T>,
+            destination_coldkey: T::AccountId,
+            hotkey: T::AccountId,
+            netuid: NetUid,
+            alpha_amount: AlphaBalance,
+            locked: bool,
+        ) -> DispatchResult {
+            Self::do_transfer_stake_lock_aware(
+                origin,
+                destination_coldkey,
+                hotkey,
+                netuid,
+                alpha_amount,
+                locked,
+            )
+        }
+
         /// Swaps a specified amount of stake from one subnet to another, while keeping the same coldkey and hotkey.
         ///
         /// # Arguments

pallets/subtensor/src/staking/lock.rs

@@ -680,6 +680,20 @@ impl<T: Config> Pallet<T> {
         }
     }
 
+    /// Returns the transferable alpha in either the locked or unlocked bucket.
+    pub fn lock_aware_transferable_alpha(
+        coldkey: &T::AccountId,
+        netuid: NetUid,
+        locked: bool,
+    ) -> AlphaBalance {
+        if locked {
+            Self::get_current_locked(coldkey, netuid)
+                .min(Self::total_coldkey_alpha_on_subnet(coldkey, netuid))
+        } else {
+            Self::available_to_unstake(coldkey, netuid)
+        }
+    }
+
     /// Ensures that the amount can be unstaked
     pub fn ensure_available_to_unstake(
         coldkey: &T::AccountId,
@@ -1671,6 +1685,7 @@ impl<T: Config> Pallet<T> {
         destination_coldkey: &T::AccountId,
         netuid: NetUid,
         amount: AlphaBalance,
+        lock_aware_transfer: Option<bool>,
     ) -> DispatchResult {
         let now = Self::get_current_block_as_u64();
 
@@ -1719,9 +1734,22 @@ impl<T: Config> Pallet<T> {
         let unavailable = source_lock.locked_mass;
         let available_stake = total_alpha.saturating_sub(unavailable);
 
-        // Reduce remaining_to_transfer by min(remaining_to_transfer, available stake)
-        let available_transfer = remaining_to_transfer.min(available_stake);
-        remaining_to_transfer = remaining_to_transfer.saturating_sub(available_transfer);
+        // In the default mode, transfers consume unlocked stake first. Any amount
+        // left in `remaining_to_transfer` after this subtraction must come from
+        // the lock and needs lock state moved with it. In locked-only mode, skip
+        // this subtraction so the whole capped amount is treated as locked so
+        // effectively we start the transfer from the locked portion without accounting
+        // for unlocked alpha.
+        if lock_aware_transfer != Some(true) {
+            let available_transfer = remaining_to_transfer.min(available_stake);
+            remaining_to_transfer = remaining_to_transfer.saturating_sub(available_transfer);
+        }
+
+        // In unlocked-only mode, the capped amount has already been fully covered
+        // by unlocked stake, so no lock state should be transferred.
+        if lock_aware_transfer == Some(false) {
+            remaining_to_transfer = AlphaBalance::ZERO;
+        }
 
         // If result is non-zero, check the hotkey match between source and destination coldkey locks
         // (if destination coldkey lock exists). If no match, error out with LockHotkeyMismatch, otherwise,

pallets/subtensor/src/staking/move_stake.rs

@@ -50,6 +50,7 @@ impl<T: Config> Pallet<T> {
             None,
             None,
             false,
+            None,
         )?;
 
         // Log the event.
@@ -140,6 +141,7 @@ impl<T: Config> Pallet<T> {
             None,
             None,
             true,
+            None,
         )?;
 
         // 9. Emit an event for logging/monitoring.
@@ -159,6 +161,51 @@ impl<T: Config> Pallet<T> {
         Ok(())
     }
 
+    /// Transfers either locked or unlocked stake from one coldkey to another within one subnet.
+    ///
+    /// This follows `do_transfer_stake`, but caps the requested amount to the
+    /// selected lock bucket. If `locked` is true, only locked alpha can move and
+    /// the matching lock state follows the stake. If `locked` is false, only
+    /// unlocked alpha can move.
+    pub fn do_transfer_stake_lock_aware(
+        origin: OriginFor<T>,
+        destination_coldkey: T::AccountId,
+        hotkey: T::AccountId,
+        netuid: NetUid,
+        alpha_amount: AlphaBalance,
+        locked: bool,
+    ) -> dispatch::DispatchResult {
+        let coldkey = ensure_signed(origin)?;
+
+        let tao_moved = Self::transition_stake_internal(
+            &coldkey,
+            &destination_coldkey,
+            &hotkey,
+            &hotkey,
+            netuid,
+            netuid,
+            alpha_amount,
+            None,
+            None,
+            true,
+            Some(locked),
+        )?;
+
+        log::debug!(
+            "StakeTransferred(origin_coldkey: {coldkey:?}, destination_coldkey: {destination_coldkey:?}, hotkey: {hotkey:?}, netuid: {netuid:?}, amount: {tao_moved:?})"
+        );
+        Self::deposit_event(Event::StakeTransferred(
+            coldkey,
+            destination_coldkey,
+            hotkey,
+            netuid,
+            netuid,
+            tao_moved,
+        ));
+
+        Ok(())
+    }
+
     /// Swaps a specified amount of stake for the same `(coldkey, hotkey)` pair from one subnet
     /// (`origin_netuid`) to another (`destination_netuid`).
     ///
@@ -204,6 +251,7 @@ impl<T: Config> Pallet<T> {
             None,
             None,
             false,
+            None,
         )?;
 
         // Emit an event for logging.
@@ -271,6 +319,7 @@ impl<T: Config> Pallet<T> {
             Some(limit_price),
             Some(allow_partial),
             false,
+            None,
         )?;
 
         // Emit an event for logging.
@@ -302,6 +351,7 @@ impl<T: Config> Pallet<T> {
         maybe_limit_price: Option<TaoBalance>,
         maybe_allow_partial: Option<bool>,
         check_transfer_toggle: bool,
+        lock_aware_transfer: Option<bool>,
     ) -> Result<TaoBalance, DispatchError> {
         // Cap the alpha_amount at available Alpha because user might be paying transaxtion fees
         // in Alpha and their total is already reduced by now.
@@ -311,6 +361,15 @@ impl<T: Config> Pallet<T> {
             origin_netuid,
         );
         let alpha_amount = alpha_amount.min(alpha_available);
+        let alpha_amount = lock_aware_transfer
+            .map(|locked| {
+                alpha_amount.min(Self::lock_aware_transferable_alpha(
+                    origin_coldkey,
+                    origin_netuid,
+                    locked,
+                ))
+            })
+            .unwrap_or(alpha_amount);
 
         // Calculate the maximum amount that can be executed
         let max_amount = if origin_netuid != destination_netuid {
@@ -393,6 +452,7 @@ impl<T: Config> Pallet<T> {
                 destination_hotkey,
                 origin_netuid,
                 move_amount,
+                lock_aware_transfer,
             )
         }
     }

pallets/subtensor/src/staking/stake_utils.rs

@@ -954,9 +954,31 @@ impl<T: Config> Pallet<T> {
         destination_hotkey: &T::AccountId,
         netuid: NetUid,
         alpha: AlphaBalance,
+        lock_aware_transfer: Option<bool>,
     ) -> Result<TaoBalance, DispatchError> {
+        // Calculate TAO equivalent based on current price (it is accurate because
+        // there's no slippage in this move) and validate it before mutating lock
+        // or stake storage.
+        let current_price =
+            <T as pallet::Config>::SwapInterface::current_alpha_price(netuid.into());
+        let tao_equivalent: TaoBalance = current_price
+            .saturating_mul(U96F32::saturating_from_num(alpha))
+            .saturating_to_num::<u64>()
+            .into();
+
+        ensure!(
+            tao_equivalent >= DefaultMinStake::<T>::get(),
+            Error::<T>::AmountTooLow
+        );
+
         // Transfer lock (may fail if destination coldkey has a conflicting lock)
-        Self::transfer_lock(origin_coldkey, destination_coldkey, netuid, alpha)?;
+        Self::transfer_lock(
+            origin_coldkey,
+            destination_coldkey,
+            netuid,
+            alpha,
+            lock_aware_transfer,
+        )?;
 
         // Decrease alpha on origin keys
         Self::decrease_stake_for_hotkey_and_coldkey_on_subnet(
@@ -988,21 +1010,6 @@ impl<T: Config> Pallet<T> {
             );
         }
 
-        // Calculate TAO equivalent based on current price (it is accurate because
-        // there's no slippage in this move)
-        let current_price =
-            <T as pallet::Config>::SwapInterface::current_alpha_price(netuid.into());
-        let tao_equivalent: TaoBalance = current_price
-            .saturating_mul(U96F32::saturating_from_num(alpha))
-            .saturating_to_num::<u64>()
-            .into();
-
-        // Ensure tao_equivalent is above DefaultMinStake
-        ensure!(
-            tao_equivalent >= DefaultMinStake::<T>::get(),
-            Error::<T>::AmountTooLow
-        );
-
         // Step 3: Update StakingHotkeys if the hotkey's total alpha, across all subnets, is zero
         // TODO: fix.
         // if Self::get_stake(hotkey, coldkey) == 0 {

pallets/subtensor/src/subnets/leasing.rs

@@ -312,6 +312,7 @@ impl<T: Config> Pallet<T> {
                     &lease.hotkey,
                     lease.netuid,
                     alpha_for_contributor.into(),
+                    None,
                 )?;
                 alpha_distributed = alpha_distributed.saturating_add(alpha_for_contributor.into());
 
@@ -332,6 +333,7 @@ impl<T: Config> Pallet<T> {
                 &lease.hotkey,
                 lease.netuid,
                 beneficiary_cut_alpha.into(),
+                None,
             )?;
             Self::deposit_event(Event::SubnetLeaseDividendsDistributed {
                 lease_id,