Merge pull request #611 from opentensor/feat/cold_key_swap

Feat/cold key swap

justfile

@@ -31,15 +31,11 @@ clippy:
 
 clippy-fix:
     @echo "Running cargo clippy with automatic fixes on potentially dirty code..."
-    cargo +{{RUSTV}} clippy --fix --allow-dirty --workspace --all-targets -- \
-        -A clippy::todo \
-        -A clippy::unimplemented \
-        -A clippy::indexing_slicing
-    @echo "Running cargo clippy with automatic fixes on potentially dirty code..."
-    cargo +{{RUSTV}} clippy --fix --allow-dirty --workspace --all-targets -- \
+    cargo +{{RUSTV}} clippy --fix --allow-dirty --allow-staged --workspace --all-targets -- \
         -A clippy::todo \
         -A clippy::unimplemented \
         -A clippy::indexing_slicing
+
 fix:
   @echo "Running cargo fix..."
   cargo +{{RUSTV}} fix --workspace

pallets/subtensor/src/errors.rs

@@ -132,5 +132,19 @@ mod errors {
         AlphaHighTooLow,
         /// Alpha low is out of range: alpha_low > 0 && alpha_low < 0.8
         AlphaLowOutOfRange,
+        /// The coldkey has already been swapped
+        ColdKeyAlreadyAssociated,
+        /// The coldkey swap transaction rate limit exceeded
+        ColdKeySwapTxRateLimitExceeded,
+        /// The new coldkey is the same as the old coldkey
+        NewColdKeyIsSameWithOld,
+        /// The coldkey does not exist
+        NotExistColdkey,
+        /// The coldkey balance is not enough to pay for the swap
+        NotEnoughBalanceToPaySwapColdKey,
+        /// No balance to transfer
+        NoBalanceToTransfer,
+        /// Same coldkey
+        SameColdkey,
     }
 }

pallets/subtensor/src/events.rs

@@ -132,5 +132,27 @@ mod events {
         MinDelegateTakeSet(u16),
         /// the target stakes per interval is set by sudo/admin transaction
         TargetStakesPerIntervalSet(u64),
+        /// A coldkey has been swapped
+        ColdkeySwapped {
+            /// the account ID of old coldkey
+            old_coldkey: T::AccountId,
+            /// the account ID of new coldkey
+            new_coldkey: T::AccountId,
+        },
+        /// All balance of a hotkey has been unstaked and transferred to a new coldkey
+        AllBalanceUnstakedAndTransferredToNewColdkey {
+            /// The account ID of the current coldkey
+            current_coldkey: T::AccountId,
+            /// The account ID of the new coldkey
+            new_coldkey: T::AccountId,
+            /// The account ID of the hotkey
+            hotkey: T::AccountId,
+            /// The current stake of the hotkey
+            current_stake: u64,
+            /// The total balance of the hotkey
+            total_balance: <<T as Config>::Currency as fungible::Inspect<
+                <T as frame_system::Config>::AccountId,
+            >>::Balance,
+        },
     }
 }

pallets/subtensor/src/lib.rs

