Feat/ Move Stake Limit

chain-extensions/src/lib.rs

@@ -223,6 +223,56 @@ where
         }
     }
 
+    fn dispatch_move_stake_limit_v1<Env>(
+        env: &mut Env,
+        origin: RawOrigin<T::AccountId>,
+    ) -> Result<RetVal, DispatchError>
+    where
+        Env: SubtensorExtensionEnv<T>,
+        <<T as SysConfig>::Lookup as StaticLookup>::Source: From<<T as SysConfig>::AccountId>,
+    {
+        let (
+            origin_hotkey,
+            destination_hotkey,
+            origin_netuid,
+            destination_netuid,
+            alpha_amount,
+            limit_price,
+            allow_partial,
+        ): (
+            T::AccountId,
+            T::AccountId,
+            NetUid,
+            NetUid,
+            AlphaBalance,
+            TaoBalance,
+            bool,
+        ) = env
+            .read_as()
+            .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?;
+
+        env.charge_weight(
+            <<T as pallet_subtensor::Config>::WeightInfo as SubtensorWeightInfo>::move_stake_limit(
+            ),
+        )?;
+
+        let call_result = pallet_subtensor::Pallet::<T>::move_stake_limit(
+            origin.into(),
+            origin_hotkey,
+            destination_hotkey,
+            origin_netuid,
+            destination_netuid,
+            alpha_amount,
+            limit_price,
+            allow_partial,
+        );
+
+        match call_result {
+            Ok(_) => Ok(RetVal::Converging(Output::Success as u32)),
+            Err(e) => Ok(RetVal::Converging(Output::from(e) as u32)),
+        }
+    }
+
     fn dispatch_transfer_stake_v1<Env>(
         env: &mut Env,
         origin: RawOrigin<T::AccountId>,
@@ -692,6 +742,14 @@ where
                 let origin = convert_origin(env.origin());
                 Self::dispatch_move_stake_v1(env, origin)
             }
+            FunctionId::MoveStakeLimitV1 => {
+                let origin = RawOrigin::Signed(env.caller());
+                Self::dispatch_move_stake_limit_v1(env, origin)
+            }
+            FunctionId::CallerMoveStakeLimitV1 => {
+                let origin = convert_origin(env.origin());
+                Self::dispatch_move_stake_limit_v1(env, origin)
+            }
             FunctionId::TransferStakeV1 => {
                 let origin = RawOrigin::Signed(env.caller());
                 Self::dispatch_transfer_stake_v1(env, origin)

chain-extensions/src/mock.rs

@@ -160,6 +160,7 @@ impl frame_support::traits::InstanceFilter<RuntimeCall> for subtensor_runtime_co
                     | RuntimeCall::SubtensorModule(pallet_subtensor::Call::swap_stake { .. })
                     | RuntimeCall::SubtensorModule(pallet_subtensor::Call::swap_stake_limit { .. })
                     | RuntimeCall::SubtensorModule(pallet_subtensor::Call::move_stake { .. })
+                    | RuntimeCall::SubtensorModule(pallet_subtensor::Call::move_stake_limit { .. })
                     | RuntimeCall::SubtensorModule(pallet_subtensor::Call::transfer_stake { .. })
             ),
             _ => false,

chain-extensions/src/tests.rs

@@ -569,6 +569,79 @@ fn move_stake_success_moves_alpha_between_hotkeys() {
     });
 }
 
