Refactor staking and neuron precompile with precompile-utils

ales-otf · ales-otf · commit 126f78049d88 · 2025-02-12T14:38:07.000+01:00
diff --git a/runtime/src/precompiles/balance_transfer.rs b/runtime/src/precompiles/balance_transfer.rs
@@ -48,8 +48,6 @@ impl BalanceTransferPrecompile {
             value: amount_sub.unique_saturated_into(),
         };
 
-        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)?;
-
-        Ok(())
+        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
     }
 }
diff --git a/runtime/src/precompiles/mod.rs b/runtime/src/precompiles/mod.rs
@@ -12,7 +12,8 @@ use pallet_evm::{
 use pallet_evm_precompile_modexp::Modexp;
 use pallet_evm_precompile_sha3fips::Sha3FIPS256;
 use pallet_evm_precompile_simple::{ECRecover, ECRecoverPublicKey, Identity, Ripemd160, Sha256};
-use sp_core::{hashing::keccak_256, H160};
+use precompile_utils::EvmResult;
+use sp_core::{hashing::keccak_256, H160, U256};
 use sp_runtime::{traits::Dispatchable, AccountId32};
 use sp_std::vec;
 
@@ -151,6 +152,12 @@ pub fn get_pubkey(data: &[u8]) -> Result<(AccountId32, vec::Vec<u8>), Precompile
     ))
 }
 
