Lock transfers

pallets/subtensor/src/staking/lock.rs

@@ -15,6 +15,7 @@ impl<T: Config> Pallet<T> {
         if !lock_state.locked_mass.is_zero() || !lock_state.unlocked_mass.is_zero() {
             Lock::<T>::insert((coldkey, netuid, hotkey), lock_state);
         } else {
+            // If there is no record previously, this is a no-op
             Lock::<T>::remove((coldkey, netuid, hotkey));
         }
     }
@@ -610,4 +611,138 @@ impl<T: Config> Pallet<T> {
         // needs to happen in that case
         let _ = Self::do_lock_stake(&subnet_owner_coldkey, netuid, &lock_hotkey, amount);
     }
+
+    /// When locked stake is transfered, the lock should follow the stake
+    ///
+    /// First, this function rolls the lock forward and checks if amount is over available
+    /// stake and if it is, the stake that's over the available amount on the destination
+    /// coldkey is locked in the same way as the original stake:
+    ///
+    ///  - If original stake is actively locked to a hotkey, it remains actively locked to
+    ///    the same hotkey
+    ///  - If original stake is being unlocked, the lock is created on the destination coldkey
+    ///    with this amount of unlocked_mass
+    pub fn transfer_lock(
+        origin_coldkey: &T::AccountId,
+        destination_coldkey: &T::AccountId,
+        netuid: NetUid,
+        amount: AlphaBalance,
+    ) -> DispatchResult {
+        let now = Self::get_current_block_as_u64();
+
+        // If no actual transfer happens, this is ok
+        if origin_coldkey == destination_coldkey || amount.is_zero() {
+            return Ok(());
+        }
+
+        // Read total alpha of the coldkey on this netuid. Do not check if total alpha is
+        // lower than amount transferred, this is responsibility of a higher level, this
+        // function needs to act protectively.
+        let total_alpha = Self::total_coldkey_alpha_on_subnet(origin_coldkey, netuid);
+        let mut remaining_to_transfer = amount;
+
+        // Read the locks for source and destination coldkey (if exist) and roll forward
+        let Some((source_hotkey, source_lock)) =
+            Lock::<T>::iter_prefix((origin_coldkey, netuid)).next()
+        else {
+            return Ok(());
+        };
+
+        let mut source_lock = Self::roll_forward_lock(source_lock, now);
+        let maybe_destination_lock = Lock::<T>::iter_prefix((destination_coldkey, netuid))
+            .next()
+            .map(|(hotkey, lock)| (hotkey, Self::roll_forward_lock(lock, now)));
+
+        let mut destination_hotkey = maybe_destination_lock
+            .as_ref()
+            .map(|(hotkey, _)| hotkey.clone())
+            .unwrap_or_else(|| source_hotkey.clone());
+        let mut destination_lock = maybe_destination_lock
+            .as_ref()
+            .map(|(_, lock)| lock.clone())
+            .unwrap_or(LockState {
+                locked_mass: AlphaBalance::ZERO,
+                unlocked_mass: AlphaBalance::ZERO,
+                conviction: U64F64::saturating_from_num(0),
+                last_update: now,
+            });
+
+        // Calculate available stake by subtracting locked_mass and unlocked_mass from total alpha.
+        let unavailable = source_lock
+            .locked_mass
+            .saturating_add(source_lock.unlocked_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);
+
+        // If result is non-zero, reduce remaining_to_transfer by min(unlocked_mass, remaining_to_transfer),
+        // reduce unlocked_mass on the source coldkey by the same amount, increase unlocked_mass on the
+        // destination coldkey by the same amount.
+        if !remaining_to_transfer.is_zero() {
+            let unlocked_transfer = source_lock.unlocked_mass.min(remaining_to_transfer);
+            remaining_to_transfer = remaining_to_transfer.saturating_sub(unlocked_transfer);
+            source_lock.unlocked_mass = source_lock.unlocked_mass.saturating_sub(unlocked_transfer);
+            destination_lock.unlocked_mass = destination_lock
+                .unlocked_mass
+                .saturating_add(unlocked_transfer);
+        }
+
+        // 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,
+        // reduce remaining_to_transfer by min(remaining_to_transfer, locked_mass), reduce locked_mass on
+        // the source coldkey by the same amount, increase locked_mass on the destination coldkey by the
+        // same amount, reduce conviction on the source coldkey proportionally, and increase conviction
+        // on the destination coldkey proportionally.
+        if !remaining_to_transfer.is_zero() {
+            if let Some((existing_hotkey, _)) = maybe_destination_lock.as_ref() {
+                ensure!(
+                    existing_hotkey == &source_hotkey,
+                    Error::<T>::LockHotkeyMismatch
+                );
+                destination_hotkey = existing_hotkey.clone();
+            }
+
+            let locked_transfer = remaining_to_transfer.min(source_lock.locked_mass);
+            let conviction_transfer =
+                if locked_transfer.is_zero() || source_lock.locked_mass.is_zero() {
+                    U64F64::saturating_from_num(0)
+                } else {
+                    // Conviction never exceeds locked_mass, so we can scale it proportionally
+                    // using integer arithmetic without overflowing fixed-point multiplication.
+                    let conviction_u128 = source_lock.conviction.saturating_to_num::<u128>();
+                    let locked_transfer_u128 = locked_transfer.to_u64() as u128;
+                    let source_locked_u128 = source_lock.locked_mass.to_u64() as u128;
+                    let transferred_conviction_u128 = conviction_u128
+                        .saturating_mul(locked_transfer_u128)
+                        .checked_div(source_locked_u128)
+                        .unwrap_or(0);
+                    U64F64::saturating_from_num(transferred_conviction_u128)
+                };
+
+            source_lock.locked_mass = source_lock.locked_mass.saturating_sub(locked_transfer);
+            source_lock.conviction = source_lock.conviction.saturating_sub(conviction_transfer);
+            destination_lock.locked_mass =
+                destination_lock.locked_mass.saturating_add(locked_transfer);
+            destination_lock.conviction = destination_lock
+                .conviction
+                .saturating_add(conviction_transfer);
+        }
+
+        source_lock.last_update = now;
+        destination_lock.last_update = now;
+
+        // Upsert updated locks (only once per this fn) even if there were no updates because
+        // of roll-forward
+        Self::insert_lock_state(origin_coldkey, netuid, &source_hotkey, source_lock);
+        Self::insert_lock_state(
+            destination_coldkey,
+            netuid,
+            &destination_hotkey,
+            destination_lock,
+        );
+
+        Ok(())
+    }
 }

