Account flags to control locked alpha transfers

pallets/subtensor/src/lib.rs

@@ -69,6 +69,9 @@ pub const MAX_SUBNET_CLAIMS: usize = 5;
 
 pub const MAX_ROOT_CLAIM_THRESHOLD: u64 = 10_000_000;
 
+/// Account flag bit that rejects incoming locked alpha transfers.
+pub const ACCOUNT_FLAGS_REJECT_LOCKED_ALPHA: u128 = 1u128 << 0;
+
 #[allow(deprecated)]
 #[deny(missing_docs)]
 #[import_section(errors::errors)]
@@ -1186,6 +1189,11 @@ pub mod pallet {
     pub type Owner<T: Config> =
         StorageMap<_, Blake2_128Concat, T::AccountId, T::AccountId, ValueQuery, DefaultAccount<T>>;
 
+    /// MAP ( coldkey ) --> flags | Account-level flags. Defaults to zero.
+    #[pallet::storage]
+    pub type AccountFlags<T: Config> =
+        StorageMap<_, Blake2_128Concat, T::AccountId, u128, ValueQuery>;
+
     /// MAP ( hot ) --> take | Returns the hotkey delegation take. And signals that this key is open for delegation
     #[pallet::storage]
     pub type Delegates<T: Config> =

pallets/subtensor/src/macros/dispatches.rs

@@ -2593,5 +2593,30 @@ mod dispatches {
             let coldkey = ensure_signed(origin)?;
             Self::do_set_perpetual_lock(&coldkey, netuid, enabled)
         }
+
+        /// Sets or clears the caller's "reject locked alpha" account flag.
+        ///
+        /// When enabled, this coldkey cannot receive locked alpha from stake
+        /// transfers or coldkey swaps.
+        #[pallet::call_index(139)]
+        #[pallet::weight((
+            <T as frame_system::Config>::DbWeight::get().reads_writes(1, 1),
+            DispatchClass::Normal,
+            Pays::Yes
+        ))]
+        pub fn set_reject_locked_alpha(origin: OriginFor<T>, enabled: bool) -> DispatchResult {
+            let coldkey = ensure_signed(origin)?;
+            AccountFlags::<T>::mutate_exists(&coldkey, |maybe_flags| {
+                let mut flags = maybe_flags.unwrap_or_default();
+                if enabled {
+                    flags |= crate::ACCOUNT_FLAGS_REJECT_LOCKED_ALPHA;
+                } else {
+                    flags &= !crate::ACCOUNT_FLAGS_REJECT_LOCKED_ALPHA;
+                }
+                *maybe_flags = if flags == 0 { None } else { Some(flags) };
+            });
+            Self::deposit_event(Event::RejectLockedAlphaUpdated { coldkey, enabled });
+            Ok(())
+        }
     }
 }

pallets/subtensor/src/macros/errors.rs

@@ -305,5 +305,7 @@ mod errors {
         CannotUseSystemAccount,
         /// Trying to unlock more than locked
         UnlockAmountTooHigh,
+        /// The destination coldkey rejects incoming locked alpha.
+        AccountRejectsLockedAlpha,
     }
 }

pallets/subtensor/src/macros/events.rs

@@ -631,5 +631,13 @@ mod events {
             /// Whether this coldkey's locks are now perpetual.
             enabled: bool,
         },
+
+        /// A coldkey's reject locked alpha account flag was updated.
+        RejectLockedAlphaUpdated {
+            /// The coldkey whose flag changed.
+            coldkey: T::AccountId,
+            /// Whether this coldkey rejects incoming locked alpha.
+            enabled: bool,
+        },
     }
 }

pallets/subtensor/src/staking/lock.rs

