diff --git a/pallets/subtensor/src/swap/swap_hotkey.rs b/pallets/subtensor/src/swap/swap_hotkey.rs index 0c7eff911..b92cbc06a 100644 --- a/pallets/subtensor/src/swap/swap_hotkey.rs +++ b/pallets/subtensor/src/swap/swap_hotkey.rs @@ -303,6 +303,36 @@ impl Pallet { } }); + // 11. Swap Alpha + // Alpha( hotkey, coldkey, netuid ) -> alpha + let old_alpha_values: Vec<((T::AccountId, u16), U64F64)> = + Alpha::::iter_prefix((old_hotkey,)).collect(); + // Clear the entire old prefix here. + let _ = Alpha::::clear_prefix((old_hotkey,), old_alpha_values.len() as u32, None); + weight.saturating_accrue(T::DbWeight::get().reads(old_alpha_values.len() as u64)); + weight.saturating_accrue(T::DbWeight::get().writes(old_alpha_values.len() as u64)); + + // Insert the new alpha values. + for ((coldkey, netuid), alpha) in old_alpha_values { + let new_alpha = Alpha::::get((new_hotkey, &coldkey, netuid)); + Alpha::::insert( + (new_hotkey, &coldkey, netuid), + new_alpha.saturating_add(alpha), + ); + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1)); + + // Swap StakingHotkeys. + // StakingHotkeys( coldkey ) --> Vec -- the hotkeys that the coldkey stakes. + let mut staking_hotkeys = StakingHotkeys::::get(&coldkey); + weight.saturating_accrue(T::DbWeight::get().reads(1)); + if staking_hotkeys.contains(&old_hotkey) { + staking_hotkeys.retain(|hk| *hk != *old_hotkey && *hk != *new_hotkey); + staking_hotkeys.push(new_hotkey.clone()); + StakingHotkeys::::insert(&coldkey, staking_hotkeys); + weight.saturating_accrue(T::DbWeight::get().writes(1)); + } + } + // 11. Swap Stake. // Stake( hotkey, coldkey ) -> stake -- the stake that the hotkey controls on behalf of the coldkey. let stakes: Vec<(T::AccountId, u64)> = Stake::::iter_prefix(old_hotkey).collect(); @@ -322,30 +352,7 @@ impl Pallet { &coldkey, new_stake_value.saturating_add(old_stake_amount), ); - weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1)); - - // 3.1 Swap Alpha - for netuid in Self::get_all_subnet_netuids() { - // Get the stake on the old (hot,coldkey) account. - let old_alpha: U64F64 = Alpha::::get((&old_hotkey, coldkey.clone(), netuid)); - // Get the stake on the new (hot,coldkey) account. - let new_alpha: U64F64 = Alpha::::get((&new_hotkey, coldkey.clone(), netuid)); - // Add the stake to new account. - Alpha::::insert( - (&new_hotkey, coldkey.clone(), netuid), - new_alpha.saturating_add(old_alpha), - ); - // Remove the value from the old account. - Alpha::::remove((&old_hotkey, coldkey.clone(), netuid)); - } - - // Swap StakingHotkeys. - // StakingHotkeys( coldkey ) --> Vec -- the hotkeys that the coldkey stakes. - let mut staking_hotkeys = StakingHotkeys::::get(coldkey.clone()); - staking_hotkeys.retain(|hk| *hk != *old_hotkey && *hk != *new_hotkey); - staking_hotkeys.push(new_hotkey.clone()); - StakingHotkeys::::insert(coldkey.clone(), staking_hotkeys); - weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1)); + weight.saturating_accrue(T::DbWeight::get().writes(1)); } // 12. Swap ChildKeys. diff --git a/pallets/subtensor/src/tests/swap_hotkey.rs b/pallets/subtensor/src/tests/swap_hotkey.rs index ddedc3fc1..1c03b838c 100644 --- a/pallets/subtensor/src/tests/swap_hotkey.rs +++ b/pallets/subtensor/src/tests/swap_hotkey.rs @@ -613,56 +613,132 @@ fn test_swap_hotkey_with_no_stake() { }); } -// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_hotkey_with_multiple_coldkeys_and_subnets --exact --nocapture +// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --package pallet-subtensor --lib -- tests::swap_hotkey::test_swap_hotkey_with_multiple_coldkeys_and_subnets --exact --show-output #[test] fn test_swap_hotkey_with_multiple_coldkeys_and_subnets() { new_test_ext(1).execute_with(|| { - assert!(false); - - // let old_hotkey = U256::from(1); - // let new_hotkey = U256::from(2); - // let coldkey1 = U256::from(3); - // let coldkey2 = U256::from(4); - // let netuid1 = 0; - // let netuid2 = 1; - // let mut weight = Weight::zero(); - - // // Set up initial state - // add_network(netuid1, 0, 1); - // add_network(netuid2, 0, 1); - // Owner::::insert(old_hotkey, coldkey1); - // Stake::::insert(old_hotkey, coldkey1, 100); - // Stake::::insert(old_hotkey, coldkey2, 200); - // IsNetworkMember::::insert(old_hotkey, netuid1, true); - // IsNetworkMember::::insert(old_hotkey, netuid2, true); - // TotalHotkeyStake::::insert(old_hotkey, 300); - - // assert_ok!(SubtensorModule::perform_hotkey_swap( - // &old_hotkey, - // &new_hotkey, - // &coldkey1, - // &mut weight - // )); - - // // Check ownership transfer - // assert!(!Owner::::contains_key(old_hotkey)); - // assert_eq!(Owner::::get(new_hotkey), coldkey1); - - // // Check stake transfer - // assert_eq!(Stake::::get(new_hotkey, coldkey1), 100); - // assert_eq!(Stake::::get(new_hotkey, coldkey2), 200); - // assert!(!Stake::::contains_key(old_hotkey, coldkey1)); - // assert!(!Stake::::contains_key(old_hotkey, coldkey2)); - - // // Check subnet membership transfer - // assert!(IsNetworkMember::::get(new_hotkey, netuid1)); - // assert!(IsNetworkMember::::get(new_hotkey, netuid2)); - // assert!(!IsNetworkMember::::get(old_hotkey, netuid1)); - // assert!(!IsNetworkMember::::get(old_hotkey, netuid2)); - - // // Check total stake transfer - // assert_eq!(TotalHotkeyStake::::get(new_hotkey), 300); - // assert!(!TotalHotkeyStake::::contains_key(old_hotkey)); + let old_hotkey = U256::from(1); + let new_hotkey = U256::from(2); + let coldkey1 = U256::from(3); + let coldkey2 = U256::from(4); + let netuid1 = 1; + let netuid2 = 2; + let stake = 1_000_000_u64; + let mut weight = Weight::zero(); + + // Set up initial state + add_network(netuid1, 0, 1); + add_network(netuid2, 0, 1); + register_ok_neuron(netuid1, old_hotkey, coldkey1, 1234); + register_ok_neuron(netuid2, old_hotkey, coldkey1, 1234); + + // Add balance to both coldkeys + SubtensorModule::add_balance_to_coldkey_account(&coldkey1, stake + 1_000); + SubtensorModule::add_balance_to_coldkey_account(&coldkey2, stake + 1_000); + + // Stake with coldkey1 + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(coldkey1), + old_hotkey, + netuid1, + stake + )); + + // Stake with coldkey2 also + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(coldkey2), + old_hotkey, + netuid2, + stake + )); + + let ck1_stake = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &old_hotkey, + &coldkey1, + netuid1, + ); + let ck2_stake = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &old_hotkey, + &coldkey2, + netuid2, + ); + assert!(ck1_stake > 0); + assert!(ck2_stake > 0); + let total_hk_stake = SubtensorModule::get_total_stake_for_hotkey(&old_hotkey); + assert!(total_hk_stake > 0); + + assert_ok!(SubtensorModule::perform_hotkey_swap( + &old_hotkey, + &new_hotkey, + &coldkey1, + &mut weight + )); + + // Check ownership transfer + assert_eq!( + SubtensorModule::get_owning_coldkey_for_hotkey(&new_hotkey), + coldkey1 + ); + assert!(!SubtensorModule::get_owned_hotkeys(&coldkey2).contains(&new_hotkey)); + + // Check stake transfer + assert_eq!( + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &new_hotkey, + &coldkey1, + netuid1 + ), + ck1_stake + ); + assert_eq!( + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &new_hotkey, + &coldkey2, + netuid2 + ), + ck2_stake + ); + assert_eq!( + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &old_hotkey, + &coldkey1, + netuid1 + ), + 0 + ); + assert_eq!( + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &old_hotkey, + &coldkey2, + netuid2 + ), + 0 + ); + + // Check subnet membership transfer + assert!(SubtensorModule::is_hotkey_registered_on_network( + netuid1, + &new_hotkey + )); + assert!(SubtensorModule::is_hotkey_registered_on_network( + netuid2, + &new_hotkey + )); + assert!(!SubtensorModule::is_hotkey_registered_on_network( + netuid1, + &old_hotkey + )); + assert!(!SubtensorModule::is_hotkey_registered_on_network( + netuid2, + &old_hotkey + )); + + // Check total stake transfer + assert_eq!( + SubtensorModule::get_total_stake_for_hotkey(&new_hotkey), + total_hk_stake + ); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&old_hotkey), 0); }); }