pallets/subtensor/src/staking/stake_utils.rs

@@ -960,6 +960,9 @@ impl<T: Config> Pallet<T> {
         netuid: NetUid,
         alpha: AlphaBalance,
     ) -> Result<TaoBalance, DispatchError> {
+        // Transfer lock (may fail if destination coldkey has a conflicting lock)
+        Self::transfer_lock(origin_coldkey, destination_coldkey, netuid, alpha)?;
+
         // Decrease alpha on origin keys
         Self::decrease_stake_for_hotkey_and_coldkey_on_subnet(
             origin_hotkey,
@@ -1327,9 +1330,9 @@ impl<T: Config> Pallet<T> {
             }
         }
 
-        // Enforce lock invariant: if the operation reduces total coldkey alpha on origin subnet
-        // (cross-coldkey transfer or cross-subnet move), the remaining amount must cover the lock.
-        if origin_coldkey != destination_coldkey || origin_netuid != destination_netuid {
+        // Enforce lock invariant: if the is cross-subnet move, the remaining amount must
+        // cover the lock.
+        if origin_netuid != destination_netuid {
             Self::ensure_available_stake(origin_coldkey, origin_netuid, alpha_amount)?;
         }
 

pallets/subtensor/src/tests/locks.rs

@@ -678,6 +678,66 @@ fn test_move_stake_same_coldkey_same_subnet_allowed() {
     });
 }
 
