diff --git a/justfile b/justfile index a75b0052b..e33fdf685 100644 --- a/justfile +++ b/justfile @@ -31,15 +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 -- \ - -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 -- \ + 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 57ba91d31..3e30c094c 100644 --- a/pallets/subtensor/src/errors.rs +++ b/pallets/subtensor/src/errors.rs @@ -132,5 +132,19 @@ 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, + /// No balance to transfer + NoBalanceToTransfer, + /// Same coldkey + SameColdkey, } } diff --git a/pallets/subtensor/src/events.rs b/pallets/subtensor/src/events.rs index 47cc9973b..167f10170 100644 --- a/pallets/subtensor/src/events.rs +++ b/pallets/subtensor/src/events.rs @@ -132,5 +132,27 @@ 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, + }, + /// 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 64aefcfee..873168d65 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 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 = StorageMap<_, Blake2_128Concat, T::AccountId, u16, ValueQuery, DefaultDefaultTake>; @@ -1204,6 +1207,13 @@ pub mod pallet { // Fill stake information. Owner::::insert(hotkey.clone(), coldkey.clone()); + // Update OwnedHotkeys map + let mut hotkeys = OwnedHotkeys::::get(coldkey); + if !hotkeys.contains(hotkey) { + hotkeys.push(hotkey.clone()); + OwnedHotkeys::::insert(coldkey, hotkeys); + } + TotalHotkeyStake::::insert(hotkey.clone(), stake); TotalColdkeyStake::::insert( coldkey.clone(), @@ -1325,7 +1335,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 OwnedHotkeys map for coldkey swap. Doesn't update storage vesion. + .saturating_add(migration::migrate_populate_owned::()); weight } @@ -1970,6 +1982,61 @@ pub mod pallet { Self::do_swap_hotkey(origin, &hotkey, &new_hotkey) } + /// 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)) + .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) + } + + /// 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/migration.rs b/pallets/subtensor/src/migration.rs index 86fc02054..5dbdd5f42 100644 --- a/pallets/subtensor/src/migration.rs +++ b/pallets/subtensor/src/migration.rs @@ -477,3 +477,65 @@ pub fn migrate_to_v2_fixed_total_stake() -> Weight { Weight::zero() } } + +/// 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); + let migration_name = "Populate OwnedHotkeys map"; + + // Check if this migration is needed (if OwnedHotkeys map is empty) + let migrate = OwnedHotkeys::::iter().next().is_none(); + + // Only runs if the migration is needed + if migrate { + info!(target: LOG_TARGET_1, ">>> Starting Migration: {}", migration_name); + + 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()); + } + + // Update the OwnedHotkeys storage + OwnedHotkeys::::insert(&coldkey, hotkeys); + storage_writes = storage_writes.saturating_add(1); // Write to OwnedHotkeys storage + } + + // Accrue weight for reads and writes + weight = weight.saturating_add(T::DbWeight::get().reads_writes(2, 1)); + }); + + // 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); + } + + 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..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. @@ -560,6 +562,13 @@ impl Pallet { if !Self::hotkey_account_exists(hotkey) { Stake::::insert(hotkey, coldkey, 0); Owner::::insert(hotkey, coldkey); + + // Update OwnedHotkeys map + let mut hotkeys = OwnedHotkeys::::get(coldkey); + if !hotkeys.contains(hotkey) { + hotkeys.push(hotkey.clone()); + OwnedHotkeys::::insert(coldkey, hotkeys); + } } } @@ -781,6 +790,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 @@ -795,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/src/swap.rs b/pallets/subtensor/src/swap.rs index ce090d736..49af72a3f 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); @@ -92,6 +82,82 @@ 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 { + ensure_signed(origin)?; + + let mut weight = T::DbWeight::get().reads(2); + + // 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(); + + // 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); + Self::swap_total_hotkey_coldkey_stakes_this_interval_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::kill_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 @@ -127,6 +193,15 @@ impl Pallet { ) { Owner::::remove(old_hotkey); Owner::::insert(new_hotkey, coldkey.clone()); + + // 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); + OwnedHotkeys::::insert(coldkey, hotkeys); + weight.saturating_accrue(T::DbWeight::get().writes(2)); } @@ -396,4 +471,174 @@ 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, + ) { + // Find all hotkeys for this 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); + Stake::::remove(&hotkey, old_coldkey); + Stake::::insert(&hotkey, new_coldkey, 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, + ) { + 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); + weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 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, + ) { + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0)); + for hotkey in OwnedHotkeys::::get(old_coldkey).iter() { + let (stake, block) = + TotalHotkeyColdkeyStakesThisInterval::::get(&hotkey, old_coldkey); + TotalHotkeyColdkeyStakesThisInterval::::remove(&hotkey, old_coldkey); + TotalHotkeyColdkeyStakesThisInterval::::insert(&hotkey, new_coldkey, (stake, block)); + weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2)); + } + } + + /// 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)); + } + + /// 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 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/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(), + ); + }); +} diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index af7d19d2d..3d7454fb7 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1048,3 +1048,335 @@ 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 hotkey1 = U256::from(3); + let hotkey2 = U256::from(4); + let netuid = 1u16; + 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, hotkey1, old_coldkey, 0); + register_ok_neuron(netuid, hotkey2, old_coldkey, 0); + + // 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), + hotkey2, + stake_amount2 + )); + + // Verify initial stakes and balances + assert_eq!( + TotalColdkeyStake::::get(old_coldkey), + stake_amount1 + stake_amount2 + ); + 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 + ); + + // Perform the swap + assert_ok!(SubtensorModule::do_swap_coldkey( + <::RuntimeOrigin>::signed(old_coldkey), + &old_coldkey, + &new_coldkey + )); + + // Verify the swap + 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(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), + free_balance_old + ); + assert_eq!(SubtensorModule::get_coldkey_balance(&old_coldkey), 0); + + // Verify event emission + System::assert_last_event( + Event::ColdkeySwapped { + old_coldkey, + new_coldkey, + } + .into(), + ); + }); +} + +#[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(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]); + + // Perform the swap + SubtensorModule::swap_stake_for_coldkey(&old_coldkey, &new_coldkey, &mut weight); + + // Verify the swap + 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 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); + }); +} + +#[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); + + // Initialize OwnedHotkeys map + OwnedHotkeys::::insert(old_coldkey, vec![hotkey1, hotkey2]); + + // 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(hotkey1, old_coldkey, stake1); + TotalHotkeyColdkeyStakesThisInterval::::insert(hotkey2, old_coldkey, stake2); + + // Populate OwnedHotkeys map + OwnedHotkeys::::insert(old_coldkey, vec![hotkey1, hotkey2]); + + // 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(hotkey1, new_coldkey), + stake1 + ); + assert_eq!( + TotalHotkeyColdkeyStakesThisInterval::::get(hotkey2, new_coldkey), + 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(5, 4); + 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: u64 = 1000u64; + + // 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); + SubnetOwner::::insert(netuid, old_coldkey); + + // Populate OwnedHotkeys map + OwnedHotkeys::::insert(old_coldkey, vec![hotkey]); + + // 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_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/runtime/src/lib.rs b/runtime/src/lib.rs index d3ba9c24d..44b4d9b00 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: 155, + spec_version: 156, impl_version: 1, apis: RUNTIME_API_VERSIONS, transaction_version: 1, 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