@@ -363,6 +363,9 @@ pub mod pallet {
     #[pallet::storage] // --- MAP ( hot ) --> cold | Returns the controlling coldkey for a hotkey.
     pub type Owner<T: Config> =
         StorageMap<_, Blake2_128Concat, T::AccountId, T::AccountId, ValueQuery, DefaultAccount<T>>;
+    #[pallet::storage] // --- MAP ( cold ) --> Vec<hot> | Returns the vector of hotkeys controlled by this coldkey.
+    pub type OwnedHotkeys<T: Config> =
+        StorageMap<_, Blake2_128Concat, T::AccountId, Vec<T::AccountId>, ValueQuery>;
     #[pallet::storage] // --- MAP ( hot ) --> take | Returns the hotkey delegation take. And signals that this key is open for delegation.
     pub type Delegates<T: Config> =
         StorageMap<_, Blake2_128Concat, T::AccountId, u16, ValueQuery, DefaultDefaultTake<T>>;
@@ -1204,6 +1207,13 @@ pub mod pallet {
                     // Fill stake information.
                     Owner::<T>::insert(hotkey.clone(), coldkey.clone());
 
+                    // Update OwnedHotkeys map
+                    let mut hotkeys = OwnedHotkeys::<T>::get(coldkey);
+                    if !hotkeys.contains(hotkey) {
+                        hotkeys.push(hotkey.clone());
+                        OwnedHotkeys::<T>::insert(coldkey, hotkeys);
+                    }
+
                     TotalHotkeyStake::<T>::insert(hotkey.clone(), stake);
                     TotalColdkeyStake::<T>::insert(
                         coldkey.clone(),
@@ -1325,7 +1335,9 @@ pub mod pallet {
                 // Storage version v4 -> v5
                 .saturating_add(migration::migrate_delete_subnet_3::<T>())
                 // Doesn't check storage version. TODO: Remove after upgrade
-                .saturating_add(migration::migration5_total_issuance::<T>(false));
+                .saturating_add(migration::migration5_total_issuance::<T>(false))
+                // Populate OwnedHotkeys map for coldkey swap. Doesn't update storage vesion.
+                .saturating_add(migration::migrate_populate_owned::<T>());
 
             weight
         }
@@ -1970,6 +1982,61 @@ pub mod pallet {
             Self::do_swap_hotkey(origin, &hotkey, &new_hotkey)
         }
 
+        /// The extrinsic for user to change the coldkey associated with their account.
+        ///
+        /// # Arguments
+        ///
+        /// * `origin` - The origin of the call, must be signed by the old coldkey.
+        /// * `old_coldkey` - The current coldkey associated with the account.
+        /// * `new_coldkey` - The new coldkey to be associated with the account.
+        ///
+        /// # Returns
+        ///
+        /// Returns a `DispatchResultWithPostInfo` indicating success or failure of the operation.
+        ///
+        /// # Weight
+        ///
+        /// Weight is calculated based on the number of database reads and writes.
+        #[pallet::call_index(71)]
+        #[pallet::weight((Weight::from_parts(1_940_000_000, 0)
+		.saturating_add(T::DbWeight::get().reads(272))
+		.saturating_add(T::DbWeight::get().writes(527)), DispatchClass::Operational, Pays::No))]
+        pub fn swap_coldkey(
+            origin: OriginFor<T>,
+            old_coldkey: T::AccountId,
+            new_coldkey: T::AccountId,
+        ) -> DispatchResultWithPostInfo {
+            Self::do_swap_coldkey(origin, &old_coldkey, &new_coldkey)
+        }
+
+        /// Unstakes all tokens associated with a hotkey and transfers them to a new coldkey.
+        ///
+        /// # Arguments
+        ///
+        /// * `origin` - The origin of the call, must be signed by the current coldkey.
+        /// * `hotkey` - The hotkey associated with the stakes to be unstaked.
+        /// * `new_coldkey` - The new coldkey to receive the unstaked tokens.
+        ///
+        /// # Returns
+        ///
+        /// Returns a `DispatchResult` indicating success or failure of the operation.
+        ///
+        /// # Weight
+        ///
+        /// Weight is calculated based on the number of database reads and writes.
+        #[pallet::call_index(72)]
+        #[pallet::weight((Weight::from_parts(1_940_000_000, 0)
+		.saturating_add(T::DbWeight::get().reads(272))
+		.saturating_add(T::DbWeight::get().writes(527)), DispatchClass::Operational, Pays::No))]
+        pub fn unstake_all_and_transfer_to_new_coldkey(
+            origin: OriginFor<T>,
+            hotkey: T::AccountId,
+            new_coldkey: T::AccountId,
+        ) -> DispatchResult {
+            let current_coldkey = ensure_signed(origin)?;
+            Self::do_unstake_all_and_transfer_to_new_coldkey(current_coldkey, hotkey, new_coldkey)
+        }
+
         // ---- SUDO ONLY FUNCTIONS ------------------------------------------------------------
 
         // ==================================

pallets/subtensor/src/migration.rs

@@ -477,3 +477,65 @@ pub fn migrate_to_v2_fixed_total_stake<T: Config>() -> Weight {
         Weight::zero()
     }
 }