+#[test]
+fn test_do_transfer_stake_same_subnet_transfers_lock_to_destination_hotkey() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey_sender = U256::from(1);
+        let coldkey_receiver = U256::from(5);
+        let hotkey = U256::from(2);
+        let netuid = setup_subnet_with_stake(coldkey_sender, hotkey, 100_000_000_000);
+
+        let total = SubtensorModule::total_coldkey_alpha_on_subnet(&coldkey_sender, netuid);
+        let lock_half = total / 2.into();
+        assert_ok!(SubtensorModule::do_lock_stake(
+            &coldkey_sender,
+            netuid,
+            &hotkey,
+            lock_half,
+        ));
+
+        let sender_lock_before =
+            Lock::<Test>::get((coldkey_sender, netuid, hotkey)).expect("sender lock should exist");
+        let hotkey_lock_before =
+            HotkeyLock::<Test>::get(netuid, hotkey).expect("hotkey lock should exist");
+
+        step_block(1);
+
+        let transfer_amount = total;
+        assert_ok!(SubtensorModule::do_transfer_stake(
+            RuntimeOrigin::signed(coldkey_sender),
+            coldkey_receiver,
+            hotkey,
+            netuid,
+            netuid,
+            transfer_amount,
+        ));
+
+        let expected_sender_lock = SubtensorModule::roll_forward_lock(
+            sender_lock_before,
+            SubtensorModule::get_current_block_as_u64(),
+        );
+
+        assert!(Lock::<Test>::get((coldkey_sender, netuid, hotkey)).is_none());
+
+        let receiver_lock = Lock::<Test>::get((coldkey_receiver, netuid, hotkey))
+            .expect("receiver lock should exist after transfer");
+        assert_eq!(receiver_lock.locked_mass, expected_sender_lock.locked_mass);
+        assert_eq!(
+            receiver_lock.unlocked_mass,
+            expected_sender_lock.unlocked_mass
+        );
+        assert!(receiver_lock.conviction > U64F64::from_num(0));
+        assert!(receiver_lock.conviction <= expected_sender_lock.conviction);
+
+        let hotkey_lock_after =
+            HotkeyLock::<Test>::get(netuid, hotkey).expect("hotkey lock should remain");
+        assert_eq!(
+            hotkey_lock_after.locked_mass,
+            hotkey_lock_before.locked_mass
+        );
+    });
+}
+
 #[test]
 fn test_move_stake_cross_subnet_blocked_by_lock() {
     new_test_ext(1).execute_with(|| {
@@ -716,39 +776,6 @@ fn test_move_stake_cross_subnet_blocked_by_lock() {
     });
 }
 
-#[test]
-fn test_transfer_stake_cross_coldkey_blocked_by_lock() {
-    new_test_ext(1).execute_with(|| {
-        let coldkey_sender = U256::from(1);
-        let coldkey_receiver = U256::from(5);
-        let hotkey = U256::from(2);
-        let netuid = setup_subnet_with_stake(coldkey_sender, hotkey, 100_000_000_000);
-
-        let total = SubtensorModule::total_coldkey_alpha_on_subnet(&coldkey_sender, netuid);
-        assert_ok!(SubtensorModule::do_lock_stake(
-            &coldkey_sender,
-            netuid,
-            &hotkey,
-            total,
-        ));
-
-        step_block(1);
-
-        let alpha = get_alpha(&hotkey, &coldkey_sender, netuid);
-        assert_noop!(
-            SubtensorModule::do_transfer_stake(
-                RuntimeOrigin::signed(coldkey_sender),
-                coldkey_receiver,
-                hotkey,
-                netuid,
-                netuid,
-                alpha,
-            ),
-            Error::<Test>::StakeUnavailable
-        );
-    });
-}
-
 #[test]
 fn test_transfer_stake_cross_coldkey_allowed_partial() {
     new_test_ext(1).execute_with(|| {
@@ -766,6 +793,9 @@ fn test_transfer_stake_cross_coldkey_allowed_partial() {
             lock_half,
         ));
 
+        let sender_lock_before =
+            Lock::<Test>::get((coldkey_sender, netuid, hotkey)).expect("sender lock should exist");
+
         step_block(1);
 
         // Transfer the unlocked portion
@@ -779,6 +809,22 @@ fn test_transfer_stake_cross_coldkey_allowed_partial() {
             netuid,
             transfer_amount,
         ));
+
+        let sender_lock_after =
+            Lock::<Test>::get((coldkey_sender, netuid, hotkey)).expect("sender lock should remain");
+        assert_eq!(
+            sender_lock_after.locked_mass,
+            sender_lock_before.locked_mass
+        );
+        assert_eq!(
+            sender_lock_after.unlocked_mass,
+            SubtensorModule::roll_forward_lock(
+                sender_lock_before,
+                SubtensorModule::get_current_block_as_u64()
+            )
+            .unlocked_mass
+        );
+        assert!(Lock::<Test>::get((coldkey_receiver, netuid, hotkey)).is_none());
     });
 }