code/pallets/subtensor/src/staking/claim_root.rs
use super::*;
use frame_support::dispatch::DispatchResult;
use frame_support::storage::{TransactionOutcome, with_transaction};
use frame_support::weights::Weight;
use sp_core::Get;
use sp_runtime::DispatchError;
use sp_std::collections::btree_map::BTreeMap;
use sp_std::collections::btree_set::BTreeSet;
use substrate_fixed::types::I96F32;
use subtensor_runtime_common::clear_prefix_with_meter;
use subtensor_swap_interface::SwapHandler;
impl<T: Config> Pallet<T> {
pub fn block_hash_to_indices(block_hash: T::Hash, k: u64, n: u64) -> Vec<u64> {
let block_hash_bytes = block_hash.as_ref();
let mut indices: BTreeSet<u64> = BTreeSet::new();
// k < n
let start_index: u64 = u64::from_be_bytes(
block_hash_bytes
.get(0..8)
.unwrap_or(&[0; 8])
.try_into()
.unwrap_or([0; 8]),
);
let mut last_idx = start_index;
for i in 0..k {
let bh_idx: usize = ((i.saturating_mul(8)) % 32) as usize;
let idx_step = u64::from_be_bytes(
block_hash_bytes
.get(bh_idx..(bh_idx.saturating_add(8)))
.unwrap_or(&[0; 8])
.try_into()
.unwrap_or([0; 8]),
);
let idx = last_idx
.saturating_add(idx_step)
.checked_rem(n)
.unwrap_or(0);
indices.insert(idx);
last_idx = idx;
}
indices.into_iter().collect()
}
pub fn increase_root_claimable_for_hotkey_and_subnet(
hotkey: &T::AccountId,
netuid: NetUid,
amount: AlphaBalance,
) {
// Get total stake on this hotkey on root.
let total: I96F32 =
I96F32::saturating_from_num(Self::get_stake_for_hotkey_on_subnet(hotkey, NetUid::ROOT));
// Get increment
let increment: I96F32 = I96F32::saturating_from_num(amount)
.checked_div(total)
.unwrap_or(I96F32::saturating_from_num(0.0));
// Unlikely to happen. This is mostly for test environment sanity checks.
if u64::from(amount) > total.saturating_to_num::<u64>() {
log::warn!("Not enough root stake. NetUID = {netuid}");
let owner = Owner::<T>::get(hotkey);
Self::increase_stake_for_hotkey_and_coldkey_on_subnet(hotkey, &owner, netuid, amount);
return;
}
// Increment claimable for this subnet.
RootClaimable::<T>::mutate(hotkey, |claimable| {
claimable
.entry(netuid)
.and_modify(|claim_total| *claim_total = claim_total.saturating_add(increment))
.or_insert(increment);
});
}
pub fn get_root_claimable_for_hotkey_coldkey(
hotkey: &T::AccountId,
coldkey: &T::AccountId,
netuid: NetUid,
) -> I96F32 {
// Get this keys stake balance on root.
let root_stake: I96F32 = I96F32::saturating_from_num(
Self::get_stake_for_hotkey_and_coldkey_on_subnet(hotkey, coldkey, NetUid::ROOT),
);
// Get the total claimable_rate for this hotkey and this network
let claimable_rate: I96F32 = *RootClaimable::<T>::get(hotkey)
.get(&netuid)
.unwrap_or(&I96F32::from(0));
// Compute the proportion owed to this coldkey via balance.
let claimable: I96F32 = claimable_rate.saturating_mul(root_stake);
claimable
}
pub fn get_root_owed_for_hotkey_coldkey_float(
hotkey: &T::AccountId,
coldkey: &T::AccountId,
netuid: NetUid,
) -> I96F32 {
let claimable = Self::get_root_claimable_for_hotkey_coldkey(hotkey, coldkey, netuid);
// Attain the root claimed to avoid overclaiming.
let root_claimed: I96F32 =
I96F32::saturating_from_num(RootClaimed::<T>::get((netuid, hotkey, coldkey)));
// Subtract the already claimed alpha.
let owed: I96F32 = claimable.saturating_sub(root_claimed);
owed
}
pub fn get_root_owed_for_hotkey_coldkey(
hotkey: &T::AccountId,
coldkey: &T::AccountId,
netuid: NetUid,
) -> u64 {
let owed = Self::get_root_owed_for_hotkey_coldkey_float(hotkey, coldkey, netuid);
// Convert owed to u64, mapping negative values to 0
let owed_u64: u64 = if owed.is_negative() {
0
} else {
owed.saturating_to_num::<u64>()
};
owed_u64
}
pub fn root_claim_on_subnet(
hotkey: &T::AccountId,
coldkey: &T::AccountId,
netuid: NetUid,
root_claim_type: RootClaimTypeEnum,
ignore_minimum_condition: bool,
) -> DispatchResult {
if DissolveCleanupQueue::<T>::get().contains(&netuid) {
log::debug!("root claim on subnet {netuid} is skipped, network is dissolved");
return Ok(());
}
// Subtract the root claimed.
let owed: I96F32 = Self::get_root_owed_for_hotkey_coldkey_float(hotkey, coldkey, netuid);
if !ignore_minimum_condition
&& owed < I96F32::saturating_from_num(RootClaimableThreshold::<T>::get(&netuid))
{
log::debug!(
"root claim on subnet {netuid} is skipped: {owed:?} for h={hotkey:?},c={coldkey:?} "
);
return Ok(()); // no-op
}
// Convert owed to u64, mapping negative values to 0
let owed_u64: u64 = if owed.is_negative() {
0
} else {
owed.saturating_to_num::<u64>()
};
if owed_u64 == 0 {
log::debug!(
"root claim on subnet {netuid} is skipped: {owed:?} for h={hotkey:?},c={coldkey:?}"
);
return Ok(()); // no-op
}
let swap = match root_claim_type {
RootClaimTypeEnum::Swap => true,
RootClaimTypeEnum::Keep => false,
RootClaimTypeEnum::KeepSubnets { subnets } => !subnets.contains(&netuid),
};
if swap {
with_transaction(|| {
// Increase stake on root. Swap the alpha owed to TAO.
let owed_tao = match Self::swap_alpha_for_tao(
netuid,
owed_u64.into(),
T::SwapInterface::min_price::<TaoBalance>(),
true,
) {
Ok(owed_tao) => owed_tao,
Err(err) => {
log::error!("Error swapping alpha for TAO: {err:?}");
return TransactionOutcome::Rollback(Err(err));
}
};
let root_subnet_account_id = match Self::get_subnet_account_id(NetUid::ROOT) {
Some(account_id) => account_id,
None => {
return TransactionOutcome::Rollback(Err(
Error::<T>::RootNetworkDoesNotExist.into(),
));
}
};
if let Err(err) = Self::transfer_tao_from_subnet(
netuid,
&root_subnet_account_id,
owed_tao.amount_paid_out.into(),
) {
log::error!("Error transferring root claim TAO from subnet: {err:?}");
return TransactionOutcome::Rollback(Err(err));
}
// Record root sell as protocol outflow (reduces protocol cost).
let root_sell_tao: TaoBalance = owed_tao.amount_paid_out;
SubnetRootSellTao::<T>::mutate(netuid, |total| {
*total = total.saturating_add(root_sell_tao);
});
Self::record_protocol_outflow(netuid, root_sell_tao);
Self::increase_stake_for_hotkey_and_coldkey_on_subnet(
hotkey,
coldkey,
NetUid::ROOT,
owed_tao.amount_paid_out.to_u64().into(),
);
// Increase root subnet SubnetTAO
SubnetTAO::<T>::mutate(NetUid::ROOT, |total| {
*total = total.saturating_add(owed_tao.amount_paid_out.into());
});
// Increase root SubnetAlphaOut
SubnetAlphaOut::<T>::mutate(NetUid::ROOT, |total| {
*total = total.saturating_add(u64::from(owed_tao.amount_paid_out).into());
});
// Increase Total Stake
TotalStake::<T>::mutate(|total| {
*total = total.saturating_add(owed_tao.amount_paid_out.into());
});
Self::add_stake_adjust_root_claimed_for_hotkey_and_coldkey(
hotkey,
coldkey,
owed_tao.amount_paid_out.into(),
);
TransactionOutcome::Commit(Ok(()))
})?;
} else
/* Keep */
{
// Increase the stake with the alpha owned
Self::increase_stake_for_hotkey_and_coldkey_on_subnet(
hotkey,
coldkey,
netuid,
owed_u64.into(),
);
}
// Increase root claimed by owed amount.
RootClaimed::<T>::mutate((netuid, hotkey, coldkey), |root_claimed| {
*root_claimed = root_claimed.saturating_add(owed_u64.into());
});
Ok(())
}
fn root_claim_on_subnet_weight(_root_claim_type: RootClaimTypeEnum) -> Weight {
Weight::from_parts(60_000_000, 6987)
.saturating_add(T::DbWeight::get().reads(7_u64))
.saturating_add(T::DbWeight::get().writes(5_u64))
}
pub fn root_claim_all(
hotkey: &T::AccountId,
coldkey: &T::AccountId,
subnets: Option<BTreeSet<NetUid>>,
) -> Result<Weight, DispatchError> {
let mut weight = Weight::default();
let root_claim_type = RootClaimType::<T>::get(coldkey);
weight.saturating_accrue(T::DbWeight::get().reads(1));
// Iterate over all the subnets this hotkey has claimable for root.
let root_claimable = RootClaimable::<T>::get(hotkey);
weight.saturating_accrue(T::DbWeight::get().reads(1));
for (netuid, _) in root_claimable.iter() {
let skip = subnets
.as_ref()
.map(|subnets| !subnets.contains(netuid))
.unwrap_or(false);
if skip {
continue;
}
Self::root_claim_on_subnet(hotkey, coldkey, *netuid, root_claim_type.clone(), false)?;
weight.saturating_accrue(Self::root_claim_on_subnet_weight(root_claim_type.clone()));
}
Ok(weight)
}
pub fn add_stake_adjust_root_claimed_for_hotkey_and_coldkey(
hotkey: &T::AccountId,
coldkey: &T::AccountId,
amount: u64,
) {
// Iterate over all the subnets this hotkey is staked on for root.
let root_claimable = RootClaimable::<T>::get(hotkey);
for (netuid, claimable_rate) in root_claimable.iter() {
// Get current staker root claimed value.
let root_claimed: u128 = RootClaimed::<T>::get((netuid, hotkey, coldkey));
// Increase root claimed based on the claimable rate.
let new_root_claimed = root_claimed.saturating_add(
claimable_rate
.saturating_mul(I96F32::from(u64::from(amount)))
.saturating_to_num(),
);
// Set the new root claimed value.
RootClaimed::<T>::insert((netuid, hotkey, coldkey), new_root_claimed);
}
}
pub fn remove_stake_adjust_root_claimed_for_hotkey_and_coldkey(
hotkey: &T::AccountId,
coldkey: &T::AccountId,
amount: AlphaBalance,
) {
// Iterate over all the subnets this hotkey is staked on for root.
let root_claimable = RootClaimable::<T>::get(hotkey);
for (netuid, claimable_rate) in root_claimable.iter() {
if *netuid == NetUid::ROOT.into() {
continue; // Skip the root netuid.
}
// Get current staker root claimed value.
let root_claimed: u128 = RootClaimed::<T>::get((netuid, hotkey, coldkey));
// Decrease root claimed based on the claimable rate.
let new_root_claimed = root_claimed.saturating_sub(
claimable_rate
.saturating_mul(I96F32::from(u64::from(amount)))
.saturating_to_num(),
);
// Set the new root_claimed value.
RootClaimed::<T>::insert((netuid, hotkey, coldkey), new_root_claimed);
}
}
pub fn do_root_claim(
coldkey: T::AccountId,
subnets: Option<BTreeSet<NetUid>>,
) -> Result<Weight, DispatchError> {
with_transaction(|| match Self::try_do_root_claim(coldkey, subnets) {
Ok(weight) => TransactionOutcome::Commit(Ok(weight)),
Err(err) => TransactionOutcome::Rollback(Err(err)),
})
}
fn try_do_root_claim(
coldkey: T::AccountId,
subnets: Option<BTreeSet<NetUid>>,
) -> Result<Weight, DispatchError> {
let mut weight = Weight::default();
let hotkeys = StakingHotkeys::<T>::get(&coldkey);
weight.saturating_accrue(T::DbWeight::get().reads(1));
for hotkey in hotkeys.iter() {
weight.saturating_accrue(T::DbWeight::get().reads(1));
weight.saturating_accrue(Self::root_claim_all(hotkey, &coldkey, subnets.clone())?);
}
Self::deposit_event(Event::RootClaimed { coldkey });
Ok(weight)
}
fn block_hash_to_indices_weight(k: u64, _n: u64) -> Weight {
Weight::from_parts(3_000_000, 1517)
.saturating_add(Weight::from_parts(100_412, 0).saturating_mul(k.into()))
}
pub fn maybe_add_coldkey_index(coldkey: &T::AccountId) {
if !StakingColdkeys::<T>::contains_key(coldkey) {
let n = NumStakingColdkeys::<T>::get();
StakingColdkeysByIndex::<T>::insert(n, coldkey.clone());
StakingColdkeys::<T>::insert(coldkey.clone(), n);
NumStakingColdkeys::<T>::mutate(|n| *n = n.saturating_add(1));
}
}
/// Returns true if `coldkey` still holds any root (netuid 0) stake on any of its
/// staking hotkeys. Used to decide whether the coldkey should remain indexed in the
/// auto-claim staking-coldkey index.
pub fn coldkey_has_root_stake(coldkey: &T::AccountId) -> bool {
StakingHotkeys::<T>::get(coldkey).iter().any(|hotkey| {
!Self::get_stake_for_hotkey_and_coldkey_on_subnet(hotkey, coldkey, NetUid::ROOT)
.is_zero()
})
}
/// Remove `coldkey` from the staking-coldkey index, compacting by moving the last
/// entry into the freed slot so the index stays dense in `[0, n)`. This is the inverse
/// of `maybe_add_coldkey_index` and keeps the
/// `StakingColdkeys[c] == i <=> StakingColdkeysByIndex[i] == c` bijection consistent.
pub fn maybe_remove_coldkey_index(coldkey: &T::AccountId) {
if let Some(idx) = StakingColdkeys::<T>::take(coldkey) {
let last = NumStakingColdkeys::<T>::get().saturating_sub(1);
if idx != last
&& let Some(moved) = StakingColdkeysByIndex::<T>::take(last)
{
StakingColdkeysByIndex::<T>::insert(idx, moved.clone());
StakingColdkeys::<T>::insert(moved, idx);
} else {
StakingColdkeysByIndex::<T>::remove(idx);
}
NumStakingColdkeys::<T>::put(last);
}
}
pub fn run_auto_claim_root_divs(last_block_hash: T::Hash) -> Weight {
let mut weight: Weight = Weight::default();
let n = NumStakingColdkeys::<T>::get();
let k = NumRootClaim::<T>::get();
weight.saturating_accrue(T::DbWeight::get().reads(2));
let coldkeys_to_claim: Vec<u64> = Self::block_hash_to_indices(last_block_hash, k, n);
weight.saturating_accrue(Self::block_hash_to_indices_weight(k, n));
for i in coldkeys_to_claim.iter() {
weight.saturating_accrue(T::DbWeight::get().reads(1));
if let Ok(coldkey) = StakingColdkeysByIndex::<T>::try_get(i) {
match Self::do_root_claim(coldkey.clone(), None) {
Ok(claim_weight) => weight.saturating_accrue(claim_weight),
Err(err) => log::error!("Error auto-claiming root dividends: {err:?}"),
}
}
}
weight
}
pub fn change_root_claim_type(coldkey: &T::AccountId, new_type: RootClaimTypeEnum) {
RootClaimType::<T>::insert(coldkey.clone(), new_type.clone());
Self::deposit_event(Event::RootClaimTypeSet {
coldkey: coldkey.clone(),
root_claim_type: new_type,
});
}
pub fn transfer_root_claimed_for_new_keys(
netuid: NetUid,
old_hotkey: &T::AccountId,
new_hotkey: &T::AccountId,
old_coldkey: &T::AccountId,
new_coldkey: &T::AccountId,
) {
let old_root_claimed = RootClaimed::<T>::get((netuid, old_hotkey, old_coldkey));
RootClaimed::<T>::remove((netuid, old_hotkey, old_coldkey));
RootClaimed::<T>::mutate((netuid, new_hotkey, new_coldkey), |new_root_claimed| {
// Sum the two already-claimed watermarks. When BOTH the source and the
// destination hold a legitimate RootClaimed — e.g. a coldkey swap onto a
// hotkey the new coldkey has already staked to, or a hotkey swap that merges
// two real positions — the merged "already claimed" total is old + new. Taking
// the max would drop one side, under-count what has already been claimed, and
// cause a future over-payment / double-claim of root dividends.
//
// GHSA-2026-010 (a *stale residual* watermark on new_hotkey inflating this sum
// in the hotkey-swap path) is prevented upstream by the root-swap cleanliness
// gate in `do_swap_hotkey`, which now also requires RootClaimed to be empty on
// new_hotkey (see `test_do_swap_hotkey_err_new_hotkey_not_clean_for_root`). With
// that gate the destination is always clean (new == 0) in the swap path, so the
// sum cannot be inflated there.
*new_root_claimed = old_root_claimed.saturating_add(*new_root_claimed);
});
}
pub fn transfer_root_claimable_for_new_hotkey(
old_hotkey: &T::AccountId,
new_hotkey: &T::AccountId,
) {
let src_root_claimable = RootClaimable::<T>::get(old_hotkey);
let mut dst_root_claimable = RootClaimable::<T>::get(new_hotkey);
RootClaimable::<T>::remove(old_hotkey);
for (netuid, claimable_rate) in src_root_claimable.into_iter() {
dst_root_claimable
.entry(netuid)
.and_modify(|total| *total = total.saturating_add(claimable_rate))
.or_insert(claimable_rate);
}
RootClaimable::<T>::insert(new_hotkey, dst_root_claimable);
}
/// Claim all root dividends for subnet and remove all associated data.
pub fn clean_up_root_claimable_for_subnet(
netuid: NetUid,
weight_meter: &mut WeightMeter,
last_key: Option<Vec<u8>>,
) -> (bool, Option<Vec<u8>>) {
// let mut to_remove_map = BTreeMap::<T::AccountId, BTreeMap<NetUid, I96F32>>::new();
// let mut read_all = true;
let iter = match last_key {
Some(raw_key) => RootClaimable::<T>::iter_from(raw_key),
None => RootClaimable::<T>::iter(),
};
fn filter_claimable(
claimable: &BTreeMap<NetUid, I96F32>,
netuid: NetUid,
) -> BTreeMap<NetUid, I96F32> {
let mut result = claimable.clone();
if result.contains_key(&netuid) {
result.remove(&netuid);
}
result
}
let (read_all, last_item) = Self::remove_storage_entries_for_netuid(
weight_meter,
iter,
|(_, _)| true,
|(hotkey, claimable)| (hotkey.clone(), claimable.clone()),
|(hotkey, claimable)| {
RootClaimable::<T>::insert(hotkey, filter_claimable(claimable, netuid))
},
1,
);
(
read_all,
last_item.map(|(hotkey, _)| RootClaimable::<T>::hashed_key_for(&hotkey)),
)
}
pub fn clean_up_root_claimed_for_subnet(
netuid: NetUid,
weight_meter: &mut WeightMeter,
) -> bool {
clear_prefix_with_meter(weight_meter, T::DbWeight::get().writes(1), |limit| {
RootClaimed::<T>::clear_prefix((netuid,), limit, None)
})
}
}