fix: match remove-side NetUid::ROOT guard in add_stake_adjust_root_cl…

pallets/subtensor/src/staking/claim_root.rs

@@ -245,12 +245,22 @@ impl<T: Config> Pallet<T> {
         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::<T>::get(hotkey);
         for (netuid, claimable_rate) in root_claimable.iter() {
             // Get current staker root claimed value.
@@ -273,13 +283,8 @@ impl<T: Config> Pallet<T> {
         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));
 

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::<Test>::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::<u128>();
+
+        // Both entries got bumped the same way by the add path.
+        assert_eq!(
+            RootClaimed::<Test>::get((NetUid::ROOT, hotkey, coldkey)),
+            expected
+        );
+        assert_eq!(
+            RootClaimed::<Test>::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::<Test>::get((NetUid::ROOT, hotkey, coldkey)),
+            0u128
+        );
+        assert_eq!(
+            RootClaimed::<Test>::get((other_netuid, hotkey, coldkey)),
+            0u128
+        );
+    });
+}