@@ -444,6 +444,29 @@ impl ConvictionModel {
 }
 
 impl<T: Config> Pallet<T> {
+    pub fn account_rejects_locked_alpha(coldkey: &T::AccountId) -> bool {
+        AccountFlags::<T>::get(coldkey) & crate::ACCOUNT_FLAGS_REJECT_LOCKED_ALPHA != 0
+    }
+
+    pub fn ensure_can_receive_locked_alpha(
+        coldkey: &T::AccountId,
+        amount: AlphaBalance,
+    ) -> DispatchResult {
+        let rejects_locked_alpha = Self::account_rejects_locked_alpha(coldkey);
+        Self::ensure_can_receive_locked_alpha_with_flag(rejects_locked_alpha, amount)
+    }
+
+    fn ensure_can_receive_locked_alpha_with_flag(
+        rejects_locked_alpha: bool,
+        amount: AlphaBalance,
+    ) -> DispatchResult {
+        if amount.is_zero() {
+            return Ok(());
+        }
+        ensure!(!rejects_locked_alpha, Error::<T>::AccountRejectsLockedAlpha);
+        Ok(())
+    }
+
     pub fn insert_lock_state(
         coldkey: &T::AccountId,
         netuid: NetUid,
@@ -1331,17 +1354,13 @@ impl<T: Config> Pallet<T> {
         Self::ensure_no_active_locks(new_coldkey)?;
 
         let mut locks_to_transfer: Vec<(NetUid, T::AccountId, LockState)> = Vec::new();
+        let now = Self::get_current_block_as_u64();
+        let unlock_rate = UnlockRate::<T>::get();
+        let maturity_rate = MaturityRate::<T>::get();
+        let new_coldkey_rejects_locked_alpha = Self::account_rejects_locked_alpha(new_coldkey);
 
         // Gather locks for old coldkey
         for ((netuid, hotkey), lock) in Lock::<T>::iter_prefix((old_coldkey,)) {
-            locks_to_transfer.push((netuid, hotkey, lock));
-        }
-
-        // Remove locks for old coldkey and insert for new
-        for (netuid, hotkey, lock) in locks_to_transfer {
-            let now = Self::get_current_block_as_u64();
-            let unlock_rate = UnlockRate::<T>::get();
-            let maturity_rate = MaturityRate::<T>::get();
             let old_lock = ConvictionModel::roll_forward_lock(
                 lock,
                 now,
@@ -1350,6 +1369,15 @@ impl<T: Config> Pallet<T> {
                 Self::is_subnet_owner_hotkey(netuid, &hotkey),
                 Self::is_perpetual_lock(old_coldkey, netuid),
             );
+            Self::ensure_can_receive_locked_alpha_with_flag(
+                new_coldkey_rejects_locked_alpha,
+                old_lock.locked_mass,
+            )?;
+            locks_to_transfer.push((netuid, hotkey, old_lock));
+        }
+
+        // Remove locks for old coldkey and insert for new
+        for (netuid, hotkey, old_lock) in locks_to_transfer {
             let new_lock = ConvictionModel::roll_forward_lock(
                 old_lock.clone(),
                 now,
@@ -1780,6 +1808,7 @@ impl<T: Config> Pallet<T> {
                 .conviction
                 .saturating_add(conviction_transfer);
         }
+        Self::ensure_can_receive_locked_alpha(destination_coldkey, locked_transfer)?;
 
         source_lock = ConvictionModel::roll_forward_lock(
             source_lock,

pallets/subtensor/src/tests/locks.rs

@@ -6,6 +6,7 @@
 )]
 
 use approx::assert_abs_diff_eq;
+use frame_support::dispatch::{GetDispatchInfo, Pays};
 use frame_support::weights::Weight;
 use frame_support::{assert_noop, assert_ok};
 use safe_math::FixedExt;
@@ -96,6 +97,40 @@ fn roll_forward_individual_lock(
     )
 }
 
+#[test]
+fn test_account_flags_default_to_zero_and_reject_locked_alpha_setter_pays_fee() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(1);
+
+        assert_eq!(AccountFlags::<Test>::get(coldkey), 0);
+        assert!(!AccountFlags::<Test>::contains_key(coldkey));
+        assert!(!SubtensorModule::account_rejects_locked_alpha(&coldkey));
+
+        let call =
+            RuntimeCall::SubtensorModule(crate::Call::set_reject_locked_alpha { enabled: true });
+        assert_eq!(call.get_dispatch_info().pays_fee, Pays::Yes);
+
+        assert_ok!(SubtensorModule::set_reject_locked_alpha(
+            RuntimeOrigin::signed(coldkey),
+            true,
+        ));
+        assert_eq!(
+            AccountFlags::<Test>::get(coldkey),
+            ACCOUNT_FLAGS_REJECT_LOCKED_ALPHA
+        );
+        assert!(AccountFlags::<Test>::contains_key(coldkey));
+        assert!(SubtensorModule::account_rejects_locked_alpha(&coldkey));
+
+        assert_ok!(SubtensorModule::set_reject_locked_alpha(
+            RuntimeOrigin::signed(coldkey),
+            false,
+        ));
+        assert_eq!(AccountFlags::<Test>::get(coldkey), 0);
+        assert!(!AccountFlags::<Test>::contains_key(coldkey));
+        assert!(!SubtensorModule::account_rejects_locked_alpha(&coldkey));
+    });
+}
+
 fn roll_forward_hotkey_lock(lock: LockState, now: u64) -> LockState {
     roll_forward_lock(lock, now, false, true)
 }
@@ -1694,6 +1729,101 @@ fn test_move_stake_cross_subnet_blocked_by_lock() {
     });
 }
 
