diff --git a/node/src/rpc.rs b/node/src/rpc.rs index 54f82447f..c169b8530 100644 --- a/node/src/rpc.rs +++ b/node/src/rpc.rs @@ -60,6 +60,7 @@ where C::Api: subtensor_custom_rpc_runtime_api::NeuronInfoRuntimeApi, C::Api: subtensor_custom_rpc_runtime_api::SubnetInfoRuntimeApi, C::Api: subtensor_custom_rpc_runtime_api::SubnetRegistrationRuntimeApi, + C::Api: subtensor_custom_rpc_runtime_api::ColdkeySwapRuntimeApi, B: sc_client_api::Backend + Send + Sync + 'static, P: TransactionPool + 'static, { diff --git a/pallets/admin-utils/tests/mock.rs b/pallets/admin-utils/tests/mock.rs index a78eb6d3d..5f71626ad 100644 --- a/pallets/admin-utils/tests/mock.rs +++ b/pallets/admin-utils/tests/mock.rs @@ -114,6 +114,7 @@ parameter_types! { pub const InitialAlphaHigh: u16 = 58982; // Represents 0.9 as per the production default pub const InitialAlphaLow: u16 = 45875; // Represents 0.7 as per the production default pub const InitialLiquidAlphaOn: bool = false; // Default value for LiquidAlphaOn + pub const InitialBaseDifficulty: u64 = 10_000; // Base difficulty } impl pallet_subtensor::Config for Test { @@ -169,6 +170,7 @@ impl pallet_subtensor::Config for Test { type AlphaHigh = InitialAlphaHigh; type AlphaLow = InitialAlphaLow; type LiquidAlphaOn = InitialLiquidAlphaOn; + type InitialBaseDifficulty = InitialBaseDifficulty; } #[derive_impl(frame_system::config_preludes::TestDefaultConfig)] diff --git a/pallets/subtensor/rpc/src/lib.rs b/pallets/subtensor/rpc/src/lib.rs index 2f71e9c21..15157be2f 100644 --- a/pallets/subtensor/rpc/src/lib.rs +++ b/pallets/subtensor/rpc/src/lib.rs @@ -12,7 +12,7 @@ use std::sync::Arc; use sp_api::ProvideRuntimeApi; pub use subtensor_custom_rpc_runtime_api::{ - DelegateInfoRuntimeApi, NeuronInfoRuntimeApi, SubnetInfoRuntimeApi, + ColdkeySwapRuntimeApi, DelegateInfoRuntimeApi, NeuronInfoRuntimeApi, SubnetInfoRuntimeApi, SubnetRegistrationRuntimeApi, }; @@ -51,6 +51,24 @@ pub trait SubtensorCustomApi { #[method(name = "subnetInfo_getLockCost")] fn get_network_lock_cost(&self, at: Option) -> RpcResult; + #[method(name = "coldkeySwap_getScheduledColdkeySwap")] + fn get_scheduled_coldkey_swap( + &self, + coldkey_account_vec: Vec, + at: Option, + ) -> RpcResult>; + #[method(name = "coldkeySwap_getRemainingArbitrationPeriod")] + fn get_remaining_arbitration_period( + &self, + coldkey_account_vec: Vec, + at: Option, + ) -> RpcResult>; + #[method(name = "coldkeySwap_getColdkeySwapDestinations")] + fn get_coldkey_swap_destinations( + &self, + coldkey_account_vec: Vec, + at: Option, + ) -> RpcResult>; } pub struct SubtensorCustom { @@ -99,6 +117,7 @@ where C::Api: NeuronInfoRuntimeApi, C::Api: SubnetInfoRuntimeApi, C::Api: SubnetRegistrationRuntimeApi, + C::Api: ColdkeySwapRuntimeApi, { fn get_delegates(&self, at: Option<::Hash>) -> RpcResult> { let api = self.client.runtime_api(); @@ -223,4 +242,51 @@ where Error::RuntimeError(format!("Unable to get subnet lock cost: {:?}", e)).into() }) } + + fn get_scheduled_coldkey_swap( + &self, + coldkey_account_vec: Vec, + at: Option<::Hash>, + ) -> RpcResult> { + let api = self.client.runtime_api(); + let at = at.unwrap_or_else(|| self.client.info().best_hash); + + api.get_scheduled_coldkey_swap(at, coldkey_account_vec) + .map_err(|e| { + Error::RuntimeError(format!("Unable to get scheduled coldkey swap: {:?}", e)).into() + }) + } + + fn get_remaining_arbitration_period( + &self, + coldkey_account_vec: Vec, + at: Option<::Hash>, + ) -> RpcResult> { + let api = self.client.runtime_api(); + let at = at.unwrap_or_else(|| self.client.info().best_hash); + + api.get_remaining_arbitration_period(at, coldkey_account_vec) + .map_err(|e| { + Error::RuntimeError(format!( + "Unable to get remaining arbitration period: {:?}", + e + )) + .into() + }) + } + + fn get_coldkey_swap_destinations( + &self, + coldkey_account_vec: Vec, + at: Option<::Hash>, + ) -> RpcResult> { + let api = self.client.runtime_api(); + let at = at.unwrap_or_else(|| self.client.info().best_hash); + + api.get_coldkey_swap_destinations(at, coldkey_account_vec) + .map_err(|e| { + Error::RuntimeError(format!("Unable to get coldkey swap destinations: {:?}", e)) + .into() + }) + } } diff --git a/pallets/subtensor/runtime-api/src/lib.rs b/pallets/subtensor/runtime-api/src/lib.rs index 9095ad54a..a647d3619 100644 --- a/pallets/subtensor/runtime-api/src/lib.rs +++ b/pallets/subtensor/runtime-api/src/lib.rs @@ -32,4 +32,10 @@ sp_api::decl_runtime_apis! { pub trait SubnetRegistrationRuntimeApi { fn get_network_registration_cost() -> u64; } + + pub trait ColdkeySwapRuntimeApi { + fn get_scheduled_coldkey_swap( coldkey_account_vec: Vec ) -> Vec; + fn get_remaining_arbitration_period( coldkey_account_vec: Vec ) -> Vec; + fn get_coldkey_swap_destinations( coldkey_account_vec: Vec ) -> Vec; + } } diff --git a/pallets/subtensor/src/benchmarks.rs b/pallets/subtensor/src/benchmarks.rs index 403ba413e..6da83c5af 100644 --- a/pallets/subtensor/src/benchmarks.rs +++ b/pallets/subtensor/src/benchmarks.rs @@ -428,4 +428,31 @@ reveal_weights { let _ = Subtensor::::commit_weights(::RuntimeOrigin::from(RawOrigin::Signed(hotkey.clone())), netuid, commit_hash); }: reveal_weights(RawOrigin::Signed(hotkey.clone()), netuid, uids, weight_values, salt, version_key) + + schedule_coldkey_swap { + let seed: u32 = 1; + let old_coldkey: T::AccountId = account("OldColdkey", 0, seed); + let new_coldkey: T::AccountId = account("NewColdkey", 0, seed + 1); + let hotkey: T::AccountId = account("Hotkey", 0, seed); + + let netuid = 1u16; + let tempo = 1u16; + let block_number: u64 = Subtensor::::get_current_block_as_u64(); + let nonce = 0; + + // Initialize the network + Subtensor::::init_new_network(netuid, tempo); + Subtensor::::set_network_registration_allowed(netuid, true); + + // Add balance to the old coldkey account + let amount_to_be_staked: u64 = 1000000u32.into(); + Subtensor::::add_balance_to_coldkey_account(&old_coldkey.clone(), amount_to_be_staked+1000000000); + // Burned register the hotkey with the old coldkey + assert_ok!(Subtensor::::burned_register( + RawOrigin::Signed(old_coldkey.clone()).into(), + netuid, + hotkey.clone() + )); + + }: schedule_coldkey_swap(RawOrigin::Signed(old_coldkey.clone()), new_coldkey.clone(), vec![], block_number, nonce) } diff --git a/pallets/subtensor/src/delegate_info.rs b/pallets/subtensor/src/delegate_info.rs index 1f8b06b69..e1eab1210 100644 --- a/pallets/subtensor/src/delegate_info.rs +++ b/pallets/subtensor/src/delegate_info.rs @@ -40,11 +40,7 @@ impl Pallet { let mut emissions_per_day: U64F64 = U64F64::from_num(0); for netuid in registrations.iter() { - let _uid = Self::get_uid_for_net_and_hotkey(*netuid, &delegate.clone()); - if _uid.is_err() { - continue; // this should never happen - } else { - let uid = _uid.expect("Delegate's UID should be ok"); + if let Ok(uid) = Self::get_uid_for_net_and_hotkey(*netuid, &delegate.clone()) { let validator_permit = Self::get_validator_permit_for_uid(*netuid, uid); if validator_permit { validator_permits.push((*netuid).into()); @@ -52,9 +48,11 @@ impl Pallet { let emission: U64F64 = Self::get_emission_for_uid(*netuid, uid).into(); let tempo: U64F64 = Self::get_tempo(*netuid).into(); - let epochs_per_day: U64F64 = U64F64::from_num(7200).saturating_div(tempo); - emissions_per_day = - emissions_per_day.saturating_add(emission.saturating_mul(epochs_per_day)); + if tempo > U64F64::from_num(0) { + let epochs_per_day: U64F64 = U64F64::from_num(7200).saturating_div(tempo); + emissions_per_day = + emissions_per_day.saturating_add(emission.saturating_mul(epochs_per_day)); + } } } @@ -63,15 +61,15 @@ impl Pallet { let total_stake: U64F64 = Self::get_total_stake_for_hotkey(&delegate.clone()).into(); - let mut return_per_1000: U64F64 = U64F64::from_num(0); - - if total_stake > U64F64::from_num(0) { - return_per_1000 = emissions_per_day + let return_per_1000: U64F64 = if total_stake > U64F64::from_num(0) { + emissions_per_day .saturating_mul(U64F64::from_num(0.82)) - .saturating_div(total_stake.saturating_div(U64F64::from_num(1000))); - } + .saturating_div(total_stake.saturating_div(U64F64::from_num(1000))) + } else { + U64F64::from_num(0) + }; - return DelegateInfo { + DelegateInfo { delegate_ss58: delegate.clone(), take, nominators, @@ -80,7 +78,7 @@ impl Pallet { validator_permits, return_per_1000: U64F64::to_num::(return_per_1000).into(), total_daily_return: U64F64::to_num::(emissions_per_day).into(), - }; + } } pub fn get_delegate(delegate_account_vec: Vec) -> Option> { @@ -132,4 +130,47 @@ impl Pallet { delegates } + + pub fn get_total_delegated_stake(coldkey: &T::AccountId) -> u64 { + let mut total_delegated = 0u64; + + // Get all hotkeys associated with this coldkey + let hotkeys = StakingHotkeys::::get(coldkey); + + for hotkey in hotkeys { + let owner = Owner::::get(&hotkey); + + for (delegator, stake) in Stake::::iter_prefix(&hotkey) { + if delegator != owner { + total_delegated = total_delegated.saturating_add(stake); + } + } + } + + log::info!( + "Total delegated stake for coldkey {:?}: {}", + coldkey, + total_delegated + ); + total_delegated + } + + // Helper function to get total delegated stake for a hotkey + pub fn get_total_hotkey_delegated_stake(hotkey: &T::AccountId) -> u64 { + let mut total_delegated = 0u64; + + // Iterate through all delegators for this hotkey + for (delegator, stake) in Stake::::iter_prefix(hotkey) { + if delegator != Self::get_coldkey_for_hotkey(hotkey) { + total_delegated = total_delegated.saturating_add(stake); + } + } + + total_delegated + } + + // Helper function to get the coldkey associated with a hotkey + pub fn get_coldkey_for_hotkey(hotkey: &T::AccountId) -> T::AccountId { + Owner::::get(hotkey) + } } diff --git a/pallets/subtensor/src/errors.rs b/pallets/subtensor/src/errors.rs index 3e30c094c..100be89b7 100644 --- a/pallets/subtensor/src/errors.rs +++ b/pallets/subtensor/src/errors.rs @@ -146,5 +146,15 @@ mod errors { NoBalanceToTransfer, /// Same coldkey SameColdkey, + /// The coldkey is in arbitration + ColdkeyIsInArbitration, + /// The new coldkey is already registered for the drain + DuplicateColdkey, + /// Error thrown on a coldkey swap. + ColdkeySwapError, + /// Insufficient Balance to Schedule coldkey swap + InsufficientBalanceToPerformColdkeySwap, + /// The maximum number of coldkey destinations has been reached + MaxColdkeyDestinationsReached, } } diff --git a/pallets/subtensor/src/events.rs b/pallets/subtensor/src/events.rs index d401eebe1..2ed0d621d 100644 --- a/pallets/subtensor/src/events.rs +++ b/pallets/subtensor/src/events.rs @@ -150,5 +150,19 @@ mod events { ::AccountId, >>::Balance, }, + /// A coldkey swap has been scheduled + ColdkeySwapScheduled { + /// The account ID of the old coldkey + old_coldkey: T::AccountId, + /// The account ID of the new coldkey + new_coldkey: T::AccountId, + /// The arbitration block for the coldkey swap + arbitration_block: u64, + }, + /// The arbitration period has been extended + ArbitrationPeriodExtended { + /// The account ID of the coldkey + coldkey: T::AccountId, + }, } } diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index c4b0c74cd..0560ca72c 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -18,6 +18,7 @@ use frame_support::{ use codec::{Decode, Encode}; use frame_support::sp_runtime::transaction_validity::InvalidTransaction; use frame_support::sp_runtime::transaction_validity::ValidTransaction; +use pallet_balances::Call as BalancesCall; use scale_info::TypeInfo; use sp_runtime::{ traits::{DispatchInfoOf, Dispatchable, PostDispatchInfoOf, SignedExtension}, @@ -50,6 +51,7 @@ mod weights; pub mod delegate_info; pub mod neuron_info; +pub mod schedule_coldkey_swap_info; pub mod stake_info; pub mod subnet_info; @@ -83,6 +85,9 @@ pub mod pallet { /// order of migrations. (i.e. always increasing) const STORAGE_VERSION: StorageVersion = StorageVersion::new(6); + /// Minimum balance required to perform a coldkey swap + pub const MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP: u64 = 100_000_000; // 0.1 TAO in RAO + #[pallet::pallet] #[pallet::without_storage_info] #[pallet::storage_version(STORAGE_VERSION)] @@ -251,6 +256,9 @@ pub mod pallet { /// A flag to indicate if Liquid Alpha is enabled. #[pallet::constant] type LiquidAlphaOn: Get; + /// The base difficulty for proof of work for coldkey swaps + #[pallet::constant] + type InitialBaseDifficulty: Get; } /// Alias for the account ID. @@ -327,6 +335,12 @@ pub mod pallet { 360 } + /// Default base difficulty for proof of work for coldkey swaps + #[pallet::type_value] + pub fn DefaultBaseDifficulty() -> u64 { + T::InitialBaseDifficulty::get() + } + #[pallet::storage] // --- ITEM ( total_stake ) pub type TotalStake = StorageValue<_, u64, ValueQuery>; #[pallet::storage] // --- ITEM ( default_take ) @@ -340,6 +354,9 @@ pub mod pallet { #[pallet::storage] // --- ITEM (target_stakes_per_interval) pub type TargetStakesPerInterval = StorageValue<_, u64, ValueQuery, DefaultTargetStakesPerInterval>; + + #[pallet::storage] // --- ITEM ( base_difficulty ) + pub type BaseDifficulty = StorageValue<_, u64, ValueQuery, DefaultBaseDifficulty>; #[pallet::storage] // --- ITEM (default_stake_interval) pub type StakeInterval = StorageValue<_, u64, ValueQuery, DefaultStakeInterval>; #[pallet::storage] // --- MAP ( hot ) --> stake | Returns the total amount of stake under a hotkey. @@ -366,6 +383,9 @@ pub mod pallet { #[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] // --- DMAP ( cold ) --> Vec | Maps coldkey to hotkeys that stake to it + pub type StakingHotkeys = + 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>; @@ -380,9 +400,37 @@ pub mod pallet { ValueQuery, DefaultAccountTake, >; - #[pallet::storage] // --- DMAP ( cold ) --> Vec | Maps coldkey to hotkeys that stake to it - pub type StakingHotkeys = - StorageMap<_, Blake2_128Concat, T::AccountId, Vec, ValueQuery>; + + #[pallet::type_value] + /// Default value for hotkeys. + pub fn EmptyAccounts() -> Vec { + vec![] + } + #[pallet::type_value] + /// Default arbitration period. + /// This value represents the default arbitration period in blocks. + /// The period is set to 18 hours, assuming a block time of 12 seconds. + pub fn DefaultArbitrationPeriod() -> u64 { + 7200 * 3 // 3 days + } + #[pallet::storage] // ---- StorageItem Global Used Work. + pub type ArbitrationPeriod = + StorageValue<_, u64, ValueQuery, DefaultArbitrationPeriod>; + #[pallet::storage] // --- MAP ( cold ) --> Vec | Returns a list of keys to drain to, if there are two, we extend the period. + pub type ColdkeySwapDestinations = StorageMap< + _, + Blake2_128Concat, + T::AccountId, + Vec, + ValueQuery, + EmptyAccounts, + >; + #[pallet::storage] // --- MAP ( cold ) --> u64 | Block when the coldkey will be arbitrated. + pub type ColdkeyArbitrationBlock = + StorageMap<_, Blake2_128Concat, T::AccountId, u64, ValueQuery>; + #[pallet::storage] // --- MAP ( u64 ) --> Vec | Coldkeys to drain on the specific block. + pub type ColdkeysToSwapAtBlock = + StorageMap<_, Identity, u64, Vec, ValueQuery, EmptyAccounts>; /// -- ITEM (switches liquid alpha on) #[pallet::type_value] pub fn DefaultLiquidAlpha() -> bool { @@ -1295,29 +1343,48 @@ pub mod pallet { #[pallet::hooks] impl Hooks> for Pallet { - // ---- Called on the initialization of this pallet. (the order of on_finalize calls is determined in the runtime) - // - // # Args: - // * 'n': (BlockNumberFor): - // - The number of the block we are initializing. + fn on_idle(_n: BlockNumberFor, _remaining_weight: Weight) -> Weight { + Weight::zero() + } + fn on_initialize(_block_number: BlockNumberFor) -> Weight { + let mut total_weight = Weight::zero(); + + // Create a Weight::MAX value to pass to swap_coldkeys_this_block + let max_weight = Weight::MAX; + + // Perform coldkey swapping + let swap_weight = match Self::swap_coldkeys_this_block(&max_weight) { + Ok(weight_used) => weight_used, + Err(e) => { + log::error!("Error while swapping coldkeys: {:?}", e); + Weight::zero() + } + }; + total_weight = total_weight.saturating_add(swap_weight); + + // Perform block step let block_step_result = Self::block_step(); match block_step_result { Ok(_) => { - // --- If the block step was successful, return the weight. - log::info!("Successfully ran block step."); - Weight::from_parts(110_634_229_000_u64, 0) - .saturating_add(T::DbWeight::get().reads(8304_u64)) - .saturating_add(T::DbWeight::get().writes(110_u64)) + log::debug!("Successfully ran block step."); + total_weight = total_weight.saturating_add( + Weight::from_parts(110_634_229_000_u64, 0) + .saturating_add(T::DbWeight::get().reads(8304_u64)) + .saturating_add(T::DbWeight::get().writes(110_u64)), + ); } Err(e) => { - // --- If the block step was unsuccessful, return the weight anyway. log::error!("Error while stepping block: {:?}", e); - Weight::from_parts(110_634_229_000_u64, 0) - .saturating_add(T::DbWeight::get().reads(8304_u64)) - .saturating_add(T::DbWeight::get().writes(110_u64)) + total_weight = total_weight.saturating_add( + Weight::from_parts(110_634_229_000_u64, 0) + .saturating_add(T::DbWeight::get().reads(8304_u64)) + .saturating_add(T::DbWeight::get().writes(110_u64)), + ); } } + + total_weight } fn on_runtime_upgrade() -> frame_support::weights::Weight { @@ -2015,10 +2082,9 @@ pub mod pallet { .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) + Self::do_swap_coldkey(origin, &new_coldkey) } /// Unstakes all tokens associated with a hotkey and transfers them to a new coldkey. @@ -2037,15 +2103,19 @@ pub mod pallet { /// /// 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( + #[pallet::weight((Weight::from_parts(21_000_000, 0) + .saturating_add(T::DbWeight::get().reads(3)) + .saturating_add(T::DbWeight::get().writes(3)), DispatchClass::Operational, Pays::No))] + pub fn schedule_coldkey_swap( origin: OriginFor, new_coldkey: T::AccountId, + work: Vec, + block_number: u64, + nonce: u64, ) -> DispatchResult { - let current_coldkey = ensure_signed(origin)?; - Self::do_unstake_all_and_transfer_to_new_coldkey(current_coldkey, new_coldkey) + // Attain the calling coldkey from the origin. + let old_coldkey: T::AccountId = ensure_signed(origin)?; + Self::do_schedule_coldkey_swap(&old_coldkey, &new_coldkey, work, block_number, nonce) } // ---- SUDO ONLY FUNCTIONS ------------------------------------------------------------ @@ -2278,10 +2348,12 @@ impl sp_std::fmt::Debug for SubtensorSignedE } } -impl SignedExtension for SubtensorSignedExtension +impl SignedExtension + for SubtensorSignedExtension where T::RuntimeCall: Dispatchable, ::RuntimeCall: IsSubType>, + ::RuntimeCall: IsSubType>, { const IDENTIFIER: &'static str = "SubtensorSignedExtension"; @@ -2301,6 +2373,19 @@ where _info: &DispatchInfoOf, _len: usize, ) -> TransactionValidity { + // Check if the call is one of the balance transfer types we want to reject + if let Some(balances_call) = call.is_sub_type() { + match balances_call { + BalancesCall::transfer_allow_death { .. } + | BalancesCall::transfer_keep_alive { .. } + | BalancesCall::transfer_all { .. } => { + if Pallet::::coldkey_in_arbitration(who) { + return Err(TransactionValidityError::Invalid(InvalidTransaction::Call)); + } + } + _ => {} // Other Balances calls are allowed + } + } match call.is_sub_type() { Some(Call::commit_weights { netuid, .. }) => { if Self::check_weights_min_stake(who) { @@ -2377,6 +2462,16 @@ where priority: Self::get_priority_vanilla(), ..Default::default() }), + Some(Call::dissolve_network { .. }) => { + if Pallet::::coldkey_in_arbitration(who) { + Err(TransactionValidityError::Invalid(InvalidTransaction::Call)) + } else { + Ok(ValidTransaction { + priority: Self::get_priority_vanilla(), + ..Default::default() + }) + } + } _ => Ok(ValidTransaction { priority: Self::get_priority_vanilla(), ..Default::default() diff --git a/pallets/subtensor/src/registration.rs b/pallets/subtensor/src/registration.rs index 4688bcbb5..6b73f2fc3 100644 --- a/pallets/subtensor/src/registration.rs +++ b/pallets/subtensor/src/registration.rs @@ -41,6 +41,10 @@ impl Pallet { ) -> DispatchResult { // --- 1. Check that the caller has signed the transaction. (the coldkey of the pairing) let coldkey = ensure_signed(origin)?; + ensure!( + !Self::coldkey_in_arbitration(&coldkey), + Error::::ColdkeyIsInArbitration + ); log::info!( "do_registration( coldkey:{:?} netuid:{:?} hotkey:{:?} )", coldkey, diff --git a/pallets/subtensor/src/root.rs b/pallets/subtensor/src/root.rs index 9e0327fb3..f39ae97e0 100644 --- a/pallets/subtensor/src/root.rs +++ b/pallets/subtensor/src/root.rs @@ -483,6 +483,10 @@ impl Pallet { // --- 1. Ensure that the call originates from a signed source and retrieve the caller's account ID (coldkey). let coldkey = ensure_signed(origin)?; + ensure!( + !Self::coldkey_in_arbitration(&coldkey), + Error::::ColdkeyIsInArbitration + ); log::info!( "do_root_register( coldkey: {:?}, hotkey: {:?} )", coldkey, @@ -623,6 +627,10 @@ impl Pallet { ) -> dispatch::DispatchResult { // Check the caller's signature. This is the coldkey of a registered account. let coldkey = ensure_signed(origin)?; + ensure!( + !Self::coldkey_in_arbitration(&coldkey), + Error::::ColdkeyIsInArbitration + ); log::info!( "do_set_root_weights( origin:{:?} netuid:{:?}, uids:{:?}, values:{:?})", coldkey, @@ -744,6 +752,10 @@ impl Pallet { ) -> DispatchResultWithPostInfo { // --- 1. Ensure that the caller has signed with their coldkey. let coldkey = ensure_signed(origin.clone())?; + ensure!( + !Self::coldkey_in_arbitration(&coldkey), + Error::::ColdkeyIsInArbitration + ); // --- 2. Ensure that the calling coldkey owns the associated hotkey. ensure!( @@ -797,6 +809,10 @@ impl Pallet { pub fn user_add_network(origin: T::RuntimeOrigin) -> dispatch::DispatchResult { // --- 0. Ensure the caller is a signed user. let coldkey = ensure_signed(origin)?; + ensure!( + !Self::coldkey_in_arbitration(&coldkey), + Error::::ColdkeyIsInArbitration + ); // --- 1. Rate limit for network registrations. let current_block = Self::get_current_block_as_u64(); @@ -885,6 +901,10 @@ impl Pallet { pub fn user_remove_network(origin: T::RuntimeOrigin, netuid: u16) -> dispatch::DispatchResult { // --- 1. Ensure the function caller is a signed user. let coldkey = ensure_signed(origin)?; + ensure!( + !Self::coldkey_in_arbitration(&coldkey), + Error::::ColdkeyIsInArbitration + ); // --- 2. Ensure this subnet exists. ensure!( diff --git a/pallets/subtensor/src/schedule_coldkey_swap_info.rs b/pallets/subtensor/src/schedule_coldkey_swap_info.rs new file mode 100644 index 000000000..e23ad23b6 --- /dev/null +++ b/pallets/subtensor/src/schedule_coldkey_swap_info.rs @@ -0,0 +1,176 @@ +use super::*; +use codec::Compact; +use frame_support::pallet_prelude::{Decode, Encode}; + +use sp_core::hexdisplay::AsBytesRef; + +#[derive(Decode, Encode, PartialEq, Eq, Clone, Debug)] +pub struct ScheduledColdkeySwapInfo { + old_coldkey: T::AccountId, + new_coldkey: T::AccountId, + arbitration_block: Compact, +} + +impl Pallet { + /// Retrieves the scheduled coldkey swap information for an existing account. + /// + /// # Arguments + /// + /// * `coldkey` - The account ID of the coldkey to check. + /// + /// # Returns + /// + /// * `Option>` - The scheduled coldkey swap information if it exists, otherwise `None`. + /// + /// # Notes + /// + /// This function checks if there are any destination coldkeys associated with the given coldkey. + /// If there are, it retrieves the arbitration block and constructs the `ScheduledColdkeySwapInfo` struct. + fn get_scheduled_coldkey_swap_by_existing_account( + coldkey: AccountIdOf, + ) -> Option> { + let destinations: Vec = ColdkeySwapDestinations::::get(&coldkey); + if destinations.is_empty() { + return None; + } + + let arbitration_block: u64 = ColdkeyArbitrationBlock::::get(&coldkey); + + Some(ScheduledColdkeySwapInfo { + old_coldkey: coldkey, + new_coldkey: destinations.first().cloned().unwrap_or_else(|| { + T::AccountId::decode(&mut sp_runtime::traits::TrailingZeroInput::zeroes()) + .expect("Infinite length input; no invalid inputs for type; qed") + }), + arbitration_block: arbitration_block.into(), + }) + } + + /// Retrieves the scheduled coldkey swap information for a given coldkey account vector. + /// + /// # Arguments + /// + /// * `coldkey_account_vec` - The vector of bytes representing the coldkey account. + /// + /// # Returns + /// + /// * `Option>` - The scheduled coldkey swap information if it exists, otherwise `None`. + /// + /// # Notes + /// + /// This function decodes the coldkey account vector into an account ID and then calls + /// `get_scheduled_coldkey_swap_by_existing_account` to retrieve the swap information. + pub fn get_scheduled_coldkey_swap( + coldkey_account_vec: Vec, + ) -> Option> { + if coldkey_account_vec.len() != 32 { + return None; + } + + let coldkey: AccountIdOf = + T::AccountId::decode(&mut coldkey_account_vec.as_bytes_ref()).ok()?; + Self::get_scheduled_coldkey_swap_by_existing_account(coldkey) + } + + /// Retrieves all scheduled coldkey swaps from storage. + /// + /// # Returns + /// + /// * `Vec>` - A vector containing all scheduled coldkey swap information. + /// + /// # Notes + /// + /// This function iterates over all coldkeys in `ColdkeySwapDestinations` and retrieves their swap information + /// using `get_scheduled_coldkey_swap_by_existing_account`. + pub fn get_all_scheduled_coldkey_swaps() -> Vec> { + let mut scheduled_swaps: Vec> = Vec::new(); + for coldkey in ColdkeySwapDestinations::::iter_keys() { + if let Some(swap_info) = Self::get_scheduled_coldkey_swap_by_existing_account(coldkey) { + scheduled_swaps.push(swap_info); + } + } + scheduled_swaps + } + + /// Retrieves the scheduled coldkey swaps for a given block. + /// + /// # Arguments + /// + /// * `block` - The block number to check for scheduled coldkey swaps. + /// + /// # Returns + /// + /// * `Vec>` - A vector containing the scheduled coldkey swap information for the given block. + /// + /// # Notes + /// + /// This function retrieves the coldkeys to swap at the given block and then retrieves their swap information + /// using `get_scheduled_coldkey_swap_by_existing_account`. + pub fn get_scheduled_coldkey_swaps_at_block(block: u64) -> Vec> { + let coldkeys_to_swap: Vec = ColdkeysToSwapAtBlock::::get(block); + let mut scheduled_swaps: Vec> = Vec::new(); + for coldkey in coldkeys_to_swap { + if let Some(swap_info) = Self::get_scheduled_coldkey_swap_by_existing_account(coldkey) { + scheduled_swaps.push(swap_info); + } + } + scheduled_swaps + } + + /// Retrieves the remaining arbitration period for a given coldkey account vector. + /// + /// # Arguments + /// + /// * `coldkey_account_vec` - The vector of bytes representing the coldkey account. + /// + /// # Returns + /// + /// * `Option` - The remaining arbitration period in blocks if it exists, otherwise `None`. + /// + /// # Notes + /// + /// This function decodes the coldkey account vector into an account ID and calculates the remaining arbitration period + /// by subtracting the current block number from the arbitration block number. + pub fn get_remaining_arbitration_period(coldkey_account_vec: Vec) -> Option { + if coldkey_account_vec.len() != 32 { + return None; + } + + let coldkey: AccountIdOf = + T::AccountId::decode(&mut coldkey_account_vec.as_bytes_ref()).ok()?; + let current_block: u64 = Self::get_current_block_as_u64(); + let arbitration_block: u64 = ColdkeyArbitrationBlock::::get(&coldkey); + + if arbitration_block > current_block { + Some(arbitration_block.saturating_sub(current_block)) + } else { + Some(0) + } + } + + /// Retrieves the destination coldkeys for a given coldkey account vector. + /// + /// # Arguments + /// + /// * `coldkey_account_vec` - The vector of bytes representing the coldkey account. + /// + /// # Returns + /// + /// * `Option>` - A vector containing the destination coldkeys if they exist, otherwise `None`. + /// + /// # Notes + /// + /// This function decodes the coldkey account vector into an account ID and retrieves the destination coldkeys + /// from `ColdkeySwapDestinations`. + pub fn get_coldkey_swap_destinations( + coldkey_account_vec: Vec, + ) -> Option> { + if coldkey_account_vec.len() != 32 { + return None; + } + + let coldkey: AccountIdOf = + T::AccountId::decode(&mut coldkey_account_vec.as_bytes_ref()).ok()?; + Some(ColdkeySwapDestinations::::get(&coldkey)) + } +} diff --git a/pallets/subtensor/src/staking.rs b/pallets/subtensor/src/staking.rs index 59f8ca4de..199234a30 100644 --- a/pallets/subtensor/src/staking.rs +++ b/pallets/subtensor/src/staking.rs @@ -1,5 +1,4 @@ use super::*; -use dispatch::RawOrigin; use frame_support::{ storage::IterableStorageDoubleMap, traits::{ @@ -10,7 +9,6 @@ use frame_support::{ Imbalance, }, }; -use num_traits::Zero; impl Pallet { /// ---- The implementation for the extrinsic become_delegate: signals that this hotkey allows delegated stake. @@ -46,6 +44,10 @@ impl Pallet { ) -> dispatch::DispatchResult { // --- 1. We check the coldkey signuture. let coldkey = ensure_signed(origin)?; + ensure!( + !Self::coldkey_in_arbitration(&coldkey), + Error::::ColdkeyIsInArbitration + ); log::info!( "do_become_delegate( origin:{:?} hotkey:{:?}, take:{:?} )", coldkey, @@ -135,6 +137,10 @@ impl Pallet { hotkey, take ); + ensure!( + !Self::coldkey_in_arbitration(&coldkey), + Error::::ColdkeyIsInArbitration + ); // --- 2. Ensure we are delegating a known key. // Ensure that the coldkey is the owner. @@ -207,6 +213,10 @@ impl Pallet { hotkey, take ); + ensure!( + !Self::coldkey_in_arbitration(&coldkey), + Error::::ColdkeyIsInArbitration + ); // --- 2. Ensure we are delegating a known key. // Ensure that the coldkey is the owner. @@ -292,6 +302,10 @@ impl Pallet { hotkey, stake_to_be_added ); + ensure!( + !Self::coldkey_in_arbitration(&coldkey), + Error::::ColdkeyIsInArbitration + ); // Ensure the callers coldkey has enough stake to perform the transaction. ensure!( @@ -404,6 +418,10 @@ impl Pallet { hotkey, stake_to_be_removed ); + ensure!( + !Self::coldkey_in_arbitration(&coldkey), + Error::::ColdkeyIsInArbitration + ); // Ensure that the hotkey account exists this is only possible through registration. ensure!( @@ -850,114 +868,4 @@ 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. - /// * `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, - new_coldkey: T::AccountId, - ) -> DispatchResult { - // Ensure the new coldkey is different from the current one - ensure!(current_coldkey != new_coldkey, Error::::SameColdkey); - - // Get all the hotkeys associated with this coldkey - let hotkeys: Vec = OwnedHotkeys::::get(¤t_coldkey); - - // iterate over all hotkeys. - for next_hotkey in hotkeys { - ensure!( - Self::hotkey_account_exists(&next_hotkey), - Error::::HotKeyAccountNotExists - ); - ensure!( - Self::coldkey_owns_hotkey(¤t_coldkey, &next_hotkey), - Error::::NonAssociatedColdKey - ); - - // Get the current stake - let current_stake: u64 = - Self::get_stake_for_coldkey_and_hotkey(¤t_coldkey, &next_hotkey); - - // Unstake all balance if there's any stake - if current_stake > 0 { - Self::do_remove_stake( - RawOrigin::Signed(current_coldkey.clone()).into(), - next_hotkey.clone(), - current_stake, - )?; - } - } - - // Unstake all delegate stake make by this coldkey to non-owned hotkeys - let staking_hotkeys = StakingHotkeys::::get(¤t_coldkey); - - // iterate over all staking hotkeys. - for hotkey in staking_hotkeys { - // 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, - )?; - } - } - - 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(), - total_balance, - }); - - Ok(()) - } } diff --git a/pallets/subtensor/src/swap.rs b/pallets/subtensor/src/swap.rs index 6ea38a94c..24d67aeb4 100644 --- a/pallets/subtensor/src/swap.rs +++ b/pallets/subtensor/src/swap.rs @@ -1,6 +1,9 @@ use super::*; +use crate::MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP; +use frame_support::traits::fungible::Mutate; +use frame_support::traits::tokens::Preservation; use frame_support::{storage::IterableStorageDoubleMap, weights::Weight}; -use sp_core::Get; +use sp_core::{Get, U256}; impl Pallet { /// Swaps the hotkey of a coldkey account. @@ -28,6 +31,10 @@ impl Pallet { new_hotkey: &T::AccountId, ) -> DispatchResultWithPostInfo { let coldkey = ensure_signed(origin)?; + ensure!( + !Self::coldkey_in_arbitration(&coldkey), + Error::::ColdkeyIsInArbitration + ); let mut weight = T::DbWeight::get().reads(2); @@ -112,50 +119,313 @@ impl Pallet { /// 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 old_coldkey = ensure_signed(origin)?; + ensure!( + !Self::coldkey_in_arbitration(&old_coldkey), + Error::::ColdkeyIsInArbitration + ); - let mut weight = T::DbWeight::get().reads(2); + let mut weight: Weight = T::DbWeight::get().reads(2); - // Check if the new coldkey is already associated with any hotkeys + // Check that the coldkey is a new key (does not exist elsewhere.) ensure!( !Self::coldkey_has_associated_hotkeys(new_coldkey), Error::::ColdKeyAlreadyAssociated ); + // Check that the new coldkey is not a hotkey. + ensure!( + !Self::hotkey_account_exists(new_coldkey), + Error::::ColdKeyAlreadyAssociated + ); - let block: u64 = Self::get_current_block_as_u64(); + // Actually do the swap. + weight = weight.saturating_add( + Self::perform_swap_coldkey(&old_coldkey, new_coldkey) + .map_err(|_| Error::::ColdkeySwapError)?, + ); + + Self::set_last_tx_block(new_coldkey, Self::get_current_block_as_u64()); + 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()) + } + + /// Checks if a coldkey is currently in arbitration. + /// + /// # Arguments + /// + /// * `coldkey` - The account ID of the coldkey to check. + /// + /// # Returns + /// + /// * `bool` - True if the coldkey is in arbitration, false otherwise. + /// + /// # Notes + /// + /// This function compares the arbitration block number of the coldkey with the current block number. + pub fn coldkey_in_arbitration(coldkey: &T::AccountId) -> bool { + ColdkeyArbitrationBlock::::get(coldkey) > Self::get_current_block_as_u64() + } + + /// Returns the remaining arbitration period for a given coldkey. + /// + /// # Arguments + /// + /// * `coldkey` - The account ID of the coldkey to check. + /// + /// # Returns + /// + /// * `u64` - The remaining arbitration period in blocks. + /// + /// + /// # Notes + /// + /// This function calculates the remaining arbitration period by subtracting the current block number + /// from the arbitration block number of the coldkey. + // pub fn get_remaining_arbitration_period(coldkey: &T::AccountId) -> u64 { + // let current_block: u64 = Self::get_current_block_as_u64(); + // let arbitration_block: u64 = ColdkeyArbitrationBlock::::get(coldkey); + // if arbitration_block > current_block { + // arbitration_block.saturating_sub(current_block) + // } else { + // 0 + // } + // } + + pub fn meets_min_allowed_coldkey_balance(coldkey: &T::AccountId) -> bool { + let all_staked_keys: Vec = StakingHotkeys::::get(coldkey); + let mut total_staking_balance: u64 = 0; + for hotkey in all_staked_keys { + total_staking_balance = total_staking_balance + .saturating_add(Self::get_stake_for_coldkey_and_hotkey(coldkey, &hotkey)); + } + total_staking_balance = + total_staking_balance.saturating_add(Self::get_coldkey_balance(coldkey)); + total_staking_balance >= MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP + } + + /// Schedules a coldkey swap to a new coldkey with arbitration. + /// + /// # Arguments + /// + /// * `old_coldkey` - The account ID of the old coldkey. + /// * `new_coldkey` - The account ID of the new coldkey. + /// * `work` - The proof of work submitted by the caller. + /// * `block_number` - The block number at which the work was performed. + /// * `nonce` - The nonce used for the proof of work. + /// + /// # Returns + /// + /// * `DispatchResult` - The result of the dispatch. + /// + /// # Errors + /// + + /// - `SameColdkey`: The old coldkey is the same as the new coldkey. + /// - `DuplicateColdkey`: The new coldkey is already in the list of destination coldkeys. + /// - `MaxColdkeyDestinationsReached`: There are already the maximum allowed destination coldkeys for the old coldkey. + /// - `InsufficientBalanceToPerformColdkeySwap`: The old coldkey doesn't have the minimum required TAO balance. + /// - `InvalidDifficulty`: The proof of work is invalid or doesn't meet the required difficulty. + /// + /// # Notes + /// + /// This function ensures that the new coldkey is not already in the list of destination coldkeys. + /// It also checks for a minimum TAO balance and verifies the proof of work. + /// The difficulty of the proof of work increases exponentially with each subsequent call. + pub fn do_schedule_coldkey_swap( + old_coldkey: &T::AccountId, + new_coldkey: &T::AccountId, + work: Vec, + block_number: u64, + nonce: u64, + ) -> DispatchResult { + ensure!(old_coldkey != new_coldkey, Error::::SameColdkey); + + // Check if the old_coldkey is a subnet owner for any network + let is_subnet_owner = (0..=TotalNetworks::::get()) + .any(|netuid| SubnetOwner::::get(netuid) == *old_coldkey); + + // Check if the old_coldkey has more than 500 TAO delegated + let total_delegated = Self::get_total_delegated_stake(old_coldkey); + let has_sufficient_delegation = total_delegated > 500_000_000_000; // 500 TAO in RAO + + // Only check the minimum balance if the old_coldkey is not a subnet owner + // and doesn't have sufficient delegation + if !(is_subnet_owner || has_sufficient_delegation) { + ensure!( + Self::meets_min_allowed_coldkey_balance(old_coldkey), + Error::::InsufficientBalanceToPerformColdkeySwap + ); + } + + // Get current destination coldkeys + let mut destination_coldkeys: Vec = + ColdkeySwapDestinations::::get(old_coldkey.clone()); + + // Calculate difficulty based on the number of existing destination coldkeys + let difficulty = Self::calculate_pow_difficulty(destination_coldkeys.len() as u32); + let work_hash = Self::vec_to_hash(work.clone()); + ensure!( + Self::hash_meets_difficulty(&work_hash, difficulty), + Error::::InvalidDifficulty + ); + + // Verify work is the product of the nonce, the block number, and coldkey + let seal = Self::create_seal_hash(block_number, nonce, old_coldkey); + ensure!(seal == work_hash, Error::::InvalidSeal); + + // Check if the new coldkey is already in the swap wallets list + ensure!( + !destination_coldkeys.contains(new_coldkey), + Error::::DuplicateColdkey + ); + + // If the destinations keys are empty or have less than the maximum allowed, we will add the new coldkey to the list + const MAX_COLDKEY_DESTINATIONS: usize = 10; + + if destination_coldkeys.len() < MAX_COLDKEY_DESTINATIONS { + destination_coldkeys.push(new_coldkey.clone()); + ColdkeySwapDestinations::::insert(old_coldkey.clone(), destination_coldkeys.clone()); + } else { + return Err(Error::::MaxColdkeyDestinationsReached.into()); + } + + // It is the first time we have seen this key + if destination_coldkeys.len() == 1_usize { + // Set the arbitration block for this coldkey + let arbitration_block: u64 = + Self::get_current_block_as_u64().saturating_add(ArbitrationPeriod::::get()); + ColdkeyArbitrationBlock::::insert(old_coldkey.clone(), arbitration_block); + + // Update the list of coldkeys to arbitrate on this block + let mut key_to_arbitrate_on_this_block: Vec = + ColdkeysToSwapAtBlock::::get(arbitration_block); + if !key_to_arbitrate_on_this_block.contains(old_coldkey) { + key_to_arbitrate_on_this_block.push(old_coldkey.clone()); + } + ColdkeysToSwapAtBlock::::insert(arbitration_block, key_to_arbitrate_on_this_block); + } + + // Emit an event indicating that a coldkey swap has been scheduled + Self::deposit_event(Event::ColdkeySwapScheduled { + old_coldkey: old_coldkey.clone(), + new_coldkey: new_coldkey.clone(), + arbitration_block: ColdkeyArbitrationBlock::::get(old_coldkey), + }); + + Ok(()) + } + + /// Calculate the proof of work difficulty based on the number of swap attempts + #[allow(clippy::arithmetic_side_effects)] + pub fn calculate_pow_difficulty(swap_attempts: u32) -> U256 { + let base_difficulty: U256 = U256::from(BaseDifficulty::::get()); // Base difficulty + base_difficulty.saturating_mul(U256::from(2).pow(U256::from(swap_attempts))) + } + + /// Arbitrates coldkeys that are scheduled to be swapped on this block. + /// + /// This function retrieves the list of coldkeys scheduled to be swapped on the current block, + /// and processes each coldkey by either extending the arbitration period or performing the swap + /// to the new coldkey. + /// + /// # Returns + /// + /// * `Weight` - The total weight consumed by this operation + pub fn swap_coldkeys_this_block(_weight_limit: &Weight) -> Result { + let mut weight_used = frame_support::weights::Weight::from_parts(0, 0); + + let current_block: u64 = Self::get_current_block_as_u64(); + log::debug!("Swapping coldkeys for block: {:?}", current_block); + + let source_coldkeys: Vec = ColdkeysToSwapAtBlock::::get(current_block); + ColdkeysToSwapAtBlock::::remove(current_block); + weight_used = weight_used.saturating_add(T::DbWeight::get().reads_writes(1, 1)); + + let mut keys_swapped = 0u64; + for coldkey_i in source_coldkeys.iter() { + // TODO: need a sane way to terminate early without locking users in. + // we should update the swap time + // if weight_used.ref_time() > weight_limit.ref_time() { + // log::warn!("Could not finish swapping all coldkeys this block due to weight limit, breaking after swapping {} keys.", keys_swapped); + // break; + // } + + let destinations_coldkeys: Vec = + ColdkeySwapDestinations::::get(coldkey_i); + weight_used = weight_used.saturating_add(T::DbWeight::get().reads(1)); + + if destinations_coldkeys.len() > 1 { + // Do not remove ColdkeySwapDestinations if there are multiple destinations + ColdkeyArbitrationBlock::::insert(coldkey_i.clone(), u64::MAX); + Self::deposit_event(Event::ArbitrationPeriodExtended { + coldkey: coldkey_i.clone(), + }); + } else if let Some(new_coldkey) = destinations_coldkeys.first() { + // Only remove ColdkeySwapDestinations if there's a single destination + ColdkeySwapDestinations::::remove(&coldkey_i); + weight_used = weight_used.saturating_add(T::DbWeight::get().writes(1)); + Self::perform_swap_coldkey(coldkey_i, new_coldkey).map(|weight| { + weight_used = weight_used.saturating_add(weight); + keys_swapped = keys_swapped.saturating_add(1); + })?; + } + } + + Ok(weight_used) + } + + pub fn perform_swap_coldkey( + old_coldkey: &T::AccountId, + new_coldkey: &T::AccountId, + ) -> Result { + log::info!( + "Performing swap for coldkey: {:?} to {:?}", + old_coldkey, + new_coldkey + ); + // Init the weight. + let mut weight = frame_support::weights::Weight::from_parts(0, 0); // 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)?; + if let Err(e) = Self::kill_coldkey_account(old_coldkey, remaining_balance) { + return Err(e.into()); + } 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(), - }); + // Swap the coldkey. + let total_balance: u64 = Self::get_coldkey_balance(old_coldkey); + if total_balance > 0 { + // Attempt to transfer the entire total balance to new_coldkey. + T::Currency::transfer( + old_coldkey, + new_coldkey, + total_balance, + Preservation::Expendable, + )?; + } - Ok(Some(weight).into()) + Ok(weight) } /// Retrieves the network membership status for a given hotkey. @@ -277,9 +547,14 @@ impl Pallet { let mut staking_hotkeys = StakingHotkeys::::get(&coldkey); if !staking_hotkeys.contains(new_hotkey) { staking_hotkeys.push(new_hotkey.clone()); - StakingHotkeys::::insert(coldkey.clone(), staking_hotkeys); writes = writes.saturating_add(1u64); // One write for insert } + if let Some(pos) = staking_hotkeys.iter().position(|x| x == old_hotkey) { + staking_hotkeys.remove(pos); + writes = writes.saturating_add(1u64); // One write for remove + } + StakingHotkeys::::insert(coldkey.clone(), staking_hotkeys); + writes = writes.saturating_add(1u64); // One write for insert } // Clear the prefix for the old hotkey after transferring all stakes @@ -506,7 +781,7 @@ impl Pallet { 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. + /// Swaps the stake associated with a coldkey from the old coldkey to the new coldkey. /// /// # Arguments /// @@ -516,55 +791,96 @@ impl Pallet { /// /// # Effects /// - /// * Removes all stakes associated with the old coldkey. - /// * Inserts all stakes for the new coldkey. + /// * Transfers all stakes from the old coldkey to the new coldkey. + /// * Updates the ownership of hotkeys. + /// * Updates the total stake for both old and new coldkeys. /// * 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); + // Retrieve the list of hotkeys owned by the old coldkey + let old_owned_hotkeys: Vec = OwnedHotkeys::::get(old_coldkey); - // Update StakingHotkeys map - let staking_hotkeys = StakingHotkeys::::get(old_coldkey); - StakingHotkeys::::insert(new_coldkey.clone(), staking_hotkeys); + // Initialize the total transferred stake to zero + let mut total_transferred_stake: u64 = 0u64; + + // Log the total stake of old and new coldkeys before the swap + log::info!( + "Before swap - Old coldkey total stake: {}", + TotalColdkeyStake::::get(old_coldkey) + ); + log::info!( + "Before swap - New coldkey total stake: {}", + TotalColdkeyStake::::get(new_coldkey) + ); - weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 3)); + // Iterate over each hotkey owned by the old coldkey + for hotkey in old_owned_hotkeys.iter() { + // Retrieve and remove the stake associated with the hotkey and old coldkey + let stake: u64 = Stake::::take(hotkey, old_coldkey); + log::info!("Transferring stake for hotkey {:?}: {}", hotkey, stake); + if stake > 0 { + // Insert the stake for the hotkey and new coldkey + Stake::::insert(hotkey, new_coldkey, stake); + total_transferred_stake = total_transferred_stake.saturating_add(stake); + + // Update the owner of the hotkey to the new coldkey + Owner::::insert(hotkey, new_coldkey); + + // Update the transaction weight + weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 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)); + // Log the total transferred stake + log::info!("Total transferred stake: {}", total_transferred_stake); + + // Update the total stake for both old and new coldkeys if any stake was transferred + if total_transferred_stake > 0 { + let old_coldkey_stake: u64 = TotalColdkeyStake::::take(old_coldkey); // Remove it here. + let new_coldkey_stake: u64 = TotalColdkeyStake::::get(new_coldkey); + + TotalColdkeyStake::::insert(old_coldkey, 0); + TotalColdkeyStake::::insert( + new_coldkey, + new_coldkey_stake.saturating_add(old_coldkey_stake), + ); + + log::info!("Updated old coldkey stake from {} to 0", old_coldkey_stake); + log::info!( + "Updated new coldkey stake from {} to {}", + new_coldkey_stake, + new_coldkey_stake.saturating_add(old_coldkey_stake) + ); + + // Update the transaction weight + weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2)); } - } + // Update the list of owned hotkeys for both old and new coldkeys + OwnedHotkeys::::remove(old_coldkey); + OwnedHotkeys::::insert(new_coldkey, old_owned_hotkeys); + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2)); + + // Update the staking hotkeys for both old and new coldkeys + let staking_hotkeys: Vec = StakingHotkeys::::take(old_coldkey); + StakingHotkeys::::insert(new_coldkey, staking_hotkeys); + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1)); + + // Log the total stake of old and new coldkeys after the swap + log::info!( + "After swap - Old coldkey total stake: {}", + TotalColdkeyStake::::get(old_coldkey) + ); + log::info!( + "After swap - New coldkey total stake: {}", + TotalColdkeyStake::::get(new_coldkey) + ); + } /// Swaps the total hotkey-coldkey stakes for the current interval from the old coldkey to the new coldkey. /// /// # Arguments @@ -603,7 +919,7 @@ impl Pallet { /// /// * `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) + !StakingHotkeys::::get(coldkey).is_empty() } /// Swaps the subnet owner from the old coldkey to the new coldkey for all networks where the old coldkey is the owner. @@ -632,28 +948,4 @@ 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 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/src/utils.rs b/pallets/subtensor/src/utils.rs index faaaecebb..94f18dfbf 100644 --- a/pallets/subtensor/src/utils.rs +++ b/pallets/subtensor/src/utils.rs @@ -616,6 +616,13 @@ impl Pallet { Self::deposit_event(Event::SubnetOwnerCutSet(subnet_owner_cut)); } + pub fn get_owned_hotkeys(coldkey: &T::AccountId) -> Vec { + OwnedHotkeys::::get(coldkey) + } + pub fn get_all_staked_hotkeys(coldkey: &T::AccountId) -> Vec { + StakingHotkeys::::get(coldkey) + } + pub fn set_total_issuance(total_issuance: u64) { TotalIssuance::::put(total_issuance); } diff --git a/pallets/subtensor/tests/mock.rs b/pallets/subtensor/tests/mock.rs index 3e7e74f95..c3a864075 100644 --- a/pallets/subtensor/tests/mock.rs +++ b/pallets/subtensor/tests/mock.rs @@ -2,21 +2,22 @@ use frame_support::derive_impl; use frame_support::dispatch::DispatchResultWithPostInfo; use frame_support::weights::constants::RocksDbWeight; +// use frame_support::weights::constants::WEIGHT_PER_SECOND; +use frame_support::weights::Weight; use frame_support::{ assert_ok, parameter_types, traits::{Everything, Hooks}, - weights, }; use frame_system as system; use frame_system::{limits, EnsureNever, EnsureRoot, RawOrigin}; +use pallet_collective::MemberCount; use sp_core::{Get, H256, U256}; +use sp_runtime::Perbill; use sp_runtime::{ traits::{BlakeTwo256, IdentityLookup}, BuildStorage, }; -use pallet_collective::MemberCount; - type Block = frame_system::mocking::MockBlock; // Configure a mock runtime to test the pallet. @@ -86,7 +87,7 @@ impl pallet_balances::Config for Test { #[derive_impl(frame_system::config_preludes::TestDefaultConfig)] impl system::Config for Test { type BaseCallFilter = Everything; - type BlockWeights = (); + type BlockWeights = BlockWeights; type BlockLength = (); type DbWeight = RocksDbWeight; type RuntimeOrigin = RuntimeOrigin; @@ -114,7 +115,10 @@ parameter_types! { pub const InitialMinAllowedWeights: u16 = 0; pub const InitialEmissionValue: u16 = 0; pub const InitialMaxWeightsLimit: u16 = u16::MAX; - pub BlockWeights: limits::BlockWeights = limits::BlockWeights::simple_max(weights::Weight::from_parts(1024, 0)); + pub BlockWeights: limits::BlockWeights = limits::BlockWeights::with_sensible_defaults( + Weight::from_parts(2_000_000_000_000, u64::MAX), + Perbill::from_percent(75), + ); pub const ExistentialDeposit: Balance = 1; pub const TransactionByteFee: Balance = 100; pub const SDebug:u64 = 1; @@ -164,6 +168,7 @@ parameter_types! { pub const InitialAlphaHigh: u16 = 58982; // Represents 0.9 as per the production default pub const InitialAlphaLow: u16 = 45875; // Represents 0.7 as per the production default pub const InitialLiquidAlphaOn: bool = false; // Default value for LiquidAlphaOn + pub const SubtensorInitialBaseDifficulty: u64 = 10_000; // Base difficulty } // Configure collective pallet for council @@ -368,6 +373,7 @@ impl pallet_subtensor::Config for Test { type AlphaHigh = InitialAlphaHigh; type AlphaLow = InitialAlphaLow; type LiquidAlphaOn = InitialLiquidAlphaOn; + type InitialBaseDifficulty = SubtensorInitialBaseDifficulty; } impl pallet_utility::Config for Test { diff --git a/pallets/subtensor/tests/staking.rs b/pallets/subtensor/tests/staking.rs index 4030a0f5c..5db439e5b 100644 --- a/pallets/subtensor/tests/staking.rs +++ b/pallets/subtensor/tests/staking.rs @@ -1,13 +1,21 @@ #![allow(clippy::unwrap_used)] +#![allow(clippy::arithmetic_side_effects)] +use frame_support::pallet_prelude::{ + InvalidTransaction, TransactionValidity, TransactionValidityError, +}; +use frame_support::traits::{OnFinalize, OnIdle, OnInitialize}; +use frame_support::weights::Weight; use frame_support::{assert_err, assert_noop, assert_ok, traits::Currency}; use frame_system::Config; mod mock; use frame_support::dispatch::{DispatchClass, DispatchInfo, GetDispatchInfo, Pays}; use frame_support::sp_runtime::DispatchError; use mock::*; +use pallet_balances::Call as BalancesCall; use pallet_subtensor::*; use sp_core::{H256, U256}; +use sp_runtime::traits::SignedExtension; /*********************************************************** staking::add_stake() tests @@ -3153,51 +3161,110 @@ fn setup_test_environment() -> (AccountId, AccountId, AccountId) { (current_coldkey, hotkey, new_coldkey) } +/// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test staking -- test_arbitrated_coldkey_swap_success --exact --nocapture #[test] -fn test_do_unstake_all_and_transfer_to_new_coldkey_success() { +fn test_arbitrated_coldkey_swap_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, - new_coldkey + let current_block = SubtensorModule::get_current_block_as_u64(); + let (work, nonce) = generate_valid_pow( + ¤t_coldkey, + current_block, + U256::from(BaseDifficulty::::get()), + ); + SubtensorModule::add_balance_to_coldkey_account( + ¤t_coldkey, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + ); + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + ¤t_coldkey.clone(), + &new_coldkey, + work.to_fixed_bytes().to_vec(), + current_block, + nonce )); - // Check that the stake has been removed - assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey), 0); + // Check that ColdkeySwapDestinations is populated correctly + assert_eq!( + pallet_subtensor::ColdkeySwapDestinations::::get(current_coldkey), + vec![new_coldkey] + ); - // 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, - total_balance: 1000, - } - .into(), + // Check that drain block is set correctly + let drain_block: u64 = 7200 * 3 + 1; + + log::info!( + "ColdkeysToSwapAtBlock before scheduling: {:?}", + pallet_subtensor::ColdkeysToSwapAtBlock::::get(drain_block) + ); + + assert_eq!( + pallet_subtensor::ColdkeysToSwapAtBlock::::get(drain_block), + vec![current_coldkey] + ); + log::info!("Drain block set correctly: {:?}", drain_block); + log::info!( + "Drain block {:?}", + pallet_subtensor::ColdkeysToSwapAtBlock::::get(drain_block) + ); + + // Make 5400 blocks pass + run_to_block(drain_block); + + // Run unstaking + SubtensorModule::swap_coldkeys_this_block(&BlockWeights::get().max_block).unwrap(); + log::info!( + "Arbitrated coldkeys for block: {:?}", + SubtensorModule::get_current_block_as_u64() + ); + + // Check the hotkey stake. + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey), 500); + + // Get the owner of the hotkey now new key. + assert_eq!( + SubtensorModule::get_owning_coldkey_for_hotkey(&hotkey), + new_coldkey ); + + // Check that the balance has been transferred to the new coldkey + assert_eq!( + SubtensorModule::get_coldkey_balance(&new_coldkey), + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP + 500 + ); // The new key as the 500 }); } +/// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test staking -- test_arbitrated_coldkey_swap_same_coldkey --exact --nocapture #[test] -fn test_do_unstake_all_and_transfer_to_new_coldkey_same_coldkey() { +fn test_arbitrated_coldkey_swap_same_coldkey() { new_test_ext(1).execute_with(|| { let (current_coldkey, _hotkey, _) = setup_test_environment(); + let current_block = SubtensorModule::get_current_block_as_u64(); + let (work, nonce) = generate_valid_pow( + ¤t_coldkey, + current_block, + U256::from(BaseDifficulty::::get()), + ); + assert_noop!( - SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey( - current_coldkey, - current_coldkey + SubtensorModule::do_schedule_coldkey_swap( + ¤t_coldkey.clone(), + ¤t_coldkey, + work.to_fixed_bytes().to_vec(), + current_block, + nonce ), Error::::SameColdkey ); }); } +/// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test staking -- test_arbitrated_coldkey_swap_no_balance --exact --nocapture #[test] -fn test_do_unstake_all_and_transfer_to_new_coldkey_no_balance() { +fn test_arbitrated_coldkey_swap_no_balance() { new_test_ext(1).execute_with(|| { // Create accounts manually let current_coldkey: AccountId = U256::from(1); @@ -3228,16 +3295,30 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_no_balance() { 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, - new_coldkey, + // Generate valid PoW + let current_block = SubtensorModule::get_current_block_as_u64(); + let (work, nonce) = generate_valid_pow( + ¤t_coldkey, + current_block, + U256::from(BaseDifficulty::::get()), + ); + + // Try to schedule coldkey swap + let result = SubtensorModule::do_schedule_coldkey_swap( + ¤t_coldkey.clone(), + &new_coldkey, + work.to_fixed_bytes().to_vec(), + current_block, + nonce, ); // Print the result - log::info!( - "Result of do_unstake_all_and_transfer_to_new_coldkey: {:?}", - result + log::info!("Result of arbitrated_coldkey_swap: {:?}", result); + + // Verify that the operation failed due to insufficient balance + assert_noop!( + result, + Error::::InsufficientBalanceToPerformColdkeySwap ); // Print final balances @@ -3254,9 +3335,6 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_no_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); @@ -3264,8 +3342,10 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_no_balance() { }); } +// To run this test, use the following command: +// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test staking -- test_arbitrated_coldkey_swap_with_no_stake --exact --nocapture #[test] -fn test_do_unstake_all_and_transfer_to_new_coldkey_with_no_stake() { +fn test_arbitrated_coldkey_swap_with_no_stake() { new_test_ext(1).execute_with(|| { // Create accounts manually let current_coldkey: AccountId = U256::from(1); @@ -3278,8 +3358,7 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_with_no_stake() { 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); + Balances::make_free_balance_be(¤t_coldkey, MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP); // Print initial balances log::info!( @@ -3296,16 +3375,36 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_with_no_stake() { ); // Ensure initial balances are correct - assert_eq!(Balances::total_balance(¤t_coldkey), initial_balance); + assert_eq!( + Balances::total_balance(¤t_coldkey), + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP + ); 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, - new_coldkey + let current_block = SubtensorModule::get_current_block_as_u64(); + let (work, nonce) = generate_valid_pow( + ¤t_coldkey, + current_block, + U256::from(BaseDifficulty::::get()), + ); + + // Schedule coldkey swap + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + ¤t_coldkey.clone(), + &new_coldkey, + work.to_fixed_bytes().to_vec(), + current_block, + nonce )); + // Make 5400 blocks pass, simulating on_idle for each block + let drain_block: u64 = 7200 * 3 + 1; + for _ in 0..drain_block { + next_block(); + SubtensorModule::on_idle(System::block_number(), Weight::MAX); + } + // Print final balances log::info!( "Final current_coldkey balance: {:?}", @@ -3321,26 +3420,25 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_with_no_stake() { ); // 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, - total_balance: initial_balance, - } - .into(), + assert_eq!( + Balances::total_balance(&new_coldkey), + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP ); + assert_eq!(Balances::total_balance(¤t_coldkey), 0); }); } + +// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test staking -- test_arbitrated_coldkey_swap_with_multiple_stakes --exact --nocapture #[test] -fn test_do_unstake_all_and_transfer_to_new_coldkey_with_multiple_stakes() { +fn test_arbitrated_coldkey_swap_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); + SubtensorModule::add_balance_to_coldkey_account( + ¤t_coldkey, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + ); // Add more stake assert_ok!(SubtensorModule::add_stake( @@ -3349,67 +3447,1275 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_with_multiple_stakes() { 300 )); - assert_ok!(SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey( - current_coldkey, - new_coldkey + let current_block = SubtensorModule::get_current_block_as_u64(); + let (work, nonce) = generate_valid_pow( + ¤t_coldkey, + current_block, + U256::from(BaseDifficulty::::get()), + ); + + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + ¤t_coldkey.clone(), + &new_coldkey, + work.to_fixed_bytes().to_vec(), + current_block, + nonce )); + // Make 5400 blocks pass, simulating on_idle for each block + let drain_block: u64 = 7200 * 3 + 1; + for _ in 0..drain_block { + next_block(); + SubtensorModule::on_idle(System::block_number(), Weight::MAX); + } + // Check that all stake has been removed - assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey), 0); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey), 800); + + // Owner has changed + assert_eq!( + SubtensorModule::get_owning_coldkey_for_hotkey(&hotkey), + new_coldkey + ); // 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, - total_balance: 1000, + assert_eq!( + SubtensorModule::get_coldkey_balance(&new_coldkey), + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP + 200 + ); + + // Check that the full balance has been transferred to the new coldkey + assert_eq!(SubtensorModule::get_coldkey_balance(¤t_coldkey), 0); + }); +} +// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test staking -- test_arbitrated_coldkey_swap_multiple_arbitrations --exact --nocapture +#[test] +fn test_arbitrated_coldkey_swap_multiple_arbitrations() { + new_test_ext(1).execute_with(|| { + // Set a very low base difficulty for testing + BaseDifficulty::::put(1); + + // Create coldkey with three choices. + let coldkey: AccountId = U256::from(1); + let new_coldkey1: AccountId = U256::from(2); + let new_coldkey2: AccountId = U256::from(3); + let new_coldkey3: AccountId = U256::from(4); + let hotkey: AccountId = U256::from(5); + + // Setup network state. + add_network(1, 0, 0); + SubtensorModule::add_balance_to_coldkey_account( + &coldkey, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + ); + ArbitrationPeriod::::put(5); // Set arbitration period to 5 blocks + register_ok_neuron(1, hotkey, coldkey, 0); + + let current_block = SubtensorModule::get_current_block_as_u64(); + + // Generate valid PoW for each swap attempt + let (work1, nonce1) = generate_valid_pow(&coldkey, current_block, U256::from(1)); + let (work2, nonce2) = generate_valid_pow(&coldkey, current_block, U256::from(2)); + let (work3, nonce3) = generate_valid_pow(&coldkey, current_block, U256::from(4)); + + // Schedule three swaps + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + &coldkey.clone(), + &new_coldkey1, + work1.to_fixed_bytes().to_vec(), + current_block, + nonce1 + )); + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + &coldkey.clone(), + &new_coldkey2, + work2.to_fixed_bytes().to_vec(), + current_block, + nonce2 + )); + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + &coldkey.clone(), + &new_coldkey3, + work3.to_fixed_bytes().to_vec(), + current_block, + nonce3 + )); + + // All three keys are added in swap destinations. + assert_eq!( + pallet_subtensor::ColdkeySwapDestinations::::get(coldkey), + vec![new_coldkey1, new_coldkey2, new_coldkey3] + ); + + // Simulate the passage of blocks and on_idle calls + for i in 0..(7200 * 3 + 1) { + next_block(); + SubtensorModule::on_idle(System::block_number(), Weight::MAX); + + log::info!( + "Block {}: Coldkey in arbitration: {}, Swap destinations: {:?}", + i + 1, + SubtensorModule::coldkey_in_arbitration(&coldkey), + pallet_subtensor::ColdkeySwapDestinations::::get(coldkey) + ); + } + + // Check that the swap destinations remain unchanged due to multiple (>2) swap calls + assert_eq!( + pallet_subtensor::ColdkeySwapDestinations::::get(coldkey), + vec![new_coldkey1, new_coldkey2, new_coldkey3], + "ColdkeySwapDestinations should remain unchanged with more than two swap calls" + ); + + // Key remains in arbitration due to multiple (>2) swap calls + assert!( + SubtensorModule::coldkey_in_arbitration(&coldkey), + "Coldkey should remain in arbitration with more than two swap calls" + ); + + // Check that no balance has been transferred + assert_eq!( + SubtensorModule::get_coldkey_balance(&coldkey), + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + "Original coldkey balance should remain unchanged" + ); + assert_eq!( + SubtensorModule::get_coldkey_balance(&new_coldkey1), + 0, + "New coldkey1 should not receive any balance" + ); + assert_eq!( + SubtensorModule::get_coldkey_balance(&new_coldkey2), + 0, + "New coldkey2 should not receive any balance" + ); + assert_eq!( + SubtensorModule::get_coldkey_balance(&new_coldkey3), + 0, + "New coldkey3 should not receive any balance" + ); + }); +} + +// TODO: Verify that we never want more than 2 destinations for a coldkey +#[test] +fn test_arbitrated_coldkey_swap_existing_destination() { + new_test_ext(1).execute_with(|| { + let (current_coldkey, _hotkey, new_coldkey) = setup_test_environment(); + let another_coldkey = U256::from(4); + let third_coldkey = U256::from(5); + + let current_block = SubtensorModule::get_current_block_as_u64(); + + SubtensorModule::add_balance_to_coldkey_account( + ¤t_coldkey, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + ); + + // First swap attempt (0 existing destinations) + let difficulty1 = SubtensorModule::calculate_pow_difficulty(0); + let (work1, nonce1) = generate_valid_pow(¤t_coldkey, current_block, difficulty1); + + // Schedule a swap to new_coldkey + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + ¤t_coldkey, + &new_coldkey, + work1.to_fixed_bytes().to_vec(), + current_block, + nonce1 + )); + + // Second swap attempt (1 existing destination) + let difficulty2 = SubtensorModule::calculate_pow_difficulty(1); + let (work2, nonce2) = generate_valid_pow(¤t_coldkey, current_block, difficulty2); + + // Attempt to schedule a swap to the same new_coldkey again + assert_noop!( + SubtensorModule::do_schedule_coldkey_swap( + ¤t_coldkey.clone(), + &new_coldkey, + work2.to_fixed_bytes().to_vec(), + current_block, + nonce2 + ), + Error::::DuplicateColdkey + ); + + // Schedule a swap to another_coldkey (still 1 existing destination) + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + ¤t_coldkey.clone(), + &another_coldkey, + work2.to_fixed_bytes().to_vec(), + current_block, + nonce2 + )); + + // Third swap attempt (2 existing destinations) + let difficulty3 = SubtensorModule::calculate_pow_difficulty(2); + let (work3, nonce3) = generate_valid_pow(¤t_coldkey, current_block, difficulty3); + + // Attempt to schedule a third swap + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + ¤t_coldkey.clone(), + &third_coldkey, + work3.to_fixed_bytes().to_vec(), + current_block, + nonce3 + )); + }); +} + +#[test] +fn test_arbitration_period_extension() { + new_test_ext(1).execute_with(|| { + let (current_coldkey, _hotkey, new_coldkey) = setup_test_environment(); + let another_coldkey = U256::from(4); + + let current_block = SubtensorModule::get_current_block_as_u64(); + let (work1, nonce1) = generate_valid_pow( + ¤t_coldkey, + current_block, + U256::from(BaseDifficulty::::get()), + ); + let (work2, nonce2) = + generate_valid_pow(¤t_coldkey, current_block, U256::from(20_000_000u64)); + SubtensorModule::add_balance_to_coldkey_account( + ¤t_coldkey, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + ); + + // Schedule a swap to new_coldkey + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + ¤t_coldkey.clone(), + &new_coldkey, + work1.to_fixed_bytes().to_vec(), + current_block, + nonce1 + )); + + // Schedule a swap to another_coldkey + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + ¤t_coldkey.clone(), + &another_coldkey, + work2.to_fixed_bytes().to_vec(), + current_block, + nonce2 + )); + + // Check that the arbitration period is extended + let arbitration_block = + SubtensorModule::get_current_block_as_u64() + ArbitrationPeriod::::get(); + assert_eq!( + pallet_subtensor::ColdkeyArbitrationBlock::::get(current_coldkey), + arbitration_block + ); + }); +} + +#[test] +fn test_concurrent_arbitrated_coldkey_swaps() { + new_test_ext(1).execute_with(|| { + // Manually create accounts + let coldkey1: AccountId = U256::from(1); + let hotkey1: AccountId = U256::from(2); + let new_coldkey1: AccountId = U256::from(3); + + let coldkey2: AccountId = U256::from(4); + let hotkey2: AccountId = U256::from(5); + let new_coldkey2: AccountId = U256::from(6); + + // Add networks + let netuid1: u16 = 1; + let netuid2: u16 = 2; + add_network(netuid1, 13, 0); + add_network(netuid2, 13, 0); + + // Register neurons in different networks + register_ok_neuron(netuid1, hotkey1, coldkey1, 0); + register_ok_neuron(netuid2, hotkey2, coldkey2, 0); + + // Add balance to coldkeys + SubtensorModule::add_balance_to_coldkey_account( + &coldkey1, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + ); + SubtensorModule::add_balance_to_coldkey_account( + &coldkey2, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + ); + + let current_block = SubtensorModule::get_current_block_as_u64(); + let (work1, nonce1) = generate_valid_pow( + &coldkey1, + current_block, + U256::from(BaseDifficulty::::get()), + ); + let (work2, nonce2) = generate_valid_pow( + &coldkey2, + current_block, + U256::from(BaseDifficulty::::get()), + ); + // Schedule swaps for both coldkeys + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + &coldkey1.clone(), + &new_coldkey1, + work1.to_fixed_bytes().to_vec(), + current_block, + nonce1 + )); + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + &coldkey2.clone(), + &new_coldkey2, + work2.to_fixed_bytes().to_vec(), + current_block, + nonce2 + )); + // Make 5400 blocks pass + let drain_block: u64 = 7200 * 3 + 1; + run_to_block(drain_block); + + // Run arbitration + SubtensorModule::swap_coldkeys_this_block(&BlockWeights::get().max_block).unwrap(); + + // Check that the balances have been transferred correctly + assert_eq!( + SubtensorModule::get_coldkey_balance(&new_coldkey1), + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP + ); + assert_eq!( + SubtensorModule::get_coldkey_balance(&new_coldkey2), + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP + ); + }); +} + +// #[test] +// fn test_get_remaining_arbitration_period() { +// new_test_ext(1).execute_with(|| { +// let coldkey_account_id = U256::from(12345); // arbitrary coldkey +// let new_coldkey_account_id = U256::from(54321); // arbitrary new coldkey + +// let current_block = SubtensorModule::get_current_block_as_u64(); +// let (work, nonce) = generate_valid_pow( +// &coldkey_account_id, +// current_block, +// U256::from(BaseDifficulty::::get()), +// ); + +// SubtensorModule::add_balance_to_coldkey_account( +// &coldkey_account_id, +// MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, +// ); + +// // Schedule a coldkey swap to set the arbitration block +// assert_ok!(SubtensorModule::do_schedule_coldkey_swap( +// &coldkey_account_id.clone(), +// &new_coldkey_account_id, +// work.to_fixed_bytes().to_vec(), +// current_block, +// nonce +// )); + +// // Get the current block number and arbitration period +// let current_block: u64 = SubtensorModule::get_current_block_as_u64(); +// let arbitration_period: u64 = ArbitrationPeriod::::get(); +// log::info!("arbitration_period: {:?}", arbitration_period); +// let arbitration_block: u64 = current_block + arbitration_period; +// log::info!("arbitration_block: {:?}", arbitration_block); + +// // Check if the remaining arbitration period is correct +// let remaining_period = +// SubtensorModule::get_remaining_arbitration_period(&coldkey_account_id); +// assert_eq!(remaining_period, arbitration_period); + +// // Move the current block forward and check again +// step_block(50); +// let remaining_period = +// SubtensorModule::get_remaining_arbitration_period(&coldkey_account_id); +// assert_eq!(remaining_period, arbitration_period - 50); + +// // Move the current block beyond the arbitration block and check again +// step_block((arbitration_period as u16) - 50 + 1); +// let remaining_period = +// SubtensorModule::get_remaining_arbitration_period(&coldkey_account_id); +// assert_eq!(remaining_period, 0); +// }); +// } + +#[test] +fn test_transfer_coldkey_in_arbitration() { + new_test_ext(1).execute_with(|| { + let coldkey_account_id = U256::from(1); + let recipient_account_id = U256::from(2); + let new_coldkey_account_id = U256::from(3); + + // Add balance to coldkey + SubtensorModule::add_balance_to_coldkey_account( + &coldkey_account_id, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + ); + + let current_block = SubtensorModule::get_current_block_as_u64(); + let (work, nonce) = generate_valid_pow( + &coldkey_account_id, + current_block, + U256::from(BaseDifficulty::::get()), + ); + + // Schedule a coldkey swap to put the coldkey in arbitration + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + &coldkey_account_id.clone(), + &new_coldkey_account_id, + work.to_fixed_bytes().to_vec(), + current_block, + nonce + )); + + // Try to transfer balance + let call = RuntimeCall::Balances(BalancesCall::transfer_allow_death { + dest: recipient_account_id, + value: 1000, + }); + + assert_eq!( + validate_transaction(&coldkey_account_id, &call), + Err(TransactionValidityError::Invalid(InvalidTransaction::Call)) + ); + }); +} + +#[test] +fn test_add_stake_coldkey_in_arbitration() { + new_test_ext(1).execute_with(|| { + let hotkey_account_id = U256::from(561337); + let coldkey_account_id = U256::from(61337); + let new_coldkey_account_id = U256::from(71337); + let netuid: u16 = 1; + let start_nonce: u64 = 0; + let tempo: u16 = 13; + + add_network(netuid, tempo, 0); + register_ok_neuron(netuid, hotkey_account_id, coldkey_account_id, start_nonce); + SubtensorModule::add_balance_to_coldkey_account( + &coldkey_account_id, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + ); + + let current_block = SubtensorModule::get_current_block_as_u64(); + let (work, nonce) = generate_valid_pow( + &coldkey_account_id, + current_block, + U256::from(BaseDifficulty::::get()), + ); + + // Schedule a coldkey swap to put the coldkey in arbitration + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + &coldkey_account_id.clone(), + &new_coldkey_account_id, + work.to_fixed_bytes().to_vec(), + current_block, + nonce + )); + let call = RuntimeCall::SubtensorModule(crate::Call::add_stake { + hotkey: hotkey_account_id, + amount_staked: 1000, + }); + + // This should now be Ok + assert!(validate_transaction(&coldkey_account_id, &call).is_ok()); + }) +} + +#[test] +fn test_remove_stake_coldkey_in_arbitration() { + new_test_ext(1).execute_with(|| { + let hotkey_account_id = U256::from(561337); + let coldkey_account_id = U256::from(61337); + let new_coldkey_account_id = U256::from(71337); + let netuid: u16 = 1; + let start_nonce: u64 = 0; + let tempo: u16 = 13; + + add_network(netuid, tempo, 0); + register_ok_neuron(netuid, hotkey_account_id, coldkey_account_id, start_nonce); + SubtensorModule::add_balance_to_coldkey_account( + &coldkey_account_id, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + ); + SubtensorModule::increase_stake_on_hotkey_account(&hotkey_account_id, 1000); + + let current_block = SubtensorModule::get_current_block_as_u64(); + let (work, nonce) = generate_valid_pow( + &coldkey_account_id, + current_block, + U256::from(BaseDifficulty::::get()), + ); + + // Schedule a coldkey swap to put the coldkey in arbitration + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + &coldkey_account_id.clone(), + &new_coldkey_account_id, + work.to_fixed_bytes().to_vec(), + current_block, + nonce + )); + + let call = RuntimeCall::SubtensorModule(crate::Call::remove_stake { + hotkey: hotkey_account_id, + amount_unstaked: 500, + }); + + // This should now be Ok + assert!(validate_transaction(&coldkey_account_id, &call).is_ok()); + }); +} + +#[test] +fn test_transfer_coldkey_not_in_arbitration() { + new_test_ext(1).execute_with(|| { + let coldkey_account_id = U256::from(61337); + let recipient_account_id = U256::from(71337); + + SubtensorModule::add_balance_to_coldkey_account(&coldkey_account_id, 60000); + + let call = RuntimeCall::Balances(BalancesCall::transfer_allow_death { + dest: recipient_account_id, + value: 1000, + }); + + // This should be Ok + assert!(validate_transaction(&coldkey_account_id, &call).is_ok()); + }); +} + +fn validate_transaction(who: &AccountId, call: &RuntimeCall) -> TransactionValidity { + SubtensorSignedExtension::::new().validate(who, call, &DispatchInfo::default(), 0) +} + +// Helper function to generate valid PoW +fn generate_valid_pow(coldkey: &U256, block_number: u64, difficulty: U256) -> (H256, u64) { + let mut nonce: u64 = 0; + loop { + let work = SubtensorModule::create_seal_hash(block_number, nonce, coldkey); + if SubtensorModule::hash_meets_difficulty(&work, difficulty) { + return (work, nonce); + } + nonce += 1; + } +} + +// Helper function to advance to the next block and run hooks +fn next_block() { + let current_block = System::block_number(); + System::on_finalize(current_block); + System::set_block_number(current_block + 1); + System::on_initialize(System::block_number()); + SubtensorModule::on_initialize(System::block_number()); +} + +// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test staking -- test_coldkey_meets_enough --exact --nocapture +#[test] +fn test_coldkey_meets_enough() { + new_test_ext(1).execute_with(|| { + let coldkey = U256::from(1); + let new_coldkey = U256::from(2); + let hotkey = U256::from(2); + let netuid = 1u16; + add_network(netuid, 13, 0); + register_ok_neuron(netuid, hotkey, coldkey, 0); + let current_block = SubtensorModule::get_current_block_as_u64(); + let (work1, nonce1) = generate_valid_pow( + &coldkey, + current_block, + U256::from(BaseDifficulty::::get()), + ); + assert_err!( + SubtensorModule::do_schedule_coldkey_swap( + &coldkey.clone(), + &new_coldkey, + work1.to_fixed_bytes().to_vec(), + current_block, + nonce1 + ), + Error::::InsufficientBalanceToPerformColdkeySwap + ); + SubtensorModule::add_balance_to_coldkey_account( + &coldkey, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + ); + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + &coldkey.clone(), + &new_coldkey, + work1.to_fixed_bytes().to_vec(), + current_block, + nonce1 + )); + }); +} + +#[test] +fn test_comprehensive_coldkey_swap_scenarios() { + new_test_ext(1).execute_with(|| { + // Set arbitration period to 5 blocks + ArbitrationPeriod::::put(5); + + let subnet_owner1 = U256::from(1); + let subnet_owner2 = U256::from(2); + let regular_user = U256::from(3); + let new_coldkey1 = U256::from(4); + let new_coldkey2 = U256::from(5); + let new_coldkey3 = U256::from(6); + let netuid1 = 1; + let netuid2 = 2; + + // Add networks and register subnet owners + add_network(netuid1, 13, 0); + add_network(netuid2, 13, 0); + SubnetOwner::::insert(netuid1, subnet_owner1); + SubnetOwner::::insert(netuid2, subnet_owner2); + + // Add balance to subnet owners and regular user + SubtensorModule::add_balance_to_coldkey_account( + &subnet_owner1, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + ); + SubtensorModule::add_balance_to_coldkey_account( + &subnet_owner2, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + ); + SubtensorModule::add_balance_to_coldkey_account( + ®ular_user, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP * 2, + ); + + // Set a very low base difficulty for testing + BaseDifficulty::::put(1); + + let current_block = SubtensorModule::get_current_block_as_u64(); + + // Schedule swaps for subnet owners and regular user + let (work1, nonce1) = generate_valid_pow(&subnet_owner1, current_block, U256::from(BaseDifficulty::::get())); + let (work2, nonce2) = generate_valid_pow(&subnet_owner2, current_block, U256::from(BaseDifficulty::::get())); + let (work3, nonce3) = generate_valid_pow(®ular_user, current_block, U256::from(BaseDifficulty::::get())); + + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + &subnet_owner1, + &new_coldkey1, + work1.to_fixed_bytes().to_vec(), + current_block, + nonce1 + )); + + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + &subnet_owner2, + &new_coldkey2, + work2.to_fixed_bytes().to_vec(), + current_block, + nonce2 + )); + + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + ®ular_user, + &new_coldkey3, + work3.to_fixed_bytes().to_vec(), + current_block, + nonce3 + )); + + // Check if swaps were scheduled correctly + assert_eq!( + ColdkeySwapDestinations::::get(subnet_owner1), + vec![new_coldkey1] + ); + assert_eq!( + ColdkeySwapDestinations::::get(subnet_owner2), + vec![new_coldkey2] + ); + assert_eq!( + ColdkeySwapDestinations::::get(regular_user), + vec![new_coldkey3] + ); + + // Run through the arbitration period plus one block + for i in 0..6 { + next_block(); + SubtensorModule::on_idle(System::block_number(), Weight::MAX); + + log::info!( + "Block {}: Coldkey in arbitration: {}, Swap destinations: {:?}", + i + 1, + SubtensorModule::coldkey_in_arbitration(&subnet_owner1), + ColdkeySwapDestinations::::get(subnet_owner1) + ); + + // Test edge case: try to schedule another swap during arbitration + if i == 2 { + let (work4, nonce4) = generate_valid_pow( + &subnet_owner1, + current_block + i as u64, + U256::from(4) * U256::from(BaseDifficulty::::get()), + ); + assert_ok!(SubtensorModule::do_schedule_coldkey_swap( + &subnet_owner1, + &new_coldkey2, + work4.to_fixed_bytes().to_vec(), + current_block + i as u64, + nonce4 + )); + // This should add new_coldkey2 to subnet_owner1's destinations + assert_eq!( + ColdkeySwapDestinations::::get(subnet_owner1), + vec![new_coldkey1, new_coldkey2] + ); } - .into(), + } + + // Check if swaps have been executed + log::info!( + "After arbitration period - Swap destinations for subnet_owner1: {:?}", + ColdkeySwapDestinations::::get(subnet_owner1) + ); + assert_eq!( + ColdkeySwapDestinations::::get(subnet_owner1), + vec![new_coldkey1, new_coldkey2], + "ColdkeySwapDestinations for subnet_owner1 should still contain two destinations after arbitration period" + ); + assert!(ColdkeySwapDestinations::::get(subnet_owner2).is_empty()); + assert!(ColdkeySwapDestinations::::get(regular_user).is_empty()); + + // Verify that subnet ownerships have NOT been transferred for subnet_owner1 + assert_eq!(SubnetOwner::::get(netuid1), subnet_owner1); + // But subnet_owner2's ownership should have been transferred + assert_eq!(SubnetOwner::::get(netuid2), new_coldkey2); + + // Verify regular user's balance has been transferred + assert_eq!( + SubtensorModule::get_coldkey_balance(&new_coldkey3), + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP * 2 ); + assert_eq!(SubtensorModule::get_coldkey_balance(®ular_user), 0); }); } #[test] -fn test_do_unstake_all_and_transfer_to_new_coldkey_with_multiple_stakes_multiple() { +fn test_get_total_delegated_stake_after_unstaking() { new_test_ext(1).execute_with(|| { - // Register the neuron to a new network - let netuid = 1; - let hotkey0 = U256::from(1); - let hotkey2 = U256::from(2); - let current_coldkey = U256::from(3); - let new_coldkey = U256::from(4); + let netuid = 1u16; + let delegate_coldkey = U256::from(1); + let delegate_hotkey = U256::from(2); + let delegator = U256::from(3); + let initial_stake = 2000; + let unstake_amount = 500; + let existential_deposit = 1; // Account for the existential deposit + add_network(netuid, 0, 0); - register_ok_neuron(1, hotkey0, current_coldkey, 0); - register_ok_neuron(1, hotkey2, current_coldkey, 0); - SubtensorModule::set_target_stakes_per_interval(10); - SubtensorModule::add_balance_to_coldkey_account(¤t_coldkey, 1000); + register_ok_neuron(netuid, delegate_hotkey, delegate_coldkey, 0); + + // Make the account a delegate + assert_ok!(SubtensorModule::become_delegate( + RuntimeOrigin::signed(delegate_coldkey), + delegate_hotkey + )); + + // Add balance to delegator + SubtensorModule::add_balance_to_coldkey_account(&delegator, initial_stake); + + // Delegate stake assert_ok!(SubtensorModule::add_stake( - RuntimeOrigin::signed(current_coldkey), - hotkey0, - 500 + RuntimeOrigin::signed(delegator), + delegate_hotkey, + initial_stake + )); + + // Check initial delegated stake + assert_eq!( + SubtensorModule::get_total_delegated_stake(&delegate_coldkey), + initial_stake - existential_deposit, + "Initial delegated stake is incorrect" + ); + + // Unstake part of the delegation + assert_ok!(SubtensorModule::remove_stake( + RuntimeOrigin::signed(delegator), + delegate_hotkey, + unstake_amount + )); + + // Calculate the expected delegated stake + let expected_delegated_stake = initial_stake - unstake_amount - existential_deposit; + + // Debug prints + println!("Initial stake: {}", initial_stake); + println!("Unstake amount: {}", unstake_amount); + println!("Existential deposit: {}", existential_deposit); + println!("Expected delegated stake: {}", expected_delegated_stake); + println!( + "Actual delegated stake: {}", + SubtensorModule::get_total_delegated_stake(&delegate_coldkey) + ); + + // Check the total delegated stake after unstaking + assert_eq!( + SubtensorModule::get_total_delegated_stake(&delegate_coldkey), + expected_delegated_stake, + "Delegated stake mismatch after unstaking" + ); + }); +} + +#[test] +fn test_get_total_delegated_stake_no_delegations() { + new_test_ext(1).execute_with(|| { + let delegate = U256::from(1); + let coldkey = U256::from(2); + let netuid = 1u16; + + add_network(netuid, 0, 0); + register_ok_neuron(netuid, delegate, coldkey, 0); + + // Make the delegate a delegate + assert_ok!(SubtensorModule::become_delegate( + RuntimeOrigin::signed(coldkey), + delegate + )); + + // Check that there's no delegated stake + assert_eq!(SubtensorModule::get_total_delegated_stake(&delegate), 0); + }); +} + +#[test] +fn test_get_total_delegated_stake_single_delegator() { + new_test_ext(1).execute_with(|| { + let netuid = 1u16; + let delegate_coldkey = U256::from(1); + let delegate_hotkey = U256::from(2); + let delegator = U256::from(3); + let stake_amount = 999; + let existential_deposit = 1; // Account for the existential deposit + + add_network(netuid, 0, 0); + register_ok_neuron(netuid, delegate_hotkey, delegate_coldkey, 0); + + // Make the account a delegate + assert_ok!(SubtensorModule::become_delegate( + RuntimeOrigin::signed(delegate_coldkey), + delegate_hotkey )); + + // Add stake from delegator + SubtensorModule::add_balance_to_coldkey_account(&delegator, stake_amount); assert_ok!(SubtensorModule::add_stake( - RuntimeOrigin::signed(current_coldkey), - hotkey2, - 300 + RuntimeOrigin::signed(delegator), + delegate_hotkey, + stake_amount )); - assert_ok!(SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey( - current_coldkey, - new_coldkey + + // Debug prints + println!("Delegate coldkey: {:?}", delegate_coldkey); + println!("Delegate hotkey: {:?}", delegate_hotkey); + println!("Delegator: {:?}", delegator); + println!("Stake amount: {}", stake_amount); + println!("Existential deposit: {}", existential_deposit); + println!("Total stake for hotkey: {}", SubtensorModule::get_total_stake_for_hotkey(&delegate_hotkey)); + println!("Delegated stake for coldkey: {}", SubtensorModule::get_total_delegated_stake(&delegate_coldkey)); + + // Calculate expected delegated stake + let expected_delegated_stake = stake_amount - existential_deposit; + let actual_delegated_stake = SubtensorModule::get_total_delegated_stake(&delegate_coldkey); + + assert_eq!( + actual_delegated_stake, + expected_delegated_stake, + "Total delegated stake should match the delegator's stake minus existential deposit. Expected: {}, Actual: {}", + expected_delegated_stake, + actual_delegated_stake + ); + }); +} + +#[test] +fn test_get_total_delegated_stake_multiple_delegators() { + new_test_ext(1).execute_with(|| { + let netuid = 1u16; + let delegate_coldkey = U256::from(1); + let delegate_hotkey = U256::from(2); + let delegator1 = U256::from(3); + let delegator2 = U256::from(4); + let stake1 = 1000; + let stake2 = 1999; + let existential_deposit = 1; // Account for the existential deposit + + add_network(netuid, 0, 0); + register_ok_neuron(netuid, delegate_hotkey, delegate_coldkey, 0); + + // Make the account a delegate + assert_ok!(SubtensorModule::become_delegate( + RuntimeOrigin::signed(delegate_coldkey), + delegate_hotkey )); - assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey0), 0); - assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey2), 0); - assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), 1000); - System::assert_last_event( - Event::AllBalanceUnstakedAndTransferredToNewColdkey { - current_coldkey, - new_coldkey, - total_balance: 1000, - } - .into(), + + // Add stake from delegator1 + SubtensorModule::add_balance_to_coldkey_account(&delegator1, stake1); + assert_ok!(SubtensorModule::add_stake( + RuntimeOrigin::signed(delegator1), + delegate_hotkey, + stake1 + )); + + // Add stake from delegator2 + SubtensorModule::add_balance_to_coldkey_account(&delegator2, stake2); + assert_ok!(SubtensorModule::add_stake( + RuntimeOrigin::signed(delegator2), + delegate_hotkey, + stake2 + )); + + // Debug prints + println!("Delegator1 stake: {}", stake1); + println!("Delegator2 stake: {}", stake2); + println!("Existential deposit: {}", existential_deposit); + println!("Total stake for hotkey: {}", SubtensorModule::get_total_stake_for_hotkey(&delegate_hotkey)); + println!("Delegated stake for coldkey: {}", SubtensorModule::get_total_delegated_stake(&delegate_coldkey)); + + // Calculate expected total delegated stake + let expected_total_delegated = stake1 + stake2 - 2 * existential_deposit; + let actual_total_delegated = SubtensorModule::get_total_delegated_stake(&delegate_coldkey); + + assert_eq!( + actual_total_delegated, + expected_total_delegated, + "Total delegated stake should match the sum of delegators' stakes minus existential deposits. Expected: {}, Actual: {}", + expected_total_delegated, + actual_total_delegated + ); + }); +} + +#[test] +fn test_get_total_delegated_stake_exclude_owner_stake() { + new_test_ext(1).execute_with(|| { + let netuid = 1u16; + let delegate_coldkey = U256::from(1); + let delegate_hotkey = U256::from(2); + let delegator = U256::from(3); + let owner_stake = 1000; + let delegator_stake = 999; + let existential_deposit = 1; // Account for the existential deposit + + add_network(netuid, 0, 0); + register_ok_neuron(netuid, delegate_hotkey, delegate_coldkey, 0); + + // Make the account a delegate + assert_ok!(SubtensorModule::become_delegate( + RuntimeOrigin::signed(delegate_coldkey), + delegate_hotkey + )); + + // Add owner stake + SubtensorModule::add_balance_to_coldkey_account(&delegate_coldkey, owner_stake); + assert_ok!(SubtensorModule::add_stake( + RuntimeOrigin::signed(delegate_coldkey), + delegate_hotkey, + owner_stake + )); + + // Add delegator stake + SubtensorModule::add_balance_to_coldkey_account(&delegator, delegator_stake); + assert_ok!(SubtensorModule::add_stake( + RuntimeOrigin::signed(delegator), + delegate_hotkey, + delegator_stake + )); + + // Debug prints + println!("Owner stake: {}", owner_stake); + println!("Delegator stake: {}", delegator_stake); + println!("Existential deposit: {}", existential_deposit); + println!( + "Total stake for hotkey: {}", + SubtensorModule::get_total_stake_for_hotkey(&delegate_hotkey) + ); + println!( + "Delegated stake for coldkey: {}", + SubtensorModule::get_total_delegated_stake(&delegate_coldkey) + ); + + // Check the total delegated stake (should exclude owner's stake) + let expected_delegated_stake = delegator_stake - existential_deposit; + let actual_delegated_stake = SubtensorModule::get_total_delegated_stake(&delegate_coldkey); + + assert_eq!( + actual_delegated_stake, expected_delegated_stake, + "Delegated stake should exclude owner's stake. Expected: {}, Actual: {}", + expected_delegated_stake, actual_delegated_stake + ); + }); +} + +#[test] +fn test_do_schedule_coldkey_swap_subnet_owner_skips_min_balance() { + new_test_ext(1).execute_with(|| { + let netuid = 1u16; + let subnet_owner = U256::from(1); + let new_coldkey = U256::from(2); + let hotkey = U256::from(3); + let current_block = 0u64; + + add_network(netuid, 0, 0); + register_ok_neuron(netuid, hotkey, subnet_owner, 0); + + // Make subnet_owner the owner of the subnet + SubnetOwner::::insert(netuid, subnet_owner); + + // Ensure subnet_owner has less than minimum balance + assert!( + SubtensorModule::get_coldkey_balance(&subnet_owner) + < MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP + ); + + // Generate valid PoW + let difficulty = U256::from(4) * U256::from(BaseDifficulty::::get()); + let (work, nonce) = generate_valid_pow(&subnet_owner, current_block, difficulty); + + // Debug prints + println!("Subnet owner: {:?}", subnet_owner); + println!("New coldkey: {:?}", new_coldkey); + println!("Current block: {}", current_block); + println!("Difficulty: {:?}", difficulty); + println!("Work: {:?}", work); + println!("Nonce: {}", nonce); + + // Verify the PoW + let seal = SubtensorModule::create_seal_hash(current_block, nonce, &subnet_owner); + println!("Calculated seal: {:?}", seal); + println!("Work matches seal: {}", work == seal); + println!( + "Seal meets difficulty: {}", + SubtensorModule::hash_meets_difficulty(&seal, difficulty) + ); + + // Attempt to schedule coldkey swap + let result = SubtensorModule::do_schedule_coldkey_swap( + &subnet_owner, + &new_coldkey, + work.to_fixed_bytes().to_vec(), + current_block, + nonce, + ); + + // Print the result + println!("Swap result: {:?}", result); + + assert_ok!(result); + + // Verify that the swap was scheduled + assert_eq!( + ColdkeySwapDestinations::::get(subnet_owner), + vec![new_coldkey] + ); + }); +} + +#[test] +fn test_do_schedule_coldkey_swap_delegate_with_500_tao_skips_min_balance() { + new_test_ext(1).execute_with(|| { + let netuid = 1u16; + let delegate_coldkey = U256::from(1); + let delegate_hotkey = U256::from(2); + let new_coldkey = U256::from(3); + let delegator = U256::from(4); + let current_block = 0u64; + + add_network(netuid, 0, 0); + register_ok_neuron(netuid, delegate_hotkey, delegate_coldkey, 0); + + // Make delegate a delegate + assert_ok!(SubtensorModule::become_delegate( + RuntimeOrigin::signed(delegate_coldkey), + delegate_hotkey + )); + + // Add more than 500 TAO of stake to the delegate's hotkey + let stake_amount = 501_000_000_000; // 501 TAO in RAO + SubtensorModule::add_balance_to_coldkey_account(&delegator, stake_amount); + assert_ok!(SubtensorModule::add_stake( + RuntimeOrigin::signed(delegator), + delegate_hotkey, + stake_amount + )); + + // Debug prints + println!( + "Delegator balance: {}", + SubtensorModule::get_coldkey_balance(&delegator) + ); + println!( + "Delegate coldkey balance: {}", + SubtensorModule::get_coldkey_balance(&delegate_coldkey) + ); + println!("Stake amount: {}", stake_amount); + println!( + "Delegate hotkey total stake: {}", + SubtensorModule::get_total_stake_for_hotkey(&delegate_hotkey) + ); + println!( + "Delegate coldkey delegated stake: {}", + SubtensorModule::get_total_delegated_stake(&delegate_coldkey) + ); + + // Ensure delegate's coldkey has less than minimum balance + assert!( + SubtensorModule::get_coldkey_balance(&delegate_coldkey) + < MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP, + "Delegate coldkey balance should be less than minimum required" + ); + + // Ensure the delegate's hotkey has more than 500 TAO delegated + assert!( + SubtensorModule::get_total_delegated_stake(&delegate_coldkey) >= 500_000_000_000, + "Delegate hotkey should have at least 500 TAO delegated" + ); + + // Generate valid PoW + let (work, nonce) = generate_valid_pow( + &delegate_coldkey, + current_block, + U256::from(4) * U256::from(BaseDifficulty::::get()), + ); + + // Debug prints + println!("Work: {:?}", work); + println!("Nonce: {}", nonce); + + // Attempt to schedule coldkey swap + let result = SubtensorModule::do_schedule_coldkey_swap( + &delegate_coldkey, + &new_coldkey, + work.to_fixed_bytes().to_vec(), + current_block, + nonce, + ); + + // Print the result + println!("Swap result: {:?}", result); + + assert_ok!(result); + + // Verify that the swap was scheduled + assert_eq!( + ColdkeySwapDestinations::::get(delegate_coldkey), + vec![new_coldkey] + ); + + // Additional debug prints after swap + println!( + "Coldkey swap destinations: {:?}", + ColdkeySwapDestinations::::get(delegate_coldkey) + ); + println!( + "Is coldkey in arbitration: {}", + SubtensorModule::coldkey_in_arbitration(&delegate_coldkey) + ); + }); +} + +#[test] +fn test_do_schedule_coldkey_swap_regular_user_fails_min_balance() { + new_test_ext(1).execute_with(|| { + let netuid = 1u16; + let regular_user = U256::from(1); + let new_coldkey = U256::from(2); + let hotkey = U256::from(3); + let current_block = 0u64; + let nonce = 0u64; + + add_network(netuid, 0, 0); + register_ok_neuron(netuid, hotkey, regular_user, 0); + + // Ensure regular_user has less than minimum balance + assert!( + SubtensorModule::get_coldkey_balance(®ular_user) + < MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP + ); + + let (work, _) = generate_valid_pow( + ®ular_user, + current_block, + U256::from(4) * U256::from(BaseDifficulty::::get()), + ); + + // Attempt to schedule coldkey swap + assert_noop!( + SubtensorModule::do_schedule_coldkey_swap( + ®ular_user, + &new_coldkey, + work.to_fixed_bytes().to_vec(), + current_block, + nonce + ), + Error::::InsufficientBalanceToPerformColdkeySwap + ); + + // Verify that the swap was not scheduled + assert!(ColdkeySwapDestinations::::get(regular_user).is_empty()); + }); +} + +#[test] +fn test_do_schedule_coldkey_swap_regular_user_passes_min_balance() { + new_test_ext(1).execute_with(|| { + let netuid = 1u16; + let regular_user = U256::from(1); + let new_coldkey = U256::from(2); + let hotkey = U256::from(3); + let current_block = 0u64; + + add_network(netuid, 0, 0); + register_ok_neuron(netuid, hotkey, regular_user, 0); + + // Ensure regular_user has more than minimum balance + SubtensorModule::add_balance_to_coldkey_account( + ®ular_user, + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP + 1, + ); + assert!( + SubtensorModule::get_coldkey_balance(®ular_user) + > MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP + ); + + // Generate valid PoW + let (work, nonce) = generate_valid_pow( + ®ular_user, + current_block, + U256::from(4) * U256::from(BaseDifficulty::::get()), + ); + + // Debug prints + println!("Regular user: {:?}", regular_user); + println!("New coldkey: {:?}", new_coldkey); + println!("Current block: {}", current_block); + println!("Work: {:?}", work); + println!("Nonce: {}", nonce); + + // Attempt to schedule coldkey swap + let result = SubtensorModule::do_schedule_coldkey_swap( + ®ular_user, + &new_coldkey, + work.to_fixed_bytes().to_vec(), + current_block, + nonce, + ); + + // Print the result + println!("Swap result: {:?}", result); + + assert_ok!(result); + + // Verify that the swap was scheduled + assert_eq!( + ColdkeySwapDestinations::::get(regular_user), + vec![new_coldkey] ); }); } diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index d527c44e8..90ebdfdcb 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -625,7 +625,7 @@ fn test_swap_stake_weight_update() { SubtensorModule::swap_stake(&old_hotkey, &new_hotkey, &mut weight); // Verify the weight update - let expected_weight = ::DbWeight::get().writes(3); + let expected_weight = ::DbWeight::get().writes(4); assert_eq!(weight, expected_weight); }); } @@ -1059,7 +1059,7 @@ fn test_do_swap_coldkey_success() { let netuid = 1u16; let stake_amount1 = 1000u64; let stake_amount2 = 2000u64; - let free_balance_old = 12345u64; + let free_balance_old = 12345u64 + MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP; // Setup initial state add_network(netuid, 13, 0); @@ -1072,6 +1072,20 @@ fn test_do_swap_coldkey_success() { stake_amount1 + stake_amount2 + free_balance_old, ); + // Log initial state + log::info!( + "Initial total stake: {}", + SubtensorModule::get_total_stake() + ); + log::info!( + "Initial old coldkey stake: {}", + SubtensorModule::get_total_stake_for_coldkey(&old_coldkey) + ); + log::info!( + "Initial new coldkey stake: {}", + SubtensorModule::get_total_stake_for_coldkey(&new_coldkey) + ); + // Add stake to the neurons assert_ok!(SubtensorModule::add_stake( <::RuntimeOrigin>::signed(old_coldkey), @@ -1084,29 +1098,43 @@ fn test_do_swap_coldkey_success() { stake_amount2 )); - // Verify initial stakes and balances - assert_eq!( - TotalColdkeyStake::::get(old_coldkey), - stake_amount1 + stake_amount2 + // Log state after adding stake + log::info!( + "Total stake after adding: {}", + SubtensorModule::get_total_stake() ); - 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] + log::info!( + "Old coldkey stake after adding: {}", + SubtensorModule::get_total_stake_for_coldkey(&old_coldkey) ); - assert_eq!( - SubtensorModule::get_coldkey_balance(&old_coldkey), - free_balance_old + log::info!( + "New coldkey stake after adding: {}", + SubtensorModule::get_total_stake_for_coldkey(&new_coldkey) ); + // Record total stake before swap + let total_stake_before_swap = SubtensorModule::get_total_stake(); + // Perform the swap assert_ok!(SubtensorModule::do_swap_coldkey( <::RuntimeOrigin>::signed(old_coldkey), - &old_coldkey, &new_coldkey )); + // Log state after swap + log::info!( + "Total stake after swap: {}", + SubtensorModule::get_total_stake() + ); + log::info!( + "Old coldkey stake after swap: {}", + SubtensorModule::get_total_stake_for_coldkey(&old_coldkey) + ); + log::info!( + "New coldkey stake after swap: {}", + SubtensorModule::get_total_stake_for_coldkey(&new_coldkey) + ); + // Verify the swap assert_eq!(Owner::::get(hotkey1), new_coldkey); assert_eq!(Owner::::get(hotkey2), new_coldkey); @@ -1114,7 +1142,7 @@ fn test_do_swap_coldkey_success() { TotalColdkeyStake::::get(new_coldkey), stake_amount1 + stake_amount2 ); - assert!(!TotalColdkeyStake::::contains_key(old_coldkey)); + assert_eq!(TotalColdkeyStake::::get(old_coldkey), 0); 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)); @@ -1134,6 +1162,13 @@ fn test_do_swap_coldkey_success() { ); assert_eq!(SubtensorModule::get_coldkey_balance(&old_coldkey), 0); + // Verify total stake remains unchanged + assert_eq!( + SubtensorModule::get_total_stake(), + total_stake_before_swap, + "Total stake changed unexpectedly" + ); + // Verify event emission System::assert_last_event( Event::ColdkeySwapped { @@ -1145,30 +1180,6 @@ fn test_do_swap_coldkey_success() { }); } -#[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(|| { @@ -1178,72 +1189,60 @@ fn test_swap_stake_for_coldkey() { let hotkey2 = U256::from(4); let stake_amount1 = 1000u64; let stake_amount2 = 2000u64; + let total_stake = stake_amount1 + stake_amount2; let mut weight = Weight::zero(); - // Initialize Stake for old_coldkey + // Setup initial state + OwnedHotkeys::::insert(old_coldkey, vec![hotkey1, hotkey2]); 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); + TotalColdkeyStake::::insert(old_coldkey, total_stake); - // Initialize TotalStake and TotalIssuance - TotalStake::::put(stake_amount1 + stake_amount2); - TotalIssuance::::put(stake_amount1 + stake_amount2); + // Set up total issuance + TotalIssuance::::put(total_stake); + TotalStake::::put(total_stake); - // Populate OwnedHotkeys map - OwnedHotkeys::::insert(old_coldkey, vec![hotkey1, hotkey2]); + // Record initial values + let initial_total_issuance = SubtensorModule::get_total_issuance(); + let initial_total_stake = SubtensorModule::get_total_stake(); // Perform the swap SubtensorModule::swap_stake_for_coldkey(&old_coldkey, &new_coldkey, &mut weight); - // Verify the swap + // Verify ownership transfer + assert_eq!( + SubtensorModule::get_owned_hotkeys(&new_coldkey), + vec![hotkey1, hotkey2] + ); + assert_eq!(SubtensorModule::get_owned_hotkeys(&old_coldkey), vec![]); + + // Verify stake transfer 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)); + assert_eq!(Stake::::get(hotkey1, old_coldkey), 0); + assert_eq!(Stake::::get(hotkey2, old_coldkey), 0); + + // Verify TotalColdkeyStake + assert_eq!(TotalColdkeyStake::::get(new_coldkey), total_stake); + assert_eq!(TotalColdkeyStake::::get(old_coldkey), 0); // 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(5, 6); - 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); + // Verify total stake and issuance remain unchanged + assert_eq!( + SubtensorModule::get_total_stake(), + initial_total_stake, + "Total stake changed unexpectedly" + ); + assert_eq!( + SubtensorModule::get_total_issuance(), + initial_total_issuance, + "Total issuance changed unexpectedly" + ); }); } @@ -1350,7 +1349,6 @@ fn test_do_swap_coldkey_with_subnet_ownership() { // Perform the swap assert_ok!(SubtensorModule::do_swap_coldkey( <::RuntimeOrigin>::signed(old_coldkey), - &old_coldkey, &new_coldkey )); @@ -1369,14 +1367,323 @@ fn test_coldkey_has_associated_hotkeys() { // Setup initial state add_network(netuid, 13, 0); register_ok_neuron(netuid, hotkey, coldkey, 0); + SubtensorModule::add_balance_to_coldkey_account(&coldkey, 1000); + }); +} + +// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test swap -- test_coldkey_swap_total --exact --nocapture +#[test] +fn test_coldkey_swap_total() { + new_test_ext(1).execute_with(|| { + let coldkey = U256::from(1); + let nominator1 = U256::from(2); + let nominator2 = U256::from(3); + let nominator3 = U256::from(4); + let delegate1 = U256::from(5); + let delegate2 = U256::from(6); + let delegate3 = U256::from(7); + let hotkey1 = U256::from(2); + let hotkey2 = U256::from(3); + let hotkey3 = U256::from(4); + let netuid1 = 1u16; + let netuid2 = 2u16; + let netuid3 = 3u16; + SubtensorModule::add_balance_to_coldkey_account(&coldkey, 1000); + SubtensorModule::add_balance_to_coldkey_account(&delegate1, 1000); + SubtensorModule::add_balance_to_coldkey_account(&delegate2, 1000); + SubtensorModule::add_balance_to_coldkey_account(&delegate3, 1000); + SubtensorModule::add_balance_to_coldkey_account(&nominator1, 1000); + SubtensorModule::add_balance_to_coldkey_account(&nominator2, 1000); + SubtensorModule::add_balance_to_coldkey_account(&nominator3, 1000); + + // Setup initial state + add_network(netuid1, 13, 0); + add_network(netuid2, 14, 0); + add_network(netuid3, 15, 0); + register_ok_neuron(netuid1, hotkey1, coldkey, 0); + register_ok_neuron(netuid2, hotkey2, coldkey, 0); + register_ok_neuron(netuid3, hotkey3, coldkey, 0); + register_ok_neuron(netuid1, delegate1, delegate1, 0); + register_ok_neuron(netuid2, delegate2, delegate2, 0); + register_ok_neuron(netuid3, delegate3, delegate3, 0); + assert_ok!(SubtensorModule::do_become_delegate( + <::RuntimeOrigin>::signed(coldkey), + hotkey1, + u16::MAX / 10 + )); + assert_ok!(SubtensorModule::do_become_delegate( + <::RuntimeOrigin>::signed(coldkey), + hotkey2, + u16::MAX / 10 + )); + assert_ok!(SubtensorModule::do_become_delegate( + <::RuntimeOrigin>::signed(coldkey), + hotkey3, + u16::MAX / 10 + )); + assert_ok!(SubtensorModule::do_become_delegate( + <::RuntimeOrigin>::signed(delegate1), + delegate1, + u16::MAX / 10 + )); + assert_ok!(SubtensorModule::do_become_delegate( + <::RuntimeOrigin>::signed(delegate2), + delegate2, + u16::MAX / 10 + )); + assert_ok!(SubtensorModule::do_become_delegate( + <::RuntimeOrigin>::signed(delegate3), + delegate3, + u16::MAX / 10 + )); + + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(coldkey), + hotkey1, + 100 + )); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(coldkey), + hotkey2, + 100 + )); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(coldkey), + hotkey3, + 100 + )); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(coldkey), + delegate1, + 100 + )); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(coldkey), + delegate2, + 100 + )); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(coldkey), + delegate3, + 100 + )); + + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(delegate1), + hotkey1, + 100 + )); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(delegate2), + hotkey2, + 100 + )); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(delegate3), + hotkey3, + 100 + )); + + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(delegate1), + delegate1, + 100 + )); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(delegate2), + delegate2, + 100 + )); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(delegate3), + delegate3, + 100 + )); - // Check if coldkey has associated hotkeys - assert!(SubtensorModule::coldkey_has_associated_hotkeys(&coldkey)); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(nominator1), + hotkey1, + 100 + )); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(nominator2), + hotkey2, + 100 + )); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(nominator3), + hotkey3, + 100 + )); - // Check for a coldkey without associated hotkeys - let unassociated_coldkey = U256::from(3); - assert!(!SubtensorModule::coldkey_has_associated_hotkeys( - &unassociated_coldkey + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(nominator1), + delegate1, + 100 + )); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(nominator2), + delegate2, + 100 )); + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(nominator3), + delegate3, + 100 + )); + + assert_eq!( + SubtensorModule::get_owned_hotkeys(&coldkey), + vec![hotkey1, hotkey2, hotkey3] + ); + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&coldkey), + vec![hotkey1, hotkey2, hotkey3, delegate1, delegate2, delegate3] + ); + assert_eq!(SubtensorModule::get_total_stake_for_coldkey(&coldkey), 600); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey1), 300); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey2), 300); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey3), 300); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&delegate1), 300); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&delegate2), 300); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&delegate3), 300); + + assert_eq!( + SubtensorModule::get_owned_hotkeys(&delegate1), + vec![delegate1] + ); + assert_eq!( + SubtensorModule::get_owned_hotkeys(&delegate2), + vec![delegate2] + ); + assert_eq!( + SubtensorModule::get_owned_hotkeys(&delegate3), + vec![delegate3] + ); + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&delegate1), + vec![delegate1, hotkey1] + ); + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&delegate2), + vec![delegate2, hotkey2] + ); + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&delegate3), + vec![delegate3, hotkey3] + ); + + assert_eq!(SubtensorModule::get_owned_hotkeys(&nominator1), vec![]); + assert_eq!(SubtensorModule::get_owned_hotkeys(&nominator2), vec![]); + assert_eq!(SubtensorModule::get_owned_hotkeys(&nominator3), vec![]); + + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&nominator1), + vec![hotkey1, delegate1] + ); + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&nominator2), + vec![hotkey2, delegate2] + ); + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&nominator3), + vec![hotkey3, delegate3] + ); + + // Perform the swap + let new_coldkey = U256::from(1100); + assert_eq!(SubtensorModule::get_total_stake_for_coldkey(&coldkey), 600); + assert_ok!(SubtensorModule::perform_swap_coldkey( + &coldkey, + &new_coldkey + )); + assert_eq!( + SubtensorModule::get_total_stake_for_coldkey(&new_coldkey), + 600 + ); + + // Check everything is swapped. + assert_eq!( + SubtensorModule::get_owned_hotkeys(&new_coldkey), + vec![hotkey1, hotkey2, hotkey3] + ); + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&new_coldkey), + vec![hotkey1, hotkey2, hotkey3, delegate1, delegate2, delegate3] + ); + assert_eq!( + SubtensorModule::get_total_stake_for_coldkey(&new_coldkey), + 600 + ); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey1), 300); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey2), 300); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey3), 300); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&delegate1), 300); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&delegate2), 300); + assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&delegate3), 300); + + assert_eq!( + SubtensorModule::get_owned_hotkeys(&delegate1), + vec![delegate1] + ); + assert_eq!( + SubtensorModule::get_owned_hotkeys(&delegate2), + vec![delegate2] + ); + assert_eq!( + SubtensorModule::get_owned_hotkeys(&delegate3), + vec![delegate3] + ); + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&delegate1), + vec![delegate1, hotkey1] + ); + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&delegate2), + vec![delegate2, hotkey2] + ); + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&delegate3), + vec![delegate3, hotkey3] + ); + + assert_eq!(SubtensorModule::get_owned_hotkeys(&nominator1), vec![]); + assert_eq!(SubtensorModule::get_owned_hotkeys(&nominator2), vec![]); + assert_eq!(SubtensorModule::get_owned_hotkeys(&nominator3), vec![]); + + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&nominator1), + vec![hotkey1, delegate1] + ); + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&nominator2), + vec![hotkey2, delegate2] + ); + assert_eq!( + SubtensorModule::get_all_staked_hotkeys(&nominator3), + vec![hotkey3, delegate3] + ); }); } + +// #[test] +// fn test_coldkey_arbitrated_sw() { +// 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 45096d33b..3e5af1b43 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: 160, + spec_version: 194, impl_version: 1, apis: RUNTIME_API_VERSIONS, transaction_version: 1, @@ -312,27 +312,12 @@ impl Contains for SafeModeWhitelistedCalls { | RuntimeCall::SafeMode(_) | RuntimeCall::Timestamp(_) | RuntimeCall::SubtensorModule( - pallet_subtensor::Call::add_stake { .. } - | pallet_subtensor::Call::become_delegate { .. } - | pallet_subtensor::Call::burned_register { .. } - | pallet_subtensor::Call::commit_weights { .. } - | pallet_subtensor::Call::decrease_take { .. } - | pallet_subtensor::Call::faucet { .. } - | pallet_subtensor::Call::increase_take { .. } - | pallet_subtensor::Call::register { .. } - | pallet_subtensor::Call::register_network { .. } - | pallet_subtensor::Call::remove_stake { .. } - | pallet_subtensor::Call::reveal_weights { .. } - | pallet_subtensor::Call::root_register { .. } - | pallet_subtensor::Call::serve_axon { .. } - | pallet_subtensor::Call::serve_prometheus { .. } - | pallet_subtensor::Call::set_root_weights { .. } + pallet_subtensor::Call::schedule_coldkey_swap { .. } | pallet_subtensor::Call::set_weights { .. } - | pallet_subtensor::Call::sudo { .. } - | pallet_subtensor::Call::sudo_unchecked_weight { .. } - | pallet_subtensor::Call::swap_hotkey { .. } - | pallet_subtensor::Call::vote { .. } + | pallet_subtensor::Call::set_root_weights { .. } + | pallet_subtensor::Call::serve_axon { .. } ) + | RuntimeCall::Commitments(pallet_commitments::Call::set_commitment { .. }) ) } } @@ -896,6 +881,7 @@ parameter_types! { pub const InitialAlphaHigh: u16 = 58982; // Represents 0.9 as per the production default pub const InitialAlphaLow: u16 = 45875; // Represents 0.7 as per the production default pub const InitialLiquidAlphaOn: bool = false; // Default value for LiquidAlphaOn + pub const SubtensorInitialBaseDifficulty: u64 = 10_000_000; // Base difficulty } impl pallet_subtensor::Config for Runtime { @@ -951,6 +937,7 @@ impl pallet_subtensor::Config for Runtime { type AlphaHigh = InitialAlphaHigh; type AlphaLow = InitialAlphaLow; type LiquidAlphaOn = InitialLiquidAlphaOn; + type InitialBaseDifficulty = SubtensorInitialBaseDifficulty; } use sp_runtime::BoundedVec; @@ -1677,6 +1664,23 @@ impl_runtime_apis! { SubtensorModule::get_network_lock_cost() } } + + impl subtensor_custom_rpc_runtime_api::ColdkeySwapRuntimeApi for Runtime { + fn get_scheduled_coldkey_swap( coldkey_account_vec: Vec ) -> Vec { + let result = SubtensorModule::get_scheduled_coldkey_swap( coldkey_account_vec ); + result.encode() + } + + fn get_remaining_arbitration_period( coldkey_account_vec: Vec ) -> Vec { + let result = SubtensorModule::get_remaining_arbitration_period( coldkey_account_vec ); + result.encode() + } + + fn get_coldkey_swap_destinations( coldkey_account_vec: Vec ) -> Vec { + let result = SubtensorModule::get_coldkey_swap_destinations( coldkey_account_vec ); + result.encode() + } + } } // #[cfg(test)] diff --git a/scripts/benchmark.sh b/scripts/benchmark.sh index 52bdaf2c5..8f54fa54a 100755 --- a/scripts/benchmark.sh +++ b/scripts/benchmark.sh @@ -7,31 +7,31 @@ OUTPUT_FILE='benchmarking.txt' # Getting arguments from user while [[ $# -gt 0 ]]; do - case $1 in - -p | --bin-path) - BIN_PATH="$2" - shift - shift - ;; - -* | --*) - echo "Unknown option $1" - exit 1 - ;; - *) - POSITIONAL_ARGS+=("$1") - shift - ;; - esac + case $1 in + -p | --bin-path) + BIN_PATH="$2" + shift + shift + ;; + -* | --*) + echo "Unknown option $1" + exit 1 + ;; + *) + POSITIONAL_ARGS+=("$1") + shift + ;; + esac done # Ensure binary exists before node-subtensor executions if [ ! -f $BIN_PATH ]; then - if [[ "$DEFAULT_BIN_PATH" == "$BIN_PATH" ]]; then - cargo build --profile production --features runtime-benchmarks - else - echo "Binary '$BIN_PATH' does not exist. You can use -p or --bin-path to specify a different location." - exit 1 - fi + if [[ "$DEFAULT_BIN_PATH" == "$BIN_PATH" ]]; then + cargo build --profile production --features runtime-benchmarks + else + echo "Binary '$BIN_PATH' does not exist. You can use -p or --bin-path to specify a different location." + exit 1 + fi fi # Build Temporary Spec @@ -39,8 +39,8 @@ $BIN_PATH build-spec --disable-default-bootnode --raw --chain local >$TMP_SPEC # Run benchmark $BIN_PATH benchmark pallet \ - --chain=$TMP_SPEC \ - --pallet pallet-subtensor --extrinsic 'benchmark_dissolve_network' \ - --output $OUTPUT_FILE +--chain=$TMP_SPEC \ +--pallet pallet-subtensor --extrinsic 'schedule_coldkey_swap' \ +--output $OUTPUT_FILE rm $TMP_SPEC