+#[test]
+fn move_stake_limit_success_moves_alpha_between_hotkeys() {
+    mock::new_test_ext(1).execute_with(|| {
+        let owner_hotkey = U256::from(6201);
+        let owner_coldkey = U256::from(6202);
+        let coldkey = U256::from(7201);
+        let origin_hotkey = U256::from(7202);
+        let destination_hotkey = U256::from(7203);
+        let min_stake = DefaultMinStake::<mock::Test>::get();
+        let stake_amount_raw = min_stake.to_u64().saturating_mul(240);
+        let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
+
+        mock::setup_reserves(
+            netuid,
+            stake_amount_raw.saturating_mul(15).into(),
+            AlphaBalance::from(stake_amount_raw.saturating_mul(25)),
+        );
+        mock::register_ok_neuron(netuid, origin_hotkey, coldkey, 0);
+        mock::register_ok_neuron(netuid, destination_hotkey, coldkey, 1);
+        add_balance_to_coldkey_account(&coldkey, (stake_amount_raw + 1_000_000_000).into());
+        assert_ok!(pallet_subtensor::Pallet::<mock::Test>::add_stake(
+            RawOrigin::Signed(coldkey).into(),
+            origin_hotkey,
+            netuid,
+            stake_amount_raw.into(),
+        ));
+
+        let origin_before =
+            pallet_subtensor::Pallet::<mock::Test>::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &origin_hotkey,
+                &coldkey,
+                netuid,
+            );
+        let amount: AlphaBalance = (origin_before.to_u64() / 2).into();
+        let expected_weight = <<mock::Test as pallet_subtensor::Config>::WeightInfo as SubtensorWeightInfo>::move_stake_limit();
+        let mut env = MockEnv::new(
+            FunctionId::MoveStakeLimitV1,
+            coldkey,
+            (
+                origin_hotkey,
+                destination_hotkey,
+                netuid,
+                netuid,
+                amount,
+                TaoBalance::MAX,
+                false,
+            )
+                .encode(),
+        )
+        .with_expected_weight(expected_weight);
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        assert_success(ret);
+        assert_eq!(env.charged_weight(), Some(expected_weight));
+        assert_eq!(
+            pallet_subtensor::Pallet::<mock::Test>::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &origin_hotkey,
+                &coldkey,
+                netuid,
+            ),
+            origin_before - amount
+        );
+        assert_eq!(
+            pallet_subtensor::Pallet::<mock::Test>::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &destination_hotkey,
+                &coldkey,
+                netuid,
+            ),
+            amount
+        );
+    });
+}
+
 #[test]
 fn unstake_all_alpha_success_moves_stake_to_root() {
     mock::new_test_ext(1).execute_with(|| {

chain-extensions/src/types.rs

@@ -44,6 +44,8 @@ pub enum FunctionId {
     GetSubnetRegistrationStateV1 = 34,
     GetColdkeyLockV1 = 35,
     GetStakeAvailabilityV1 = 36,
+    MoveStakeLimitV1 = 37,
+    CallerMoveStakeLimitV1 = 38,
 }
 
 #[freeze_struct("5dc33d60abed5c08")]
@@ -188,11 +190,13 @@ mod function_id_tests {
         assert_eq!(FunctionId::GetSubnetRegistrationStateV1 as u16, 34);
         assert_eq!(FunctionId::GetColdkeyLockV1 as u16, 35);
         assert_eq!(FunctionId::GetStakeAvailabilityV1 as u16, 36);
+        assert_eq!(FunctionId::MoveStakeLimitV1 as u16, 37);
+        assert_eq!(FunctionId::CallerMoveStakeLimitV1 as u16, 38);
     }
 
     #[test]
     fn caller_ids_roundtrip_try_from_primitive() {
-        for id in 16u16..=36u16 {
+        for id in 16u16..=38u16 {
             let v = FunctionId::try_from_primitive(id)
                 .unwrap_or_else(|_| panic!("try_from_primitive failed for {id}"));
             assert_eq!(v as u16, id);

eco-tests/src/tests_mentat_indexer.rs

@@ -339,6 +339,28 @@ fn indexer_extrinsic_move_stake() {
     });
 }
 
+#[test]
+fn indexer_extrinsic_move_stake_limit() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(1);
+        let origin_hotkey = U256::from(2);
+        let destination_hotkey = U256::from(3);
+        let origin_netuid = NetUid::from(1u16);
+        let destination_netuid = NetUid::from(2u16);
+
+        let _ = SubtensorModule::move_stake_limit(
+            RuntimeOrigin::signed(coldkey),
+            origin_hotkey,
+            destination_hotkey,
+            origin_netuid,
+            destination_netuid,
+            AlphaBalance::from(1_000_000_000u64),
+            TaoBalance::from(1_000_000_000u64),
+            false,
+        );
+    });
+}
+
 #[test]
 fn indexer_extrinsic_set_children() {
     new_test_ext(1).execute_with(|| {

ink-contract/lib.rs

@@ -43,6 +43,8 @@ pub enum FunctionId {
     GetSubnetRegistrationStateV1 = 34,
     GetColdkeyLockV1 = 35,
     GetStakeAvailabilityV1 = 36,
+    MoveStakeLimitV1 = 37,
+    CallerMoveStakeLimitV1 = 38,
 }
 
 #[ink::chain_extension(extension = 0x1000)]
@@ -287,6 +289,28 @@ pub trait RuntimeReadWrite {
         coldkey: <CustomEnvironment as ink::env::Environment>::AccountId,
         netuid: u16,
     ) -> StakeAvailability;
+
+    #[ink(function = 37)]
+    fn move_stake_limit(
+        origin_hotkey: <CustomEnvironment as ink::env::Environment>::AccountId,
+        destination_hotkey: <CustomEnvironment as ink::env::Environment>::AccountId,
+        origin_netuid: u16,
+        destination_netuid: u16,
+        amount: u64,
+        limit_price: u64,
+        allow_partial: bool,
+    );
+
+    #[ink(function = 38)]
+    fn caller_move_stake_limit(
+        origin_hotkey: <CustomEnvironment as ink::env::Environment>::AccountId,
+        destination_hotkey: <CustomEnvironment as ink::env::Environment>::AccountId,
+        origin_netuid: u16,
+        destination_netuid: u16,
+        amount: u64,
+        limit_price: u64,
+        allow_partial: bool,
+    );
 }
 
 #[ink::scale_derive(Encode, Decode, TypeInfo)]
@@ -543,6 +567,31 @@ mod bittensor {
                 .map_err(|_e| ReadWriteErrorCode::WriteFailed)
         }
 
+        #[ink(message)]
+        pub fn move_stake_limit(
+            &self,
+            origin_hotkey: [u8; 32],
+            destination_hotkey: [u8; 32],
+            origin_netuid: u16,
+            destination_netuid: u16,
+            amount: u64,
+            limit_price: u64,
+            allow_partial: bool,
+        ) -> Result<(), ReadWriteErrorCode> {
+            self.env()
+                .extension()
+                .move_stake_limit(
+                    origin_hotkey.into(),
+                    destination_hotkey.into(),
+                    origin_netuid,
+                    destination_netuid,
+                    amount,
+                    limit_price,
+                    allow_partial,
+                )
+                .map_err(|_e| ReadWriteErrorCode::WriteFailed)
+        }
+
         #[ink(message)]
         pub fn remove_stake_full_limit(
             &self,
@@ -801,6 +850,31 @@ mod bittensor {
                 .map_err(|_e| ReadWriteErrorCode::WriteFailed)
         }
 
+        #[ink(message)]
+        pub fn caller_move_stake_limit(
+            &self,
+            origin_hotkey: [u8; 32],
+            destination_hotkey: [u8; 32],
+            origin_netuid: u16,
+            destination_netuid: u16,
+            amount: u64,
+            limit_price: u64,
+            allow_partial: bool,
+        ) -> Result<(), ReadWriteErrorCode> {
+            self.env()
+                .extension()
+                .caller_move_stake_limit(
+                    origin_hotkey.into(),
+                    destination_hotkey.into(),
+                    origin_netuid,
+                    destination_netuid,
+                    amount,
+                    limit_price,
+                    allow_partial,
+                )
+                .map_err(|_e| ReadWriteErrorCode::WriteFailed)
+        }
+
         #[ink(message)]
         pub fn caller_remove_stake_full_limit(
             &self,

pallets/subtensor/src/benchmarks/benchmarks.rs

@@ -1089,6 +1089,66 @@ mod pallet_benchmarks {
         );
     }
 
+    #[benchmark]
+    fn move_stake_limit() {
+        let coldkey: T::AccountId = whitelisted_caller::<AccountIdOf<T>>();
+        let origin_hotkey: T::AccountId = account("A", 0, 1);
+        let destination_hotkey: T::AccountId = account("B", 0, 2);
+        let origin_netuid = NetUid::from(1);
+        let destination_netuid = NetUid::from(2);
+
+        for netuid in [origin_netuid, destination_netuid] {
+            SubtokenEnabled::<T>::insert(netuid, true);
+            Subtensor::<T>::init_new_network(netuid, 1);
+            Subtensor::<T>::set_network_registration_allowed(netuid, true);
+            add_lock::<T>(&coldkey, netuid);
+        }
+
+        let tao_reserve = TaoBalance::from(150_000_000_000_u64);
+        let alpha_in = AlphaBalance::from(100_000_000_000_u64);
+        set_reserves::<T>(origin_netuid, tao_reserve, alpha_in);
+        SubnetTAO::<T>::insert(destination_netuid, tao_reserve);
+        Subtensor::<T>::increase_total_stake(1_000_000_000_000_u64.into());
+
+        let balance = TaoBalance::from(900_000_000_000_u64);
+        let stake_limit = TaoBalance::from(6_000_000_000_u64);
+        let move_limit = TaoBalance::from(1_000_000_000_u64);
+        let amount_to_stake = TaoBalance::from(440_000_000_000_u64);
+        let amount_to_move = AlphaBalance::from(30_000_000_000_u64);
+        add_balance_to_coldkey_account::<T>(&coldkey, balance);
+
+        assert_ok!(Subtensor::<T>::burned_register(
+            RawOrigin::Signed(coldkey.clone()).into(),
+            origin_netuid,
+            origin_hotkey.clone()
+        ));
+        assert_ok!(Subtensor::<T>::burned_register(
+            RawOrigin::Signed(coldkey.clone()).into(),
+            destination_netuid,
+            destination_hotkey.clone()
+        ));
+        assert_ok!(Subtensor::<T>::add_stake_limit(
+            RawOrigin::Signed(coldkey.clone()).into(),
+            origin_hotkey.clone(),
+            origin_netuid,
+            amount_to_stake,
+            stake_limit,
+            true,
+        ));
+
+        #[extrinsic_call]
+        _(
+            RawOrigin::Signed(coldkey),
+            origin_hotkey,
+            destination_hotkey,
+            origin_netuid,
+            destination_netuid,
+            amount_to_move,
+            move_limit,
+            true,
+        );
+    }
+
     #[benchmark]
     fn transfer_stake() {
         let coldkey: T::AccountId = whitelisted_caller();