+fn try_u16_from_u256(value: U256) -> Result<u16, PrecompileFailure> {
+    u16::try_from(value.as_u32()).map_err(|_| PrecompileFailure::Error {
+        exit_status: ExitError::Other("the value is outside of u16 bounds".into()),
+    })
+}
+
 fn parse_netuid(data: &[u8], offset: usize) -> Result<u16, PrecompileFailure> {
     if data.len() < offset + 2 {
         return Err(PrecompileFailure::Error {
@@ -175,7 +182,7 @@ fn try_dispatch_runtime_call(
     handle: &mut impl PrecompileHandle,
     call: impl Into<RuntimeCall>,
     origin: RawOrigin<AccountId32>,
-) -> PrecompileResult {
+) -> EvmResult<()> {
     let call = Into::<RuntimeCall>::into(call);
     let info = call.get_dispatch_info();
 
@@ -220,10 +227,7 @@ fn try_dispatch_runtime_call(
 
             log::info!("Dispatch succeeded. Post info: {:?}", post_info);
 
-            Ok(PrecompileOutput {
-                exit_status: ExitSucceed::Returned,
-                output: Default::default(),
-            })
+            Ok(())
         }
         Err(e) => {
             log::error!("Dispatch failed. Error: {:?}", e);
diff --git a/runtime/src/precompiles/neuron.rs b/runtime/src/precompiles/neuron.rs
@@ -1,62 +1,39 @@
-use pallet_evm::{ExitError, PrecompileFailure, PrecompileHandle, PrecompileResult};
-
-use crate::precompiles::{
-    contract_to_origin, get_method_id, get_pubkey, get_slice, parse_netuid,
-    try_dispatch_runtime_call,
-};
+use frame_system::RawOrigin;
+use pallet_evm::{AddressMapping, HashedAddressMapping, PrecompileHandle};
+use precompile_utils::EvmResult;
+use sp_core::H256;
+use sp_runtime::traits::BlakeTwo256;
 use sp_runtime::AccountId32;
-use sp_std::vec;
 
+use crate::precompiles::{get_pubkey, try_dispatch_runtime_call};
 use crate::{Runtime, RuntimeCall};
+
 pub const NEURON_PRECOMPILE_INDEX: u64 = 2052;
 
-// ss58 public key i.e., the contract sends funds it received to the destination address from the
-// method parameter.
-const CONTRACT_ADDRESS_SS58: [u8; 32] = [
-    0xbc, 0x46, 0x35, 0x79, 0xbc, 0x99, 0xf9, 0xee, 0x7c, 0x59, 0xed, 0xee, 0x20, 0x61, 0xa3, 0x09,
-    0xd2, 0x1e, 0x68, 0xd5, 0x39, 0xb6, 0x40, 0xec, 0x66, 0x46, 0x90, 0x30, 0xab, 0x74, 0xc1, 0xdb,
-];
 pub struct NeuronPrecompile;
 
+#[precompile_utils::precompile]
 impl NeuronPrecompile {
-    pub fn execute(handle: &mut impl PrecompileHandle) -> PrecompileResult {
-        let txdata = handle.input();
-        let method_id = get_slice(txdata, 0, 4)?;
-        let method_input = txdata
-            .get(4..)
-            .map_or_else(vec::Vec::new, |slice| slice.to_vec()); // Avoiding borrowing conflicts
-
-        match method_id {
-            id if id == get_method_id("burnedRegister(uint16,bytes32)") => {
-                Self::burned_register(handle, &method_input)
-            }
-
-            _ => Err(PrecompileFailure::Error {
-                exit_status: ExitError::InvalidRange,
-            }),
-        }
-    }
-
-    pub fn burned_register(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
-        let (netuid, hotkey) = Self::parse_netuid_hotkey_parameter(data)?;
+    #[precompile::public("burnedRegister(uint16,bytes32)")]
+    #[precompile::payable]
+    fn burned_register(
+        handle: &mut impl PrecompileHandle,
+        netuid: u16,
+        hotkey: H256,
+    ) -> EvmResult<()> {
+        let coldkey =
+            <HashedAddressMapping<BlakeTwo256> as AddressMapping<AccountId32>>::into_account_id(
+                handle.context().caller,
+            );
+        let (hotkey, _) = get_pubkey(hotkey.as_bytes())?;
         let call =
             RuntimeCall::SubtensorModule(pallet_subtensor::Call::<Runtime>::burned_register {
                 netuid,
                 hotkey,
             });
-        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
-    }
-
-    fn parse_netuid_hotkey_parameter(data: &[u8]) -> Result<(u16, AccountId32), PrecompileFailure> {
-        if data.len() < 64 {
-            return Err(PrecompileFailure::Error {
-                exit_status: ExitError::InvalidRange,
-            });
-        }
-        let netuid = parse_netuid(data, 30)?;
 
-        let (hotkey, _) = get_pubkey(get_slice(data, 32, 64)?)?;
+        try_dispatch_runtime_call(handle, call, RawOrigin::Signed(coldkey))?;
 
-        Ok((netuid, hotkey))
+        Ok(())
     }
 }
diff --git a/runtime/src/precompiles/staking.rs b/runtime/src/precompiles/staking.rs
@@ -25,20 +25,23 @@
 //   - Precompile checks the result of do_remove_stake and, in case of a failure, reverts the transaction.
 //
 
-use crate::precompiles::{
-    get_method_id, get_pubkey, get_slice, parse_netuid, try_dispatch_runtime_call,
-};
-use crate::{ProxyType, Runtime, RuntimeCall};
 use frame_system::RawOrigin;
 use pallet_evm::{
     AddressMapping, BalanceConverter, ExitError, ExitSucceed, HashedAddressMapping,
     PrecompileFailure, PrecompileHandle, PrecompileOutput, PrecompileResult,
 };
-use sp_core::U256;
+use precompile_utils::EvmResult;
+use sp_core::{H256, U256};
 use sp_runtime::traits::{BlakeTwo256, Dispatchable, StaticLookup, UniqueSaturatedInto};
 use sp_runtime::AccountId32;
 use sp_std::vec;
 
+use crate::precompiles::{
+    get_method_id, get_pubkey, get_slice, parse_netuid, try_dispatch_runtime_call,
+    try_u16_from_u256,
+};
+use crate::{ProxyType, Runtime, RuntimeCall};
+
 pub const STAKING_PRECOMPILE_INDEX: u64 = 2049;
 
 // ss58 public key i.e., the contract sends funds it received to the destination address from the
@@ -50,40 +53,17 @@ const CONTRACT_ADDRESS_SS58: [u8; 32] = [
 
 pub struct StakingPrecompile;
 
+#[precompile_utils::precompile]
 impl StakingPrecompile {
-    pub fn execute(handle: &mut impl PrecompileHandle) -> PrecompileResult {
-        let txdata = handle.input();
-        let method_id = get_slice(txdata, 0, 4)?;
-        let method_input = txdata
-            .get(4..)
-            .map_or_else(vec::Vec::new, |slice| slice.to_vec()); // Avoiding borrowing conflicts
-
-        if method_id == get_method_id("addStake(bytes32,uint256)") {
-            Self::add_stake(handle, &method_input)
-        } else if method_id == get_method_id("removeStake(bytes32,uint256,uint256)") {
-            Self::remove_stake(handle, &method_input)
-        } else if method_id == get_method_id("getStake(bytes32,bytes32,uint256)") {
-            Self::get_stake(&method_input)
-        } else if method_id == get_method_id("addProxy(bytes32)") {
-            Self::add_proxy(handle, &method_input)
-        } else if method_id == get_method_id("removeProxy(bytes32)") {
-            Self::remove_proxy(handle, &method_input)
-        } else {
-            Err(PrecompileFailure::Error {
-                exit_status: ExitError::InvalidRange,
-            })
-        }
-    }
-
-    fn add_stake(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
+    #[precompile::public("addStake(bytes32,uint256)")]
+    #[precompile::payable]
+    fn add_stake(handle: &mut impl PrecompileHandle, address: H256, netuid: U256) -> EvmResult<()> {
         let account_id =
             <HashedAddressMapping<BlakeTwo256> as AddressMapping<AccountId32>>::into_account_id(
                 handle.context().caller,
             );
 
-        let (hotkey, _) = get_pubkey(data)?;
-        let amount: U256 = handle.context().apparent_value;
-        let netuid = parse_netuid(data, 0x3E)?;
+        let amount = handle.context().apparent_value;
 
         if !amount.is_zero() {
             Self::transfer_back_to_caller(&account_id, amount)?;
@@ -93,6 +73,8 @@ impl StakingPrecompile {
             <Runtime as pallet_evm::Config>::BalanceConverter::into_substrate_balance(amount)
                 .ok_or(ExitError::OutOfFund)?;
 
+        let (hotkey, _) = get_pubkey(address.as_bytes())?;
+        let netuid = try_u16_from_u256(netuid)?;
         let call = RuntimeCall::SubtensorModule(pallet_subtensor::Call::<Runtime>::add_stake {
             hotkey,
             netuid,
@@ -102,39 +84,40 @@ impl StakingPrecompile {
         try_dispatch_runtime_call(handle, call, RawOrigin::Signed(account_id))
     }
 
-    fn remove_stake(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
+    #[precompile::public("removeStake(bytes32,uint256,uint256)")]
+    fn remove_stake(
+        handle: &mut impl PrecompileHandle,
+        address: H256,
+        amount: U256,
+        netuid: U256,
+    ) -> EvmResult<()> {
         let account_id =
             <HashedAddressMapping<BlakeTwo256> as AddressMapping<AccountId32>>::into_account_id(
                 handle.context().caller,
             );
-        let (hotkey, _) = get_pubkey(data)?;
-        let netuid = parse_netuid(data, 0x5E)?;
-
-        // We have to treat this as uint256 (because of Solidity ABI encoding rules, it pads uint64),
-        // but this will never exceed 8 bytes, se we will ignore higher bytes and will only use lower
-        // 8 bytes.
-        let amount = data
-            .get(56..64)
-            .map(U256::from_big_endian)
-            .ok_or(ExitError::OutOfFund)?;
+
         let amount_sub =
             <Runtime as pallet_evm::Config>::BalanceConverter::into_substrate_balance(amount)
                 .ok_or(ExitError::OutOfFund)?;
 
+        let (hotkey, _) = get_pubkey(address.as_bytes())?;
+        let netuid = try_u16_from_u256(netuid)?;
         let call = RuntimeCall::SubtensorModule(pallet_subtensor::Call::<Runtime>::remove_stake {
             hotkey,
             netuid,
             amount_unstaked: amount_sub.unique_saturated_into(),
         });
+
         try_dispatch_runtime_call(handle, call, RawOrigin::Signed(account_id))
     }
 
-    fn add_proxy(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
+    #[precompile::public("addProxy(bytes32)")]
+    fn add_proxy(handle: &mut impl PrecompileHandle, delegate: H256) -> EvmResult<()> {
         let account_id =
             <HashedAddressMapping<BlakeTwo256> as AddressMapping<AccountId32>>::into_account_id(
                 handle.context().caller,
             );
-        let (delegate, _) = get_pubkey(data)?;
+        let (delegate, _) = get_pubkey(delegate.as_bytes())?;
         let delegate = <Runtime as frame_system::Config>::Lookup::unlookup(delegate);
         let call = RuntimeCall::Proxy(pallet_proxy::Call::<Runtime>::add_proxy {
             delegate,
@@ -145,12 +128,13 @@ impl StakingPrecompile {
         try_dispatch_runtime_call(handle, call, RawOrigin::Signed(account_id))
     }
 
-    fn remove_proxy(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
+    #[precompile::public("removeProxy(bytes32)")]
+    fn remove_proxy(handle: &mut impl PrecompileHandle, delegate: H256) -> EvmResult<()> {
         let account_id =
             <HashedAddressMapping<BlakeTwo256> as AddressMapping<AccountId32>>::into_account_id(
                 handle.context().caller,
             );
-        let (delegate, _) = get_pubkey(data)?;
+        let (delegate, _) = get_pubkey(delegate.as_bytes())?;
         let delegate = <Runtime as frame_system::Config>::Lookup::unlookup(delegate);
         let call = RuntimeCall::Proxy(pallet_proxy::Call::<Runtime>::remove_proxy {
             delegate,
@@ -161,29 +145,25 @@ impl StakingPrecompile {
         try_dispatch_runtime_call(handle, call, RawOrigin::Signed(account_id))
     }
 
-    fn get_stake(data: &[u8]) -> PrecompileResult {
-        let (hotkey, left_data) = get_pubkey(data)?;
-        let (coldkey, _) = get_pubkey(&left_data)?;
-        let netuid = parse_netuid(data, 0x5E)?;
+    #[precompile::public("getStake(bytes32,bytes32,uint256)")]
+    #[precompile::view]
+    fn get_stake(
+        _: &mut impl PrecompileHandle,
+        hotkey: H256,
+        coldkey: H256,
+        netuid: U256,
+    ) -> EvmResult<U256> {
+        let (hotkey, _) = get_pubkey(hotkey.as_bytes())?;
+        let (coldkey, _) = get_pubkey(coldkey.as_bytes())?;
+        let netuid = try_u16_from_u256(netuid)?;
 
         let stake = pallet_subtensor::Pallet::<Runtime>::get_stake_for_hotkey_and_coldkey_on_subnet(
             &hotkey, &coldkey, netuid,
         );
 
         // Convert to EVM decimals
-        let stake_u256 = U256::from(stake);
-        let stake_eth =
-            <Runtime as pallet_evm::Config>::BalanceConverter::into_evm_balance(stake_u256)
-                .ok_or(ExitError::InvalidRange)?;
-
-        // Format output
-        let mut result = [0_u8; 32];
-        U256::to_big_endian(&stake_eth, &mut result);
-
-        Ok(PrecompileOutput {
-            exit_status: ExitSucceed::Returned,
-            output: result.into(),
-        })
+        <Runtime as pallet_evm::Config>::BalanceConverter::into_evm_balance(stake.into())
+            .ok_or(ExitError::InvalidRange.into())
     }
 
     fn transfer_back_to_caller(
diff --git a/runtime/src/precompiles/subnet.rs b/runtime/src/precompiles/subnet.rs
@@ -1,13 +1,15 @@
-use crate::precompiles::{get_method_id, get_pubkey, get_slice, try_dispatch_runtime_call};
-use crate::{Runtime, RuntimeCall};
 use frame_system::RawOrigin;
 use pallet_evm::{
-    AddressMapping, ExitError, HashedAddressMapping, PrecompileFailure, PrecompileHandle,
-    PrecompileResult,
+    AddressMapping, ExitError, ExitSucceed, HashedAddressMapping, PrecompileFailure,
+    PrecompileHandle, PrecompileOutput, PrecompileResult,
 };
 use sp_runtime::traits::BlakeTwo256;
 use sp_runtime::AccountId32;
 use sp_std::vec;
+
+use crate::precompiles::{get_method_id, get_pubkey, get_slice, try_dispatch_runtime_call};
+use crate::{Runtime, RuntimeCall};
+
 pub const SUBNET_PRECOMPILE_INDEX: u64 = 2051;
 // bytes with max lenght 1K
 pub const MAX_SINGLE_PARAMETER_SIZE: usize = 1024;
@@ -21,6 +23,7 @@ const CONTRACT_ADDRESS_SS58: [u8; 32] = [
     0x3a, 0x86, 0x18, 0xfb, 0xbb, 0x1b, 0xbc, 0x47, 0x86, 0x64, 0xff, 0x53, 0x46, 0x18, 0x0c, 0x35,
     0xd0, 0x9f, 0xac, 0x26, 0xf2, 0x02, 0x70, 0x85, 0xb3, 0x1c, 0x56, 0xc1, 0x06, 0x3c, 0x1c, 0xd3,
 ];
+
 pub struct SubnetPrecompile;
 
 impl SubnetPrecompile {
@@ -93,7 +96,11 @@ impl SubnetPrecompile {
             );
 
         // Dispatch the register_network call
-        try_dispatch_runtime_call(handle, call, RawOrigin::Signed(account_id))
+        try_dispatch_runtime_call(handle, call, RawOrigin::Signed(account_id))?;
+        Ok(PrecompileOutput {
+            exit_status: ExitSucceed::Returned,
+            output: Default::default(),
+        })
     }
 
     fn parse_register_network_parameters(

Original file line number	Diff line number	Diff line change
`@@ -48,8 +48,6 @@ impl BalanceTransferPrecompile {`
`48`	`48`	`value: amount_sub.unique_saturated_into(),`
`49`	`49`	`};`
`50`	`50`
`51`		`- try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)?;`
`52`		`-`
`53`		`- Ok(())`
	`51`	`+ try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)`
`54`	`52`	`}`
`55`	`53`	`}`