+
+/// Migrate the OwnedHotkeys map to the new storage format
+pub fn migrate_populate_owned<T: Config>() -> Weight {
+    // Setup migration weight
+    let mut weight = T::DbWeight::get().reads(1);
+    let migration_name = "Populate OwnedHotkeys map";
+
+    // Check if this migration is needed (if OwnedHotkeys map is empty)
+    let migrate = OwnedHotkeys::<T>::iter().next().is_none();
+
+    // Only runs if the migration is needed
+    if migrate {
+        info!(target: LOG_TARGET_1, ">>> Starting Migration: {}", migration_name);
+
+        let mut longest_hotkey_vector: usize = 0;
+        let mut longest_coldkey: Option<T::AccountId> = None;
+        let mut keys_touched: u64 = 0;
+        let mut storage_reads: u64 = 0;
+        let mut storage_writes: u64 = 0;
+
+        // Iterate through all Owner entries
+        Owner::<T>::iter().for_each(|(hotkey, coldkey)| {
+            storage_reads = storage_reads.saturating_add(1); // Read from Owner storage
+            let mut hotkeys = OwnedHotkeys::<T>::get(&coldkey);
+            storage_reads = storage_reads.saturating_add(1); // Read from OwnedHotkeys storage
+
+            // Add the hotkey if it's not already in the vector
+            if !hotkeys.contains(&hotkey) {
+                hotkeys.push(hotkey);
+                keys_touched = keys_touched.saturating_add(1);
+
+                // Update longest hotkey vector info
+                if longest_hotkey_vector < hotkeys.len() {
+                    longest_hotkey_vector = hotkeys.len();
+                    longest_coldkey = Some(coldkey.clone());
+                }
+
+                // Update the OwnedHotkeys storage
+                OwnedHotkeys::<T>::insert(&coldkey, hotkeys);
+                storage_writes = storage_writes.saturating_add(1); // Write to OwnedHotkeys storage
+            }
+
+            // Accrue weight for reads and writes
+            weight = weight.saturating_add(T::DbWeight::get().reads_writes(2, 1));
+        });
+
+        // Log migration results
+        info!(
+            target: LOG_TARGET_1,
+            "Migration {} finished. Keys touched: {}, Longest hotkey vector: {}, Storage reads: {}, Storage writes: {}",
+            migration_name, keys_touched, longest_hotkey_vector, storage_reads, storage_writes
+        );
+        if let Some(c) = longest_coldkey {
+            info!(target: LOG_TARGET_1, "Longest hotkey vector is controlled by: {:?}", c);
+        }
+
+        weight
+    } else {
+        info!(target: LOG_TARGET_1, "Migration {} already done!", migration_name);
+        Weight::zero()
+    }
+}

pallets/subtensor/src/staking.rs

@@ -1,4 +1,5 @@
 use super::*;
