diff --git a/pallets/subtensor/src/staking/claim_root.rs b/pallets/subtensor/src/staking/claim_root.rs index 58babb79a6..e80177e04b 100644 --- a/pallets/subtensor/src/staking/claim_root.rs +++ b/pallets/subtensor/src/staking/claim_root.rs @@ -245,12 +245,22 @@ impl Pallet { weight } + // Invariant: `RootClaimable[hotkey]` never carries a `NetUid::ROOT` key. + // `increase_root_claimable_for_hotkey_and_subnet` is the only writer, and + // its sole caller `run_coinbase` filters ROOT out of the distribution loop + // before populating the map (see `run_coinbase.rs:31`). + // `transfer_root_claimable_for_new_hotkey` and + // `finalize_all_subnet_root_dividends` only propagate or remove existing + // keys, so they can't introduce a ROOT key either. The two functions + // below therefore iterate the full map without a ROOT special-case; if a + // ROOT entry ever does appear through a future code path, add and remove + // will stay symmetric rather than drifting. + 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::::get(hotkey); for (netuid, claimable_rate) in root_claimable.iter() { // Get current staker root claimed value. @@ -273,13 +283,8 @@ impl Pallet { coldkey: &T::AccountId, amount: AlphaBalance, ) { - // Iterate over all the subnets this hotkey is staked on for root. let root_claimable = RootClaimable::::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::::get((netuid, hotkey, coldkey)); diff --git a/pallets/subtensor/src/tests/claim_root.rs b/pallets/subtensor/src/tests/claim_root.rs index f86ca0b1e6..71d099b84c 100644 --- a/pallets/subtensor/src/tests/claim_root.rs +++ b/pallets/subtensor/src/tests/claim_root.rs @@ -2059,3 +2059,63 @@ fn test_claim_root_with_moved_stake() { assert_abs_diff_eq!(bob_stake_diff2, estimated_stake as u64, epsilon = 100u64,); }); } + +// RootClaimable should never carry a NetUid::ROOT key under current code +// paths, so this scenario isn't reachable today. The test pins the +// symmetric-iteration property: if a ROOT entry ever does sneak in through +// a future code path, add and remove both process it the same way so the +// pair stays balanced rather than drifting as it would if only one side +// handled it specially. +#[test] +fn test_add_and_remove_stake_adjust_root_claimed_are_symmetric() { + new_test_ext(1).execute_with(|| { + let hotkey = U256::from(1); + let coldkey = U256::from(2); + let other_netuid = NetUid::from(7); + + let rate = I96F32::saturating_from_num(3); + + // Seed RootClaimable with both a ROOT entry and a non-root entry. + // ROOT should not normally be here; we seed it to verify the two + // functions stay symmetric regardless. + let mut claimable = std::collections::BTreeMap::new(); + claimable.insert(NetUid::ROOT, rate); + claimable.insert(other_netuid, rate); + RootClaimable::::insert(hotkey, claimable); + + let amount: u64 = 100; + SubtensorModule::add_stake_adjust_root_claimed_for_hotkey_and_coldkey( + &hotkey, &coldkey, amount, + ); + + let expected: u128 = rate + .saturating_mul(I96F32::from(amount)) + .saturating_to_num::(); + + // Both entries got bumped the same way by the add path. + assert_eq!( + RootClaimed::::get((NetUid::ROOT, hotkey, coldkey)), + expected + ); + assert_eq!( + RootClaimed::::get((other_netuid, hotkey, coldkey)), + expected + ); + + // Remove the same amount and check both entries drain back to 0. + SubtensorModule::remove_stake_adjust_root_claimed_for_hotkey_and_coldkey( + &hotkey, + &coldkey, + AlphaBalance::from(amount), + ); + + assert_eq!( + RootClaimed::::get((NetUid::ROOT, hotkey, coldkey)), + 0u128 + ); + assert_eq!( + RootClaimed::::get((other_netuid, hotkey, coldkey)), + 0u128 + ); + }); +}