bittensor_rs/queries/
account.rs

1//! # Account Queries
2//!
3//! Query account balances and stake information.
4
5use crate::api::api;
6use crate::error::BittensorError;
7use crate::types::Balance;
8use crate::AccountId;
9use std::str::FromStr;
10use subxt::OnlineClient;
11use subxt::PolkadotConfig;
12
13/// Get the free balance of an account
14///
15/// # Arguments
16///
17/// * `client` - The subxt client
18/// * `address` - The SS58 address to query
19///
20/// # Returns
21///
22/// The free balance in RAO
23///
24/// # Example
25///
26/// ```rust,no_run
27/// use bittensor_rs::queries::get_balance;
28///
29/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
30/// # let client: subxt::OnlineClient<subxt::PolkadotConfig> = todo!();
31/// let balance = get_balance(&client, "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY").await?;
32/// println!("Balance: {}", balance);
33/// # Ok(())
34/// # }
35/// ```
36pub async fn get_balance(
37    client: &OnlineClient<PolkadotConfig>,
38    address: &str,
39) -> Result<Balance, BittensorError> {
40    let account_id = AccountId::from_str(address).map_err(|_| BittensorError::InvalidHotkey {
41        hotkey: address.to_string(),
42    })?;
43
44    let storage = api::storage().system().account(account_id);
45
46    let account_info = client
47        .storage()
48        .at_latest()
49        .await
50        .map_err(|e| BittensorError::RpcError {
51            message: format!("Failed to get storage: {}", e),
52        })?
53        .fetch(&storage)
54        .await
55        .map_err(|e| BittensorError::StorageQueryError {
56            key: "account".to_string(),
57            message: format!("Failed to fetch account: {}", e),
58        })?;
59
60    match account_info {
61        Some(info) => Ok(Balance::from_rao(info.data.free)),
62        None => Ok(Balance::zero()),
63    }
64}
65
66/// Get the stake of a hotkey on a specific coldkey for a subnet
67///
68/// # Arguments
69///
70/// * `client` - The subxt client
71/// * `hotkey` - The hotkey SS58 address
72/// * `coldkey` - The coldkey SS58 address
73/// * `netuid` - The subnet netuid
74pub async fn get_stake(
75    client: &OnlineClient<PolkadotConfig>,
76    hotkey: &str,
77    coldkey: &str,
78    netuid: u16,
79) -> Result<Balance, BittensorError> {
80    let hotkey_id = AccountId::from_str(hotkey).map_err(|_| BittensorError::InvalidHotkey {
81        hotkey: hotkey.to_string(),
82    })?;
83
84    let coldkey_id = AccountId::from_str(coldkey).map_err(|_| BittensorError::InvalidHotkey {
85        hotkey: coldkey.to_string(),
86    })?;
87
88    let storage = api::storage()
89        .subtensor_module()
90        .alpha(hotkey_id, coldkey_id, netuid);
91
92    let stake = client
93        .storage()
94        .at_latest()
95        .await
96        .map_err(|e| BittensorError::RpcError {
97            message: format!("Failed to get storage: {}", e),
98        })?
99        .fetch(&storage)
100        .await
101        .map_err(|e| BittensorError::StorageQueryError {
102            key: "stake".to_string(),
103            message: format!("Failed to fetch stake: {}", e),
104        })?;
105
106    match stake {
107        // alpha returns FixedU128 with a `bits` field
108        // The fixed-point representation stores the value shifted by 64 bits
109        Some(amount) => {
110            // Get the raw bits and convert to u64 representing RAO
111            let raw_bits: u128 = amount.bits;
112            // Shift right by 64 to get the integer part
113            let rao = (raw_bits >> 64) as u64;
114            Ok(Balance::from_rao(rao))
115        }
116        None => Ok(Balance::zero()),
117    }
118}
119
120/// Get the global total stake on the network
121///
122/// This returns the total stake across all coldkeys and hotkeys.
123///
124/// # Arguments
125///
126/// * `client` - The subxt client
127pub async fn get_total_network_stake(
128    client: &OnlineClient<PolkadotConfig>,
129) -> Result<Balance, BittensorError> {
130    let storage = api::storage().subtensor_module().total_stake();
131
132    let stake = client
133        .storage()
134        .at_latest()
135        .await
136        .map_err(|e| BittensorError::RpcError {
137            message: format!("Failed to get storage: {}", e),
138        })?
139        .fetch(&storage)
140        .await
141        .map_err(|e| BittensorError::StorageQueryError {
142            key: "total_stake".to_string(),
143            message: format!("Failed to fetch total stake: {}", e),
144        })?;
145
146    match stake {
147        Some(amount) => Ok(Balance::from_rao(amount)),
148        None => Ok(Balance::zero()),
149    }
150}
151
152/// Stake information returned from runtime API
153#[derive(Debug, Clone)]
154pub struct StakeInfo {
155    /// The hotkey address
156    pub hotkey: AccountId,
157    /// The coldkey address
158    pub coldkey: AccountId,
159    /// The subnet netuid
160    pub netuid: u16,
161    /// The stake amount
162    pub stake: Balance,
163    /// Locked stake amount
164    pub locked: Balance,
165    /// Emission earned
166    pub emission: Balance,
167    /// TAO emission earned
168    pub tao_emission: Balance,
169    /// Drain amount
170    pub drain: Balance,
171    /// Whether this is registered
172    pub is_registered: bool,
173}
174
175/// Get stake information for a coldkey using runtime API
176///
177/// # Arguments
178///
179/// * `client` - The subxt client
180/// * `coldkey` - The coldkey SS58 address
181pub async fn get_stake_info_for_coldkey(
182    client: &OnlineClient<PolkadotConfig>,
183    coldkey: &str,
184) -> Result<Vec<StakeInfo>, BittensorError> {
185    let coldkey_id = AccountId::from_str(coldkey).map_err(|_| BittensorError::InvalidHotkey {
186        hotkey: coldkey.to_string(),
187    })?;
188
189    let runtime_api =
190        client
191            .runtime_api()
192            .at_latest()
193            .await
194            .map_err(|e| BittensorError::RpcError {
195                message: format!("Failed to get runtime API: {}", e),
196            })?;
197
198    let stake_infos = runtime_api
199        .call(
200            api::runtime_apis::stake_info_runtime_api::StakeInfoRuntimeApi
201                .get_stake_info_for_coldkey(coldkey_id.clone()),
202        )
203        .await
204        .map_err(|e| BittensorError::RpcMethodError {
205            method: "get_stake_info_for_coldkey".to_string(),
206            message: e.to_string(),
207        })?;
208
209    let result = stake_infos
210        .into_iter()
211        .map(|info| StakeInfo {
212            hotkey: info.hotkey,
213            coldkey: info.coldkey,
214            netuid: info.netuid,
215            stake: Balance::from_rao(info.stake),
216            locked: Balance::from_rao(info.locked),
217            emission: Balance::from_rao(info.emission),
218            tao_emission: Balance::from_rao(info.tao_emission),
219            drain: Balance::from_rao(info.drain),
220            is_registered: info.is_registered,
221        })
222        .collect();
223
224    Ok(result)
225}
226
227#[cfg(test)]
228mod tests {
229    #[test]
230    fn test_balance_type() {
231        use crate::types::Balance;
232        let balance = Balance::from_rao(1_000_000_000);
233        assert_eq!(balance.as_tao(), 1.0);
234    }
235
236    #[test]
237    fn test_stake_info_struct() {
238        use super::*;
239        use subxt::utils::AccountId32;
240
241        let _info = StakeInfo {
242            hotkey: AccountId32::from([0u8; 32]),
243            coldkey: AccountId32::from([1u8; 32]),
244            netuid: 1,
245            stake: Balance::from_rao(1000),
246            locked: Balance::from_rao(0),
247            emission: Balance::from_rao(100),
248            tao_emission: Balance::from_rao(50),
249            drain: Balance::from_rao(0),
250            is_registered: true,
251        };
252    }
253}