+use dispatch::RawOrigin;
 use frame_support::{
     storage::IterableStorageDoubleMap,
     traits::{
@@ -9,6 +10,7 @@ use frame_support::{
         Imbalance,
     },
 };
+use num_traits::Zero;
 
 impl<T: Config> Pallet<T> {
     /// ---- The implementation for the extrinsic become_delegate: signals that this hotkey allows delegated stake.
@@ -560,6 +562,13 @@ impl<T: Config> Pallet<T> {
         if !Self::hotkey_account_exists(hotkey) {
             Stake::<T>::insert(hotkey, coldkey, 0);
             Owner::<T>::insert(hotkey, coldkey);
+
+            // Update OwnedHotkeys map
+            let mut hotkeys = OwnedHotkeys::<T>::get(coldkey);
+            if !hotkeys.contains(hotkey) {
+                hotkeys.push(hotkey.clone());
+                OwnedHotkeys::<T>::insert(coldkey, hotkeys);
+            }
         }
     }
 
@@ -781,6 +790,31 @@ impl<T: Config> Pallet<T> {
         Ok(credit)
     }
 
+    pub fn kill_coldkey_account(
+        coldkey: &T::AccountId,
+        amount: <<T as Config>::Currency as fungible::Inspect<<T as system::Config>::AccountId>>::Balance,
+    ) -> Result<u64, DispatchError> {
+        if amount == 0 {
+            return Ok(0);
+        }
+
+        let credit = T::Currency::withdraw(
+            coldkey,
+            amount,
+            Precision::Exact,
+            Preservation::Expendable,
+            Fortitude::Force,
+        )
+        .map_err(|_| Error::<T>::BalanceWithdrawalError)?
+        .peek();
+
+        if credit == 0 {
+            return Err(Error::<T>::ZeroBalanceAfterWithdrawn.into());
+        }
+
+        Ok(credit)
+    }
+
     pub fn unstake_all_coldkeys_from_hotkey_account(hotkey: &T::AccountId) {
         // Iterate through all coldkeys that have a stake on this hotkey account.
         for (delegate_coldkey_i, stake_i) in
@@ -795,4 +829,96 @@ impl<T: Config> Pallet<T> {
             Self::add_balance_to_coldkey_account(&delegate_coldkey_i, stake_i);
         }
     }
+
+    /// Unstakes all tokens associated with a hotkey and transfers them to a new coldkey.
+    ///
+    /// This function performs the following operations:
+    /// 1. Verifies that the hotkey exists and is owned by the current coldkey.
+    /// 2. Ensures that the new coldkey is different from the current one.
+    /// 3. Unstakes all balance if there's any stake.
+    /// 4. Transfers the entire balance of the hotkey to the new coldkey.
+    /// 5. Verifies the success of the transfer and handles partial transfers if necessary.
+    ///
+    /// # Arguments
+    ///
+    /// * `current_coldkey` - The AccountId of the current coldkey.
+    /// * `hotkey` - The AccountId of the hotkey whose balance is being unstaked and transferred.
+    /// * `new_coldkey` - The AccountId of the new coldkey to receive the unstaked tokens.
+    ///
+    /// # Returns
+    ///
+    /// Returns a `DispatchResult` indicating success or failure of the operation.
+    ///
+    /// # Errors
+    ///
+    /// This function will return an error if:
+    /// * The hotkey account does not exist.
+    /// * The current coldkey does not own the hotkey.
+    /// * The new coldkey is the same as the current coldkey.
+    /// * There is no balance to transfer.
+    /// * The transfer fails or is only partially successful.
+    ///
+    /// # Events
+    ///
+    /// Emits an `AllBalanceUnstakedAndTransferredToNewColdkey` event upon successful execution.
+    /// Emits a `PartialBalanceTransferredToNewColdkey` event if only a partial transfer is successful.
+    ///
+    pub fn do_unstake_all_and_transfer_to_new_coldkey(
+        current_coldkey: T::AccountId,
+        hotkey: T::AccountId,
+        new_coldkey: T::AccountId,
+    ) -> DispatchResult {
+        // Ensure the hotkey exists and is owned by the current coldkey
+        ensure!(
+            Self::hotkey_account_exists(&hotkey),
+            Error::<T>::HotKeyAccountNotExists
+        );
+        ensure!(
+            Self::coldkey_owns_hotkey(&current_coldkey, &hotkey),
+            Error::<T>::NonAssociatedColdKey
+        );
+
+        // Ensure the new coldkey is different from the current one
+        ensure!(current_coldkey != new_coldkey, Error::<T>::SameColdkey);
+
+        // Get the current stake
+        let current_stake: u64 = Self::get_stake_for_coldkey_and_hotkey(&current_coldkey, &hotkey);
+
+        // Unstake all balance if there's any stake
+        if current_stake > 0 {
+            Self::do_remove_stake(
+                RawOrigin::Signed(current_coldkey.clone()).into(),
+                hotkey.clone(),
+                current_stake,
+            )?;
+        }
+
+        // Get the total balance of the current coldkey account
+        // let total_balance: <<T as Config>::Currency as fungible::Inspect<<T as system::Config>::AccountId>>::Balance = T::Currency::total_balance(&current_coldkey);
+
+        let total_balance = Self::get_coldkey_balance(&current_coldkey);
+        log::info!("Total Bank Balance: {:?}", total_balance);
+
+        // Ensure there's a balance to transfer
+        ensure!(!total_balance.is_zero(), Error::<T>::NoBalanceToTransfer);
+
+        // Attempt to transfer the entire total balance to the new coldkey
+        T::Currency::transfer(
+            &current_coldkey,
+            &new_coldkey,
+            total_balance,
+            Preservation::Expendable,
+        )?;
+
+        // Emit the event
+        Self::deposit_event(Event::AllBalanceUnstakedAndTransferredToNewColdkey {
+            current_coldkey: current_coldkey.clone(),
+            new_coldkey: new_coldkey.clone(),
+            hotkey: hotkey.clone(),
+            current_stake,
+            total_balance,
+        });
+
+        Ok(())
+    }
 }