From 6abe023f66a895e3b7b22cb3a32774b48a302aea Mon Sep 17 00:00:00 2001
From: Tanyu <tanyu@tensorplex.ai>
Date: Thu, 20 Mar 2025 14:24:35 +0800
Subject: [PATCH] Add transferStake and swapStake precompile

---
 precompiles/src/solidity/stakingV2.sol | 52 ++++++++++++++++++++++++++
 precompiles/src/staking.rs             | 48 ++++++++++++++++++++++++
 2 files changed, 100 insertions(+)

diff --git a/precompiles/src/solidity/stakingV2.sol b/precompiles/src/solidity/stakingV2.sol
index 67ac0cb129..ec10c4a1a5 100644
--- a/precompiles/src/solidity/stakingV2.sol
+++ b/precompiles/src/solidity/stakingV2.sol
@@ -22,6 +22,30 @@ interface IStaking {
      */
     function addStake(bytes32 hotkey, uint256 amount, uint256 netuid) external payable;
 
+    /**
+     * @dev Swaps a subtensor stake from one subnet to another.
+     *
+     * This function allows external accounts and contracts to swap a subtensor stake between two different subnets.
+     * It effectively calls `swap_stake` on the subtensor pallet with specified hotkey, origin subnet, destination subnet,
+     * and alpha amount as parameters.
+     *
+     * @param hotkey The hotkey public key (32 bytes).
+     * @param originNetuid The subnet to swap from (uint256).
+     * @param destinationNetuid The subnet to swap to (uint256).
+     * @param alphaAmount The amount to swap in rao
+     *
+     * Requirements:
+     * - `hotkey` must be a valid hotkey registered on the network, ensuring that the stake is
+     *   correctly attributed.
+     * - `originNetuid` and `destinationNetuid` must be valid subnets on the network.
+     */
+    function swapStake(
+      bytes32 hotkey,
+      uint256 originNetuid,
+      uint256 destinationNetuid,
+      uint256 alphaAmount
+    ) external;
+
     /**
      * @dev Removes a subtensor stake `amount` from the specified `hotkey`.
      *
@@ -46,6 +70,34 @@ interface IStaking {
         uint256 netuid
     ) external;
 
+    /**
+     * @dev Transfers a subtensor stake from one hotkey to another.
+     *
+     * This function allows external accounts and contracts to transfer a subtensor stake from one hotkey to another.
+     * It effectively calls `transfer_stake` on the subtensor pallet with specified destination coldkey, hotkey, origin subnet,
+     * destination subnet, and alpha amount as parameters.
+     *
+     * @param destinationColdkey The coldkey public key (32 bytes) of the destination account.
+     * @param hotkey The hotkey public key (32 bytes) of the source account.
+     * @param originNetuid The subnet to transfer from (uint256).
+     * @param destinationNetuid The subnet to transfer to (uint256).
+     * @param alphaAmount The amount to transfer in rao.
+     *
+     * Requirements:
+     * - `hotkey` must be a valid hotkey registered on the network, ensuring that the stake is
+     *   correctly attributed.
+     * - `destinationColdkey` must be a valid coldkey registered on the network, ensuring that the stake is
+     *   correctly attributed.
+     * - `originNetuid` and `destinationNetuid` must be valid subnets on the network.
+     */
+    function transferStake(
+        bytes32 destinationColdkey,
+        bytes32 hotkey,
+        uint256 originNetuid,
+        uint256 destinationNetuid,
+        uint256 alphaAmount
+    ) external;
+
     /**
      * @dev Returns the amount of RAO staked by the coldkey.
      *
diff --git a/precompiles/src/staking.rs b/precompiles/src/staking.rs
index 192e55f57c..66b284735c 100644
--- a/precompiles/src/staking.rs
+++ b/precompiles/src/staking.rs
@@ -99,6 +99,28 @@ where
         handle.try_dispatch_runtime_call::<R, _>(call, RawOrigin::Signed(account_id))
     }
 
+    #[precompile::public("swapStake(bytes32,uint256,uint256,uint256)")]
+    fn swap_stake(
+        handle: &mut impl PrecompileHandle,
+        hotkey: H256,
+        origin_netuid: U256,
+        destination_netuid: U256,
+        alpha_amount: U256,
+    ) -> EvmResult<()> {
+        let account_id = handle.caller_account_id::<R>();
+        let hotkey = R::AccountId::from(hotkey.0);
+        let origin_netuid = try_u16_from_u256(origin_netuid)?;
+        let destination_netuid = try_u16_from_u256(destination_netuid)?;
+        let alpha_amount = alpha_amount.unique_saturated_into();
+        let call = pallet_subtensor::Call::<R>::swap_stake {
+            hotkey,
+            origin_netuid,
+            destination_netuid,
+            alpha_amount,
+        };
+        handle.try_dispatch_runtime_call::<R, _>(call, RawOrigin::Signed(account_id))
+    }
+
     #[precompile::public("removeStake(bytes32,uint256,uint256)")]
     fn remove_stake(
         handle: &mut impl PrecompileHandle,
@@ -119,6 +141,32 @@ where
         handle.try_dispatch_runtime_call::<R, _>(call, RawOrigin::Signed(account_id))
     }
 
+    #[precompile::public("transferStake(bytes32,bytes32,uint256,uint256,uint256)")]
+    fn transfer_stake(
+        handle: &mut impl PrecompileHandle,
+        destination_coldkey: H256,
+        hotkey: H256,
+        origin_netuid: U256,
+        destination_netuid: U256,
+        alpha_amount: U256,
+    ) -> EvmResult<()> {
+        let account_id = handle.caller_account_id::<R>();
+        let destination_coldkey = R::AccountId::from(destination_coldkey.0);
+        let hotkey = R::AccountId::from(hotkey.0);
+        let origin_netuid = try_u16_from_u256(origin_netuid)?;
+        let destination_netuid = try_u16_from_u256(destination_netuid)?;
+        let alpha_amount = alpha_amount.unique_saturated_into();
+        let call = pallet_subtensor::Call::<R>::transfer_stake {
+            destination_coldkey,
+            hotkey,
+            origin_netuid,
+            destination_netuid,
+            alpha_amount,
+        };
+
+        handle.try_dispatch_runtime_call::<R, _>(call, RawOrigin::Signed(account_id))
+    }
+
     #[precompile::public("getTotalColdkeyStake(bytes32)")]
     #[precompile::view]
     fn get_total_coldkey_stake(