From ac74347d616f0407c0146cef0d8940379bf83a80 Mon Sep 17 00:00:00 2001 From: Samuel Dare Date: Wed, 3 Jul 2024 22:11:47 +0400 Subject: [PATCH 01/14] feat: cold key swap --- pallets/subtensor/src/errors.rs | 10 + pallets/subtensor/src/events.rs | 7 + pallets/subtensor/src/swap.rs | 285 +++++++++++++++++++++++ pallets/subtensor/tests/swap.rs | 394 ++++++++++++++++++++++++++++++++ scripts/test_specific.sh | 4 +- 5 files changed, 699 insertions(+), 1 deletion(-) diff --git a/pallets/subtensor/src/errors.rs b/pallets/subtensor/src/errors.rs index 57ba91d31..9dd01fda4 100644 --- a/pallets/subtensor/src/errors.rs +++ b/pallets/subtensor/src/errors.rs @@ -132,5 +132,15 @@ mod errors { AlphaHighTooLow, /// Alpha low is out of range: alpha_low > 0 && alpha_low < 0.8 AlphaLowOutOfRange, + /// The coldkey has already been swapped + ColdKeyAlreadyAssociated, + /// The coldkey swap transaction rate limit exceeded + ColdKeySwapTxRateLimitExceeded, + /// The new coldkey is the same as the old coldkey + NewColdKeyIsSameWithOld, + /// The coldkey does not exist + NotExistColdkey, + /// The coldkey balance is not enough to pay for the swap + NotEnoughBalanceToPaySwapColdKey, } } diff --git a/pallets/subtensor/src/events.rs b/pallets/subtensor/src/events.rs index 47cc9973b..73bda500b 100644 --- a/pallets/subtensor/src/events.rs +++ b/pallets/subtensor/src/events.rs @@ -132,5 +132,12 @@ mod events { MinDelegateTakeSet(u16), /// the target stakes per interval is set by sudo/admin transaction TargetStakesPerIntervalSet(u64), + /// A coldkey has been swapped + ColdkeySwapped { + /// the account ID of old coldkey + old_coldkey: T::AccountId, + /// the account ID of new coldkey + new_coldkey: T::AccountId, + }, } } diff --git a/pallets/subtensor/src/swap.rs b/pallets/subtensor/src/swap.rs index ce090d736..fd934fa08 100644 --- a/pallets/subtensor/src/swap.rs +++ b/pallets/subtensor/src/swap.rs @@ -92,6 +92,102 @@ impl Pallet { Ok(Some(weight).into()) } + /// Swaps the coldkey associated with a set of hotkeys from an old coldkey to a new coldkey. + /// + /// # Arguments + /// + /// * `origin` - The origin of the call, which must be signed by the old coldkey. + /// * `old_coldkey` - The account ID of the old coldkey. + /// * `new_coldkey` - The account ID of the new coldkey. + /// + /// # Returns + /// + /// Returns a `DispatchResultWithPostInfo` indicating success or failure, along with the weight consumed. + /// + /// # Errors + /// + /// This function will return an error if: + /// - The caller is not the old coldkey. + /// - The new coldkey is the same as the old coldkey. + /// - The new coldkey is already associated with other hotkeys. + /// - The transaction rate limit for coldkey swaps has been exceeded. + /// - There's not enough balance to pay for the swap. + /// + /// # Events + /// + /// Emits a `ColdkeySwapped` event when successful. + /// + /// # Weight + /// + /// Weight is tracked and updated throughout the function execution. + pub fn do_swap_coldkey( + origin: T::RuntimeOrigin, + old_coldkey: &T::AccountId, + new_coldkey: &T::AccountId, + ) -> DispatchResultWithPostInfo { + let caller = ensure_signed(origin)?; + + // Ensure the caller is the old coldkey + ensure!(caller == *old_coldkey, Error::::NonAssociatedColdKey); + + let mut weight = T::DbWeight::get().reads(2); + + ensure!( + old_coldkey != new_coldkey, + Error::::NewColdKeyIsSameWithOld + ); + + // Check if the new coldkey is already associated with any hotkeys + ensure!( + !Self::coldkey_has_associated_hotkeys(new_coldkey), + Error::::ColdKeyAlreadyAssociated + ); + + let block: u64 = Self::get_current_block_as_u64(); + ensure!( + !Self::exceeds_tx_rate_limit(Self::get_last_tx_block(old_coldkey), block), + Error::::ColdKeySwapTxRateLimitExceeded + ); + + // Note: we probably want to make this free + let swap_cost = Self::get_coldkey_swap_cost(); + ensure!( + Self::can_remove_balance_from_coldkey_account(old_coldkey, swap_cost), + Error::::NotEnoughBalanceToPaySwapColdKey + ); + let actual_burn_amount = Self::remove_balance_from_coldkey_account(old_coldkey, swap_cost)?; + Self::burn_tokens(actual_burn_amount); + + // Swap coldkey references in storage maps + Self::swap_total_coldkey_stake(old_coldkey, new_coldkey, &mut weight); + Self::swap_stake_for_coldkey(old_coldkey, new_coldkey, &mut weight); + Self::swap_owner_for_coldkey(old_coldkey, new_coldkey, &mut weight); + Self::swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey( + old_coldkey, + new_coldkey, + &mut weight, + ); + Self::swap_keys_for_coldkey(old_coldkey, new_coldkey, &mut weight); + Self::swap_subnet_owner_for_coldkey(old_coldkey, new_coldkey, &mut weight); + + // Transfer any remaining balance from old_coldkey to new_coldkey + let remaining_balance = Self::get_coldkey_balance(old_coldkey); + if remaining_balance > 0 { + Self::remove_balance_from_coldkey_account(old_coldkey, remaining_balance)?; + Self::add_balance_to_coldkey_account(new_coldkey, remaining_balance); + } + + Self::set_last_tx_block(new_coldkey, block); + weight.saturating_accrue(T::DbWeight::get().writes(1)); + + Self::deposit_event(Event::ColdkeySwapped { + old_coldkey: old_coldkey.clone(), + new_coldkey: new_coldkey.clone(), + }); + + Ok(Some(weight).into()) + } + /// Retrieves the network membership status for a given hotkey. /// /// # Arguments @@ -396,4 +492,193 @@ impl Pallet { weight.saturating_accrue(T::DbWeight::get().writes(2)); // One write for insert and one for remove } } + + /// Swaps the total stake associated with a coldkey from the old coldkey to the new coldkey. + /// + /// # Arguments + /// + /// * `old_coldkey` - The AccountId of the old coldkey. + /// * `new_coldkey` - The AccountId of the new coldkey. + /// * `weight` - Mutable reference to the weight of the transaction. + /// + /// # Effects + /// + /// * Removes the total stake from the old coldkey. + /// * Inserts the total stake for the new coldkey. + /// * Updates the transaction weight. + pub fn swap_total_coldkey_stake( + old_coldkey: &T::AccountId, + new_coldkey: &T::AccountId, + weight: &mut Weight, + ) { + let stake = TotalColdkeyStake::::get(old_coldkey); + TotalColdkeyStake::::remove(old_coldkey); + TotalColdkeyStake::::insert(new_coldkey, stake); + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2)); + } + + /// Swaps all stakes associated with a coldkey from the old coldkey to the new coldkey. + /// + /// # Arguments + /// + /// * `old_coldkey` - The AccountId of the old coldkey. + /// * `new_coldkey` - The AccountId of the new coldkey. + /// * `weight` - Mutable reference to the weight of the transaction. + /// + /// # Effects + /// + /// * Removes all stakes associated with the old coldkey. + /// * Inserts all stakes for the new coldkey. + /// * Updates the transaction weight. + pub fn swap_stake_for_coldkey( + old_coldkey: &T::AccountId, + new_coldkey: &T::AccountId, + weight: &mut Weight, + ) { + for (hotkey, stake) in Stake::::iter_prefix(old_coldkey) { + Stake::::remove(old_coldkey, &hotkey); + Stake::::insert(new_coldkey, &hotkey, stake); + } + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2)); + } + + /// Swaps the owner of all hotkeys from the old coldkey to the new coldkey. + /// + /// # Arguments + /// + /// * `old_coldkey` - The AccountId of the old coldkey. + /// * `new_coldkey` - The AccountId of the new coldkey. + /// * `weight` - Mutable reference to the weight of the transaction. + /// + /// # Effects + /// + /// * Updates the owner of all hotkeys associated with the old coldkey to the new coldkey. + /// * Updates the transaction weight. + pub fn swap_owner_for_coldkey( + old_coldkey: &T::AccountId, + new_coldkey: &T::AccountId, + weight: &mut Weight, + ) { + for (hotkey, _) in Owner::::iter() { + if Owner::::get(&hotkey) == *old_coldkey { + Owner::::insert(&hotkey, new_coldkey); + } + } + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1)); + } + + /// Swaps the total hotkey-coldkey stakes for the current interval from the old coldkey to the new coldkey. + /// + /// # Arguments + /// + /// * `old_coldkey` - The AccountId of the old coldkey. + /// * `new_coldkey` - The AccountId of the new coldkey. + /// * `weight` - Mutable reference to the weight of the transaction. + /// + /// # Effects + /// + /// * Removes all total hotkey-coldkey stakes for the current interval associated with the old coldkey. + /// * Inserts all total hotkey-coldkey stakes for the current interval for the new coldkey. + /// * Updates the transaction weight. + pub fn swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey( + old_coldkey: &T::AccountId, + new_coldkey: &T::AccountId, + weight: &mut Weight, + ) { + for (hotkey, (stakes, block)) in + TotalHotkeyColdkeyStakesThisInterval::::iter_prefix(old_coldkey) + { + TotalHotkeyColdkeyStakesThisInterval::::remove(old_coldkey, &hotkey); + TotalHotkeyColdkeyStakesThisInterval::::insert( + new_coldkey, + &hotkey, + (stakes, block), + ); + } + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2)); + } + + /// Swaps all keys associated with a coldkey from the old coldkey to the new coldkey across all networks. + /// + /// # Arguments + /// + /// * `old_coldkey` - The AccountId of the old coldkey. + /// * `new_coldkey` - The AccountId of the new coldkey. + /// * `weight` - Mutable reference to the weight of the transaction. + /// + /// # Effects + /// + /// * Updates all keys associated with the old coldkey to be associated with the new coldkey. + /// * Updates the transaction weight. + pub fn swap_keys_for_coldkey( + old_coldkey: &T::AccountId, + new_coldkey: &T::AccountId, + weight: &mut Weight, + ) { + for netuid in 0..TotalNetworks::::get() { + for (uid, hotkey) in Keys::::iter_prefix(netuid) { + if Owner::::get(&hotkey) == *old_coldkey { + Keys::::remove(netuid, uid); + Keys::::insert(netuid, uid, new_coldkey); + } + } + } + weight.saturating_accrue(T::DbWeight::get().reads_writes( + TotalNetworks::::get() as u64, + TotalNetworks::::get() as u64, + )); + } + + /// Checks if a coldkey has any associated hotkeys. + /// + /// # Arguments + /// + /// * `coldkey` - The AccountId of the coldkey to check. + /// + /// # Returns + /// + /// * `bool` - True if the coldkey has any associated hotkeys, false otherwise. + pub fn coldkey_has_associated_hotkeys(coldkey: &T::AccountId) -> bool { + Owner::::iter().any(|(_, owner)| owner == *coldkey) + } + + /// Swaps the subnet owner from the old coldkey to the new coldkey for all networks where the old coldkey is the owner. + /// + /// # Arguments + /// + /// * `old_coldkey` - The AccountId of the old coldkey. + /// * `new_coldkey` - The AccountId of the new coldkey. + /// * `weight` - Mutable reference to the weight of the transaction. + /// + /// # Effects + /// + /// * Updates the subnet owner to the new coldkey for all networks where the old coldkey was the owner. + /// * Updates the transaction weight. + pub fn swap_subnet_owner_for_coldkey( + old_coldkey: &T::AccountId, + new_coldkey: &T::AccountId, + weight: &mut Weight, + ) { + for netuid in 0..TotalNetworks::::get() { + let subnet_owner = SubnetOwner::::get(netuid); + if subnet_owner == *old_coldkey { + SubnetOwner::::insert(netuid, new_coldkey.clone()); + weight.saturating_accrue(T::DbWeight::get().writes(1)); + } + } + weight.saturating_accrue(T::DbWeight::get().reads(TotalNetworks::::get() as u64)); + } + + /// Returns the cost of swapping a coldkey. + /// + /// # Returns + /// + /// * `u64` - The cost of swapping a coldkey in Rao. + /// + /// # Note + /// + /// This function returns a hardcoded value. In a production environment, this should be configurable or determined dynamically. + pub fn get_coldkey_swap_cost() -> u64 { + 1_000_000 // Example cost in Rao + } } diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index af7d19d2d..47f4fa401 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1048,3 +1048,397 @@ fn test_swap_total_hotkey_coldkey_stakes_this_interval_weight_update() { assert_eq!(weight, expected_weight); }); } + +#[test] +fn test_do_swap_coldkey_success() { + new_test_ext(1).execute_with(|| { + let old_coldkey = U256::from(1); + let new_coldkey = U256::from(2); + let hotkey = U256::from(3); + let netuid = 1u16; + let stake_amount = 1000u64; + let swap_cost = SubtensorModule::get_coldkey_swap_cost(); + + // Setup initial state + add_network(netuid, 13, 0); + register_ok_neuron(netuid, hotkey, old_coldkey, 0); + SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + swap_cost); + + // Add stake to the neuron + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(old_coldkey), + hotkey, + stake_amount + )); + + log::info!("TotalColdkeyStake::::get(old_coldkey): {:?}", TotalColdkeyStake::::get(old_coldkey)); + log::info!("Stake::::get(old_coldkey, hotkey): {:?}", Stake::::get(old_coldkey, hotkey)); + + // Verify initial stake + assert_eq!(TotalColdkeyStake::::get(old_coldkey), stake_amount); + // assert_eq!(Stake::::get(old_coldkey, hotkey), stake_amount); + + + // Perform the swap + assert_ok!(SubtensorModule::do_swap_coldkey( + <::RuntimeOrigin>::signed(old_coldkey), + &old_coldkey, + &new_coldkey + )); + + // Verify the swap + assert_eq!(Owner::::get(hotkey), new_coldkey); + assert_eq!(TotalColdkeyStake::::get(new_coldkey), stake_amount); + assert!(!TotalColdkeyStake::::contains_key(old_coldkey)); + assert_eq!(Stake::::get(new_coldkey, hotkey), stake_amount); + assert!(!Stake::::contains_key(old_coldkey, hotkey)); + + // Verify balance transfer + assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), stake_amount); + assert_eq!(SubtensorModule::get_coldkey_balance(&old_coldkey), 0); + + + + // Verify event emission + System::assert_last_event(Event::ColdkeySwapped { + old_coldkey: old_coldkey, + new_coldkey: new_coldkey, + }.into()); + }); +} + +#[test] +fn test_do_swap_coldkey_not_enough_balance() { + new_test_ext(1).execute_with(|| { + let old_coldkey = U256::from(1); + let new_coldkey = U256::from(2); + let swap_cost = SubtensorModule::get_coldkey_swap_cost(); + + // Setup initial state with insufficient balance + SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, swap_cost - 1); + + // Attempt the swap + assert_err!( + SubtensorModule::do_swap_coldkey( + <::RuntimeOrigin>::signed(old_coldkey), + &old_coldkey, + &new_coldkey + ), + Error::::NotEnoughBalanceToPaySwapColdKey + ); + }); +} + +#[test] +fn test_do_swap_coldkey_same_keys() { + new_test_ext(1).execute_with(|| { + let coldkey = U256::from(1); + + // Attempt the swap with same old and new coldkeys + assert_err!( + SubtensorModule::do_swap_coldkey( + <::RuntimeOrigin>::signed(coldkey), + &coldkey, + &coldkey + ), + Error::::NewColdKeyIsSameWithOld + ); + }); +} + +#[test] +fn test_do_swap_coldkey_new_key_already_associated() { + new_test_ext(1).execute_with(|| { + let old_coldkey = U256::from(1); + let new_coldkey = U256::from(2); + let hotkey1 = U256::from(3); + let hotkey2 = U256::from(4); + let netuid = 1u16; + let swap_cost = SubtensorModule::get_coldkey_swap_cost(); + + // Setup initial state + add_network(netuid, 13, 0); + register_ok_neuron(netuid, hotkey1, old_coldkey, 0); + register_ok_neuron(netuid, hotkey2, new_coldkey, 0); + SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, swap_cost); + + // Attempt the swap + assert_err!( + SubtensorModule::do_swap_coldkey( + <::RuntimeOrigin>::signed(old_coldkey), + &old_coldkey, + &new_coldkey + ), + Error::::ColdKeyAlreadyAssociated + ); + }); +} + +#[test] +fn test_swap_total_coldkey_stake() { + new_test_ext(1).execute_with(|| { + let old_coldkey = U256::from(1); + let new_coldkey = U256::from(2); + let stake_amount = 1000u64; + let mut weight = Weight::zero(); + + // Initialize TotalColdkeyStake for old_coldkey + TotalColdkeyStake::::insert(old_coldkey, stake_amount); + + // Perform the swap + SubtensorModule::swap_total_coldkey_stake(&old_coldkey, &new_coldkey, &mut weight); + + // Verify the swap + assert_eq!(TotalColdkeyStake::::get(new_coldkey), stake_amount); + assert!(!TotalColdkeyStake::::contains_key(old_coldkey)); + + // Verify weight update + let expected_weight = ::DbWeight::get().reads_writes(1, 2); + assert_eq!(weight, expected_weight); + }); +} + +#[test] +fn test_swap_stake_for_coldkey() { + new_test_ext(1).execute_with(|| { + let old_coldkey = U256::from(1); + let new_coldkey = U256::from(2); + let hotkey1 = U256::from(3); + let hotkey2 = U256::from(4); + let stake_amount1 = 1000u64; + let stake_amount2 = 2000u64; + let mut weight = Weight::zero(); + + // Initialize Stake for old_coldkey + Stake::::insert(old_coldkey, hotkey1, stake_amount1); + Stake::::insert(old_coldkey, hotkey2, stake_amount2); + + // Perform the swap + SubtensorModule::swap_stake_for_coldkey(&old_coldkey, &new_coldkey, &mut weight); + + // Verify the swap + assert_eq!(Stake::::get(new_coldkey, hotkey1), stake_amount1); + assert_eq!(Stake::::get(new_coldkey, hotkey2), stake_amount2); + assert!(!Stake::::contains_key(old_coldkey, hotkey1)); + assert!(!Stake::::contains_key(old_coldkey, hotkey2)); + + // Verify weight update + let expected_weight = ::DbWeight::get().reads_writes(1, 4); + assert_eq!(weight, expected_weight); + }); +} + +#[test] +fn test_swap_owner_for_coldkey() { + new_test_ext(1).execute_with(|| { + let old_coldkey = U256::from(1); + let new_coldkey = U256::from(2); + let hotkey1 = U256::from(3); + let hotkey2 = U256::from(4); + let mut weight = Weight::zero(); + + // Initialize Owner for old_coldkey + Owner::::insert(hotkey1, old_coldkey); + Owner::::insert(hotkey2, old_coldkey); + + // Perform the swap + SubtensorModule::swap_owner_for_coldkey(&old_coldkey, &new_coldkey, &mut weight); + + // Verify the swap + assert_eq!(Owner::::get(hotkey1), new_coldkey); + assert_eq!(Owner::::get(hotkey2), new_coldkey); + + // Verify weight update + let expected_weight = ::DbWeight::get().reads_writes(1, 2); + assert_eq!(weight, expected_weight); + }); +} + +#[test] +fn test_swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey() { + new_test_ext(1).execute_with(|| { + let old_coldkey = U256::from(1); + let new_coldkey = U256::from(2); + let hotkey1 = U256::from(3); + let hotkey2 = U256::from(4); + let stake1 = (1000u64, 100u64); + let stake2 = (2000u64, 200u64); + let mut weight = Weight::zero(); + + // Initialize TotalHotkeyColdkeyStakesThisInterval for old_coldkey + TotalHotkeyColdkeyStakesThisInterval::::insert(old_coldkey, hotkey1, stake1); + TotalHotkeyColdkeyStakesThisInterval::::insert(old_coldkey, hotkey2, stake2); + + // Perform the swap + SubtensorModule::swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey( + &old_coldkey, + &new_coldkey, + &mut weight, + ); + + // Verify the swap + assert_eq!( + TotalHotkeyColdkeyStakesThisInterval::::get(new_coldkey, hotkey1), + stake1 + ); + assert_eq!( + TotalHotkeyColdkeyStakesThisInterval::::get(new_coldkey, hotkey2), + stake2 + ); + assert!(!TotalHotkeyColdkeyStakesThisInterval::::contains_key( + old_coldkey, + hotkey1 + )); + assert!(!TotalHotkeyColdkeyStakesThisInterval::::contains_key( + old_coldkey, + hotkey2 + )); + + // Verify weight update + let expected_weight = ::DbWeight::get().reads_writes(1, 4); + assert_eq!(weight, expected_weight); + }); +} + +#[test] +fn test_swap_keys_for_coldkey() { + new_test_ext(1).execute_with(|| { + let old_coldkey = U256::from(1); + let new_coldkey = U256::from(2); + let hotkey1 = U256::from(3); + let hotkey2 = U256::from(4); + let netuid1 = 1u16; + let netuid2 = 2u16; + let uid1 = 10u16; + let uid2 = 20u16; + let mut weight = Weight::zero(); + + // Initialize Keys and Owner for old_coldkey + Keys::::insert(netuid1, uid1, hotkey1); + Keys::::insert(netuid2, uid2, hotkey2); + Owner::::insert(hotkey1, old_coldkey); + Owner::::insert(hotkey2, old_coldkey); + + // Set up TotalNetworks + TotalNetworks::::put(3); + + // Perform the swap + SubtensorModule::swap_keys_for_coldkey(&old_coldkey, &new_coldkey, &mut weight); + + // Verify the swap + assert_eq!(Keys::::get(netuid1, uid1), hotkey1); + assert_eq!(Keys::::get(netuid2, uid2), hotkey2); + + // Verify weight update + let expected_weight = ::DbWeight::get().reads_writes(3, 2); + assert_eq!(weight, expected_weight); + }); +} + +#[test] +fn test_swap_subnet_owner_for_coldkey() { + new_test_ext(1).execute_with(|| { + let old_coldkey = U256::from(1); + let new_coldkey = U256::from(2); + let netuid1 = 1u16; + let netuid2 = 2u16; + let mut weight = Weight::zero(); + + // Initialize SubnetOwner for old_coldkey + SubnetOwner::::insert(netuid1, old_coldkey); + SubnetOwner::::insert(netuid2, old_coldkey); + + // Set up TotalNetworks + TotalNetworks::::put(3); + + // Perform the swap + SubtensorModule::swap_subnet_owner_for_coldkey(&old_coldkey, &new_coldkey, &mut weight); + + // Verify the swap + assert_eq!(SubnetOwner::::get(netuid1), new_coldkey); + assert_eq!(SubnetOwner::::get(netuid2), new_coldkey); + + // Verify weight update + let expected_weight = ::DbWeight::get().reads_writes(3, 2); + assert_eq!(weight, expected_weight); + }); +} + +#[test] +fn test_do_swap_coldkey_with_subnet_ownership() { + new_test_ext(1).execute_with(|| { + let old_coldkey = U256::from(1); + let new_coldkey = U256::from(2); + let hotkey = U256::from(3); + let netuid = 1u16; + let stake_amount = 1000u64; + let swap_cost = SubtensorModule::get_coldkey_swap_cost(); + + // Setup initial state + add_network(netuid, 13, 0); + register_ok_neuron(netuid, hotkey, old_coldkey, 0); + SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + swap_cost); + SubnetOwner::::insert(netuid, old_coldkey); + + // Perform the swap + assert_ok!(SubtensorModule::do_swap_coldkey( + <::RuntimeOrigin>::signed(old_coldkey), + &old_coldkey, + &new_coldkey + )); + + // Verify subnet ownership transfer + assert_eq!(SubnetOwner::::get(netuid), new_coldkey); + }); +} + +#[test] +fn test_do_swap_coldkey_tx_rate_limit() { + new_test_ext(1).execute_with(|| { + let old_coldkey = U256::from(1); + let new_coldkey = U256::from(2); + let swap_cost = SubtensorModule::get_coldkey_swap_cost(); + + // Setup initial state + SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, swap_cost * 2); + + // Perform first swap + assert_ok!(SubtensorModule::do_swap_coldkey( + <::RuntimeOrigin>::signed(old_coldkey), + &old_coldkey, + &new_coldkey + )); + + // Attempt second swap immediately + assert_err!( + SubtensorModule::do_swap_coldkey( + <::RuntimeOrigin>::signed(new_coldkey), + &new_coldkey, + &old_coldkey + ), + Error::::ColdKeySwapTxRateLimitExceeded + ); + }); +} + +#[test] +fn test_coldkey_has_associated_hotkeys() { + new_test_ext(1).execute_with(|| { + let coldkey = U256::from(1); + let hotkey = U256::from(2); + let netuid = 1u16; + + // Setup initial state + add_network(netuid, 13, 0); + register_ok_neuron(netuid, hotkey, coldkey, 0); + + // Check if coldkey has associated hotkeys + assert!(SubtensorModule::coldkey_has_associated_hotkeys(&coldkey)); + + // Check for a coldkey without associated hotkeys + let unassociated_coldkey = U256::from(3); + assert!(!SubtensorModule::coldkey_has_associated_hotkeys( + &unassociated_coldkey + )); + }); +} diff --git a/scripts/test_specific.sh b/scripts/test_specific.sh index 4e413c6d1..018872d33 100755 --- a/scripts/test_specific.sh +++ b/scripts/test_specific.sh @@ -1,4 +1,6 @@ pallet="${3:-pallet-subtensor}" features="${4:-pow-faucet}" -RUST_LOG="pallet_subtensor=trace,info" cargo test --release --features=$features -p $pallet --test $1 -- $2 --nocapture --exact \ No newline at end of file +# RUST_LOG="pallet_subtensor=info" cargo test --release --features=$features -p $pallet --test $1 -- $2 --nocapture --exact + +RUST_LOG=INFO cargo test --release --features=$features -p $pallet --test $1 -- $2 --nocapture --exact \ No newline at end of file From 86e9d7b729f363ce5004b3fa881cdca775df945b Mon Sep 17 00:00:00 2001 From: Greg Zaitsev Date: Wed, 3 Jul 2024 16:16:15 -0400 Subject: [PATCH 02/14] Add Owned map and migration to populate it --- pallets/subtensor/src/lib.rs | 28 +++++++++++++++++++- pallets/subtensor/src/migration.rs | 41 ++++++++++++++++++++++++++++++ pallets/subtensor/src/staking.rs | 8 ++++++ pallets/subtensor/src/swap.rs | 36 +++++++++++++++++++------- pallets/subtensor/tests/swap.rs | 17 ++++++++----- 5 files changed, 113 insertions(+), 17 deletions(-) diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index 64aefcfee..243079528 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -363,6 +363,9 @@ pub mod pallet { #[pallet::storage] // --- MAP ( hot ) --> cold | Returns the controlling coldkey for a hotkey. pub type Owner = StorageMap<_, Blake2_128Concat, T::AccountId, T::AccountId, ValueQuery, DefaultAccount>; + #[pallet::storage] // --- MAP ( cold ) --> Vec | Returns the vector of hotkeys controlled by this coldkey. + pub type Owned = + StorageMap<_, Blake2_128Concat, T::AccountId, Vec, ValueQuery>; #[pallet::storage] // --- MAP ( hot ) --> take | Returns the hotkey delegation take. And signals that this key is open for delegation. pub type Delegates = StorageMap<_, Blake2_128Concat, T::AccountId, u16, ValueQuery, DefaultDefaultTake>; @@ -1204,6 +1207,14 @@ pub mod pallet { // Fill stake information. Owner::::insert(hotkey.clone(), coldkey.clone()); + // Update Owned map + let mut hotkeys = Owned::::get(&coldkey); + hotkeys.push(hotkey.clone()); + Owned::::insert( + &coldkey, + hotkeys, + ); + TotalHotkeyStake::::insert(hotkey.clone(), stake); TotalColdkeyStake::::insert( coldkey.clone(), @@ -1325,7 +1336,9 @@ pub mod pallet { // Storage version v4 -> v5 .saturating_add(migration::migrate_delete_subnet_3::()) // Doesn't check storage version. TODO: Remove after upgrade - .saturating_add(migration::migration5_total_issuance::(false)); + .saturating_add(migration::migration5_total_issuance::(false)) + // Populate Owned map for coldkey swap. Doesn't update storage vesion. + .saturating_add(migration::migrate_populate_owned::()); weight } @@ -1970,6 +1983,19 @@ pub mod pallet { Self::do_swap_hotkey(origin, &hotkey, &new_hotkey) } + /// The extrinsic for user to change the coldkey + #[pallet::call_index(71)] + #[pallet::weight((Weight::from_parts(1_940_000_000, 0) + .saturating_add(T::DbWeight::get().reads(272)) + .saturating_add(T::DbWeight::get().writes(527)), DispatchClass::Operational, Pays::No))] + pub fn swap_coldkey( + origin: OriginFor, + old_coldkey: T::AccountId, + new_coldkey: T::AccountId, + ) -> DispatchResultWithPostInfo { + Self::do_swap_coldkey(origin, &old_coldkey, &new_coldkey) + } + // ---- SUDO ONLY FUNCTIONS ------------------------------------------------------------ // ================================== diff --git a/pallets/subtensor/src/migration.rs b/pallets/subtensor/src/migration.rs index 86fc02054..f5a5ee7db 100644 --- a/pallets/subtensor/src/migration.rs +++ b/pallets/subtensor/src/migration.rs @@ -477,3 +477,44 @@ pub fn migrate_to_v2_fixed_total_stake() -> Weight { Weight::zero() } } + +pub fn migrate_populate_owned() -> Weight { + // Setup migration weight + let mut weight = T::DbWeight::get().reads(1); + let migration_name = "Populate Owned map"; + + // Check if this migration is needed (if Owned map is empty) + let migrate = Owned::::iter().next().is_none(); + + // Only runs if we haven't already updated version past above new_storage_version. + if migrate { + info!(target: LOG_TARGET_1, ">>> Migration: {}", migration_name); + + let mut longest_hotkey_vector = 0; + let mut longest_coldkey: Option = None; + Owner::::iter().for_each(|(hotkey, coldkey)| { + let mut hotkeys = Owned::::get(&coldkey); + hotkeys.push(hotkey); + if longest_hotkey_vector < hotkeys.len() { + longest_hotkey_vector = hotkeys.len(); + longest_coldkey = Some(coldkey.clone()); + } + + Owned::::insert( + &coldkey, + hotkeys, + ); + + weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 1)); + }); + info!(target: LOG_TARGET_1, "Migration {} finished. Longest hotkey vector: {}", migration_name, longest_hotkey_vector); + if let Some(c) = longest_coldkey { + info!(target: LOG_TARGET_1, "Longest hotkey vector is controlled by: {}", c); + } + + weight + } else { + info!(target: LOG_TARGET_1, "Migration {} already done!", migration_name); + Weight::zero() + } +} diff --git a/pallets/subtensor/src/staking.rs b/pallets/subtensor/src/staking.rs index 7c3328396..0cce2dd35 100644 --- a/pallets/subtensor/src/staking.rs +++ b/pallets/subtensor/src/staking.rs @@ -560,6 +560,14 @@ impl Pallet { if !Self::hotkey_account_exists(hotkey) { Stake::::insert(hotkey, coldkey, 0); Owner::::insert(hotkey, coldkey); + + // Update Owned map + let mut hotkeys = Owned::::get(&coldkey); + hotkeys.push(hotkey.clone()); + Owned::::insert( + &coldkey, + hotkeys, + ); } } diff --git a/pallets/subtensor/src/swap.rs b/pallets/subtensor/src/swap.rs index fd934fa08..d01be5ee4 100644 --- a/pallets/subtensor/src/swap.rs +++ b/pallets/subtensor/src/swap.rs @@ -223,6 +223,15 @@ impl Pallet { ) { Owner::::remove(old_hotkey); Owner::::insert(new_hotkey, coldkey.clone()); + + // Update Owned map + let mut hotkeys = Owned::::get(&coldkey); + hotkeys.push(new_hotkey.clone()); + Owned::::insert( + &coldkey, + hotkeys, + ); + weight.saturating_accrue(T::DbWeight::get().writes(2)); } @@ -535,11 +544,15 @@ impl Pallet { new_coldkey: &T::AccountId, weight: &mut Weight, ) { - for (hotkey, stake) in Stake::::iter_prefix(old_coldkey) { - Stake::::remove(old_coldkey, &hotkey); - Stake::::insert(new_coldkey, &hotkey, stake); + // Find all hotkeys for this coldkey + let hotkeys = Owned::::get(old_coldkey); + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0)); + for hotkey in hotkeys.iter() { + let stake = Stake::::get(&hotkey, old_coldkey); + Stake::::remove(&hotkey, old_coldkey); + Stake::::insert(&hotkey, new_coldkey, stake); + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2)); } - weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2)); } /// Swaps the owner of all hotkeys from the old coldkey to the new coldkey. @@ -559,12 +572,17 @@ impl Pallet { new_coldkey: &T::AccountId, weight: &mut Weight, ) { - for (hotkey, _) in Owner::::iter() { - if Owner::::get(&hotkey) == *old_coldkey { - Owner::::insert(&hotkey, new_coldkey); - } + let hotkeys = Owned::::get(old_coldkey); + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0)); + for hotkey in hotkeys.iter() { + Owner::::insert(&hotkey, new_coldkey); + weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 1)); } - weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1)); + + // Update Owned map with new coldkey + Owned::::remove(old_coldkey); + Owned::::insert(new_coldkey, hotkeys); + weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 2)); } /// Swaps the total hotkey-coldkey stakes for the current interval from the old coldkey to the new coldkey. diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index 47f4fa401..5d3f5e6a8 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1210,20 +1210,23 @@ fn test_swap_stake_for_coldkey() { let mut weight = Weight::zero(); // Initialize Stake for old_coldkey - Stake::::insert(old_coldkey, hotkey1, stake_amount1); - Stake::::insert(old_coldkey, hotkey2, stake_amount2); + Stake::::insert(hotkey1, old_coldkey, stake_amount1); + Stake::::insert(hotkey2, old_coldkey, stake_amount2); + + // Populate Owned map + Owned::::insert(old_coldkey, vec![hotkey1, hotkey2]); // Perform the swap SubtensorModule::swap_stake_for_coldkey(&old_coldkey, &new_coldkey, &mut weight); // Verify the swap - assert_eq!(Stake::::get(new_coldkey, hotkey1), stake_amount1); - assert_eq!(Stake::::get(new_coldkey, hotkey2), stake_amount2); - assert!(!Stake::::contains_key(old_coldkey, hotkey1)); - assert!(!Stake::::contains_key(old_coldkey, hotkey2)); + assert_eq!(Stake::::get(hotkey1, new_coldkey), stake_amount1); + assert_eq!(Stake::::get(hotkey2, new_coldkey), stake_amount2); + assert!(!Stake::::contains_key(hotkey1, old_coldkey)); + assert!(!Stake::::contains_key(hotkey2, old_coldkey)); // Verify weight update - let expected_weight = ::DbWeight::get().reads_writes(1, 4); + let expected_weight = ::DbWeight::get().reads_writes(3, 4); assert_eq!(weight, expected_weight); }); } From 347b8de6dd9543a8a3cec0c154a18d1472e5ce2d Mon Sep 17 00:00:00 2001 From: Greg Zaitsev Date: Wed, 3 Jul 2024 16:23:30 -0400 Subject: [PATCH 03/14] Swapping coldkeys: unbreak the build --- pallets/subtensor/src/migration.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pallets/subtensor/src/migration.rs b/pallets/subtensor/src/migration.rs index f5a5ee7db..a24d0d03b 100644 --- a/pallets/subtensor/src/migration.rs +++ b/pallets/subtensor/src/migration.rs @@ -509,7 +509,7 @@ pub fn migrate_populate_owned() -> Weight { }); info!(target: LOG_TARGET_1, "Migration {} finished. Longest hotkey vector: {}", migration_name, longest_hotkey_vector); if let Some(c) = longest_coldkey { - info!(target: LOG_TARGET_1, "Longest hotkey vector is controlled by: {}", c); + info!(target: LOG_TARGET_1, "Longest hotkey vector is controlled by: {:?}", c); } weight From bda0c5a31b63d0916b32022cc14629d1442e8f60 Mon Sep 17 00:00:00 2001 From: Greg Zaitsev Date: Wed, 3 Jul 2024 19:16:40 -0400 Subject: [PATCH 04/14] Fix coldkey_swap and its tests --- pallets/subtensor/src/lib.rs | 12 ++-- pallets/subtensor/src/migration.rs | 10 ++-- pallets/subtensor/src/staking.rs | 37 ++++++++++-- pallets/subtensor/src/swap.rs | 85 +++++++++++++--------------- pallets/subtensor/tests/swap.rs | 91 ++++++++++++++---------------- 5 files changed, 125 insertions(+), 110 deletions(-) diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index 243079528..2fde2f293 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -1209,11 +1209,13 @@ pub mod pallet { // Update Owned map let mut hotkeys = Owned::::get(&coldkey); - hotkeys.push(hotkey.clone()); - Owned::::insert( - &coldkey, - hotkeys, - ); + if !hotkeys.contains(&hotkey) { + hotkeys.push(hotkey.clone()); + Owned::::insert( + &coldkey, + hotkeys, + ); + } TotalHotkeyStake::::insert(hotkey.clone(), stake); TotalColdkeyStake::::insert( diff --git a/pallets/subtensor/src/migration.rs b/pallets/subtensor/src/migration.rs index a24d0d03b..0873ef034 100644 --- a/pallets/subtensor/src/migration.rs +++ b/pallets/subtensor/src/migration.rs @@ -494,10 +494,12 @@ pub fn migrate_populate_owned() -> Weight { let mut longest_coldkey: Option = None; Owner::::iter().for_each(|(hotkey, coldkey)| { let mut hotkeys = Owned::::get(&coldkey); - hotkeys.push(hotkey); - if longest_hotkey_vector < hotkeys.len() { - longest_hotkey_vector = hotkeys.len(); - longest_coldkey = Some(coldkey.clone()); + if !hotkeys.contains(&hotkey) { + hotkeys.push(hotkey); + if longest_hotkey_vector < hotkeys.len() { + longest_hotkey_vector = hotkeys.len(); + longest_coldkey = Some(coldkey.clone()); + } } Owned::::insert( diff --git a/pallets/subtensor/src/staking.rs b/pallets/subtensor/src/staking.rs index 0cce2dd35..913a7208b 100644 --- a/pallets/subtensor/src/staking.rs +++ b/pallets/subtensor/src/staking.rs @@ -563,11 +563,13 @@ impl Pallet { // Update Owned map let mut hotkeys = Owned::::get(&coldkey); - hotkeys.push(hotkey.clone()); - Owned::::insert( - &coldkey, - hotkeys, - ); + if !hotkeys.contains(&hotkey) { + hotkeys.push(hotkey.clone()); + Owned::::insert( + &coldkey, + hotkeys, + ); + } } } @@ -789,6 +791,31 @@ impl Pallet { Ok(credit) } + pub fn kill_coldkey_account( + coldkey: &T::AccountId, + amount: <::Currency as fungible::Inspect<::AccountId>>::Balance, + ) -> Result { + if amount == 0 { + return Ok(0); + } + + let credit = T::Currency::withdraw( + coldkey, + amount, + Precision::Exact, + Preservation::Expendable, + Fortitude::Force, + ) + .map_err(|_| Error::::BalanceWithdrawalError)? + .peek(); + + if credit == 0 { + return Err(Error::::ZeroBalanceAfterWithdrawn.into()); + } + + Ok(credit) + } + pub fn unstake_all_coldkeys_from_hotkey_account(hotkey: &T::AccountId) { // Iterate through all coldkeys that have a stake on this hotkey account. for (delegate_coldkey_i, stake_i) in diff --git a/pallets/subtensor/src/swap.rs b/pallets/subtensor/src/swap.rs index d01be5ee4..dd06bc1fe 100644 --- a/pallets/subtensor/src/swap.rs +++ b/pallets/subtensor/src/swap.rs @@ -167,13 +167,13 @@ impl Pallet { new_coldkey, &mut weight, ); - Self::swap_keys_for_coldkey(old_coldkey, new_coldkey, &mut weight); Self::swap_subnet_owner_for_coldkey(old_coldkey, new_coldkey, &mut weight); + Self::swap_owned_for_coldkey(old_coldkey, new_coldkey, &mut weight); // Transfer any remaining balance from old_coldkey to new_coldkey let remaining_balance = Self::get_coldkey_balance(old_coldkey); if remaining_balance > 0 { - Self::remove_balance_from_coldkey_account(old_coldkey, remaining_balance)?; + Self::kill_coldkey_account(old_coldkey, remaining_balance)?; Self::add_balance_to_coldkey_account(new_coldkey, remaining_balance); } @@ -226,7 +226,10 @@ impl Pallet { // Update Owned map let mut hotkeys = Owned::::get(&coldkey); - hotkeys.push(new_hotkey.clone()); + if !hotkeys.contains(&new_hotkey) { + hotkeys.push(new_hotkey.clone()); + } + hotkeys.retain(|hk| *hk != *old_hotkey); Owned::::insert( &coldkey, hotkeys, @@ -578,11 +581,6 @@ impl Pallet { Owner::::insert(&hotkey, new_coldkey); weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 1)); } - - // Update Owned map with new coldkey - Owned::::remove(old_coldkey); - Owned::::insert(new_coldkey, hotkeys); - weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 2)); } /// Swaps the total hotkey-coldkey stakes for the current interval from the old coldkey to the new coldkey. @@ -603,48 +601,17 @@ impl Pallet { new_coldkey: &T::AccountId, weight: &mut Weight, ) { - for (hotkey, (stakes, block)) in - TotalHotkeyColdkeyStakesThisInterval::::iter_prefix(old_coldkey) - { - TotalHotkeyColdkeyStakesThisInterval::::remove(old_coldkey, &hotkey); + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0)); + for hotkey in Owned::::get(old_coldkey).iter() { + let (stake, block) = TotalHotkeyColdkeyStakesThisInterval::::get(&hotkey, old_coldkey); + TotalHotkeyColdkeyStakesThisInterval::::remove(&hotkey, old_coldkey); TotalHotkeyColdkeyStakesThisInterval::::insert( - new_coldkey, &hotkey, - (stakes, block), + new_coldkey, + (stake, block), ); + weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2)); } - weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2)); - } - - /// Swaps all keys associated with a coldkey from the old coldkey to the new coldkey across all networks. - /// - /// # Arguments - /// - /// * `old_coldkey` - The AccountId of the old coldkey. - /// * `new_coldkey` - The AccountId of the new coldkey. - /// * `weight` - Mutable reference to the weight of the transaction. - /// - /// # Effects - /// - /// * Updates all keys associated with the old coldkey to be associated with the new coldkey. - /// * Updates the transaction weight. - pub fn swap_keys_for_coldkey( - old_coldkey: &T::AccountId, - new_coldkey: &T::AccountId, - weight: &mut Weight, - ) { - for netuid in 0..TotalNetworks::::get() { - for (uid, hotkey) in Keys::::iter_prefix(netuid) { - if Owner::::get(&hotkey) == *old_coldkey { - Keys::::remove(netuid, uid); - Keys::::insert(netuid, uid, new_coldkey); - } - } - } - weight.saturating_accrue(T::DbWeight::get().reads_writes( - TotalNetworks::::get() as u64, - TotalNetworks::::get() as u64, - )); } /// Checks if a coldkey has any associated hotkeys. @@ -677,7 +644,7 @@ impl Pallet { new_coldkey: &T::AccountId, weight: &mut Weight, ) { - for netuid in 0..TotalNetworks::::get() { + for netuid in 0..=TotalNetworks::::get() { let subnet_owner = SubnetOwner::::get(netuid); if subnet_owner == *old_coldkey { SubnetOwner::::insert(netuid, new_coldkey.clone()); @@ -687,6 +654,30 @@ impl Pallet { weight.saturating_accrue(T::DbWeight::get().reads(TotalNetworks::::get() as u64)); } + /// Swaps the owned hotkeys for the coldkey + /// + /// # Arguments + /// + /// * `old_coldkey` - The AccountId of the old coldkey. + /// * `new_coldkey` - The AccountId of the new coldkey. + /// * `weight` - Mutable reference to the weight of the transaction. + /// + /// # Effects + /// + /// * Updates the subnet owner to the new coldkey for all networks where the old coldkey was the owner. + /// * Updates the transaction weight. + pub fn swap_owned_for_coldkey( + old_coldkey: &T::AccountId, + new_coldkey: &T::AccountId, + weight: &mut Weight, + ) { + // Update Owned map with new coldkey + let hotkeys = Owned::::get(&old_coldkey); + Owned::::remove(old_coldkey); + Owned::::insert(new_coldkey, hotkeys); + weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 2)); + } + /// Returns the cost of swapping a coldkey. /// /// # Returns diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index 5d3f5e6a8..db9c55823 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1058,11 +1058,12 @@ fn test_do_swap_coldkey_success() { let netuid = 1u16; let stake_amount = 1000u64; let swap_cost = SubtensorModule::get_coldkey_swap_cost(); + let free_balance = 12345; // Setup initial state add_network(netuid, 13, 0); register_ok_neuron(netuid, hotkey, old_coldkey, 0); - SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + swap_cost); + SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + swap_cost + free_balance); // Add stake to the neuron assert_ok!(SubtensorModule::add_stake( @@ -1072,12 +1073,22 @@ fn test_do_swap_coldkey_success() { )); log::info!("TotalColdkeyStake::::get(old_coldkey): {:?}", TotalColdkeyStake::::get(old_coldkey)); - log::info!("Stake::::get(old_coldkey, hotkey): {:?}", Stake::::get(old_coldkey, hotkey)); + log::info!("Stake::::get(old_coldkey, hotkey): {:?}", Stake::::get(hotkey, old_coldkey)); // Verify initial stake assert_eq!(TotalColdkeyStake::::get(old_coldkey), stake_amount); - // assert_eq!(Stake::::get(old_coldkey, hotkey), stake_amount); + assert_eq!(Stake::::get(hotkey, old_coldkey), stake_amount); + + + assert_eq!(Owned::::get(old_coldkey), vec![hotkey]); + assert!(!Owned::::contains_key(new_coldkey)); + + + + // Get coldkey free balance before swap + let balance = SubtensorModule::get_coldkey_balance(&old_coldkey); + assert_eq!(balance, free_balance + swap_cost); // Perform the swap assert_ok!(SubtensorModule::do_swap_coldkey( @@ -1086,19 +1097,19 @@ fn test_do_swap_coldkey_success() { &new_coldkey )); - // Verify the swap + // Verify the swap assert_eq!(Owner::::get(hotkey), new_coldkey); assert_eq!(TotalColdkeyStake::::get(new_coldkey), stake_amount); assert!(!TotalColdkeyStake::::contains_key(old_coldkey)); - assert_eq!(Stake::::get(new_coldkey, hotkey), stake_amount); - assert!(!Stake::::contains_key(old_coldkey, hotkey)); + assert_eq!(Stake::::get(hotkey, new_coldkey), stake_amount); + assert!(!Stake::::contains_key(hotkey, old_coldkey)); + assert_eq!(Owned::::get(new_coldkey), vec![hotkey]); + assert!(!Owned::::contains_key(old_coldkey)); // Verify balance transfer - assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), stake_amount); + assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), balance - swap_cost); assert_eq!(SubtensorModule::get_coldkey_balance(&old_coldkey), 0); - - // Verify event emission System::assert_last_event(Event::ColdkeySwapped { old_coldkey: old_coldkey, @@ -1244,6 +1255,9 @@ fn test_swap_owner_for_coldkey() { Owner::::insert(hotkey1, old_coldkey); Owner::::insert(hotkey2, old_coldkey); + // Initialize Owned map + Owned::::insert(old_coldkey, vec![hotkey1, hotkey2]); + // Perform the swap SubtensorModule::swap_owner_for_coldkey(&old_coldkey, &new_coldkey, &mut weight); @@ -1269,8 +1283,11 @@ fn test_swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey() { let mut weight = Weight::zero(); // Initialize TotalHotkeyColdkeyStakesThisInterval for old_coldkey - TotalHotkeyColdkeyStakesThisInterval::::insert(old_coldkey, hotkey1, stake1); - TotalHotkeyColdkeyStakesThisInterval::::insert(old_coldkey, hotkey2, stake2); + TotalHotkeyColdkeyStakesThisInterval::::insert(hotkey1, old_coldkey, stake1); + TotalHotkeyColdkeyStakesThisInterval::::insert(hotkey2, old_coldkey, stake2); + + // Populate Owned map + Owned::::insert(old_coldkey, vec![hotkey1, hotkey2]); // Perform the swap SubtensorModule::swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey( @@ -1281,11 +1298,11 @@ fn test_swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey() { // Verify the swap assert_eq!( - TotalHotkeyColdkeyStakesThisInterval::::get(new_coldkey, hotkey1), + TotalHotkeyColdkeyStakesThisInterval::::get(hotkey1, new_coldkey), stake1 ); assert_eq!( - TotalHotkeyColdkeyStakesThisInterval::::get(new_coldkey, hotkey2), + TotalHotkeyColdkeyStakesThisInterval::::get(hotkey2, new_coldkey), stake2 ); assert!(!TotalHotkeyColdkeyStakesThisInterval::::contains_key( @@ -1298,42 +1315,7 @@ fn test_swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey() { )); // Verify weight update - let expected_weight = ::DbWeight::get().reads_writes(1, 4); - assert_eq!(weight, expected_weight); - }); -} - -#[test] -fn test_swap_keys_for_coldkey() { - new_test_ext(1).execute_with(|| { - let old_coldkey = U256::from(1); - let new_coldkey = U256::from(2); - let hotkey1 = U256::from(3); - let hotkey2 = U256::from(4); - let netuid1 = 1u16; - let netuid2 = 2u16; - let uid1 = 10u16; - let uid2 = 20u16; - let mut weight = Weight::zero(); - - // Initialize Keys and Owner for old_coldkey - Keys::::insert(netuid1, uid1, hotkey1); - Keys::::insert(netuid2, uid2, hotkey2); - Owner::::insert(hotkey1, old_coldkey); - Owner::::insert(hotkey2, old_coldkey); - - // Set up TotalNetworks - TotalNetworks::::put(3); - - // Perform the swap - SubtensorModule::swap_keys_for_coldkey(&old_coldkey, &new_coldkey, &mut weight); - - // Verify the swap - assert_eq!(Keys::::get(netuid1, uid1), hotkey1); - assert_eq!(Keys::::get(netuid2, uid2), hotkey2); - - // Verify weight update - let expected_weight = ::DbWeight::get().reads_writes(3, 2); + let expected_weight = ::DbWeight::get().reads_writes(5, 4); assert_eq!(weight, expected_weight); }); } @@ -1380,9 +1362,16 @@ fn test_do_swap_coldkey_with_subnet_ownership() { // Setup initial state add_network(netuid, 13, 0); register_ok_neuron(netuid, hotkey, old_coldkey, 0); + + // Set TotalNetworks because swap relies on it + pallet_subtensor::TotalNetworks::::set(1); + SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + swap_cost); SubnetOwner::::insert(netuid, old_coldkey); + // Populate Owned map + Owned::::insert(old_coldkey, vec![hotkey]); + // Perform the swap assert_ok!(SubtensorModule::do_swap_coldkey( <::RuntimeOrigin>::signed(old_coldkey), @@ -1402,8 +1391,12 @@ fn test_do_swap_coldkey_tx_rate_limit() { let new_coldkey = U256::from(2); let swap_cost = SubtensorModule::get_coldkey_swap_cost(); + // Set non-zero tx rate limit + SubtensorModule::set_tx_rate_limit(1); + // Setup initial state SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, swap_cost * 2); + SubtensorModule::add_balance_to_coldkey_account(&new_coldkey, swap_cost * 2); // Perform first swap assert_ok!(SubtensorModule::do_swap_coldkey( From 49d19b1f86bbf6dd7ab7dc75544cdedea33d4d6c Mon Sep 17 00:00:00 2001 From: Samuel Dare Date: Thu, 4 Jul 2024 14:39:57 +0400 Subject: [PATCH 05/14] chore: lints --- justfile | 6 +---- pallets/subtensor/src/lib.rs | 13 ++++------- pallets/subtensor/src/migration.rs | 5 +--- pallets/subtensor/src/staking.rs | 9 +++----- pallets/subtensor/src/swap.rs | 24 ++++++++----------- pallets/subtensor/tests/swap.rs | 37 +++++++++++++++++++----------- 6 files changed, 43 insertions(+), 51 deletions(-) diff --git a/justfile b/justfile index a75b0052b..871daa6f9 100644 --- a/justfile +++ b/justfile @@ -35,11 +35,7 @@ clippy-fix: -A clippy::todo \ -A clippy::unimplemented \ -A clippy::indexing_slicing - @echo "Running cargo clippy with automatic fixes on potentially dirty code..." - cargo +{{RUSTV}} clippy --fix --allow-dirty --workspace --all-targets -- \ - -A clippy::todo \ - -A clippy::unimplemented \ - -A clippy::indexing_slicing + fix: @echo "Running cargo fix..." cargo +{{RUSTV}} fix --workspace diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index 2fde2f293..395115935 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -1208,15 +1208,12 @@ pub mod pallet { Owner::::insert(hotkey.clone(), coldkey.clone()); // Update Owned map - let mut hotkeys = Owned::::get(&coldkey); - if !hotkeys.contains(&hotkey) { + let mut hotkeys = Owned::::get(coldkey); + if !hotkeys.contains(hotkey) { hotkeys.push(hotkey.clone()); - Owned::::insert( - &coldkey, - hotkeys, - ); + Owned::::insert(coldkey, hotkeys); } - + TotalHotkeyStake::::insert(hotkey.clone(), stake); TotalColdkeyStake::::insert( coldkey.clone(), @@ -1994,7 +1991,7 @@ pub mod pallet { origin: OriginFor, old_coldkey: T::AccountId, new_coldkey: T::AccountId, - ) -> DispatchResultWithPostInfo { + ) -> DispatchResultWithPostInfo { Self::do_swap_coldkey(origin, &old_coldkey, &new_coldkey) } diff --git a/pallets/subtensor/src/migration.rs b/pallets/subtensor/src/migration.rs index 0873ef034..8263e347b 100644 --- a/pallets/subtensor/src/migration.rs +++ b/pallets/subtensor/src/migration.rs @@ -502,10 +502,7 @@ pub fn migrate_populate_owned() -> Weight { } } - Owned::::insert( - &coldkey, - hotkeys, - ); + Owned::::insert(&coldkey, hotkeys); weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 1)); }); diff --git a/pallets/subtensor/src/staking.rs b/pallets/subtensor/src/staking.rs index 913a7208b..eac47ba21 100644 --- a/pallets/subtensor/src/staking.rs +++ b/pallets/subtensor/src/staking.rs @@ -562,13 +562,10 @@ impl Pallet { Owner::::insert(hotkey, coldkey); // Update Owned map - let mut hotkeys = Owned::::get(&coldkey); - if !hotkeys.contains(&hotkey) { + let mut hotkeys = Owned::::get(coldkey); + if !hotkeys.contains(hotkey) { hotkeys.push(hotkey.clone()); - Owned::::insert( - &coldkey, - hotkeys, - ); + Owned::::insert(coldkey, hotkeys); } } } diff --git a/pallets/subtensor/src/swap.rs b/pallets/subtensor/src/swap.rs index dd06bc1fe..d56f60897 100644 --- a/pallets/subtensor/src/swap.rs +++ b/pallets/subtensor/src/swap.rs @@ -225,15 +225,12 @@ impl Pallet { Owner::::insert(new_hotkey, coldkey.clone()); // Update Owned map - let mut hotkeys = Owned::::get(&coldkey); - if !hotkeys.contains(&new_hotkey) { + let mut hotkeys = Owned::::get(coldkey); + if !hotkeys.contains(new_hotkey) { hotkeys.push(new_hotkey.clone()); } hotkeys.retain(|hk| *hk != *old_hotkey); - Owned::::insert( - &coldkey, - hotkeys, - ); + Owned::::insert(coldkey, hotkeys); weight.saturating_accrue(T::DbWeight::get().writes(2)); } @@ -547,7 +544,7 @@ impl Pallet { new_coldkey: &T::AccountId, weight: &mut Weight, ) { - // Find all hotkeys for this coldkey + // Find all hotkeys for this coldkey let hotkeys = Owned::::get(old_coldkey); weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0)); for hotkey in hotkeys.iter() { @@ -603,13 +600,10 @@ impl Pallet { ) { weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0)); for hotkey in Owned::::get(old_coldkey).iter() { - let (stake, block) = TotalHotkeyColdkeyStakesThisInterval::::get(&hotkey, old_coldkey); + let (stake, block) = + TotalHotkeyColdkeyStakesThisInterval::::get(&hotkey, old_coldkey); TotalHotkeyColdkeyStakesThisInterval::::remove(&hotkey, old_coldkey); - TotalHotkeyColdkeyStakesThisInterval::::insert( - &hotkey, - new_coldkey, - (stake, block), - ); + TotalHotkeyColdkeyStakesThisInterval::::insert(&hotkey, new_coldkey, (stake, block)); weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2)); } } @@ -672,12 +666,12 @@ impl Pallet { weight: &mut Weight, ) { // Update Owned map with new coldkey - let hotkeys = Owned::::get(&old_coldkey); + let hotkeys = Owned::::get(old_coldkey); Owned::::remove(old_coldkey); Owned::::insert(new_coldkey, hotkeys); weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 2)); } - + /// Returns the cost of swapping a coldkey. /// /// # Returns diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index db9c55823..273a21bcf 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1063,7 +1063,10 @@ fn test_do_swap_coldkey_success() { // Setup initial state add_network(netuid, 13, 0); register_ok_neuron(netuid, hotkey, old_coldkey, 0); - SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + swap_cost + free_balance); + SubtensorModule::add_balance_to_coldkey_account( + &old_coldkey, + stake_amount + swap_cost + free_balance, + ); // Add stake to the neuron assert_ok!(SubtensorModule::add_stake( @@ -1072,20 +1075,22 @@ fn test_do_swap_coldkey_success() { stake_amount )); - log::info!("TotalColdkeyStake::::get(old_coldkey): {:?}", TotalColdkeyStake::::get(old_coldkey)); - log::info!("Stake::::get(old_coldkey, hotkey): {:?}", Stake::::get(hotkey, old_coldkey)); + log::info!( + "TotalColdkeyStake::::get(old_coldkey): {:?}", + TotalColdkeyStake::::get(old_coldkey) + ); + log::info!( + "Stake::::get(old_coldkey, hotkey): {:?}", + Stake::::get(hotkey, old_coldkey) + ); // Verify initial stake assert_eq!(TotalColdkeyStake::::get(old_coldkey), stake_amount); assert_eq!(Stake::::get(hotkey, old_coldkey), stake_amount); - assert_eq!(Owned::::get(old_coldkey), vec![hotkey]); assert!(!Owned::::contains_key(new_coldkey)); - - - // Get coldkey free balance before swap let balance = SubtensorModule::get_coldkey_balance(&old_coldkey); assert_eq!(balance, free_balance + swap_cost); @@ -1107,14 +1112,20 @@ fn test_do_swap_coldkey_success() { assert!(!Owned::::contains_key(old_coldkey)); // Verify balance transfer - assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), balance - swap_cost); + assert_eq!( + SubtensorModule::get_coldkey_balance(&new_coldkey), + balance - swap_cost + ); assert_eq!(SubtensorModule::get_coldkey_balance(&old_coldkey), 0); // Verify event emission - System::assert_last_event(Event::ColdkeySwapped { - old_coldkey: old_coldkey, - new_coldkey: new_coldkey, - }.into()); + System::assert_last_event( + Event::ColdkeySwapped { + old_coldkey, + new_coldkey, + } + .into(), + ); }); } @@ -1255,7 +1266,7 @@ fn test_swap_owner_for_coldkey() { Owner::::insert(hotkey1, old_coldkey); Owner::::insert(hotkey2, old_coldkey); - // Initialize Owned map + // Initialize Owned map Owned::::insert(old_coldkey, vec![hotkey1, hotkey2]); // Perform the swap From c82f82feb7900dde4e7eb61075066e1151c5d4c5 Mon Sep 17 00:00:00 2001 From: Samuel Dare Date: Thu, 4 Jul 2024 15:44:14 +0400 Subject: [PATCH 06/14] chore: pr comments --- pallets/subtensor/src/lib.rs | 10 ++-- pallets/subtensor/src/migration.rs | 10 ++-- pallets/subtensor/src/staking.rs | 6 +-- pallets/subtensor/src/swap.rs | 38 ++++++------- pallets/subtensor/tests/swap.rs | 85 +++++------------------------- 5 files changed, 44 insertions(+), 105 deletions(-) diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index 395115935..ba768f950 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -364,7 +364,7 @@ pub mod pallet { pub type Owner = StorageMap<_, Blake2_128Concat, T::AccountId, T::AccountId, ValueQuery, DefaultAccount>; #[pallet::storage] // --- MAP ( cold ) --> Vec | Returns the vector of hotkeys controlled by this coldkey. - pub type Owned = + pub type OwnedHotkeys = StorageMap<_, Blake2_128Concat, T::AccountId, Vec, ValueQuery>; #[pallet::storage] // --- MAP ( hot ) --> take | Returns the hotkey delegation take. And signals that this key is open for delegation. pub type Delegates = @@ -1207,11 +1207,11 @@ pub mod pallet { // Fill stake information. Owner::::insert(hotkey.clone(), coldkey.clone()); - // Update Owned map - let mut hotkeys = Owned::::get(coldkey); + // Update OwnedHotkeys map + let mut hotkeys = OwnedHotkeys::::get(coldkey); if !hotkeys.contains(hotkey) { hotkeys.push(hotkey.clone()); - Owned::::insert(coldkey, hotkeys); + OwnedHotkeys::::insert(coldkey, hotkeys); } TotalHotkeyStake::::insert(hotkey.clone(), stake); @@ -1336,7 +1336,7 @@ pub mod pallet { .saturating_add(migration::migrate_delete_subnet_3::()) // Doesn't check storage version. TODO: Remove after upgrade .saturating_add(migration::migration5_total_issuance::(false)) - // Populate Owned map for coldkey swap. Doesn't update storage vesion. + // Populate OwnedHotkeys map for coldkey swap. Doesn't update storage vesion. .saturating_add(migration::migrate_populate_owned::()); weight diff --git a/pallets/subtensor/src/migration.rs b/pallets/subtensor/src/migration.rs index 8263e347b..cc77ac978 100644 --- a/pallets/subtensor/src/migration.rs +++ b/pallets/subtensor/src/migration.rs @@ -481,10 +481,10 @@ pub fn migrate_to_v2_fixed_total_stake() -> Weight { pub fn migrate_populate_owned() -> Weight { // Setup migration weight let mut weight = T::DbWeight::get().reads(1); - let migration_name = "Populate Owned map"; + let migration_name = "Populate OwnedHotkeys map"; - // Check if this migration is needed (if Owned map is empty) - let migrate = Owned::::iter().next().is_none(); + // Check if this migration is needed (if OwnedHotkeys map is empty) + let migrate = OwnedHotkeys::::iter().next().is_none(); // Only runs if we haven't already updated version past above new_storage_version. if migrate { @@ -493,7 +493,7 @@ pub fn migrate_populate_owned() -> Weight { let mut longest_hotkey_vector = 0; let mut longest_coldkey: Option = None; Owner::::iter().for_each(|(hotkey, coldkey)| { - let mut hotkeys = Owned::::get(&coldkey); + let mut hotkeys = OwnedHotkeys::::get(&coldkey); if !hotkeys.contains(&hotkey) { hotkeys.push(hotkey); if longest_hotkey_vector < hotkeys.len() { @@ -502,7 +502,7 @@ pub fn migrate_populate_owned() -> Weight { } } - Owned::::insert(&coldkey, hotkeys); + OwnedHotkeys::::insert(&coldkey, hotkeys); weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 1)); }); diff --git a/pallets/subtensor/src/staking.rs b/pallets/subtensor/src/staking.rs index eac47ba21..de2a54e4f 100644 --- a/pallets/subtensor/src/staking.rs +++ b/pallets/subtensor/src/staking.rs @@ -561,11 +561,11 @@ impl Pallet { Stake::::insert(hotkey, coldkey, 0); Owner::::insert(hotkey, coldkey); - // Update Owned map - let mut hotkeys = Owned::::get(coldkey); + // Update OwnedHotkeys map + let mut hotkeys = OwnedHotkeys::::get(coldkey); if !hotkeys.contains(hotkey) { hotkeys.push(hotkey.clone()); - Owned::::insert(coldkey, hotkeys); + OwnedHotkeys::::insert(coldkey, hotkeys); } } } diff --git a/pallets/subtensor/src/swap.rs b/pallets/subtensor/src/swap.rs index d56f60897..2971fbfe2 100644 --- a/pallets/subtensor/src/swap.rs +++ b/pallets/subtensor/src/swap.rs @@ -137,17 +137,17 @@ impl Pallet { Error::::NewColdKeyIsSameWithOld ); - // Check if the new coldkey is already associated with any hotkeys - ensure!( - !Self::coldkey_has_associated_hotkeys(new_coldkey), - Error::::ColdKeyAlreadyAssociated - ); + // // Check if the new coldkey is already associated with any hotkeys + // ensure!( + // !Self::coldkey_has_associated_hotkeys(new_coldkey), + // Error::::ColdKeyAlreadyAssociated + // ); let block: u64 = Self::get_current_block_as_u64(); - ensure!( - !Self::exceeds_tx_rate_limit(Self::get_last_tx_block(old_coldkey), block), - Error::::ColdKeySwapTxRateLimitExceeded - ); + // ensure!( + // !Self::exceeds_tx_rate_limit(Self::get_last_tx_block(old_coldkey), block), + // Error::::ColdKeySwapTxRateLimitExceeded + // ); // Note: we probably want to make this free let swap_cost = Self::get_coldkey_swap_cost(); @@ -224,13 +224,13 @@ impl Pallet { Owner::::remove(old_hotkey); Owner::::insert(new_hotkey, coldkey.clone()); - // Update Owned map - let mut hotkeys = Owned::::get(coldkey); + // Update OwnedHotkeys map + let mut hotkeys = OwnedHotkeys::::get(coldkey); if !hotkeys.contains(new_hotkey) { hotkeys.push(new_hotkey.clone()); } hotkeys.retain(|hk| *hk != *old_hotkey); - Owned::::insert(coldkey, hotkeys); + OwnedHotkeys::::insert(coldkey, hotkeys); weight.saturating_accrue(T::DbWeight::get().writes(2)); } @@ -545,7 +545,7 @@ impl Pallet { weight: &mut Weight, ) { // Find all hotkeys for this coldkey - let hotkeys = Owned::::get(old_coldkey); + let hotkeys = OwnedHotkeys::::get(old_coldkey); weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0)); for hotkey in hotkeys.iter() { let stake = Stake::::get(&hotkey, old_coldkey); @@ -572,7 +572,7 @@ impl Pallet { new_coldkey: &T::AccountId, weight: &mut Weight, ) { - let hotkeys = Owned::::get(old_coldkey); + let hotkeys = OwnedHotkeys::::get(old_coldkey); weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0)); for hotkey in hotkeys.iter() { Owner::::insert(&hotkey, new_coldkey); @@ -599,7 +599,7 @@ impl Pallet { weight: &mut Weight, ) { weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0)); - for hotkey in Owned::::get(old_coldkey).iter() { + for hotkey in OwnedHotkeys::::get(old_coldkey).iter() { let (stake, block) = TotalHotkeyColdkeyStakesThisInterval::::get(&hotkey, old_coldkey); TotalHotkeyColdkeyStakesThisInterval::::remove(&hotkey, old_coldkey); @@ -665,10 +665,10 @@ impl Pallet { new_coldkey: &T::AccountId, weight: &mut Weight, ) { - // Update Owned map with new coldkey - let hotkeys = Owned::::get(old_coldkey); - Owned::::remove(old_coldkey); - Owned::::insert(new_coldkey, hotkeys); + // Update OwnedHotkeys map with new coldkey + let hotkeys = OwnedHotkeys::::get(old_coldkey); + OwnedHotkeys::::remove(old_coldkey); + OwnedHotkeys::::insert(new_coldkey, hotkeys); weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 2)); } diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index 273a21bcf..ebdb5aed7 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1088,8 +1088,8 @@ fn test_do_swap_coldkey_success() { assert_eq!(TotalColdkeyStake::::get(old_coldkey), stake_amount); assert_eq!(Stake::::get(hotkey, old_coldkey), stake_amount); - assert_eq!(Owned::::get(old_coldkey), vec![hotkey]); - assert!(!Owned::::contains_key(new_coldkey)); + assert_eq!(OwnedHotkeys::::get(old_coldkey), vec![hotkey]); + assert!(!OwnedHotkeys::::contains_key(new_coldkey)); // Get coldkey free balance before swap let balance = SubtensorModule::get_coldkey_balance(&old_coldkey); @@ -1108,8 +1108,8 @@ fn test_do_swap_coldkey_success() { assert!(!TotalColdkeyStake::::contains_key(old_coldkey)); assert_eq!(Stake::::get(hotkey, new_coldkey), stake_amount); assert!(!Stake::::contains_key(hotkey, old_coldkey)); - assert_eq!(Owned::::get(new_coldkey), vec![hotkey]); - assert!(!Owned::::contains_key(old_coldkey)); + assert_eq!(OwnedHotkeys::::get(new_coldkey), vec![hotkey]); + assert!(!OwnedHotkeys::::contains_key(old_coldkey)); // Verify balance transfer assert_eq!( @@ -1168,34 +1168,6 @@ fn test_do_swap_coldkey_same_keys() { }); } -#[test] -fn test_do_swap_coldkey_new_key_already_associated() { - new_test_ext(1).execute_with(|| { - let old_coldkey = U256::from(1); - let new_coldkey = U256::from(2); - let hotkey1 = U256::from(3); - let hotkey2 = U256::from(4); - let netuid = 1u16; - let swap_cost = SubtensorModule::get_coldkey_swap_cost(); - - // Setup initial state - add_network(netuid, 13, 0); - register_ok_neuron(netuid, hotkey1, old_coldkey, 0); - register_ok_neuron(netuid, hotkey2, new_coldkey, 0); - SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, swap_cost); - - // Attempt the swap - assert_err!( - SubtensorModule::do_swap_coldkey( - <::RuntimeOrigin>::signed(old_coldkey), - &old_coldkey, - &new_coldkey - ), - Error::::ColdKeyAlreadyAssociated - ); - }); -} - #[test] fn test_swap_total_coldkey_stake() { new_test_ext(1).execute_with(|| { @@ -1235,8 +1207,8 @@ fn test_swap_stake_for_coldkey() { Stake::::insert(hotkey1, old_coldkey, stake_amount1); Stake::::insert(hotkey2, old_coldkey, stake_amount2); - // Populate Owned map - Owned::::insert(old_coldkey, vec![hotkey1, hotkey2]); + // Populate OwnedHotkeys map + OwnedHotkeys::::insert(old_coldkey, vec![hotkey1, hotkey2]); // Perform the swap SubtensorModule::swap_stake_for_coldkey(&old_coldkey, &new_coldkey, &mut weight); @@ -1266,8 +1238,8 @@ fn test_swap_owner_for_coldkey() { Owner::::insert(hotkey1, old_coldkey); Owner::::insert(hotkey2, old_coldkey); - // Initialize Owned map - Owned::::insert(old_coldkey, vec![hotkey1, hotkey2]); + // Initialize OwnedHotkeys map + OwnedHotkeys::::insert(old_coldkey, vec![hotkey1, hotkey2]); // Perform the swap SubtensorModule::swap_owner_for_coldkey(&old_coldkey, &new_coldkey, &mut weight); @@ -1297,8 +1269,8 @@ fn test_swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey() { TotalHotkeyColdkeyStakesThisInterval::::insert(hotkey1, old_coldkey, stake1); TotalHotkeyColdkeyStakesThisInterval::::insert(hotkey2, old_coldkey, stake2); - // Populate Owned map - Owned::::insert(old_coldkey, vec![hotkey1, hotkey2]); + // Populate OwnedHotkeys map + OwnedHotkeys::::insert(old_coldkey, vec![hotkey1, hotkey2]); // Perform the swap SubtensorModule::swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey( @@ -1380,8 +1352,8 @@ fn test_do_swap_coldkey_with_subnet_ownership() { SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + swap_cost); SubnetOwner::::insert(netuid, old_coldkey); - // Populate Owned map - Owned::::insert(old_coldkey, vec![hotkey]); + // Populate OwnedHotkeys map + OwnedHotkeys::::insert(old_coldkey, vec![hotkey]); // Perform the swap assert_ok!(SubtensorModule::do_swap_coldkey( @@ -1395,39 +1367,6 @@ fn test_do_swap_coldkey_with_subnet_ownership() { }); } -#[test] -fn test_do_swap_coldkey_tx_rate_limit() { - new_test_ext(1).execute_with(|| { - let old_coldkey = U256::from(1); - let new_coldkey = U256::from(2); - let swap_cost = SubtensorModule::get_coldkey_swap_cost(); - - // Set non-zero tx rate limit - SubtensorModule::set_tx_rate_limit(1); - - // Setup initial state - SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, swap_cost * 2); - SubtensorModule::add_balance_to_coldkey_account(&new_coldkey, swap_cost * 2); - - // Perform first swap - assert_ok!(SubtensorModule::do_swap_coldkey( - <::RuntimeOrigin>::signed(old_coldkey), - &old_coldkey, - &new_coldkey - )); - - // Attempt second swap immediately - assert_err!( - SubtensorModule::do_swap_coldkey( - <::RuntimeOrigin>::signed(new_coldkey), - &new_coldkey, - &old_coldkey - ), - Error::::ColdKeySwapTxRateLimitExceeded - ); - }); -} - #[test] fn test_coldkey_has_associated_hotkeys() { new_test_ext(1).execute_with(|| { From 0855df737704acd9868f7d3688021d643657c89a Mon Sep 17 00:00:00 2001 From: Samuel Dare Date: Thu, 4 Jul 2024 19:06:06 +0400 Subject: [PATCH 07/14] chore: pr review, make swaps free --- pallets/subtensor/src/migration.rs | 36 +++++++++++++++---- pallets/subtensor/src/swap.rs | 52 +++------------------------ pallets/subtensor/tests/swap.rs | 57 +++--------------------------- runtime/src/lib.rs | 2 +- 4 files changed, 39 insertions(+), 108 deletions(-) diff --git a/pallets/subtensor/src/migration.rs b/pallets/subtensor/src/migration.rs index cc77ac978..5dbdd5f42 100644 --- a/pallets/subtensor/src/migration.rs +++ b/pallets/subtensor/src/migration.rs @@ -478,6 +478,7 @@ pub fn migrate_to_v2_fixed_total_stake() -> Weight { } } +/// Migrate the OwnedHotkeys map to the new storage format pub fn migrate_populate_owned() -> Weight { // Setup migration weight let mut weight = T::DbWeight::get().reads(1); @@ -486,27 +487,48 @@ pub fn migrate_populate_owned() -> Weight { // Check if this migration is needed (if OwnedHotkeys map is empty) let migrate = OwnedHotkeys::::iter().next().is_none(); - // Only runs if we haven't already updated version past above new_storage_version. + // Only runs if the migration is needed if migrate { - info!(target: LOG_TARGET_1, ">>> Migration: {}", migration_name); + info!(target: LOG_TARGET_1, ">>> Starting Migration: {}", migration_name); - let mut longest_hotkey_vector = 0; + let mut longest_hotkey_vector: usize = 0; let mut longest_coldkey: Option = None; + let mut keys_touched: u64 = 0; + let mut storage_reads: u64 = 0; + let mut storage_writes: u64 = 0; + + // Iterate through all Owner entries Owner::::iter().for_each(|(hotkey, coldkey)| { + storage_reads = storage_reads.saturating_add(1); // Read from Owner storage let mut hotkeys = OwnedHotkeys::::get(&coldkey); + storage_reads = storage_reads.saturating_add(1); // Read from OwnedHotkeys storage + + // Add the hotkey if it's not already in the vector if !hotkeys.contains(&hotkey) { hotkeys.push(hotkey); + keys_touched = keys_touched.saturating_add(1); + + // Update longest hotkey vector info if longest_hotkey_vector < hotkeys.len() { longest_hotkey_vector = hotkeys.len(); longest_coldkey = Some(coldkey.clone()); } - } - OwnedHotkeys::::insert(&coldkey, hotkeys); + // Update the OwnedHotkeys storage + OwnedHotkeys::::insert(&coldkey, hotkeys); + storage_writes = storage_writes.saturating_add(1); // Write to OwnedHotkeys storage + } - weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 1)); + // Accrue weight for reads and writes + weight = weight.saturating_add(T::DbWeight::get().reads_writes(2, 1)); }); - info!(target: LOG_TARGET_1, "Migration {} finished. Longest hotkey vector: {}", migration_name, longest_hotkey_vector); + + // Log migration results + info!( + target: LOG_TARGET_1, + "Migration {} finished. Keys touched: {}, Longest hotkey vector: {}, Storage reads: {}, Storage writes: {}", + migration_name, keys_touched, longest_hotkey_vector, storage_reads, storage_writes + ); if let Some(c) = longest_coldkey { info!(target: LOG_TARGET_1, "Longest hotkey vector is controlled by: {:?}", c); } diff --git a/pallets/subtensor/src/swap.rs b/pallets/subtensor/src/swap.rs index 2971fbfe2..b6aed8b72 100644 --- a/pallets/subtensor/src/swap.rs +++ b/pallets/subtensor/src/swap.rs @@ -53,16 +53,6 @@ impl Pallet { T::DbWeight::get().reads((TotalNetworks::::get().saturating_add(1u16)) as u64), ); - let swap_cost = Self::get_hotkey_swap_cost(); - log::debug!("Swap cost: {:?}", swap_cost); - - ensure!( - Self::can_remove_balance_from_coldkey_account(&coldkey, swap_cost), - Error::::NotEnoughBalanceToPaySwapHotKey - ); - let actual_burn_amount = Self::remove_balance_from_coldkey_account(&coldkey, swap_cost)?; - Self::burn_tokens(actual_burn_amount); - Self::swap_owner(old_hotkey, new_hotkey, &coldkey, &mut weight); Self::swap_total_hotkey_stake(old_hotkey, new_hotkey, &mut weight); Self::swap_delegates(old_hotkey, new_hotkey, &mut weight); @@ -125,38 +115,17 @@ impl Pallet { old_coldkey: &T::AccountId, new_coldkey: &T::AccountId, ) -> DispatchResultWithPostInfo { - let caller = ensure_signed(origin)?; - - // Ensure the caller is the old coldkey - ensure!(caller == *old_coldkey, Error::::NonAssociatedColdKey); + ensure_signed(origin)?; let mut weight = T::DbWeight::get().reads(2); + // Check if the new coldkey is already associated with any hotkeys ensure!( - old_coldkey != new_coldkey, - Error::::NewColdKeyIsSameWithOld + !Self::coldkey_has_associated_hotkeys(new_coldkey), + Error::::ColdKeyAlreadyAssociated ); - // // Check if the new coldkey is already associated with any hotkeys - // ensure!( - // !Self::coldkey_has_associated_hotkeys(new_coldkey), - // Error::::ColdKeyAlreadyAssociated - // ); - let block: u64 = Self::get_current_block_as_u64(); - // ensure!( - // !Self::exceeds_tx_rate_limit(Self::get_last_tx_block(old_coldkey), block), - // Error::::ColdKeySwapTxRateLimitExceeded - // ); - - // Note: we probably want to make this free - let swap_cost = Self::get_coldkey_swap_cost(); - ensure!( - Self::can_remove_balance_from_coldkey_account(old_coldkey, swap_cost), - Error::::NotEnoughBalanceToPaySwapColdKey - ); - let actual_burn_amount = Self::remove_balance_from_coldkey_account(old_coldkey, swap_cost)?; - Self::burn_tokens(actual_burn_amount); // Swap coldkey references in storage maps Self::swap_total_coldkey_stake(old_coldkey, new_coldkey, &mut weight); @@ -671,17 +640,4 @@ impl Pallet { OwnedHotkeys::::insert(new_coldkey, hotkeys); weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 2)); } - - /// Returns the cost of swapping a coldkey. - /// - /// # Returns - /// - /// * `u64` - The cost of swapping a coldkey in Rao. - /// - /// # Note - /// - /// This function returns a hardcoded value. In a production environment, this should be configurable or determined dynamically. - pub fn get_coldkey_swap_cost() -> u64 { - 1_000_000 // Example cost in Rao - } } diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index ebdb5aed7..65c9f2d4b 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1057,16 +1057,12 @@ fn test_do_swap_coldkey_success() { let hotkey = U256::from(3); let netuid = 1u16; let stake_amount = 1000u64; - let swap_cost = SubtensorModule::get_coldkey_swap_cost(); let free_balance = 12345; // Setup initial state add_network(netuid, 13, 0); register_ok_neuron(netuid, hotkey, old_coldkey, 0); - SubtensorModule::add_balance_to_coldkey_account( - &old_coldkey, - stake_amount + swap_cost + free_balance, - ); + SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + free_balance); // Add stake to the neuron assert_ok!(SubtensorModule::add_stake( @@ -1093,7 +1089,7 @@ fn test_do_swap_coldkey_success() { // Get coldkey free balance before swap let balance = SubtensorModule::get_coldkey_balance(&old_coldkey); - assert_eq!(balance, free_balance + swap_cost); + assert_eq!(balance, free_balance); // Perform the swap assert_ok!(SubtensorModule::do_swap_coldkey( @@ -1112,10 +1108,7 @@ fn test_do_swap_coldkey_success() { assert!(!OwnedHotkeys::::contains_key(old_coldkey)); // Verify balance transfer - assert_eq!( - SubtensorModule::get_coldkey_balance(&new_coldkey), - balance - swap_cost - ); + assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), balance); assert_eq!(SubtensorModule::get_coldkey_balance(&old_coldkey), 0); // Verify event emission @@ -1129,45 +1122,6 @@ fn test_do_swap_coldkey_success() { }); } -#[test] -fn test_do_swap_coldkey_not_enough_balance() { - new_test_ext(1).execute_with(|| { - let old_coldkey = U256::from(1); - let new_coldkey = U256::from(2); - let swap_cost = SubtensorModule::get_coldkey_swap_cost(); - - // Setup initial state with insufficient balance - SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, swap_cost - 1); - - // Attempt the swap - assert_err!( - SubtensorModule::do_swap_coldkey( - <::RuntimeOrigin>::signed(old_coldkey), - &old_coldkey, - &new_coldkey - ), - Error::::NotEnoughBalanceToPaySwapColdKey - ); - }); -} - -#[test] -fn test_do_swap_coldkey_same_keys() { - new_test_ext(1).execute_with(|| { - let coldkey = U256::from(1); - - // Attempt the swap with same old and new coldkeys - assert_err!( - SubtensorModule::do_swap_coldkey( - <::RuntimeOrigin>::signed(coldkey), - &coldkey, - &coldkey - ), - Error::::NewColdKeyIsSameWithOld - ); - }); -} - #[test] fn test_swap_total_coldkey_stake() { new_test_ext(1).execute_with(|| { @@ -1339,8 +1293,7 @@ fn test_do_swap_coldkey_with_subnet_ownership() { let new_coldkey = U256::from(2); let hotkey = U256::from(3); let netuid = 1u16; - let stake_amount = 1000u64; - let swap_cost = SubtensorModule::get_coldkey_swap_cost(); + let stake_amount: u64 = 1000u64; // Setup initial state add_network(netuid, 13, 0); @@ -1349,7 +1302,7 @@ fn test_do_swap_coldkey_with_subnet_ownership() { // Set TotalNetworks because swap relies on it pallet_subtensor::TotalNetworks::::set(1); - SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + swap_cost); + SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount); SubnetOwner::::insert(netuid, old_coldkey); // Populate OwnedHotkeys map diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs index 9a43fdf93..fe9ea0bd8 100644 --- a/runtime/src/lib.rs +++ b/runtime/src/lib.rs @@ -139,7 +139,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion { // `spec_version`, and `authoring_version` are the same between Wasm and native. // This value is set to 100 to notify Polkadot-JS App (https://polkadot.js.org/apps) to use // the compatible custom types. - spec_version: 154, + spec_version: 156, impl_version: 1, apis: RUNTIME_API_VERSIONS, transaction_version: 1, From 621725f462a149b94dbcc63b1ba6f628287de4df Mon Sep 17 00:00:00 2001 From: Samuel Dare Date: Thu, 4 Jul 2024 21:36:37 +0400 Subject: [PATCH 08/14] chore: unstake + transfer --- justfile | 4 +- pallets/subtensor/src/errors.rs | 4 + pallets/subtensor/src/events.rs | 15 ++ pallets/subtensor/src/lib.rs | 44 ++++- pallets/subtensor/src/staking.rs | 94 ++++++++++ pallets/subtensor/tests/staking.rs | 290 +++++++++++++++++++++++++++++ 6 files changed, 448 insertions(+), 3 deletions(-) diff --git a/justfile b/justfile index 871daa6f9..e33fdf685 100644 --- a/justfile +++ b/justfile @@ -31,11 +31,11 @@ clippy: clippy-fix: @echo "Running cargo clippy with automatic fixes on potentially dirty code..." - cargo +{{RUSTV}} clippy --fix --allow-dirty --workspace --all-targets -- \ + cargo +{{RUSTV}} clippy --fix --allow-dirty --allow-staged --workspace --all-targets -- \ -A clippy::todo \ -A clippy::unimplemented \ -A clippy::indexing_slicing - + fix: @echo "Running cargo fix..." cargo +{{RUSTV}} fix --workspace diff --git a/pallets/subtensor/src/errors.rs b/pallets/subtensor/src/errors.rs index 9dd01fda4..3e30c094c 100644 --- a/pallets/subtensor/src/errors.rs +++ b/pallets/subtensor/src/errors.rs @@ -142,5 +142,9 @@ mod errors { NotExistColdkey, /// The coldkey balance is not enough to pay for the swap NotEnoughBalanceToPaySwapColdKey, + /// No balance to transfer + NoBalanceToTransfer, + /// Same coldkey + SameColdkey, } } diff --git a/pallets/subtensor/src/events.rs b/pallets/subtensor/src/events.rs index 73bda500b..167f10170 100644 --- a/pallets/subtensor/src/events.rs +++ b/pallets/subtensor/src/events.rs @@ -139,5 +139,20 @@ mod events { /// the account ID of new coldkey new_coldkey: T::AccountId, }, + /// All balance of a hotkey has been unstaked and transferred to a new coldkey + AllBalanceUnstakedAndTransferredToNewColdkey { + /// The account ID of the current coldkey + current_coldkey: T::AccountId, + /// The account ID of the new coldkey + new_coldkey: T::AccountId, + /// The account ID of the hotkey + hotkey: T::AccountId, + /// The current stake of the hotkey + current_stake: u64, + /// The total balance of the hotkey + total_balance: <::Currency as fungible::Inspect< + ::AccountId, + >>::Balance, + }, } } diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index ba768f950..873168d65 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -1982,7 +1982,21 @@ pub mod pallet { Self::do_swap_hotkey(origin, &hotkey, &new_hotkey) } - /// The extrinsic for user to change the coldkey + /// The extrinsic for user to change the coldkey associated with their account. + /// + /// # Arguments + /// + /// * `origin` - The origin of the call, must be signed by the old coldkey. + /// * `old_coldkey` - The current coldkey associated with the account. + /// * `new_coldkey` - The new coldkey to be associated with the account. + /// + /// # Returns + /// + /// Returns a `DispatchResultWithPostInfo` indicating success or failure of the operation. + /// + /// # Weight + /// + /// Weight is calculated based on the number of database reads and writes. #[pallet::call_index(71)] #[pallet::weight((Weight::from_parts(1_940_000_000, 0) .saturating_add(T::DbWeight::get().reads(272)) @@ -1995,6 +2009,34 @@ pub mod pallet { Self::do_swap_coldkey(origin, &old_coldkey, &new_coldkey) } + /// Unstakes all tokens associated with a hotkey and transfers them to a new coldkey. + /// + /// # Arguments + /// + /// * `origin` - The origin of the call, must be signed by the current coldkey. + /// * `hotkey` - The hotkey associated with the stakes to be unstaked. + /// * `new_coldkey` - The new coldkey to receive the unstaked tokens. + /// + /// # Returns + /// + /// Returns a `DispatchResult` indicating success or failure of the operation. + /// + /// # Weight + /// + /// Weight is calculated based on the number of database reads and writes. + #[pallet::call_index(72)] + #[pallet::weight((Weight::from_parts(1_940_000_000, 0) + .saturating_add(T::DbWeight::get().reads(272)) + .saturating_add(T::DbWeight::get().writes(527)), DispatchClass::Operational, Pays::No))] + pub fn unstake_all_and_transfer_to_new_coldkey( + origin: OriginFor, + hotkey: T::AccountId, + new_coldkey: T::AccountId, + ) -> DispatchResult { + let current_coldkey = ensure_signed(origin)?; + Self::do_unstake_all_and_transfer_to_new_coldkey(current_coldkey, hotkey, new_coldkey) + } + // ---- SUDO ONLY FUNCTIONS ------------------------------------------------------------ // ================================== diff --git a/pallets/subtensor/src/staking.rs b/pallets/subtensor/src/staking.rs index de2a54e4f..6c87f1131 100644 --- a/pallets/subtensor/src/staking.rs +++ b/pallets/subtensor/src/staking.rs @@ -1,4 +1,5 @@ use super::*; +use dispatch::RawOrigin; use frame_support::{ storage::IterableStorageDoubleMap, traits::{ @@ -9,6 +10,7 @@ use frame_support::{ Imbalance, }, }; +use num_traits::Zero; impl Pallet { /// ---- The implementation for the extrinsic become_delegate: signals that this hotkey allows delegated stake. @@ -827,4 +829,96 @@ impl Pallet { Self::add_balance_to_coldkey_account(&delegate_coldkey_i, stake_i); } } + + /// Unstakes all tokens associated with a hotkey and transfers them to a new coldkey. + /// + /// This function performs the following operations: + /// 1. Verifies that the hotkey exists and is owned by the current coldkey. + /// 2. Ensures that the new coldkey is different from the current one. + /// 3. Unstakes all balance if there's any stake. + /// 4. Transfers the entire balance of the hotkey to the new coldkey. + /// 5. Verifies the success of the transfer and handles partial transfers if necessary. + /// + /// # Arguments + /// + /// * `current_coldkey` - The AccountId of the current coldkey. + /// * `hotkey` - The AccountId of the hotkey whose balance is being unstaked and transferred. + /// * `new_coldkey` - The AccountId of the new coldkey to receive the unstaked tokens. + /// + /// # Returns + /// + /// Returns a `DispatchResult` indicating success or failure of the operation. + /// + /// # Errors + /// + /// This function will return an error if: + /// * The hotkey account does not exist. + /// * The current coldkey does not own the hotkey. + /// * The new coldkey is the same as the current coldkey. + /// * There is no balance to transfer. + /// * The transfer fails or is only partially successful. + /// + /// # Events + /// + /// Emits an `AllBalanceUnstakedAndTransferredToNewColdkey` event upon successful execution. + /// Emits a `PartialBalanceTransferredToNewColdkey` event if only a partial transfer is successful. + /// + pub fn do_unstake_all_and_transfer_to_new_coldkey( + current_coldkey: T::AccountId, + hotkey: T::AccountId, + new_coldkey: T::AccountId, + ) -> DispatchResult { + // Ensure the hotkey exists and is owned by the current coldkey + ensure!( + Self::hotkey_account_exists(&hotkey), + Error::::HotKeyAccountNotExists + ); + ensure!( + Self::coldkey_owns_hotkey(¤t_coldkey, &hotkey), + Error::::NonAssociatedColdKey + ); + + // Ensure the new coldkey is different from the current one + ensure!(current_coldkey != new_coldkey, Error::::SameColdkey); + + // Get the current stake + let current_stake: u64 = Self::get_stake_for_coldkey_and_hotkey(¤t_coldkey, &hotkey); + + // Unstake all balance if there's any stake + if current_stake > 0 { + Self::do_remove_stake( + RawOrigin::Signed(current_coldkey.clone()).into(), + hotkey.clone(), + current_stake, + )?; + } + + // Get the total balance of the current coldkey account + // let total_balance: <::Currency as fungible::Inspect<::AccountId>>::Balance = T::Currency::total_balance(¤t_coldkey); + + let total_balance = Self::get_coldkey_balance(¤t_coldkey); + log::info!("Total Bank Balance: {:?}", total_balance); + + // Ensure there's a balance to transfer + ensure!(!total_balance.is_zero(), Error::::NoBalanceToTransfer); + + // Attempt to transfer the entire total balance to the new coldkey + T::Currency::transfer( + ¤t_coldkey, + &new_coldkey, + total_balance, + Preservation::Expendable, + )?; + + // Emit the event + Self::deposit_event(Event::AllBalanceUnstakedAndTransferredToNewColdkey { + current_coldkey: current_coldkey.clone(), + new_coldkey: new_coldkey.clone(), + hotkey: hotkey.clone(), + current_stake, + total_balance, + }); + + Ok(()) + } } diff --git a/pallets/subtensor/tests/staking.rs b/pallets/subtensor/tests/staking.rs index 529332f04..d4441c7c9 100644 --- a/pallets/subtensor/tests/staking.rs +++ b/pallets/subtensor/tests/staking.rs @@ -3127,3 +3127,293 @@ fn test_rate_limits_enforced_on_increase_take() { assert_eq!(SubtensorModule::get_hotkey_take(&hotkey0), u16::MAX / 8); }); } + +// Helper function to set up a test environment +fn setup_test_environment() -> (AccountId, AccountId, AccountId) { + let current_coldkey = U256::from(1); + let hotkey = U256::from(2); + let new_coldkey = U256::from(3); + // Register the neuron to a new network + let netuid = 1; + add_network(netuid, 0, 0); + + // Register the hotkey and associate it with the current coldkey + register_ok_neuron(1, hotkey, current_coldkey, 0); + + // Add some balance to the hotkey + SubtensorModule::add_balance_to_coldkey_account(¤t_coldkey, 1000); + + // Stake some amount + assert_ok!(SubtensorModule::add_stake( + RuntimeOrigin::signed(current_coldkey), + hotkey, + 500 + )); + + (current_coldkey, hotkey, new_coldkey) +} + +#[test] +fn test_do_unstake_all_and_transfer_to_new_coldkey_success() { + new_test_ext(1).execute_with(|| { + let (current_coldkey, hotkey, new_coldkey) = setup_test_environment(); + + assert_ok!(SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey( + current_coldkey, + hotkey, + new_coldkey + )); + + // Check that the stake has been removed + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey), 0); + + // Check that the balance has been transferred to the new coldkey + assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), 1000); + + // Check that the appropriate event was emitted + System::assert_last_event( + Event::AllBalanceUnstakedAndTransferredToNewColdkey { + current_coldkey, + new_coldkey, + hotkey, + current_stake: 500, + total_balance: 1000, + } + .into(), + ); + }); +} + +#[test] +fn test_do_unstake_all_and_transfer_to_new_coldkey_hotkey_not_exists() { + new_test_ext(1).execute_with(|| { + let current_coldkey = U256::from(1); + let hotkey = U256::from(2); + let new_coldkey = U256::from(3); + + assert_err!( + SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey( + current_coldkey, + hotkey, + new_coldkey + ), + Error::::HotKeyAccountNotExists + ); + }); +} + +#[test] +fn test_do_unstake_all_and_transfer_to_new_coldkey_non_associated_coldkey() { + new_test_ext(1).execute_with(|| { + let (_, hotkey, new_coldkey) = setup_test_environment(); + let wrong_coldkey = U256::from(4); + + assert_noop!( + SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey( + wrong_coldkey, + hotkey, + new_coldkey + ), + Error::::NonAssociatedColdKey + ); + }); +} + +#[test] +fn test_do_unstake_all_and_transfer_to_new_coldkey_same_coldkey() { + new_test_ext(1).execute_with(|| { + let (current_coldkey, hotkey, _) = setup_test_environment(); + + assert_noop!( + SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey( + current_coldkey, + hotkey, + current_coldkey + ), + Error::::SameColdkey + ); + }); +} + +#[test] +fn test_do_unstake_all_and_transfer_to_new_coldkey_no_balance() { + new_test_ext(1).execute_with(|| { + // Create accounts manually + let current_coldkey: AccountId = U256::from(1); + let hotkey: AccountId = U256::from(2); + let new_coldkey: AccountId = U256::from(3); + + add_network(1, 0, 0); + + // Register the hotkey and associate it with the current coldkey + register_ok_neuron(1, hotkey, current_coldkey, 0); + + // Print initial balances + log::info!( + "Initial current_coldkey balance: {:?}", + Balances::total_balance(¤t_coldkey) + ); + log::info!( + "Initial hotkey balance: {:?}", + Balances::total_balance(&hotkey) + ); + log::info!( + "Initial new_coldkey balance: {:?}", + Balances::total_balance(&new_coldkey) + ); + + // Ensure there's no balance in any of the accounts + assert_eq!(Balances::total_balance(¤t_coldkey), 0); + assert_eq!(Balances::total_balance(&hotkey), 0); + assert_eq!(Balances::total_balance(&new_coldkey), 0); + + // Try to unstake and transfer + let result = SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey( + current_coldkey, + hotkey, + new_coldkey, + ); + + // Print the result + log::info!( + "Result of do_unstake_all_and_transfer_to_new_coldkey: {:?}", + result + ); + + // Print final balances + log::info!( + "Final current_coldkey balance: {:?}", + Balances::total_balance(¤t_coldkey) + ); + log::info!( + "Final hotkey balance: {:?}", + Balances::total_balance(&hotkey) + ); + log::info!( + "Final new_coldkey balance: {:?}", + Balances::total_balance(&new_coldkey) + ); + + // Assert the expected error + assert_noop!(result, Error::::NoBalanceToTransfer); + + // Verify that no balance was transferred + assert_eq!(Balances::total_balance(¤t_coldkey), 0); + assert_eq!(Balances::total_balance(&hotkey), 0); + assert_eq!(Balances::total_balance(&new_coldkey), 0); + }); +} + +#[test] +fn test_do_unstake_all_and_transfer_to_new_coldkey_with_no_stake() { + new_test_ext(1).execute_with(|| { + // Create accounts manually + let current_coldkey: AccountId = U256::from(1); + let hotkey: AccountId = U256::from(2); + let new_coldkey: AccountId = U256::from(3); + + add_network(1, 0, 0); + + // Register the hotkey and associate it with the current coldkey + register_ok_neuron(1, hotkey, current_coldkey, 0); + + // Add balance to the current coldkey without staking + let initial_balance = 500; + Balances::make_free_balance_be(¤t_coldkey, initial_balance); + + // Print initial balances + log::info!( + "Initial current_coldkey balance: {:?}", + Balances::total_balance(¤t_coldkey) + ); + log::info!( + "Initial hotkey balance: {:?}", + Balances::total_balance(&hotkey) + ); + log::info!( + "Initial new_coldkey balance: {:?}", + Balances::total_balance(&new_coldkey) + ); + + // Ensure initial balances are correct + assert_eq!(Balances::total_balance(¤t_coldkey), initial_balance); + assert_eq!(Balances::total_balance(&hotkey), 0); + assert_eq!(Balances::total_balance(&new_coldkey), 0); + + // Perform unstake and transfer + assert_ok!(SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey( + current_coldkey, + hotkey, + new_coldkey + )); + + // Print final balances + log::info!( + "Final current_coldkey balance: {:?}", + Balances::total_balance(¤t_coldkey) + ); + log::info!( + "Final hotkey balance: {:?}", + Balances::total_balance(&hotkey) + ); + log::info!( + "Final new_coldkey balance: {:?}", + Balances::total_balance(&new_coldkey) + ); + + // Check that the balance has been transferred to the new coldkey + assert_eq!(Balances::total_balance(&new_coldkey), initial_balance); + assert_eq!(Balances::total_balance(¤t_coldkey), 0); + + // Check that the appropriate event was emitted + System::assert_last_event( + Event::AllBalanceUnstakedAndTransferredToNewColdkey { + current_coldkey, + new_coldkey, + hotkey, + current_stake: 0, + total_balance: initial_balance, + } + .into(), + ); + }); +} +#[test] +fn test_do_unstake_all_and_transfer_to_new_coldkey_with_multiple_stakes() { + new_test_ext(1).execute_with(|| { + let (current_coldkey, hotkey, new_coldkey) = setup_test_environment(); + + SubtensorModule::set_target_stakes_per_interval(10); + + // Add more stake + assert_ok!(SubtensorModule::add_stake( + RuntimeOrigin::signed(current_coldkey), + hotkey, + 300 + )); + + assert_ok!(SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey( + current_coldkey, + hotkey, + new_coldkey + )); + + // Check that all stake has been removed + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey), 0); + + // Check that the full balance has been transferred to the new coldkey + assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), 1000); + + // Check that the appropriate event was emitted + System::assert_last_event( + Event::AllBalanceUnstakedAndTransferredToNewColdkey { + current_coldkey, + new_coldkey, + hotkey, + current_stake: 800, + total_balance: 1000, + } + .into(), + ); + }); +} From 65861567aea0431fa8986a527aabe86ff9113486 Mon Sep 17 00:00:00 2001 From: distributedstatemachine <112424909+distributedstatemachine@users.noreply.github.com> Date: Thu, 4 Jul 2024 23:08:30 +0400 Subject: [PATCH 09/14] Update pallets/subtensor/tests/swap.rs Co-authored-by: Cameron Fairchild --- pallets/subtensor/tests/swap.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index 65c9f2d4b..1ffd61845 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1172,7 +1172,10 @@ fn test_swap_stake_for_coldkey() { assert_eq!(Stake::::get(hotkey2, new_coldkey), stake_amount2); assert!(!Stake::::contains_key(hotkey1, old_coldkey)); assert!(!Stake::::contains_key(hotkey2, old_coldkey)); - +assert_eq!(TotalHotkeyStake::::get(hotkey1), stake_amount1); +assert_eq!(TotalHotkeyStake::::get(hotkey2), stake_amount2); +assert_eq!(TotalStake::::get(), stake_amount1 + stake_amount2); +assert_eq!(TotalIssuance::::get(), stake_amount1 + stake_amount2); // Verify weight update let expected_weight = ::DbWeight::get().reads_writes(3, 4); assert_eq!(weight, expected_weight); From e6b8aabf3798dba1743d12dc58a5d62f72d5205e Mon Sep 17 00:00:00 2001 From: distributedstatemachine <112424909+distributedstatemachine@users.noreply.github.com> Date: Thu, 4 Jul 2024 23:08:58 +0400 Subject: [PATCH 10/14] Update pallets/subtensor/tests/swap.rs Co-authored-by: Cameron Fairchild --- pallets/subtensor/tests/swap.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index 1ffd61845..5a7d9b97b 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1108,7 +1108,7 @@ fn test_do_swap_coldkey_success() { assert!(!OwnedHotkeys::::contains_key(old_coldkey)); // Verify balance transfer - assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), balance); + assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), balance + balance_new_coldkey); assert_eq!(SubtensorModule::get_coldkey_balance(&old_coldkey), 0); // Verify event emission From c4a1e6eb6f2e836afe8383d40ead121a09626898 Mon Sep 17 00:00:00 2001 From: distributedstatemachine <112424909+distributedstatemachine@users.noreply.github.com> Date: Thu, 4 Jul 2024 23:13:34 +0400 Subject: [PATCH 11/14] Update pallets/subtensor/tests/swap.rs Co-authored-by: Cameron Fairchild --- pallets/subtensor/tests/swap.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index 5a7d9b97b..aecb0f7dd 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1085,7 +1085,7 @@ fn test_do_swap_coldkey_success() { assert_eq!(Stake::::get(hotkey, old_coldkey), stake_amount); assert_eq!(OwnedHotkeys::::get(old_coldkey), vec![hotkey]); - assert!(!OwnedHotkeys::::contains_key(new_coldkey)); + assert!(!OwnedHotkeys::::get(new_coldkey).contains(hotkey)); // Get coldkey free balance before swap let balance = SubtensorModule::get_coldkey_balance(&old_coldkey); From d70ef060998ab80208177b96880fabddc9a738c8 Mon Sep 17 00:00:00 2001 From: distributedstatemachine <112424909+distributedstatemachine@users.noreply.github.com> Date: Thu, 4 Jul 2024 23:18:47 +0400 Subject: [PATCH 12/14] Update pallets/subtensor/tests/swap.rs Co-authored-by: Cameron Fairchild --- pallets/subtensor/tests/swap.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index aecb0f7dd..679393b6d 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1091,6 +1091,8 @@ fn test_do_swap_coldkey_success() { let balance = SubtensorModule::get_coldkey_balance(&old_coldkey); assert_eq!(balance, free_balance); +let balance_new_coldkey = SubtensorModule::get_coldkey_balance(&new_coldkey); + // Perform the swap assert_ok!(SubtensorModule::do_swap_coldkey( <::RuntimeOrigin>::signed(old_coldkey), From 7062a3e5263128c3f8ee65e391427db0409d33d2 Mon Sep 17 00:00:00 2001 From: Samuel Dare Date: Thu, 4 Jul 2024 23:31:49 +0400 Subject: [PATCH 13/14] chore: additional tests --- pallets/subtensor/src/swap.rs | 1 + pallets/subtensor/tests/swap.rs | 95 ++++++++++++++++++++------------- 2 files changed, 59 insertions(+), 37 deletions(-) diff --git a/pallets/subtensor/src/swap.rs b/pallets/subtensor/src/swap.rs index b6aed8b72..49af72a3f 100644 --- a/pallets/subtensor/src/swap.rs +++ b/pallets/subtensor/src/swap.rs @@ -128,6 +128,7 @@ impl Pallet { let block: u64 = Self::get_current_block_as_u64(); // Swap coldkey references in storage maps + // NOTE The order of these calls is important Self::swap_total_coldkey_stake(old_coldkey, new_coldkey, &mut weight); Self::swap_stake_for_coldkey(old_coldkey, new_coldkey, &mut weight); Self::swap_owner_for_coldkey(old_coldkey, new_coldkey, &mut weight); diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index 679393b6d..324ca3f10 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1054,44 +1054,51 @@ fn test_do_swap_coldkey_success() { new_test_ext(1).execute_with(|| { let old_coldkey = U256::from(1); let new_coldkey = U256::from(2); - let hotkey = U256::from(3); + let hotkey1 = U256::from(3); + let hotkey2 = U256::from(4); let netuid = 1u16; - let stake_amount = 1000u64; - let free_balance = 12345; + let stake_amount1 = 1000u64; + let stake_amount2 = 2000u64; + let free_balance_old = 12345u64; // Setup initial state add_network(netuid, 13, 0); - register_ok_neuron(netuid, hotkey, old_coldkey, 0); - SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + free_balance); + register_ok_neuron(netuid, hotkey1, old_coldkey, 0); + register_ok_neuron(netuid, hotkey2, old_coldkey, 0); - // Add stake to the neuron + // Add balance to old coldkey + SubtensorModule::add_balance_to_coldkey_account( + &old_coldkey, + stake_amount1 + stake_amount2 + free_balance_old, + ); + + // Add stake to the neurons + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(old_coldkey), + hotkey1, + stake_amount1 + )); assert_ok!(SubtensorModule::add_stake( <::RuntimeOrigin>::signed(old_coldkey), - hotkey, - stake_amount + hotkey2, + stake_amount2 )); - log::info!( - "TotalColdkeyStake::::get(old_coldkey): {:?}", - TotalColdkeyStake::::get(old_coldkey) + // Verify initial stakes and balances + assert_eq!( + TotalColdkeyStake::::get(old_coldkey), + stake_amount1 + stake_amount2 ); - log::info!( - "Stake::::get(old_coldkey, hotkey): {:?}", - Stake::::get(hotkey, old_coldkey) + assert_eq!(Stake::::get(hotkey1, old_coldkey), stake_amount1); + assert_eq!(Stake::::get(hotkey2, old_coldkey), stake_amount2); + assert_eq!( + OwnedHotkeys::::get(old_coldkey), + vec![hotkey1, hotkey2] + ); + assert_eq!( + SubtensorModule::get_coldkey_balance(&old_coldkey), + free_balance_old ); - - // Verify initial stake - assert_eq!(TotalColdkeyStake::::get(old_coldkey), stake_amount); - assert_eq!(Stake::::get(hotkey, old_coldkey), stake_amount); - - assert_eq!(OwnedHotkeys::::get(old_coldkey), vec![hotkey]); - assert!(!OwnedHotkeys::::get(new_coldkey).contains(hotkey)); - - // Get coldkey free balance before swap - let balance = SubtensorModule::get_coldkey_balance(&old_coldkey); - assert_eq!(balance, free_balance); - -let balance_new_coldkey = SubtensorModule::get_coldkey_balance(&new_coldkey); // Perform the swap assert_ok!(SubtensorModule::do_swap_coldkey( @@ -1101,16 +1108,30 @@ let balance_new_coldkey = SubtensorModule::get_coldkey_balance(&new_coldkey); )); // Verify the swap - assert_eq!(Owner::::get(hotkey), new_coldkey); - assert_eq!(TotalColdkeyStake::::get(new_coldkey), stake_amount); + assert_eq!(Owner::::get(hotkey1), new_coldkey); + assert_eq!(Owner::::get(hotkey2), new_coldkey); + assert_eq!( + TotalColdkeyStake::::get(new_coldkey), + stake_amount1 + stake_amount2 + ); assert!(!TotalColdkeyStake::::contains_key(old_coldkey)); - assert_eq!(Stake::::get(hotkey, new_coldkey), stake_amount); - assert!(!Stake::::contains_key(hotkey, old_coldkey)); - assert_eq!(OwnedHotkeys::::get(new_coldkey), vec![hotkey]); + assert_eq!(Stake::::get(hotkey1, new_coldkey), stake_amount1); + assert_eq!(Stake::::get(hotkey2, new_coldkey), stake_amount2); + assert!(!Stake::::contains_key(hotkey1, old_coldkey)); + assert!(!Stake::::contains_key(hotkey2, old_coldkey)); + + // Verify OwnedHotkeys + let new_owned_hotkeys = OwnedHotkeys::::get(new_coldkey); + assert!(new_owned_hotkeys.contains(&hotkey1)); + assert!(new_owned_hotkeys.contains(&hotkey2)); + assert_eq!(new_owned_hotkeys.len(), 2); assert!(!OwnedHotkeys::::contains_key(old_coldkey)); // Verify balance transfer - assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), balance + balance_new_coldkey); + assert_eq!( + SubtensorModule::get_coldkey_balance(&new_coldkey), + free_balance_old + ); assert_eq!(SubtensorModule::get_coldkey_balance(&old_coldkey), 0); // Verify event emission @@ -1174,10 +1195,10 @@ fn test_swap_stake_for_coldkey() { assert_eq!(Stake::::get(hotkey2, new_coldkey), stake_amount2); assert!(!Stake::::contains_key(hotkey1, old_coldkey)); assert!(!Stake::::contains_key(hotkey2, old_coldkey)); -assert_eq!(TotalHotkeyStake::::get(hotkey1), stake_amount1); -assert_eq!(TotalHotkeyStake::::get(hotkey2), stake_amount2); -assert_eq!(TotalStake::::get(), stake_amount1 + stake_amount2); -assert_eq!(TotalIssuance::::get(), stake_amount1 + stake_amount2); + assert_eq!(TotalHotkeyStake::::get(hotkey1), stake_amount1); + assert_eq!(TotalHotkeyStake::::get(hotkey2), stake_amount2); + assert_eq!(TotalStake::::get(), stake_amount1 + stake_amount2); + assert_eq!(TotalIssuance::::get(), stake_amount1 + stake_amount2); // Verify weight update let expected_weight = ::DbWeight::get().reads_writes(3, 4); assert_eq!(weight, expected_weight); From a0fa527c4fba9b951771844e2230b5819f7710b1 Mon Sep 17 00:00:00 2001 From: Samuel Dare Date: Fri, 5 Jul 2024 01:20:40 +0400 Subject: [PATCH 14/14] chore: fix tests --- pallets/subtensor/tests/swap.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index 324ca3f10..3d7454fb7 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1184,6 +1184,14 @@ fn test_swap_stake_for_coldkey() { Stake::::insert(hotkey1, old_coldkey, stake_amount1); Stake::::insert(hotkey2, old_coldkey, stake_amount2); + // Initialize TotalHotkeyStake + TotalHotkeyStake::::insert(hotkey1, stake_amount1); + TotalHotkeyStake::::insert(hotkey2, stake_amount2); + + // Initialize TotalStake and TotalIssuance + TotalStake::::put(stake_amount1 + stake_amount2); + TotalIssuance::::put(stake_amount1 + stake_amount2); + // Populate OwnedHotkeys map OwnedHotkeys::::insert(old_coldkey, vec![hotkey1, hotkey2]); @@ -1195,10 +1203,15 @@ fn test_swap_stake_for_coldkey() { assert_eq!(Stake::::get(hotkey2, new_coldkey), stake_amount2); assert!(!Stake::::contains_key(hotkey1, old_coldkey)); assert!(!Stake::::contains_key(hotkey2, old_coldkey)); + + // Verify TotalHotkeyStake remains unchanged assert_eq!(TotalHotkeyStake::::get(hotkey1), stake_amount1); assert_eq!(TotalHotkeyStake::::get(hotkey2), stake_amount2); + + // Verify TotalStake and TotalIssuance remain unchanged assert_eq!(TotalStake::::get(), stake_amount1 + stake_amount2); assert_eq!(TotalIssuance::::get(), stake_amount1 + stake_amount2); + // Verify weight update let expected_weight = ::DbWeight::get().reads_writes(3, 4); assert_eq!(weight, expected_weight);