Custom dust collection type in subtensor

pallets/subtensor/src/coinbase/tao.rs

@@ -33,33 +33,21 @@ impl<T: Config> Pallet<T> {
         SubnetTAO::<T>::get(netuid)
     }
 
-    /// Internal function that transfers and updates subtensor pallet total issuance
-    /// in case of dust collection.
+    /// Internal function that transfers TAO and allows the origin account to be reaped.
+    ///
+    /// Dust collection is handled by the runtime's Balances `DustRemoval` implementation.
     fn transfer_allow_death_update_ti(
         origin_coldkey: &T::AccountId,
         destination_coldkey: &T::AccountId,
         amount: BalanceOf<T>,
     ) -> DispatchResult {
-        // If account balance remainder drops below ED, then account is killed, balance
-        // is lost, and we need to reduce total issuance in subtensor pallet. Measure
-        // balance TI before and after to detect the dust.
-        let balances_ti_before = <T as pallet::Config>::Currency::total_issuance();
-
         <T as pallet::Config>::Currency::transfer(
             origin_coldkey,
             destination_coldkey,
             amount,
             Preservation::Expendable,
         )?;
 
-        let balances_ti_after = <T as pallet::Config>::Currency::total_issuance();
-        if balances_ti_after < balances_ti_before {
-            let burned = balances_ti_before.saturating_sub(balances_ti_after);
-            TotalIssuance::<T>::mutate(|total| {
-                *total = total.saturating_sub(burned);
-            });
-        }
-
         Ok(())
     }
 

pallets/subtensor/src/lib.rs

@@ -69,6 +69,21 @@ pub const MAX_SUBNET_CLAIMS: usize = 5;
 
 pub const MAX_ROOT_CLAIM_THRESHOLD: u64 = 10_000_000;
 
+pub struct SubtensorDustRemoval<T>(PhantomData<T>);
+impl<T> frame_support::traits::OnUnbalanced<pallet_balances::CreditOf<T, ()>>
+    for SubtensorDustRemoval<T>
+where
+    T: Config + pallet_balances::Config,
+    <T as pallet_balances::Config>::Balance: Into<TaoBalance> + Copy,
+{
+    fn on_nonzero_unbalanced(dust: pallet_balances::CreditOf<T, ()>) {
+        let amount: TaoBalance = frame_support::traits::Imbalance::peek(&dust).into();
+        TotalIssuance::<T>::mutate(|total| {
+            *total = total.saturating_sub(amount);
+        });
+    }
+}
+
 #[allow(deprecated)]
 #[deny(missing_docs)]
 #[import_section(errors::errors)]

pallets/subtensor/src/migrations/migrate_fix_total_issuance_evm_fees.rs

