diff --git a/Cargo.lock b/Cargo.lock index c9ba8aede2..aec9116f1f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9708,6 +9708,7 @@ dependencies = [ "fp-storage", "frame-support", "frame-system", + "log", "pallet-evm", "parity-scale-codec", "scale-info", diff --git a/pallets/subtensor/src/coinbase/run_coinbase.rs b/pallets/subtensor/src/coinbase/run_coinbase.rs index 2091946598..e39f67b2b3 100644 --- a/pallets/subtensor/src/coinbase/run_coinbase.rs +++ b/pallets/subtensor/src/coinbase/run_coinbase.rs @@ -631,6 +631,10 @@ impl Pallet { tou64!(root_alpha).into(), ); + PendingRootAlpha::::mutate(&hotkey, |alpha| { + *alpha = alpha.saturating_add(tou64!(root_alpha).into()) + }); + // Record root alpha dividends for this validator on this subnet. RootAlphaDividendsPerSubnet::::mutate(netuid, &hotkey, |divs| { *divs = divs.saturating_add(tou64!(root_alpha).into()); diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index ef2d44e68b..68431e12c9 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -1644,6 +1644,11 @@ pub mod pallet { pub type PendingRootAlphaDivs = StorageMap<_, Identity, NetUid, AlphaCurrency, ValueQuery, DefaultZeroAlpha>; + /// --- MAP ( hotkey ) --> pending_root_alpha + #[pallet::storage] + pub type PendingRootAlpha = + StorageMap<_, Blake2_128Concat, T::AccountId, u128, ValueQuery, DefaultZeroU128>; + /// --- MAP ( netuid ) --> pending_owner_cut #[pallet::storage] pub type PendingOwnerCut = diff --git a/pallets/subtensor/src/migrations/migrate_init_pending_root_alpha.rs b/pallets/subtensor/src/migrations/migrate_init_pending_root_alpha.rs new file mode 100644 index 0000000000..a0d57479fb --- /dev/null +++ b/pallets/subtensor/src/migrations/migrate_init_pending_root_alpha.rs @@ -0,0 +1,79 @@ +use super::*; +use alloc::{collections::BTreeMap, string::String}; +use frame_support::{traits::Get, weights::Weight}; +use substrate_fixed::types::I96F32; + +/// Migration to initialize PendingRootAlpha storage based on RootClaimed storage. +/// This aggregates all RootClaimed values across all netuids and coldkeys for each hotkey. +pub fn migrate_init_pending_root_alpha() -> Weight { + let migration_name = b"migrate_init_pending_root_alpha".to_vec(); + let mut weight = T::DbWeight::get().reads(1); + + // Check if the migration has already run + if HasMigrationRun::::get(&migration_name) { + log::info!( + "Migration '{:?}' has already run. Skipping.", + String::from_utf8_lossy(&migration_name) + ); + return weight; + } + + log::info!( + "Running migration '{}'", + String::from_utf8_lossy(&migration_name) + ); + + // Aggregate RootClaimable values by hotkey + let mut root_claimable_alpha_map: BTreeMap = BTreeMap::new(); + for (hotkey, root_claimable) in RootClaimable::::iter() { + // Sum rates as I96F32 (not converting to u128 first, which would truncate) + let claimable_rate: I96F32 = root_claimable + .values() + .fold(I96F32::from(0), |acc, x| acc.saturating_add(*x)); + + let root_stake = Pallet::::get_stake_for_hotkey_on_subnet(&hotkey, NetUid::ROOT); + let total = claimable_rate.saturating_mul(I96F32::saturating_from_num(root_stake)); + + root_claimable_alpha_map.insert(hotkey, total.saturating_to_num::()); + weight = weight.saturating_add(T::DbWeight::get().reads(1)); + } + + // Aggregate RootClaimed values by hotkey + // Key: hotkey, Value: sum of all RootClaimed values for that hotkey + let mut root_claimed_alpha_map: BTreeMap = BTreeMap::new(); + + // Iterate over all RootClaimed entries: (netuid, hotkey, coldkey) -> claimed_value + for ((_netuid, hotkey, _coldkey), claimed_value) in RootClaimed::::iter() { + // Aggregate the claimed value for this hotkey + root_claimed_alpha_map + .entry(hotkey.clone()) + .and_modify(|total| *total = total.saturating_add(claimed_value)) + .or_insert(claimed_value); + + // Account for read operation + weight = weight.saturating_add(T::DbWeight::get().reads(1)); + } + + // Set PendingRootAlpha for each hotkey + let mut migrated_count = 0u64; + for (hotkey, claimable) in root_claimable_alpha_map { + let claimed = root_claimed_alpha_map.get(&hotkey).unwrap_or(&0); + let pending = claimable.saturating_sub(*claimed); + PendingRootAlpha::::insert(&hotkey, pending); + migrated_count = migrated_count.saturating_add(1); + } + + weight = weight.saturating_add(T::DbWeight::get().writes(migrated_count)); + + log::info!( + "Migration '{}' completed successfully. Initialized PendingRootAlpha for {} hotkeys.", + String::from_utf8_lossy(&migration_name), + migrated_count + ); + + // Mark the migration as completed + HasMigrationRun::::insert(&migration_name, true); + weight = weight.saturating_add(T::DbWeight::get().writes(1)); + + weight +} diff --git a/pallets/subtensor/src/migrations/mod.rs b/pallets/subtensor/src/migrations/mod.rs index a03da9289e..46b6da4154 100644 --- a/pallets/subtensor/src/migrations/mod.rs +++ b/pallets/subtensor/src/migrations/mod.rs @@ -20,6 +20,7 @@ pub mod migrate_fix_is_network_member; pub mod migrate_fix_root_subnet_tao; pub mod migrate_fix_root_tao_and_alpha_in; pub mod migrate_fix_staking_hot_keys; +pub mod migrate_init_pending_root_alpha; pub mod migrate_init_tao_flow; pub mod migrate_init_total_issuance; pub mod migrate_kappa_map_to_default; diff --git a/pallets/subtensor/src/staking/claim_root.rs b/pallets/subtensor/src/staking/claim_root.rs index 24a26d154c..fca80168a1 100644 --- a/pallets/subtensor/src/staking/claim_root.rs +++ b/pallets/subtensor/src/staking/claim_root.rs @@ -1,7 +1,7 @@ use super::*; use frame_support::weights::Weight; use sp_core::Get; -use sp_std::collections::btree_set::BTreeSet; +use sp_std::collections::{btree_map::BTreeMap, btree_set::BTreeSet}; use substrate_fixed::types::I96F32; use subtensor_swap_interface::SwapHandler; @@ -207,6 +207,10 @@ impl Pallet { RootClaimed::::mutate((netuid, hotkey, coldkey), |root_claimed| { *root_claimed = root_claimed.saturating_add(owed_u64.into()); }); + + PendingRootAlpha::::mutate(hotkey, |value| { + *value = value.saturating_sub(owed_u64.into()); + }); } fn root_claim_on_subnet_weight(_root_claim_type: RootClaimTypeEnum) -> Weight { @@ -257,11 +261,14 @@ impl Pallet { let root_claimed: u128 = RootClaimed::::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(), - ); + let added_amount = claimable_rate + .saturating_mul(I96F32::from(u64::from(amount))) + .saturating_to_num(); + let new_root_claimed = root_claimed.saturating_add(added_amount); + + PendingRootAlpha::::mutate(hotkey, |value| { + *value = value.saturating_sub(added_amount); + }); // Set the new root claimed value. RootClaimed::::insert((netuid, hotkey, coldkey), new_root_claimed); @@ -284,11 +291,14 @@ impl Pallet { let root_claimed: u128 = RootClaimed::::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(), - ); + let subed_amount = claimable_rate + .saturating_mul(I96F32::from(u64::from(amount))) + .saturating_to_num(); + let new_root_claimed = root_claimed.saturating_sub(subed_amount); + + PendingRootAlpha::::mutate(hotkey, |value| { + *value = value.saturating_add(subed_amount); + }); // Set the new root_claimed value. RootClaimed::::insert((netuid, hotkey, coldkey), new_root_claimed); @@ -388,16 +398,62 @@ impl Pallet { RootClaimable::::insert(new_hotkey, dst_root_claimable); } + pub fn transfer_pending_root_alpha_for_new_hotkey( + old_hotkey: &T::AccountId, + new_hotkey: &T::AccountId, + ) { + let src_pending_root_alpha = PendingRootAlpha::::get(old_hotkey); + let dst_pending_root_alpha = PendingRootAlpha::::get(new_hotkey); + PendingRootAlpha::::remove(old_hotkey); + PendingRootAlpha::::insert( + new_hotkey, + dst_pending_root_alpha.saturating_add(src_pending_root_alpha), + ); + } + /// Claim all root dividends for subnet and remove all associated data. pub fn finalize_all_subnet_root_dividends(netuid: NetUid) { let hotkeys = RootClaimable::::iter_keys().collect::>(); + let mut root_claimable_alpha_map: BTreeMap = BTreeMap::new(); + for hotkey in hotkeys.iter() { + let root_stake = Self::get_stake_for_hotkey_on_subnet(hotkey, NetUid::ROOT); + let claimable_rate = RootClaimable::::get(hotkey) + .values() + .fold(I96F32::from(0), |acc, x| acc.saturating_add(*x)); + let total = claimable_rate.saturating_mul(I96F32::saturating_from_num(root_stake)); + // let pending_root_alpha = PendingRootAlpha::::get(hotkey); + + root_claimable_alpha_map.insert(hotkey.clone(), total.saturating_to_num::()); + RootClaimable::::mutate(hotkey, |claimable| { claimable.remove(&netuid); }); } + let mut root_claimed_alpha_map: BTreeMap = BTreeMap::new(); + + for ((_netuid, hotkey, _coldkey), root_claimed) in RootClaimed::::iter() { + if !hotkeys.contains(&hotkey) { + continue; + } + + root_claimed_alpha_map + .entry(hotkey.clone()) + .and_modify(|total| *total = total.saturating_add(root_claimed)) + .or_insert(root_claimed); + } + let _ = RootClaimed::::clear_prefix((netuid,), u32::MAX, None); + + for (hotkey, claimable) in root_claimable_alpha_map { + let claimed = root_claimed_alpha_map.get(&hotkey).unwrap_or(&0); + // still some root alpha not claimed + let pending = claimable.saturating_sub(*claimed); + PendingRootAlpha::::mutate(&hotkey, |value| { + *value = value.saturating_sub(pending); + }); + } } } diff --git a/pallets/subtensor/src/swap/swap_hotkey.rs b/pallets/subtensor/src/swap/swap_hotkey.rs index 4fdf87fb7b..eaf162cbf8 100644 --- a/pallets/subtensor/src/swap/swap_hotkey.rs +++ b/pallets/subtensor/src/swap/swap_hotkey.rs @@ -531,6 +531,9 @@ impl Pallet { } } + // 9.3 update pending root alpha for the hotkeys. + Self::transfer_pending_root_alpha_for_new_hotkey(old_hotkey, new_hotkey); + Ok(()) } } diff --git a/pallets/subtensor/src/tests/migration.rs b/pallets/subtensor/src/tests/migration.rs index bed77e797f..1811a167c9 100644 --- a/pallets/subtensor/src/tests/migration.rs +++ b/pallets/subtensor/src/tests/migration.rs @@ -2968,3 +2968,599 @@ fn test_migrate_remove_unknown_neuron_axon_cert_prom() { } } } + +#[test] +fn test_migrate_init_pending_root_alpha() { + new_test_ext(1).execute_with(|| { + const MIGRATION_NAME: &str = "migrate_init_pending_root_alpha"; + + // Setup: Create networks + let netuid1 = NetUid::from(1); + let netuid2 = NetUid::from(2); + add_network(netuid1, 1, 0); + add_network(netuid2, 1, 0); + + // Setup: Create hotkeys and coldkeys + let hotkey1 = U256::from(1001); + let hotkey2 = U256::from(1002); + let coldkey1 = U256::from(2001); + let coldkey2 = U256::from(2002); + let coldkey3 = U256::from(2003); + + // Setup: Set root stake for hotkeys + let root_stake1 = 1_000_000u64; // 1M TAO + let root_stake2 = 2_000_000u64; // 2M TAO + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey1, + &coldkey1, + NetUid::ROOT, + root_stake1.into(), + ); + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey2, + &coldkey2, + NetUid::ROOT, + root_stake2.into(), + ); + + // Setup: Set RootClaimable for hotkeys + // Hotkey1: netuid1 with rate 0.1, netuid2 with rate 0.2 + // Total rate = 0.3, claimable = 0.3 * 1M = 300k + RootClaimable::::mutate(hotkey1, |claimable| { + claimable.insert(netuid1, I96F32::from_num(0.1)); + claimable.insert(netuid2, I96F32::from_num(0.2)); + }); + + // Hotkey2: netuid1 with rate 0.15 + // Total rate = 0.15, claimable = 0.15 * 2M = 300k + RootClaimable::::mutate(hotkey2, |claimable| { + claimable.insert(netuid1, I96F32::from_num(0.15)); + }); + + // Setup: Set RootClaimed entries + // Hotkey1: claimed 50k from netuid1 (coldkey1) + 30k from netuid2 (coldkey2) = 80k total + RootClaimed::::insert((netuid1, hotkey1, coldkey1), 50_000u128); + RootClaimed::::insert((netuid2, hotkey1, coldkey2), 30_000u128); + + // Hotkey2: claimed 100k from netuid1 (coldkey2) + 50k from netuid1 (coldkey3) = 150k total + RootClaimed::::insert((netuid1, hotkey2, coldkey2), 100_000u128); + RootClaimed::::insert((netuid1, hotkey2, coldkey3), 50_000u128); + + // Verify initial state: PendingRootAlpha should be empty + assert_eq!(PendingRootAlpha::::get(hotkey1), 0_u128); + assert_eq!(PendingRootAlpha::::get(hotkey2), 0_u128); + + // Verify migration hasn't run yet + assert!( + !HasMigrationRun::::get(MIGRATION_NAME.as_bytes().to_vec()), + "Migration should not have run yet." + ); + + // Run the migration + let weight = + crate::migrations::migrate_init_pending_root_alpha::migrate_init_pending_root_alpha::< + Test, + >(); + + // Verify migration has run + assert!( + HasMigrationRun::::get(MIGRATION_NAME.as_bytes().to_vec()), + "Migration should be marked as run." + ); + + // Verify weight is non-zero + assert!(!weight.is_zero(), "Migration weight should be non-zero"); + + // Verify PendingRootAlpha for hotkey1 + // Expected: claimable (300k) - claimed (80k) = 220k + let expected_pending1 = 300_000u128 - 80_000u128; // 220k + assert_eq!( + PendingRootAlpha::::get(hotkey1), + expected_pending1, + "Hotkey1 pending root alpha should be claimable - claimed" + ); + + // Verify PendingRootAlpha for hotkey2 + // Expected: claimable (300k) - claimed (150k) = 150k + // Note: Fixed-point arithmetic may cause small rounding differences + let expected_pending2 = 300_000u128 - 150_000u128; // 150k + let actual_pending2 = PendingRootAlpha::::get(hotkey2); + assert!( + actual_pending2 >= expected_pending2.saturating_sub(1) && actual_pending2 <= expected_pending2.saturating_add(1), + "Hotkey2 pending root alpha should be approximately claimable - claimed. Expected: {}, Got: {}", + expected_pending2, + actual_pending2 + ); + + // Test: Migration should be idempotent (running twice should not change values) + let actual_pending1_first = PendingRootAlpha::::get(hotkey1); + let actual_pending2_first = PendingRootAlpha::::get(hotkey2); + let weight_second_run = + crate::migrations::migrate_init_pending_root_alpha::migrate_init_pending_root_alpha::< + Test, + >(); + assert_eq!( + PendingRootAlpha::::get(hotkey1), + actual_pending1_first, + "Second migration run should not change values" + ); + assert_eq!( + PendingRootAlpha::::get(hotkey2), + actual_pending2_first, + "Second migration run should not change values" + ); + // Second run should return early with just the read weight + assert_eq!( + weight_second_run, + ::DbWeight::get().reads(1), + "Second run should only read the migration flag" + ); + + // Test: Hotkey with no RootClaimable should not have PendingRootAlpha set + let hotkey3 = U256::from(1003); + assert_eq!( + PendingRootAlpha::::get(hotkey3), + 0_u128, + "Hotkey without RootClaimable should have zero pending" + ); + + // Test: Hotkey with RootClaimable but no RootClaimed should have full claimable as pending + let hotkey4 = U256::from(1004); + let root_stake4 = 500_000u64; + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey4, + &coldkey1, + NetUid::ROOT, + root_stake4.into(), + ); + RootClaimable::::mutate(hotkey4, |claimable| { + claimable.insert(netuid1, I96F32::from_num(0.1)); + }); + // Re-run migration to pick up new hotkey + HasMigrationRun::::remove(MIGRATION_NAME.as_bytes().to_vec()); + crate::migrations::migrate_init_pending_root_alpha::migrate_init_pending_root_alpha::( + ); + // Expected: 0.1 * 500k = 50k, no claimed = 50k pending (may have small rounding differences) + let expected_hotkey4 = 50_000u128; + let actual_hotkey4 = PendingRootAlpha::::get(hotkey4); + assert!( + actual_hotkey4 >= expected_hotkey4.saturating_sub(1) && actual_hotkey4 <= expected_hotkey4.saturating_add(1), + "Hotkey with claimable but no claimed should have full claimable as pending. Expected: {}, Got: {}", + expected_hotkey4, + actual_hotkey4 + ); + }); +} + +#[test] +fn test_migrate_init_pending_root_alpha_edge_cases() { + new_test_ext(1).execute_with(|| { + const MIGRATION_NAME: &str = "migrate_init_pending_root_alpha"; + + // Test 1: Hotkey with zero root stake + let hotkey_zero_stake = U256::from(2001); + let coldkey_zero_stake = U256::from(2002); + let netuid = NetUid::from(1); + add_network(netuid, 1, 0); + + RootClaimable::::mutate(hotkey_zero_stake, |claimable| { + claimable.insert(netuid, I96F32::from_num(0.1)); + }); + // No root stake set + + HasMigrationRun::::remove(MIGRATION_NAME.as_bytes().to_vec()); + crate::migrations::migrate_init_pending_root_alpha::migrate_init_pending_root_alpha::( + ); + + // With zero stake, claimable = rate * 0 = 0 + assert_eq!( + PendingRootAlpha::::get(hotkey_zero_stake), + 0_u128, + "Hotkey with zero root stake should have zero pending" + ); + + // Test 2: Hotkey with claimed > claimable (should not go negative) + let hotkey_overclaimed = U256::from(2003); + let coldkey_overclaimed = U256::from(2004); + let root_stake_overclaimed = 1_000_000u64; + + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey_overclaimed, + &coldkey_overclaimed, + NetUid::ROOT, + root_stake_overclaimed.into(), + ); + + RootClaimable::::mutate(hotkey_overclaimed, |claimable| { + claimable.insert(netuid, I96F32::from_num(0.1)); + }); + + // Claim more than claimable: claimable = 0.1 * 1M = 100k, claimed = 150k + RootClaimed::::insert( + (netuid, hotkey_overclaimed, coldkey_overclaimed), + 150_000u128, + ); + + HasMigrationRun::::remove(MIGRATION_NAME.as_bytes().to_vec()); + crate::migrations::migrate_init_pending_root_alpha::migrate_init_pending_root_alpha::( + ); + + // Pending should be 0, not negative (100k - 150k = 0 due to saturating_sub) + assert_eq!( + PendingRootAlpha::::get(hotkey_overclaimed), + 0_u128, + "Hotkey with claimed > claimable should have zero pending (not negative)" + ); + + // Test 3: Hotkey with multiple subnets + let hotkey_multi = U256::from(2005); + let coldkey_multi = U256::from(2006); + let netuid2 = NetUid::from(2); + add_network(netuid2, 1, 0); + let root_stake_multi = 2_000_000u64; + + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey_multi, + &coldkey_multi, + NetUid::ROOT, + root_stake_multi.into(), + ); + + RootClaimable::::mutate(hotkey_multi, |claimable| { + claimable.insert(netuid, I96F32::from_num(0.1)); // 0.1 rate + claimable.insert(netuid2, I96F32::from_num(0.2)); // 0.2 rate + }); + // Total rate = 0.3, claimable = 0.3 * 2M = 600k + + RootClaimed::::insert((netuid, hotkey_multi, coldkey_multi), 100_000u128); + RootClaimed::::insert((netuid2, hotkey_multi, coldkey_multi), 150_000u128); + // Total claimed = 250k + + HasMigrationRun::::remove(MIGRATION_NAME.as_bytes().to_vec()); + crate::migrations::migrate_init_pending_root_alpha::migrate_init_pending_root_alpha::( + ); + + // Expected: 600k - 250k = 350k (may have small rounding differences) + let expected_multi = 350_000u128; + let actual_multi = PendingRootAlpha::::get(hotkey_multi); + assert!( + actual_multi >= expected_multi.saturating_sub(1) && actual_multi <= expected_multi.saturating_add(1), + "Hotkey with multiple subnets should aggregate correctly. Expected: {}, Got: {}", + expected_multi, + actual_multi + ); + + // Test 4: Hotkey with multiple coldkeys claiming + let hotkey_multi_cold = U256::from(2007); + let coldkey_multi_cold1 = U256::from(2008); + let coldkey_multi_cold2 = U256::from(2009); + let root_stake_multi_cold = 1_000_000u64; + + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey_multi_cold, + &coldkey_multi_cold1, + NetUid::ROOT, + root_stake_multi_cold.into(), + ); + + RootClaimable::::mutate(hotkey_multi_cold, |claimable| { + claimable.insert(netuid, I96F32::from_num(0.2)); + }); + // Claimable = 0.2 * 1M = 200k + + RootClaimed::::insert((netuid, hotkey_multi_cold, coldkey_multi_cold1), 80_000u128); + RootClaimed::::insert((netuid, hotkey_multi_cold, coldkey_multi_cold2), 50_000u128); + // Total claimed = 130k + + HasMigrationRun::::remove(MIGRATION_NAME.as_bytes().to_vec()); + crate::migrations::migrate_init_pending_root_alpha::migrate_init_pending_root_alpha::( + ); + + // Expected: 200k - 130k = 70k (may have small rounding differences) + let expected_multi_cold = 70_000u128; + let actual_multi_cold = PendingRootAlpha::::get(hotkey_multi_cold); + assert!( + actual_multi_cold >= expected_multi_cold.saturating_sub(1) && actual_multi_cold <= expected_multi_cold.saturating_add(1), + "Hotkey with multiple coldkeys should aggregate claimed correctly. Expected: {}, Got: {}", + expected_multi_cold, + actual_multi_cold + ); + }); +} + +#[test] +fn test_pending_root_alpha_claiming_decrements() { + new_test_ext(1).execute_with(|| { + let netuid = NetUid::from(1); + add_network(netuid, 1, 0); + + let hotkey = U256::from(3001); + let coldkey = U256::from(3002); + let root_stake = 1_000_000u64; + + // Setup root stake + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &coldkey, + NetUid::ROOT, + root_stake.into(), + ); + + // Setup RootClaimable + RootClaimable::::mutate(hotkey, |claimable| { + claimable.insert(netuid, I96F32::from_num(0.2)); + }); + + // Initialize PendingRootAlpha manually (simulating migration) + // Claimable = 0.2 * 1M = 200k, no claimed = 200k pending + PendingRootAlpha::::insert(hotkey, 200_000u128); + + // Claim 50k + let initial_pending = PendingRootAlpha::::get(hotkey); + SubtensorModule::root_claim_on_subnet( + &hotkey, + &coldkey, + netuid, + RootClaimTypeEnum::Keep, + true, // ignore minimum condition + ); + + let final_pending = PendingRootAlpha::::get(hotkey); + let claimed = RootClaimed::::get((netuid, hotkey, coldkey)); + + // Pending should decrease by the claimed amount + assert!( + final_pending < initial_pending, + "PendingRootAlpha should decrease after claiming" + ); + assert_eq!( + initial_pending.saturating_sub(final_pending), + claimed, + "PendingRootAlpha decrease should equal claimed amount" + ); + }); +} + +#[test] +fn test_pending_root_alpha_add_stake_decrements() { + new_test_ext(1).execute_with(|| { + let netuid = NetUid::from(1); + add_network(netuid, 1, 0); + + let hotkey = U256::from(3003); + let coldkey = U256::from(3004); + let initial_stake = 1_000_000u64; + let additional_stake = 500_000u64; + + // Setup initial root stake + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &coldkey, + NetUid::ROOT, + initial_stake.into(), + ); + + // Setup RootClaimable + RootClaimable::::mutate(hotkey, |claimable| { + claimable.insert(netuid, I96F32::from_num(0.1)); + }); + + // Initialize PendingRootAlpha + PendingRootAlpha::::insert(hotkey, 100_000u128); + + // Add stake - this should adjust RootClaimed and decrement PendingRootAlpha + let initial_pending = PendingRootAlpha::::get(hotkey); + SubtensorModule::add_stake_adjust_root_claimed_for_hotkey_and_coldkey( + &hotkey, + &coldkey, + additional_stake, + ); + + let final_pending = PendingRootAlpha::::get(hotkey); + let root_claimed = RootClaimed::::get((netuid, hotkey, coldkey)); + + // Pending should decrease + assert!( + final_pending < initial_pending, + "PendingRootAlpha should decrease when adding stake" + ); + // The decrease should equal the added_amount = rate * additional_stake + // rate = 0.1, additional_stake = 500k, added_amount = 0.1 * 500k = 50k + assert_eq!( + root_claimed, 50_000u128, + "RootClaimed should increase by rate * additional_stake" + ); + }); +} + +#[test] +fn test_pending_root_alpha_remove_stake_increments() { + new_test_ext(1).execute_with(|| { + let netuid = NetUid::from(1); + add_network(netuid, 1, 0); + + let hotkey = U256::from(3005); + let coldkey = U256::from(3006); + let initial_stake = 2_000_000u64; + let remove_stake = AlphaCurrency::from(500_000u64); + + // Setup initial root stake + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &coldkey, + NetUid::ROOT, + initial_stake.into(), + ); + + // Setup RootClaimable + RootClaimable::::mutate(hotkey, |claimable| { + claimable.insert(netuid, I96F32::from_num(0.15)); + }); + + // Set initial RootClaimed + RootClaimed::::insert((netuid, hotkey, coldkey), 100_000u128); + + // Initialize PendingRootAlpha + PendingRootAlpha::::insert(hotkey, 200_000u128); + + // Remove stake - this should adjust RootClaimed and increment PendingRootAlpha + let initial_pending = PendingRootAlpha::::get(hotkey); + SubtensorModule::remove_stake_adjust_root_claimed_for_hotkey_and_coldkey( + &hotkey, + &coldkey, + remove_stake, + ); + + let final_pending = PendingRootAlpha::::get(hotkey); + let root_claimed = RootClaimed::::get((netuid, hotkey, coldkey)); + + // Pending should increase + assert!( + final_pending > initial_pending, + "PendingRootAlpha should increase when removing stake" + ); + // The increase should equal the subed_amount = rate * remove_stake + // rate = 0.15, remove_stake = 500k, subed_amount = 0.15 * 500k = 75k + // RootClaimed should decrease from 100k to 25k (may have small rounding differences) + let expected_root_claimed = 25_000u128; + assert!( + root_claimed >= expected_root_claimed.saturating_sub(1) && root_claimed <= expected_root_claimed.saturating_add(1), + "RootClaimed should decrease by rate * remove_stake. Expected: {}, Got: {}", + expected_root_claimed, + root_claimed + ); + // PendingRootAlpha should increase by subed_amount (may have small rounding differences) + let expected_pending_increase = 75_000u128; + let actual_increase = final_pending.saturating_sub(initial_pending); + assert!( + actual_increase >= expected_pending_increase.saturating_sub(1) && actual_increase <= expected_pending_increase.saturating_add(1), + "PendingRootAlpha should increase by subed_amount. Expected increase: {}, Actual increase: {}", + expected_pending_increase, + actual_increase + ); + }); +} + +#[test] +fn test_transfer_pending_root_alpha_for_new_hotkey() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(3007); + let new_hotkey = U256::from(3008); + let old_pending = 150_000u128; + let new_pending = 50_000u128; + + // Setup initial pending values + PendingRootAlpha::::insert(old_hotkey, old_pending); + PendingRootAlpha::::insert(new_hotkey, new_pending); + + // Transfer pending root alpha + SubtensorModule::transfer_pending_root_alpha_for_new_hotkey(&old_hotkey, &new_hotkey); + + // Old hotkey should have zero pending + assert_eq!( + PendingRootAlpha::::get(old_hotkey), + 0_u128, + "Old hotkey should have zero pending after transfer" + ); + + // New hotkey should have sum of both + assert_eq!( + PendingRootAlpha::::get(new_hotkey), + old_pending.saturating_add(new_pending), + "New hotkey should have sum of old and new pending" + ); + + // Test with new_hotkey having zero initial pending + let old_hotkey2 = U256::from(3009); + let new_hotkey2 = U256::from(3010); + PendingRootAlpha::::insert(old_hotkey2, 100_000u128); + PendingRootAlpha::::insert(new_hotkey2, 0_u128); + + SubtensorModule::transfer_pending_root_alpha_for_new_hotkey(&old_hotkey2, &new_hotkey2); + + assert_eq!( + PendingRootAlpha::::get(old_hotkey2), + 0_u128, + "Old hotkey2 should have zero pending" + ); + assert_eq!( + PendingRootAlpha::::get(new_hotkey2), + 100_000u128, + "New hotkey2 should receive all pending" + ); + }); +} + +#[test] +fn test_finalize_all_subnet_root_dividends_decrements_pending() { + new_test_ext(1).execute_with(|| { + let netuid = NetUid::from(1); + add_network(netuid, 1, 0); + + let hotkey1 = U256::from(3011); + let hotkey2 = U256::from(3012); + let coldkey1 = U256::from(3013); + let coldkey2 = U256::from(3014); + + // Setup root stake for both hotkeys + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey1, + &coldkey1, + NetUid::ROOT, + 1_000_000u64.into(), + ); + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey2, + &coldkey2, + NetUid::ROOT, + 2_000_000u64.into(), + ); + + // Setup RootClaimable + RootClaimable::::mutate(hotkey1, |claimable| { + claimable.insert(netuid, I96F32::from_num(0.1)); + }); + RootClaimable::::mutate(hotkey2, |claimable| { + claimable.insert(netuid, I96F32::from_num(0.15)); + }); + + // Set RootClaimed + RootClaimed::::insert((netuid, hotkey1, coldkey1), 50_000u128); + RootClaimed::::insert((netuid, hotkey2, coldkey2), 100_000u128); + + // Initialize PendingRootAlpha + // hotkey1: claimable = 0.1 * 1M = 100k, claimed = 50k, pending = 50k + // hotkey2: claimable = 0.15 * 2M = 300k, claimed = 100k, pending = 200k + PendingRootAlpha::::insert(hotkey1, 50_000u128); + PendingRootAlpha::::insert(hotkey2, 200_000u128); + + // Finalize subnet root dividends + SubtensorModule::finalize_all_subnet_root_dividends(netuid); + + // PendingRootAlpha should be decremented by the unclaimed amount + // hotkey1: pending = 50k - 50k = 0 + // hotkey2: pending = 200k - 200k = 0 (may have small rounding differences) + let pending1 = PendingRootAlpha::::get(hotkey1); + let pending2 = PendingRootAlpha::::get(hotkey2); + assert!( + pending1 <= 1, + "Hotkey1 pending should be decremented to zero (or very close). Got: {}", + pending1 + ); + assert!( + pending2 <= 1, + "Hotkey2 pending should be decremented to zero (or very close). Got: {}", + pending2 + ); + + // RootClaimable should have netuid removed + assert!( + !RootClaimable::::get(hotkey1).contains_key(&netuid), + "RootClaimable should not contain netuid after finalize" + ); + assert!( + !RootClaimable::::get(hotkey2).contains_key(&netuid), + "RootClaimable should not contain netuid after finalize" + ); + }); +}