Merge pull request #1449 from opentensor/debloat-TotalHotkeyAlpha

Debloat total hotkey alpha

Author: sam0x17

pallets/subtensor/src/macros/hooks.rs

@@ -83,7 +83,9 @@ mod hooks {
                 // Remove unused maps entries
 				.saturating_add(migrations::migrate_remove_unused_maps_and_values::migrate_remove_unused_maps_and_values::<T>())
                 // Set last emission block number for all existed subnets before start call feature applied
-                .saturating_add(migrations::migrate_set_first_emission_block_number::migrate_set_first_emission_block_number::<T>());
+                .saturating_add(migrations::migrate_set_first_emission_block_number::migrate_set_first_emission_block_number::<T>())
+                // Remove all zero value entries in TotalHotkeyAlpha
+                .saturating_add(migrations::migrate_remove_zero_total_hotkey_alpha::migrate_remove_zero_total_hotkey_alpha::<T>());
             weight
         }
 

pallets/subtensor/src/migrations/migrate_remove_zero_total_hotkey_alpha.rs

@@ -0,0 +1,60 @@
+use super::*;
+use frame_support::{traits::Get, weights::Weight};
+use log;
+use scale_info::prelude::string::String;
+
+pub fn migrate_remove_zero_total_hotkey_alpha<T: Config>() -> Weight {
+    let migration_name = b"migrate_remove_zero_total_hotkey_alpha".to_vec();
+    let mut weight = T::DbWeight::get().reads(1);
+
+    // ------------------------------
+    // Step 0: Check if already run
+    // ------------------------------
+    if HasMigrationRun::<T>::get(&migration_name) {
+        log::info!(
+            "Migration '{:?}' has already run. Skipping.",
+            migration_name
+        );
+        return weight;
+    }
+
+    log::info!(
+        "Running migration '{}'",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    // ------------------------------
+    // Step 1: Remove any zero entries in TotalHotkeyAlpha
+    // ------------------------------
+
+    let mut removed_entries_count = 0u64;
+
+    // For each (hotkey, netuid, alpha) entry, remove if alpha == 0
+    for (hotkey, netuid, alpha) in TotalHotkeyAlpha::<T>::iter() {
+        if alpha == 0 {
+            TotalHotkeyAlpha::<T>::remove(&hotkey, netuid);
+            removed_entries_count = removed_entries_count.saturating_add(1);
+        }
+    }
+
+    weight = weight.saturating_add(T::DbWeight::get().reads(removed_entries_count));
+    weight = weight.saturating_add(T::DbWeight::get().writes(removed_entries_count));
+
+    log::info!(
+        "Removed {} zero entries from TotalHotkeyAlpha.",
+        removed_entries_count
+    );
+
+    // ------------------------------
+    // Step 2: Mark Migration as Completed
+    // ------------------------------
+    HasMigrationRun::<T>::insert(&migration_name, true);
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    log::info!(
+        "Migration '{:?}' completed successfully.",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    weight
+}

pallets/subtensor/src/migrations/mod.rs

@@ -11,6 +11,7 @@ pub mod migrate_populate_owned_hotkeys;
 pub mod migrate_rao;
 pub mod migrate_remove_stake_map;
 pub mod migrate_remove_unused_maps_and_values;
+pub mod migrate_remove_zero_total_hotkey_alpha;
 pub mod migrate_set_first_emission_block_number;
 pub mod migrate_set_min_burn;
 pub mod migrate_set_min_difficulty;

pallets/subtensor/src/staking/recycle_alpha.rs

@@ -42,7 +42,14 @@ impl<T: Config> Pallet<T> {
             Error::<T>::InsufficientLiquidity
         );
 
-        TotalHotkeyAlpha::<T>::mutate(&hotkey, netuid, |v| *v = v.saturating_sub(amount));
+        if TotalHotkeyAlpha::<T>::mutate(&hotkey, netuid, |v| {
+            *v = v.saturating_sub(amount);
+            *v
+        }) == 0
+        {
+            TotalHotkeyAlpha::<T>::remove(&hotkey, netuid);
+        }
+
         SubnetAlphaOut::<T>::mutate(netuid, |total| {
             *total = total.saturating_sub(amount);
         });
@@ -92,7 +99,13 @@ impl<T: Config> Pallet<T> {
             Error::<T>::InsufficientLiquidity
         );
 
-        TotalHotkeyAlpha::<T>::mutate(&hotkey, netuid, |v| *v = v.saturating_sub(amount));
+        if TotalHotkeyAlpha::<T>::mutate(&hotkey, netuid, |v| {
+            *v = v.saturating_sub(amount);
+            *v
+        }) == 0
+        {
+            TotalHotkeyAlpha::<T>::remove(&hotkey, netuid);
+        }
 
         // Deposit event
         Self::deposit_event(Event::AlphaBurned(coldkey, hotkey, amount, netuid));

pallets/subtensor/src/tests/migration.rs

@@ -438,3 +438,45 @@ fn test_migrate_set_first_emission_block_number() {
     }
 });
 }
+
+#[test]
+fn test_migrate_remove_zero_total_hotkey_alpha() {
+    new_test_ext(1).execute_with(|| {
+        const MIGRATION_NAME: &str = "migrate_remove_zero_total_hotkey_alpha";
+        let netuid = 1u16;
+
+        let hotkey_zero = U256::from(100u64);
+        let hotkey_nonzero = U256::from(101u64);
+
+        // Insert one zero-alpha entry and one non-zero entry
+        TotalHotkeyAlpha::<Test>::insert(hotkey_zero, netuid, 0u64);
+        TotalHotkeyAlpha::<Test>::insert(hotkey_nonzero, netuid, 123u64);
+
+        assert_eq!(TotalHotkeyAlpha::<Test>::get(hotkey_zero, netuid), 0u64);
+        assert_eq!(TotalHotkeyAlpha::<Test>::get(hotkey_nonzero, netuid), 123u64);
+
+        assert!(
+            !HasMigrationRun::<Test>::get(MIGRATION_NAME.as_bytes().to_vec()),
+            "Migration should not have run yet."
+        );
+
+        let weight = crate::migrations::migrate_remove_zero_total_hotkey_alpha::migrate_remove_zero_total_hotkey_alpha::<Test>();
+
+        assert!(
+            HasMigrationRun::<Test>::get(MIGRATION_NAME.as_bytes().to_vec()),
+            "Migration should be marked as run."
+        );
+
+        assert!(
+            !TotalHotkeyAlpha::<Test>::contains_key(hotkey_zero, netuid),
+            "Zero-alpha entry should have been removed."
+        );
+
+        assert_eq!(TotalHotkeyAlpha::<Test>::get(hotkey_nonzero, netuid), 123u64);
+
+        assert!(
+            !weight.is_zero(),
+            "Migration weight should be non-zero."
+        );
+    });
+}