From 4f92d340958114ebef2320b7ea62b910354e9661 Mon Sep 17 00:00:00 2001
From: open-junius <zhou@opentensor.dev>
Date: Mon, 18 Nov 2024 23:11:13 +0800
Subject: [PATCH] add netuid argument

---
 runtime/src/precompiles/staking.rs | 35 ++++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/runtime/src/precompiles/staking.rs b/runtime/src/precompiles/staking.rs
index daa95ba93..99d392fb9 100644
--- a/runtime/src/precompiles/staking.rs
+++ b/runtime/src/precompiles/staking.rs
@@ -72,11 +72,17 @@ impl StakingPrecompile {
             <Runtime as pallet_evm::Config>::BalanceConverter::into_substrate_balance(amount)
                 .ok_or(ExitError::OutOfFund)?;
 
+        let netuid =
+            Self::parse_netuid(data.get(56..64).unwrap_or(Err(PrecompileFailure::Error {
+                exit_status: ExitError::InvalidRange,
+            })?))?;
+
+        // let netuid_u16 = netuid.as_u32();
+
         // Create the add_stake call
         let call = RuntimeCall::SubtensorModule(pallet_subtensor::Call::<Runtime>::add_stake {
             hotkey,
-            // TODO update contract to add netuid
-            netuid: 1,
+            netuid,
             amount_staked: amount_sub.unique_saturated_into(),
         });
         // Dispatch the add_stake call
@@ -96,10 +102,14 @@ impl StakingPrecompile {
             <Runtime as pallet_evm::Config>::BalanceConverter::into_substrate_balance(amount)
                 .ok_or(ExitError::OutOfFund)?;
 
+        let netuid =
+            Self::parse_netuid(data.get(64..72).unwrap_or(Err(PrecompileFailure::Error {
+                exit_status: ExitError::InvalidRange,
+            })?))?;
+
         let call = RuntimeCall::SubtensorModule(pallet_subtensor::Call::<Runtime>::remove_stake {
             hotkey,
-            // TODO update contract to add netuid
-            netuid: 1,
+            netuid,
             amount_unstaked: amount_sub.unique_saturated_into(),
         });
         Self::dispatch(handle, call)
@@ -116,6 +126,23 @@ impl StakingPrecompile {
         Ok(hotkey)
     }
 
+    fn parse_netuid(data: &[u8]) -> Result<u16, PrecompileFailure> {
+        let netuid = data
+            .get(0..8)
+            .map(U256::from_big_endian)
+            .ok_or(ExitError::InvalidRange)?;
+
+        let u16_max_u256 = U256::from(u16::MAX);
+        if netuid > u16_max_u256 {
+            // if netuid.as_u128() > u16::MAX as u128 {
+            return Err(PrecompileFailure::Error {
+                exit_status: ExitError::InvalidRange,
+            });
+        }
+
+        Ok(netuid.as_u32() as u16)
+    }
+
     fn dispatch(handle: &mut impl PrecompileHandle, call: RuntimeCall) -> PrecompileResult {
         let account_id =
             <HashedAddressMapping<BlakeTwo256> as AddressMapping<AccountId32>>::into_account_id(