code/pallets/subtensor/src/staking/helpers.rs
use alloc::collections::BTreeMap;
use safe_math::*;
use share_pool::SafeFloat;
use sp_runtime::PerU16;
use substrate_fixed::types::{U64F64, U96F32};
use subtensor_runtime_common::{NetUid, TaoBalance};
use subtensor_swap_interface::{Order, SwapHandler};
use super::*;
impl<T: Config> Pallet<T> {
// Returns true if the passed hotkey allow delegative staking.
//
pub fn hotkey_is_delegate(hotkey: &T::AccountId) -> bool {
Delegates::<T>::contains_key(hotkey)
}
// Sets the hotkey as a delegate with take.
//
pub fn delegate_hotkey(hotkey: &T::AccountId, take: u16) {
Delegates::<T>::insert(hotkey, PerU16::from_parts(take));
}
// Returns the total amount of stake in the staking table.
//
pub fn get_total_stake() -> TaoBalance {
TotalStake::<T>::get()
}
// Increases the total amount of stake by the passed amount.
//
pub fn increase_total_stake(increment: TaoBalance) {
TotalStake::<T>::put(Self::get_total_stake().saturating_add(increment));
}
// Decreases the total amount of stake by the passed amount.
//
pub fn decrease_total_stake(decrement: TaoBalance) {
TotalStake::<T>::put(Self::get_total_stake().saturating_sub(decrement));
}
/// Returns the total amount of stake (in TAO) under a hotkey (delegative or otherwise)
pub fn get_total_stake_for_hotkey(hotkey: &T::AccountId) -> TaoBalance {
Self::get_all_subnet_netuids()
.into_iter()
.map(|netuid| {
let alpha = U64F64::saturating_from_num(Self::get_stake_for_hotkey_on_subnet(
hotkey, netuid,
));
let alpha_price = T::SwapInterface::current_alpha_price(netuid.into());
alpha.saturating_mul(alpha_price)
})
.sum::<U64F64>()
.saturating_to_num::<u64>()
.into()
}
// Returns the total amount of stake under a coldkey
//
pub fn get_total_stake_for_coldkey(coldkey: &T::AccountId) -> TaoBalance {
let hotkeys = StakingHotkeys::<T>::get(coldkey);
hotkeys
.iter()
.map(|hotkey| {
Self::alpha_iter_prefix((hotkey, coldkey))
.map(|(netuid, _)| {
let alpha_stake = Self::get_stake_for_hotkey_and_coldkey_on_subnet(
hotkey, coldkey, netuid,
);
let order = GetTaoForAlpha::<T>::with_amount(alpha_stake);
T::SwapInterface::sim_swap(netuid.into(), order)
.map(|r| {
let fee: u64 = U64F64::saturating_from_num(r.fee_paid)
.saturating_mul(T::SwapInterface::current_alpha_price(
netuid.into(),
))
.saturating_to_num();
r.amount_paid_out.to_u64().saturating_add(fee)
})
.unwrap_or_default()
})
.sum::<u64>()
})
.sum::<u64>()
.into()
}
// Returns the total amount of stake under a coldkey on a subnet
//
pub fn get_total_stake_for_coldkey_on_subnet(
coldkey: &T::AccountId,
netuid: NetUid,
) -> TaoBalance {
let hotkeys = StakingHotkeys::<T>::get(coldkey);
hotkeys
.iter()
.map(|hotkey| {
Self::alpha_iter_prefix((hotkey, coldkey))
.map(|(netuid_on_storage, _)| {
if netuid_on_storage == netuid {
let alpha_stake = Self::get_stake_for_hotkey_and_coldkey_on_subnet(
hotkey, coldkey, netuid,
);
let order = GetTaoForAlpha::<T>::with_amount(alpha_stake);
T::SwapInterface::sim_swap(netuid.into(), order)
.map(|r| {
let fee: u64 = U64F64::saturating_from_num(r.fee_paid)
.saturating_mul(T::SwapInterface::current_alpha_price(
netuid.into(),
))
.saturating_to_num();
r.amount_paid_out.to_u64().saturating_add(fee)
})
.unwrap_or_default()
} else {
0
}
})
.sum::<u64>()
})
.sum::<u64>()
.into()
}
// Creates a cold - hot pairing account if the hotkey is not already an active account.
//
pub fn create_account_if_non_existent(
coldkey: &T::AccountId,
hotkey: &T::AccountId,
) -> DispatchResult {
// Only allow to register non-system hotkeys
ensure!(
Self::is_subnet_account_id(hotkey).is_none(),
Error::<T>::CannotUseSystemAccount
);
if !Self::hotkey_account_exists(hotkey) {
Owner::<T>::insert(hotkey, coldkey);
// Update OwnedHotkeys map
let mut hotkeys = OwnedHotkeys::<T>::get(coldkey);
if !hotkeys.contains(hotkey) {
hotkeys.push(hotkey.clone());
OwnedHotkeys::<T>::insert(coldkey, hotkeys);
}
// Update StakingHotkeys map
let mut staking_hotkeys = StakingHotkeys::<T>::get(coldkey);
if !staking_hotkeys.contains(hotkey) {
staking_hotkeys.push(hotkey.clone());
StakingHotkeys::<T>::insert(coldkey, staking_hotkeys);
}
}
Ok(())
}
pub fn set_hotkey_owner(coldkey: &T::AccountId, hotkey: &T::AccountId) -> DispatchResult {
// Only allow to register non-system hotkeys
ensure!(
Self::is_subnet_account_id(hotkey).is_none(),
Error::<T>::CannotUseSystemAccount
);
Owner::<T>::insert(hotkey, coldkey);
Ok(())
}
//// If the hotkey is not a delegate, make it a delegate.
pub fn maybe_become_delegate(hotkey: &T::AccountId) {
if !Self::hotkey_is_delegate(hotkey) {
Self::delegate_hotkey(hotkey, Self::get_hotkey_take(hotkey));
}
}
/// Returns the coldkey owning this hotkey. This function should only be called for active accounts.
///
/// # Arguments
/// * `hotkey`: The hotkey account ID.
///
/// # Returns
/// The coldkey account ID that owns the hotkey.
pub fn get_owning_coldkey_for_hotkey(hotkey: &T::AccountId) -> T::AccountId {
Owner::<T>::get(hotkey)
}
/// Returns the hotkey take.
///
/// # Arguments
/// * `hotkey`: The hotkey account ID.
///
/// # Returns
/// The take value of the hotkey.
pub fn get_hotkey_take(hotkey: &T::AccountId) -> u16 {
Delegates::<T>::get(hotkey).deconstruct()
}
pub fn get_hotkey_take_float(hotkey: &T::AccountId) -> U96F32 {
U96F32::saturating_from_num(Self::get_hotkey_take(hotkey))
.safe_div(U96F32::saturating_from_num(u16::MAX))
}
/// Returns true if the hotkey account has been created.
///
/// # Arguments
/// * `hotkey`: The hotkey account ID.
///
/// # Returns
/// True if the hotkey account exists, false otherwise.
pub fn hotkey_account_exists(hotkey: &T::AccountId) -> bool {
Owner::<T>::contains_key(hotkey)
}
/// Returns true if the passed coldkey owns the hotkey.
///
/// # Arguments
/// * `coldkey`: The coldkey account ID.
/// * `hotkey`: The hotkey account ID.
///
/// # Returns
/// True if the coldkey owns the hotkey, false otherwise.
pub fn coldkey_owns_hotkey(coldkey: &T::AccountId, hotkey: &T::AccountId) -> bool {
if Self::hotkey_account_exists(hotkey) {
Owner::<T>::get(hotkey) == *coldkey
} else {
false
}
}
/// Clears the nomination for an account, if it is a nominator account and the stake is below the minimum required threshold.
pub fn clear_small_nomination_if_required(
hotkey: &T::AccountId,
coldkey: &T::AccountId,
netuid: NetUid,
) {
// Verify if the account is a nominator account by checking ownership of the hotkey by the coldkey.
if !Self::coldkey_owns_hotkey(coldkey, hotkey) {
// If the stake is non-zero and below the minimum required, it's considered a small nomination and needs to be cleared.
// Log if the stake is below the minimum required
let alpha_stake =
Self::get_stake_for_hotkey_and_coldkey_on_subnet(hotkey, coldkey, netuid);
let min_alpha_stake =
U64F64::saturating_from_num(Self::get_nominator_min_required_stake())
.safe_div(T::SwapInterface::current_alpha_price(netuid))
.saturating_to_num::<u64>();
if alpha_stake > 0.into() && alpha_stake < min_alpha_stake.into() {
// Log the clearing of a small nomination
// Remove the stake from the nominator account. (this is a more forceful unstake operation which )
// Actually deletes the staking account.
// Do not apply any fees
if Self::unstake_from_subnet(
hotkey,
coldkey,
coldkey,
netuid,
alpha_stake,
T::SwapInterface::min_price(),
false,
)
.is_err()
{
// Ignore errors if unstaking fails
// Just clear small alpha
let alpha =
Self::get_stake_for_hotkey_and_coldkey_on_subnet(hotkey, coldkey, netuid);
Self::decrease_stake_for_hotkey_and_coldkey_on_subnet(
hotkey, coldkey, netuid, alpha,
);
}
// Reduce lock (if exists) by the cleaned stake amount
Self::force_reduce_lock(coldkey, netuid, alpha_stake);
}
}
}
/// Clears small nominations for all accounts.
///
/// WARN: This is an O(N) operation, where N is the number of staking accounts. It should be
/// used with caution.
pub fn clear_small_nominations() {
// Loop through all staking accounts to identify and clear nominations below the minimum stake.
for ((hotkey, coldkey, netuid), _) in Self::alpha_iter() {
Self::clear_small_nomination_if_required(&hotkey, &coldkey, netuid);
}
}
/// The function clears Alpha map in batches. Each run will check ALPHA_MAP_BATCH_SIZE
/// alphas. It keeps the alpha value stored when it's >= than MIN_ALPHA.
/// The function uses AlphaMapLastKey as a storage for key iterator between runs.
pub fn populate_root_coldkey_staking_maps() {
// Get starting key for the batch. Get the first key if we restart the process.
let mut new_starting_raw_key = AlphaMapLastKey::<T>::get();
let mut starting_key = None;
if new_starting_raw_key.is_none() {
starting_key = Alpha::<T>::iter_keys().next();
new_starting_raw_key = starting_key.as_ref().map(Alpha::<T>::hashed_key_for);
}
if let Some(starting_raw_key) = new_starting_raw_key {
// Get the key batch
let mut keys = Alpha::<T>::iter_keys_from(starting_raw_key)
.take(ALPHA_MAP_BATCH_SIZE)
.collect::<Vec<_>>();
// New iteration: insert the starting key in the batch if it's a new iteration
// iter_keys_from() skips the starting key
if let Some(starting_key) = starting_key {
if keys.len() == ALPHA_MAP_BATCH_SIZE {
keys.remove(keys.len().saturating_sub(1));
}
keys.insert(0, starting_key);
}
let mut new_starting_key = None;
let new_iteration = keys.len() < ALPHA_MAP_BATCH_SIZE;
// Check and remove alphas if necessary
for key in keys {
let (_, coldkey, netuid) = key.clone();
if netuid == NetUid::ROOT {
Self::maybe_add_coldkey_index(&coldkey);
}
new_starting_key = Some(Alpha::<T>::hashed_key_for(key));
}
// Restart the process if it's the last batch
if new_iteration {
new_starting_key = None;
}
AlphaMapLastKey::<T>::put(new_starting_key);
}
}
// Same thing as populate_root_coldkey_staking_maps, but for AlphaV2
// TODO: Remove this function and AlphaV2MapLastKey when slow migration is finished
pub fn populate_root_coldkey_staking_maps_v2() {
// Get starting key for the batch. Get the first key if we restart the process.
let mut new_starting_raw_key = AlphaV2MapLastKey::<T>::get();
let mut starting_key = None;
if new_starting_raw_key.is_none() {
starting_key = AlphaV2::<T>::iter_keys().next();
new_starting_raw_key = starting_key.as_ref().map(AlphaV2::<T>::hashed_key_for);
}
if let Some(starting_raw_key) = new_starting_raw_key {
// Get the key batch
let mut keys = AlphaV2::<T>::iter_keys_from(starting_raw_key)
.take(ALPHA_MAP_BATCH_SIZE)
.collect::<Vec<_>>();
// New iteration: insert the starting key in the batch if it's a new iteration
// iter_keys_from() skips the starting key
if let Some(starting_key) = starting_key {
if keys.len() == ALPHA_MAP_BATCH_SIZE {
keys.remove(keys.len().saturating_sub(1));
}
keys.insert(0, starting_key);
}
let mut new_starting_key = None;
let new_iteration = keys.len() < ALPHA_MAP_BATCH_SIZE;
// Check and remove alphas if necessary
for key in keys {
let (_, coldkey, netuid) = key.clone();
if netuid == NetUid::ROOT {
Self::maybe_add_coldkey_index(&coldkey);
}
new_starting_key = Some(AlphaV2::<T>::hashed_key_for(key));
}
// Restart the process if it's the last batch
if new_iteration {
new_starting_key = None;
}
AlphaV2MapLastKey::<T>::put(new_starting_key);
}
}
/// Several alpha iteration helpers that merge key space from Alpha and AlphaV2 maps
pub fn alpha_iter() -> impl Iterator<Item = ((T::AccountId, T::AccountId, NetUid), SafeFloat)> {
// Old Alpha shares format: U64F64 -> SafeFloat
let legacy = Alpha::<T>::iter().map(|(key, val_u64f64)| {
let sf: SafeFloat = val_u64f64.into();
(key, sf)
});
// New Alpha shares format
let v2 = AlphaV2::<T>::iter();
// Merge and prefer v2 on duplicates
let merged: BTreeMap<_, SafeFloat> =
legacy
.chain(v2)
.fold(BTreeMap::new(), |mut acc, (key, val)| {
acc.entry(key)
.and_modify(|existing| {
*existing = val.clone();
})
.or_insert(val);
acc
});
merged.into_iter()
}
pub fn alpha_iter_prefix(
prefix: (&T::AccountId, &T::AccountId),
) -> impl Iterator<Item = (NetUid, SafeFloat)>
where
T::AccountId: Clone,
{
// Old Alpha shares format: U64F64 -> SafeFloat
let legacy = Alpha::<T>::iter_prefix(prefix).map(|(netuid, val_u64f64)| {
let sf: SafeFloat = val_u64f64.into();
(netuid, sf)
});
// New Alpha shares format
let v2 = AlphaV2::<T>::iter_prefix(prefix);
// Merge by netuid and sum SafeFloat values
let merged: BTreeMap<NetUid, SafeFloat> =
legacy
.chain(v2)
.fold(BTreeMap::new(), |mut acc, (netuid, sf)| {
acc.entry(netuid)
.and_modify(|existing| {
*existing = sf.clone();
})
.or_insert(sf);
acc
});
merged
.into_iter()
.filter(|(_, alpha_share)| !alpha_share.is_zero())
}
pub fn alpha_iter_single_prefix(
prefix: &T::AccountId,
) -> impl Iterator<Item = (T::AccountId, NetUid, SafeFloat)>
where
T::AccountId: Clone,
{
// Old Alpha shares format: U64F64 -> SafeFloat
let legacy =
Alpha::<T>::iter_prefix((prefix.clone(),)).map(|((coldkey, netuid), val_u64f64)| {
let sf: SafeFloat = val_u64f64.into();
((coldkey, netuid), sf)
});
// New Alpha shares format
let v2 = AlphaV2::<T>::iter_prefix((prefix,));
let merged: BTreeMap<(T::AccountId, NetUid), SafeFloat> =
legacy
.chain(v2)
.fold(BTreeMap::new(), |mut acc, (key, sf)| {
acc.entry(key)
.and_modify(|existing| {
*existing = sf.clone();
})
.or_insert(sf);
acc
});
merged
.into_iter()
.map(|((coldkey, netuid), sf)| (coldkey, netuid, sf))
}
}