@@ -3,24 +3,31 @@ use frame_support::traits::fungible::Inspect;
 use frame_support::weights::Weight;
 
 pub fn migrate_fix_total_issuance_evm_fees<T: Config>() -> Weight {
-    let migration_name = b"migrate_fix_total_issuance_evm_fees".to_vec();
-    let mut weight = T::DbWeight::get().reads(1);
-
-    if HasMigrationRun::<T>::get(&migration_name) {
-        log::info!(
-            "Migration '{:?}' has already run. Skipping.",
-            String::from_utf8_lossy(&migration_name)
-        );
+    let migration_names: [&[u8]; 2] = [
+        // Fix testnet TotalIssuance after the earlier EVM fees issue caused the
+        // Subtensor pallet's accounting to diverge from the balances pallet.
+        b"migrate_fix_total_issuance_evm_fees",
+        // Fix Subtensor TotalIssuance after dust collection caused accounting drift.
+        b"migrate_fix_total_issuance_after_dust_collection",
+    ];
+    let mut weight = T::DbWeight::get().reads(migration_names.len() as u64);
+
+    let Some(migration_name) = migration_names
+        .iter()
+        .map(|name| name.to_vec())
+        .find(|name| !HasMigrationRun::<T>::get(name))
+    else {
+        log::info!("All total issuance fix migrations have already run. Skipping.");
         return weight;
-    }
+    };
 
     log::info!(
         "Running migration '{}'",
         String::from_utf8_lossy(&migration_name)
     );
 
-    // Fix testnet TotalIssuance after the earlier EVM fees issue caused the
-    // Subtensor pallet's accounting to diverge from the balances pallet.
+    // All migration instances reset Subtensor TotalIssuance to the authoritative
+    // Balances pallet total issuance.
     let balances_total_issuance = <T as Config>::Currency::total_issuance();
     let subtensor_total_issuance_before = TotalIssuance::<T>::get();
     TotalIssuance::<T>::put(balances_total_issuance);

pallets/subtensor/src/swap/swap_hotkey.rs

@@ -46,16 +46,24 @@ impl<T: Config> Pallet<T> {
             Error::<T>::NonAssociatedColdKey
         );
 
-        // 3. Initialize the weight for this operation
+        // 3. If the new hotkey already exists globally, ensure the coldkey owns it
+        if Self::hotkey_account_exists(new_hotkey) {
+            ensure!(
+                Self::coldkey_owns_hotkey(&coldkey, new_hotkey),
+                Error::<T>::NonAssociatedColdKey
+            );
+        }
+
+        // 4. Initialize the weight for this operation
         let mut weight = T::DbWeight::get().reads(2);
 
-        // 4. Ensure the new hotkey is different from the old one
+        // 5. Ensure the new hotkey is different from the old one
         ensure!(old_hotkey != new_hotkey, Error::<T>::NewHotKeyIsSameWithOld);
 
-        // 5. Get the current block number
+        // 6. Get the current block number
         let block: u64 = Self::get_current_block_as_u64();
 
-        // 6. Ensure the transaction rate limit is not exceeded
+        // 7. Ensure the transaction rate limit is not exceeded
         ensure!(
             !Self::exceeds_tx_rate_limit(Self::get_last_tx_block(&coldkey), block),
             Error::<T>::HotKeySetTxRateLimitExceeded
@@ -64,14 +72,14 @@ impl<T: Config> Pallet<T> {
         weight.saturating_accrue(T::DbWeight::get().reads(2));
 
         match netuid {
-            // 7. Ensure the hotkey is not registered on the network before, if netuid is provided
+            // 8. Ensure the hotkey is not registered on the network before, if netuid is provided
             Some(netuid) => {
                 ensure!(
                     !Self::is_hotkey_registered_on_specific_network(new_hotkey, netuid),
                     Error::<T>::HotKeyAlreadyRegisteredInSubNet
                 );
             }
-            // 7.1 Ensure the new hotkey is not already registered on any network, only if netuid is none
+            // 8.1 Ensure the new hotkey is not already registered on any network, only if netuid is none
             None => {
                 ensure!(
                     !Self::is_hotkey_registered_on_any_network(new_hotkey),
@@ -80,7 +88,7 @@ impl<T: Config> Pallet<T> {
             }
         }
 
-        // 7.2 If the swap touches the root subnet, require that new_hotkey is clean
+        // 8.2 If the swap touches the root subnet, require that new_hotkey is clean
         // on root (no outstanding claimable rate and no existing root stake). Merging
         // a non-empty rate-book would either violate total conservation or misallocate
         // dividends across coldkeys that never staked on old_hotkey.
@@ -96,46 +104,46 @@ impl<T: Config> Pallet<T> {
             );
         }
 
-        // 8. Swap LastTxBlock
+        // 9. Swap LastTxBlock
         let last_tx_block: u64 = Self::get_last_tx_block(old_hotkey);
         Self::set_last_tx_block(new_hotkey, last_tx_block);
         weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
 
-        // 9. Swap LastTxBlockDelegateTake
+        // 10. Swap LastTxBlockDelegateTake
         let last_tx_block_delegate_take: u64 = Self::get_last_tx_block_delegate_take(old_hotkey);
         Self::set_last_tx_block_delegate_take(new_hotkey, last_tx_block_delegate_take);
         weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
 
-        // 10. Swap LastTxBlockChildKeyTake
+        // 11. Swap LastTxBlockChildKeyTake
         let last_tx_block_child_key_take: u64 = Self::get_last_tx_block_childkey_take(old_hotkey);
         Self::set_last_tx_block_childkey(new_hotkey, last_tx_block_child_key_take);
         weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
 
-        // 11. fork for swap hotkey on a specific subnet case after do the common check
+        // 12. fork for swap hotkey on a specific subnet case after do the common check
         if let Some(netuid) = netuid {
             return Self::swap_hotkey_on_subnet(
                 &coldkey, old_hotkey, new_hotkey, netuid, weight, keep_stake,
             );
         };
 
         // Start to do everything for swap hotkey on all subnets case
-        // 12. Get the cost for swapping the key
+        // 13. Get the cost for swapping the key
         let swap_cost = Self::get_key_swap_cost();
         log::debug!("Swap cost: {swap_cost:?}");
 
-        // 13. Ensure the coldkey has enough balance to pay for the swap
+        // 14. Ensure the coldkey has enough balance to pay for the swap
         ensure!(
             Self::can_remove_balance_from_coldkey_account(&coldkey, swap_cost.into()),
             Error::<T>::NotEnoughBalanceToPaySwapHotKey
         );
 
         weight.saturating_accrue(T::DbWeight::get().reads_writes(3, 0));
 
-        // 14. Remove the swap cost from the coldkey's account + Recycle the tokens
+        // 15. Remove the swap cost from the coldkey's account + Recycle the tokens
         Self::recycle_tao(&coldkey, swap_cost.into())?;
         weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 2));
 
-        // 19. Perform the hotkey swap
+        // 16. Perform the hotkey swap
         Self::perform_hotkey_swap_on_all_subnets(
             old_hotkey,
             new_hotkey,
@@ -144,18 +152,18 @@ impl<T: Config> Pallet<T> {
             keep_stake,
         )?;
 
-        // 20. Update the last transaction block for the coldkey
+        // 17. Update the last transaction block for the coldkey
         Self::set_last_tx_block(&coldkey, block);
         weight.saturating_accrue(T::DbWeight::get().writes(1));
 
-        // 21. Emit an event for the hotkey swap
+        // 18. Emit an event for the hotkey swap
         Self::deposit_event(Event::HotkeySwapped {
             coldkey,
             old_hotkey: old_hotkey.clone(),
             new_hotkey: new_hotkey.clone(),
         });
 
-        // 22. Return the weight of the operation
+        // 19. Return the weight of the operation
         Ok(Some(weight).into())
     }
 

pallets/subtensor/src/tests/migration.rs

@@ -4691,6 +4691,7 @@ fn test_migrate_subnet_balances() {
 fn test_migrate_fix_total_issuance_evm_fees() {
     new_test_ext(1).execute_with(|| {
         const MIGRATION_NAME: &[u8] = b"migrate_fix_total_issuance_evm_fees";
+        const DUST_MIGRATION_NAME: &[u8] = b"migrate_fix_total_issuance_after_dust_collection";
 
         let account = U256::from(42);
         let balances_total_issuance = TaoBalance::from(123_456_789_u64);
@@ -4711,16 +4712,29 @@ fn test_migrate_fix_total_issuance_evm_fees() {
         assert!(!weight.is_zero(), "weight must be non-zero");
         assert_eq!(TotalIssuance::<Test>::get(), balances_total_issuance);
         assert!(HasMigrationRun::<Test>::get(MIGRATION_NAME.to_vec()));
+        assert!(!HasMigrationRun::<Test>::get(
+            DUST_MIGRATION_NAME.to_vec()
+        ));
 
         let second_wrong_value = TaoBalance::from(555_u64);
         TotalIssuance::<Test>::put(second_wrong_value);
 
         crate::migrations::migrate_fix_total_issuance_evm_fees::migrate_fix_total_issuance_evm_fees::<Test>();
 
+        assert_eq!(TotalIssuance::<Test>::get(), balances_total_issuance);
+        assert!(HasMigrationRun::<Test>::get(
+            DUST_MIGRATION_NAME.to_vec()
+        ));
+
+        let third_wrong_value = TaoBalance::from(777_u64);
+        TotalIssuance::<Test>::put(third_wrong_value);
+
+        crate::migrations::migrate_fix_total_issuance_evm_fees::migrate_fix_total_issuance_evm_fees::<Test>();
+
         assert_eq!(
             TotalIssuance::<Test>::get(),
-            second_wrong_value,
-            "migration must not run more than once"
+            third_wrong_value,
+            "migration must not run after all known migration keys have run"
         );
     });
 }

pallets/subtensor/src/tests/mock.rs

@@ -88,7 +88,7 @@ pub type BlockNumber = u64;
 impl pallet_balances::Config for Test {
     type Balance = Balance;
     type RuntimeEvent = RuntimeEvent;
-    type DustRemoval = ();
+    type DustRemoval = crate::SubtensorDustRemoval<Test>;
     type ExistentialDeposit = ExistentialDeposit;
     type AccountStore = System;
     type MaxLocks = ();

pallets/subtensor/src/tests/mock_high_ed.rs

@@ -70,7 +70,7 @@ pub type BlockNumber = u64;
 impl pallet_balances::Config for Test {
     type Balance = Balance;
     type RuntimeEvent = RuntimeEvent;
-    type DustRemoval = ();
+    type DustRemoval = crate::SubtensorDustRemoval<Test>;
     type ExistentialDeposit = ExistentialDeposit;
     type AccountStore = System;
     type MaxLocks = ();

pallets/subtensor/src/tests/swap_hotkey_with_subnet.rs

@@ -926,7 +926,7 @@ fn test_swap_owner_new_hotkey_already_exists() {
                 Some(netuid),
                 false
             ),
-            Error::<Test>::HotKeyAlreadyRegisteredInSubNet
+            Error::<Test>::NonAssociatedColdKey
         );
 
         // Verify the swap

runtime/src/lib.rs

@@ -234,7 +234,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: 417,
+    spec_version: 418,
     impl_version: 1,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 1,
@@ -476,7 +476,7 @@ impl pallet_balances::Config for Runtime {
     type Balance = Balance;
     // The ubiquitous event type.
     type RuntimeEvent = RuntimeEvent;
-    type DustRemoval = ();
+    type DustRemoval = pallet_subtensor::SubtensorDustRemoval<Runtime>;
     type ExistentialDeposit = ExistentialDeposit;
     type AccountStore = System;
     type WeightInfo = pallet_balances::weights::SubstrateWeight<Runtime>;

runtime/tests/balances_dust.rs

@@ -0,0 +1,60 @@
+#![allow(clippy::unwrap_used)]
+
+use frame_support::traits::fungible::{Inspect, Mutate};
+use frame_support::traits::tokens::Preservation;
+use node_subtensor_runtime::{Balances, BuildStorage, Runtime, RuntimeGenesisConfig};
+use sp_core::crypto::AccountId32;
+use subtensor_runtime_common::TaoBalance;
+
+fn new_test_ext() -> sp_io::TestExternalities {
+    let mut ext: sp_io::TestExternalities = RuntimeGenesisConfig {
+        ..Default::default()
+    }
+    .build_storage()
+    .unwrap()
+    .into();
+    ext.execute_with(|| frame_system::Pallet::<Runtime>::set_block_number(1));
+    ext
+}
+
+fn add_balance_to_account(account: &AccountId32, tao: TaoBalance) {
+    let credit = pallet_subtensor::Pallet::<Runtime>::mint_tao(tao);
+    let _ = pallet_subtensor::Pallet::<Runtime>::spend_tao(account, credit, tao).unwrap();
+}
+
+#[test]
+fn balances_dust_removal_updates_subtensor_total_issuance() {
+    new_test_ext().execute_with(|| {
+        let origin = AccountId32::new([1u8; 32]);
+        let destination = AccountId32::new([2u8; 32]);
+        let existential_deposit = TaoBalance::from(500u64);
+        let dust = TaoBalance::from(1u64);
+        let transfer_amount = existential_deposit;
+
+        add_balance_to_account(&origin, transfer_amount + dust);
+
+        let balances_issuance_before = Balances::total_issuance();
+        let subtensor_issuance_before = pallet_subtensor::Pallet::<Runtime>::get_total_issuance();
+        assert_eq!(balances_issuance_before, subtensor_issuance_before);
+
+        <Balances as Mutate<AccountId32>>::transfer(
+            &origin,
+            &destination,
+            transfer_amount,
+            Preservation::Expendable,
+        )
+        .unwrap();
+
+        assert_eq!(Balances::total_balance(&origin), 0u64.into());
+        assert_eq!(Balances::total_balance(&destination), transfer_amount);
+        assert_eq!(balances_issuance_before - Balances::total_issuance(), dust);
+        assert_eq!(
+            subtensor_issuance_before - pallet_subtensor::Pallet::<Runtime>::get_total_issuance(),
+            dust
+        );
+        assert_eq!(
+            Balances::total_issuance(),
+            pallet_subtensor::Pallet::<Runtime>::get_total_issuance()
+        );
+    });
+}