Add runtime api for stake

pallets/subtensor/runtime-api/src/lib.rs

@@ -22,4 +22,9 @@ sp_api::decl_runtime_apis! {
         fn get_subnet_info(netuid: u16) -> Vec<u8>;
         fn get_subnets_info() -> Vec<u8>;
     }
+
+	pub trait StakeInfoRuntimeApi {
+		fn get_stake_info_for_coldkey( coldkey_account_vec: Vec<u8> ) -> Vec<u8>;
+		fn get_stake_info_for_coldkeys( coldkey_account_vecs: Vec<Vec<u8>> ) -> Vec<u8>;
+	}
 }

pallets/subtensor/src/lib.rs

@@ -48,6 +48,7 @@ mod weights;
 pub mod delegate_info;
 pub mod neuron_info;
 pub mod subnet_info;
+pub mod stake_info;
 
 // apparently this is stabilized since rust 1.36
 extern crate alloc;

pallets/subtensor/src/stake_info.rs

@@ -0,0 +1,75 @@
+use super::*;
+use frame_support::pallet_prelude::{Decode, Encode};
+extern crate alloc;
+use alloc::vec::Vec;
+use codec::Compact;
+use sp_core::hexdisplay::AsBytesRef;
+
+#[derive(Decode, Encode, PartialEq, Eq, Clone, Debug)]
+pub struct StakeInfo<T: Config> {
+    hotkey: T::AccountId,
+    coldkey: T::AccountId,
+    stake: Compact<u64>,
+}
+
+impl<T: Config> Pallet<T> {
+	fn _get_stake_info_for_coldkeys(coldkeys: Vec<T::AccountId>) -> Vec<(T::AccountId, Vec<StakeInfo<T>>)> {
+		if coldkeys.len() == 0 {
+			return Vec::new(); // No coldkeys to check
+		}
+
+		let mut stake_info: Vec<(T::AccountId, Vec<StakeInfo<T>>)> = Vec::new();
+		for coldkey_ in coldkeys {
+			let mut stake_info_for_coldkey: Vec<StakeInfo<T>> = Vec::new();
+
+			for (hotkey, coldkey, stake) in <Stake<T>>::iter() {
+				if coldkey == coldkey_ {
+					stake_info_for_coldkey.push(StakeInfo {
+						hotkey,
+						coldkey,
+						stake: stake.into(),
+					});
+				}
+			}
+
+			stake_info.push((coldkey_, stake_info_for_coldkey));
+		}
+
+		return stake_info;
+	}
+
+	pub fn get_stake_info_for_coldkeys(coldkey_account_vecs: Vec<Vec<u8>>) -> Vec<(T::AccountId, Vec<StakeInfo<T>>)> {
+		let mut coldkeys: Vec<T::AccountId> = Vec::new();
+		for coldkey_account_vec in coldkey_account_vecs {
+			if coldkey_account_vec.len() != 32 {
+				continue; // Invalid coldkey
+			}
+			let coldkey: AccountIdOf<T> = T::AccountId::decode( &mut coldkey_account_vec.as_bytes_ref() ).unwrap();
+			coldkeys.push(coldkey);
+		}
+
+		if coldkeys.len() == 0 {
+			return Vec::new(); // Invalid coldkey
+		}
+
+		let stake_info = Self::_get_stake_info_for_coldkeys(coldkeys);
+
+		return stake_info;
+	}
+
+	pub fn get_stake_info_for_coldkey(coldkey_account_vec: Vec<u8>) -> Vec<StakeInfo<T>> {
+		if coldkey_account_vec.len() != 32 {
+            return Vec::new(); // Invalid coldkey
+        }
+
+		let coldkey: AccountIdOf<T> = T::AccountId::decode( &mut coldkey_account_vec.as_bytes_ref() ).unwrap();
+		let stake_info = Self::_get_stake_info_for_coldkeys(vec![coldkey]);
+
+        if stake_info.len() == 0 {
+			return Vec::new(); // Invalid coldkey
+		} else {
+			return stake_info.get(0).unwrap().1.clone();
+		}
+	}
+}
+

runtime/src/lib.rs

@@ -997,6 +997,18 @@ impl_runtime_apis! {
             result.encode()
         }
     }
+
+	impl subtensor_custom_rpc_runtime_api::StakeInfoRuntimeApi<Block> for Runtime {
+        fn get_stake_info_for_coldkey( coldkey_account_vec: Vec<u8> ) -> Vec<u8> {
+            let result = SubtensorModule::get_stake_info_for_coldkey( coldkey_account_vec );
+			result.encode()
+        }
+
+        fn get_stake_info_for_coldkeys( coldkey_account_vecs: Vec<Vec<u8>> ) -> Vec<u8> {
+            let result = SubtensorModule::get_stake_info_for_coldkeys( coldkey_account_vecs );
+            result.encode()
+        }
+    }
 }
 
 #[cfg(test)]