+#[test]
+fn test_do_transfer_stake_rejects_locked_alpha_to_flagged_destination() {
+    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,
+        ));
+        assert_ok!(SubtensorModule::set_reject_locked_alpha(
+            RuntimeOrigin::signed(coldkey_receiver),
+            true,
+        ));
+
+        let sender_lock_before =
+            Lock::<Test>::get((coldkey_sender, netuid, hotkey)).expect("sender lock should exist");
+        let sender_alpha_before =
+            SubtensorModule::total_coldkey_alpha_on_subnet(&coldkey_sender, netuid);
+        let receiver_alpha_before =
+            SubtensorModule::total_coldkey_alpha_on_subnet(&coldkey_receiver, netuid);
+
+        assert_noop!(
+            SubtensorModule::do_transfer_stake(
+                RuntimeOrigin::signed(coldkey_sender),
+                coldkey_receiver,
+                hotkey,
+                netuid,
+                netuid,
+                total,
+            ),
+            Error::<Test>::AccountRejectsLockedAlpha
+        );
+
+        assert_eq!(
+            Lock::<Test>::get((coldkey_sender, netuid, hotkey)),
+            Some(sender_lock_before)
+        );
+        assert!(Lock::<Test>::get((coldkey_receiver, netuid, hotkey)).is_none());
+        assert_eq!(
+            SubtensorModule::total_coldkey_alpha_on_subnet(&coldkey_sender, netuid),
+            sender_alpha_before
+        );
+        assert_eq!(
+            SubtensorModule::total_coldkey_alpha_on_subnet(&coldkey_receiver, netuid),
+            receiver_alpha_before
+        );
+    });
+}
+
+#[test]
+fn test_do_transfer_stake_allows_unlocked_alpha_to_flagged_destination() {
+    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,
+        ));
+        assert_ok!(SubtensorModule::set_reject_locked_alpha(
+            RuntimeOrigin::signed(coldkey_receiver),
+            true,
+        ));
+
+        let unlocked_transfer = lock_half / 2.into();
+        assert_ok!(SubtensorModule::do_transfer_stake(
+            RuntimeOrigin::signed(coldkey_sender),
+            coldkey_receiver,
+            hotkey,
+            netuid,
+            netuid,
+            unlocked_transfer,
+        ));
+
+        assert!(Lock::<Test>::get((coldkey_receiver, netuid, hotkey)).is_none());
+        assert_eq!(
+            SubtensorModule::total_coldkey_alpha_on_subnet(&coldkey_receiver, netuid),
+            unlocked_transfer
+        );
+    });
+}
+
 #[test]
 fn test_transfer_stake_cross_coldkey_allowed_partial() {
     new_test_ext(1).execute_with(|| {
@@ -2888,6 +3018,49 @@ fn test_coldkey_swap_rejects_destination_lock() {
     });
 }
 
+#[test]
+fn test_coldkey_swap_rejects_locked_alpha_to_flagged_destination() {
+    new_test_ext(1).execute_with(|| {
+        let old_coldkey = U256::from(1);
+        let new_coldkey = U256::from(10);
+        let old_hotkey = U256::from(2);
+        let netuid = subtensor_runtime_common::NetUid::from(1);
+
+        let old_locked = AlphaBalance::from(7_000u64);
+        let old_conviction = U64F64::from_num(77);
+
+        SubtensorModule::insert_lock_state(
+            &old_coldkey,
+            netuid,
+            &old_hotkey,
+            LockState {
+                locked_mass: old_locked,
+                conviction: old_conviction,
+                last_update: SubtensorModule::get_current_block_as_u64(),
+            },
+        );
+        assert_ok!(SubtensorModule::set_reject_locked_alpha(
+            RuntimeOrigin::signed(new_coldkey),
+            true,
+        ));
+
+        assert_noop!(
+            SubtensorModule::swap_coldkey_locks(&old_coldkey, &new_coldkey),
+            Error::<Test>::AccountRejectsLockedAlpha
+        );
+
+        let source_lock = Lock::<Test>::get((old_coldkey, netuid, old_hotkey))
+            .expect("source lock should remain after failed transfer");
+        assert_eq!(source_lock.locked_mass, old_locked);
+        assert_eq!(source_lock.conviction, old_conviction);
+        assert!(
+            Lock::<Test>::iter_prefix((new_coldkey, netuid))
+                .next()
+                .is_none()
+        );
+    });
+}
+
 #[test]
 // The public coldkey swap extrinsic runs inside a storage layer, so a late failure rolls back the earlier writes.
 fn test_failed_coldkey_swap_extrinsic_rolls_back_state_changes() {