From 9688e037061a66b0bfdb507081c71f76e912fd6a Mon Sep 17 00:00:00 2001
From: open-junius <zhou@opentensor.dev>
Date: Tue, 12 Nov 2024 21:59:14 +0800
Subject: [PATCH] parse keys for stake map

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

diff --git a/runtime/src/precompiles/staking.rs b/runtime/src/precompiles/staking.rs
index 104bb2116..0190db413 100644
--- a/runtime/src/precompiles/staking.rs
+++ b/runtime/src/precompiles/staking.rs
@@ -60,7 +60,7 @@ impl StakingPrecompile {
                 Self::remove_stake(handle, &method_input)
             }
             id if id == get_method_id("getStake(bytes32,bytes32)") => {
-                Self::get_stake(handle, &method_input)
+                Self::get_stake(&method_input)
             }
             _ => Err(PrecompileFailure::Error {
                 exit_status: ExitError::InvalidRange,
@@ -104,18 +104,37 @@ impl StakingPrecompile {
         Self::dispatch(handle, call)
     }
 
-    fn get_stake(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
-        let hotkey: AccountId32 = Self::parse_hotkey(data)?.into();
+    fn get_stake(data: &[u8]) -> PrecompileResult {
+        let (hotkey, coldkey) = Self::parse_hotkey_coldkey(data)?.into();
 
-        let data =
-            pallet_subtensor::Pallet::<Runtime>::get_stake_for_coldkey_and_hotkey(&hotkey, &hotkey);
+        let stake = pallet_subtensor::Pallet::<Runtime>::get_stake_for_coldkey_and_hotkey(
+            &hotkey.into(),
+            &coldkey.into(),
+        );
+
+        let stake_u256 = U256::from(stake);
+        let mut result = [0_u8; 32];
+        U256::to_big_endian(&stake_u256, &mut result);
 
         Ok(PrecompileOutput {
             exit_status: ExitSucceed::Returned,
-            output: data.to_le_bytes().to_vec(),
+            output: result.into(),
         })
     }
 
+    fn parse_hotkey_coldkey(data: &[u8]) -> Result<([u8; 32], [u8; 32]), PrecompileFailure> {
+        if data.len() < 64 {
+            return Err(PrecompileFailure::Error {
+                exit_status: ExitError::InvalidRange,
+            });
+        }
+        let mut hotkey = [0u8; 32];
+        hotkey.copy_from_slice(get_slice(data, 0, 32)?);
+        let mut coldkey = [0u8; 32];
+        coldkey.copy_from_slice(get_slice(data, 32, 64)?);
+        Ok((hotkey, coldkey))
+    }
+
     fn parse_hotkey(data: &[u8]) -> Result<[u8; 32], PrecompileFailure> {
         if data.len() < 32 {
             return Err(PrecompileFailure::Error {