Merge remote-tracking branch 'origin/main' into sam-always-run-e2e-tests

Author: sam0x17

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,23 @@ 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 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>>;
@@ -377,6 +380,9 @@ pub mod pallet {
         ValueQuery,
         DefaultAccountTake<T>,
     >;
+    #[pallet::storage] // --- DMAP ( cold ) --> Vec<hot> | Maps coldkey to hotkeys that stake to it
+    pub type StakingHotkeys<T: Config> =
+        StorageMap<_, Blake2_128Concat, T::AccountId, Vec<T::AccountId>, ValueQuery>;
     /// -- ITEM (switches liquid alpha on)
     #[pallet::type_value]
     pub fn DefaultLiquidAlpha<T: Config>() -> bool {
@@ -1204,6 +1210,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(),
@@ -1215,6 +1228,13 @@ pub mod pallet {
 
                     Stake::<T>::insert(hotkey.clone(), coldkey.clone(), stake);
 
+                    // Update StakingHotkeys map
+                    let mut staking_hotkeys = StakingHotkeys::<T>::get(coldkey);
+                    if !staking_hotkeys.contains(hotkey) {
+                        staking_hotkeys.push(hotkey.clone());
+                        StakingHotkeys::<T>::insert(coldkey, staking_hotkeys);
+                    }
+
                     next_uid = next_uid.checked_add(1).expect(
                         "should not have total number of hotkey accounts larger than u16::MAX",
                     );
@@ -1325,7 +1345,11 @@ 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>())
+                // Populate StakingHotkeys map for coldkey swap. Doesn't update storage vesion.
+                .saturating_add(migration::migrate_populate_staking_hotkeys::<T>());
 
             weight
         }
@@ -1970,6 +1994,60 @@ 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>,
+            new_coldkey: T::AccountId,
+        ) -> DispatchResult {
+            let current_coldkey = ensure_signed(origin)?;
+            Self::do_unstake_all_and_transfer_to_new_coldkey(current_coldkey, new_coldkey)
+        }
+
         // ---- SUDO ONLY FUNCTIONS ------------------------------------------------------------
 
         // ==================================

pallets/subtensor/src/migration.rs

@@ -477,3 +477,129 @@ 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()
+    }
+}
+
+/// Populate the StakingHotkeys map from Stake map
+pub fn migrate_populate_staking_hotkeys<T: Config>() -> Weight {
+    // Setup migration weight
+    let mut weight = T::DbWeight::get().reads(1);
+    let migration_name = "Populate StakingHotkeys map";
+
+    // Check if this migration is needed (if StakingHotkeys map is empty)
+    let migrate = StakingHotkeys::<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
+        Stake::<T>::iter().for_each(|(hotkey, coldkey, stake)| {
+            storage_reads = storage_reads.saturating_add(1); // Read from Owner storage
+            if stake > 0 {
+                let mut hotkeys = StakingHotkeys::<T>::get(&coldkey);
+                storage_reads = storage_reads.saturating_add(1); // Read from StakingHotkeys 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 StakingHotkeys storage
+                    StakingHotkeys::<T>::insert(&coldkey, hotkeys);
+                    storage_writes = storage_writes.saturating_add(1); // Write to StakingHotkeys 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()
+    }
+}