Fix/charge-for-swap-before-schedule (#1141)

* charge for swap before sched

* add swap cost to the events

* pass swap cost into swap_coldkey call

* fmt

* oops, add arg

* commit Cargo.lock

* bump spec

* fix tests

pallets/subtensor/src/benchmarks.rs

@@ -515,7 +515,7 @@ reveal_weights {
     Identities::<T>::insert(&old_coldkey, identity);
 
     // Benchmark setup complete, now execute the extrinsic
-}: swap_coldkey(RawOrigin::Root, old_coldkey.clone(), new_coldkey.clone())
+}: swap_coldkey(RawOrigin::Root, old_coldkey.clone(), new_coldkey.clone(), swap_cost)
 
 batch_reveal_weights {
   let tempo: u16 = 0;

pallets/subtensor/src/macros/dispatches.rs

@@ -962,12 +962,13 @@ mod dispatches {
             origin: OriginFor<T>,
             old_coldkey: T::AccountId,
             new_coldkey: T::AccountId,
+            swap_cost: u64,
         ) -> DispatchResultWithPostInfo {
             // Ensure it's called with root privileges (scheduler has root privileges)
             ensure_root(origin)?;
             log::info!("swap_coldkey: {:?} -> {:?}", old_coldkey, new_coldkey);
 
-            Self::do_swap_coldkey(&old_coldkey, &new_coldkey)
+            Self::do_swap_coldkey(&old_coldkey, &new_coldkey, swap_cost)
         }
 
         /// Sets the childkey take for a given hotkey.
@@ -1327,13 +1328,21 @@ mod dispatches {
                 Error::<T>::SwapAlreadyScheduled
             );
 
+            // Calculate the swap cost and ensure sufficient balance
+            let swap_cost = Self::get_key_swap_cost();
+            ensure!(
+                Self::can_remove_balance_from_coldkey_account(&who, swap_cost),
+                Error::<T>::NotEnoughBalanceToPaySwapColdKey
+            );
+
             let current_block: BlockNumberFor<T> = <frame_system::Pallet<T>>::block_number();
             let duration: BlockNumberFor<T> = ColdkeySwapScheduleDuration::<T>::get();
             let when: BlockNumberFor<T> = current_block.saturating_add(duration);
 
             let call = Call::<T>::swap_coldkey {
                 old_coldkey: who.clone(),
                 new_coldkey: new_coldkey.clone(),
+                swap_cost,
             };
 
             let bound_call = T::Preimages::bound(LocalCallOf::<T>::from(call.clone()))
@@ -1354,6 +1363,7 @@ mod dispatches {
                 old_coldkey: who.clone(),
                 new_coldkey: new_coldkey.clone(),
                 execution_block: when,
+                swap_cost,
             });
 
             Ok(().into())

pallets/subtensor/src/macros/events.rs

@@ -155,6 +155,8 @@ mod events {
             old_coldkey: T::AccountId,
             /// the account ID of new coldkey
             new_coldkey: T::AccountId,
+            /// the swap cost
+            swap_cost: u64,
         },
         /// All balance of a hotkey has been unstaked and transferred to a new coldkey
         AllBalanceUnstakedAndTransferredToNewColdkey {
@@ -175,6 +177,8 @@ mod events {
             new_coldkey: T::AccountId,
             /// The arbitration block for the coldkey swap
             execution_block: BlockNumberFor<T>,
+            /// The swap cost
+            swap_cost: u64,
         },
         /// The arbitration period has been extended
         ArbitrationPeriodExtended {

pallets/subtensor/src/swap/swap_coldkey.rs

@@ -33,6 +33,7 @@ impl<T: Config> Pallet<T> {
     pub fn do_swap_coldkey(
         old_coldkey: &T::AccountId,
         new_coldkey: &T::AccountId,
+        swap_cost: u64,
     ) -> DispatchResultWithPostInfo {
         // 2. Initialize the weight for this operation
         let mut weight: Weight = T::DbWeight::get().reads(2);
@@ -55,8 +56,7 @@ impl<T: Config> Pallet<T> {
             Identities::<T>::insert(new_coldkey, identity);
         }
 
-        // 6. Calculate the swap cost and ensure sufficient balance
-        let swap_cost = Self::get_key_swap_cost();
+        // 6. Ensure sufficient balance for the swap cost
         ensure!(
             Self::can_remove_balance_from_coldkey_account(old_coldkey, swap_cost),
             Error::<T>::NotEnoughBalanceToPaySwapColdKey
@@ -83,6 +83,7 @@ impl<T: Config> Pallet<T> {
         Self::deposit_event(Event::ColdkeySwapped {
             old_coldkey: old_coldkey.clone(),
             new_coldkey: new_coldkey.clone(),
+            swap_cost,
         });
 
         // 12. Return the result with the updated weight

pallets/subtensor/src/tests/swap_coldkey.rs

@@ -709,7 +709,8 @@ fn test_do_swap_coldkey_success() {
         assert_ok!(SubtensorModule::do_swap_coldkey(
             // <<Test as Config>::RuntimeOrigin>::signed(old_coldkey),
             &old_coldkey,
-            &new_coldkey
+            &new_coldkey,
+            swap_cost
         ));
 
         // Log state after swap
@@ -782,6 +783,7 @@ fn test_do_swap_coldkey_success() {
             Event::ColdkeySwapped {
                 old_coldkey,
                 new_coldkey,
+                swap_cost,
             }
             .into(),
         );
@@ -1195,7 +1197,11 @@ fn test_do_swap_coldkey_with_subnet_ownership() {
         OwnedHotkeys::<Test>::insert(old_coldkey, vec![hotkey]);
 
         // Perform the swap
-        assert_ok!(SubtensorModule::do_swap_coldkey(&old_coldkey, &new_coldkey));
+        assert_ok!(SubtensorModule::do_swap_coldkey(
+            &old_coldkey,
+            &new_coldkey,
+            swap_cost
+        ));
 
         // Verify subnet ownership transfer
         assert_eq!(SubnetOwner::<Test>::get(netuid), new_coldkey);
@@ -1652,8 +1658,10 @@ fn test_schedule_swap_coldkey_success() {
         let old_coldkey: U256 = U256::from(1);
         let new_coldkey: U256 = U256::from(2);
 
+        let swap_cost = SubtensorModule::get_key_swap_cost();
+
         // Add balance to the old coldkey account
-        SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, 1000);
+        SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, swap_cost + 1_000);
 
         // Schedule the coldkey swap
         assert_ok!(SubtensorModule::schedule_swap_coldkey(
@@ -1673,6 +1681,7 @@ fn test_schedule_swap_coldkey_success() {
                 old_coldkey,
                 new_coldkey,
                 execution_block: expected_execution_block,
+                swap_cost,
             }
             .into(),
         );
@@ -1689,7 +1698,9 @@ fn test_schedule_swap_coldkey_duplicate() {
         let old_coldkey = U256::from(1);
         let new_coldkey = U256::from(2);
 
-        SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, 2000);
+        let swap_cost = SubtensorModule::get_key_swap_cost();
+
+        SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, swap_cost + 2_000);
 
         assert_ok!(SubtensorModule::schedule_swap_coldkey(
             <<Test as Config>::RuntimeOrigin>::signed(old_coldkey),
@@ -1734,6 +1745,8 @@ fn test_schedule_swap_coldkey_execution() {
             "Initial ownership check failed"
         );
 
+        let swap_cost = SubtensorModule::get_key_swap_cost();
+
         // Schedule the swap
         assert_ok!(SubtensorModule::schedule_swap_coldkey(
             <<Test as Config>::RuntimeOrigin>::signed(old_coldkey),
@@ -1749,6 +1762,7 @@ fn test_schedule_swap_coldkey_execution() {
                 old_coldkey,
                 new_coldkey,
                 execution_block,
+                swap_cost,
             }
             .into(),
         );
@@ -1790,6 +1804,7 @@ fn test_schedule_swap_coldkey_execution() {
             Event::ColdkeySwapped {
                 old_coldkey,
                 new_coldkey,
+                swap_cost,
             }
             .into(),
         );
@@ -1807,7 +1822,8 @@ fn test_direct_swap_coldkey_call_fails() {
             SubtensorModule::swap_coldkey(
                 <<Test as Config>::RuntimeOrigin>::signed(old_coldkey),
                 old_coldkey,
-                new_coldkey
+                new_coldkey,
+                0
             ),
             BadOrigin
         );
@@ -1822,7 +1838,9 @@ fn test_schedule_swap_coldkey_with_pending_swap() {
         let new_coldkey1 = U256::from(2);
         let new_coldkey2 = U256::from(3);
 
-        SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, 2000);
+        let swap_cost = SubtensorModule::get_key_swap_cost();
+
+        SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, swap_cost + 1_000);
 
         assert_ok!(SubtensorModule::schedule_swap_coldkey(
             <<Test as Config>::RuntimeOrigin>::signed(old_coldkey),
@@ -1876,7 +1894,11 @@ fn test_coldkey_swap_delegate_identity_updated() {
         assert!(Identities::<Test>::get(old_coldkey).is_some());
         assert!(Identities::<Test>::get(new_coldkey).is_none());
 
-        assert_ok!(SubtensorModule::do_swap_coldkey(&old_coldkey, &new_coldkey));
+        assert_ok!(SubtensorModule::do_swap_coldkey(
+            &old_coldkey,
+            &new_coldkey,
+            burn_cost
+        ));
 
         assert!(Identities::<Test>::get(old_coldkey).is_none());
         assert!(Identities::<Test>::get(new_coldkey).is_some());
@@ -1912,7 +1934,11 @@ fn test_coldkey_swap_no_identity_no_changes() {
         assert!(Identities::<Test>::get(old_coldkey).is_none());
 
         // Perform the coldkey swap
-        assert_ok!(SubtensorModule::do_swap_coldkey(&old_coldkey, &new_coldkey));
+        assert_ok!(SubtensorModule::do_swap_coldkey(
+            &old_coldkey,
+            &new_coldkey,
+            burn_cost
+        ));
 
         // Ensure no identities have been changed
         assert!(Identities::<Test>::get(old_coldkey).is_none());
@@ -1956,10 +1982,33 @@ fn test_coldkey_swap_no_identity_no_changes_newcoldkey_exists() {
         assert!(Identities::<Test>::get(old_coldkey).is_none());
 
         // Perform the coldkey swap
-        assert_ok!(SubtensorModule::do_swap_coldkey(&old_coldkey, &new_coldkey));
+        assert_ok!(SubtensorModule::do_swap_coldkey(
+            &old_coldkey,
+            &new_coldkey,
+            burn_cost
+        ));
 
         // Ensure no identities have been changed
         assert!(Identities::<Test>::get(old_coldkey).is_none());
         assert!(Identities::<Test>::get(new_coldkey).is_some());
     });
 }
+
+// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test swap_coldkey -- test_cant_schedule_swap_without_enough_to_burn --exact --nocapture
+#[test]
+fn test_cant_schedule_swap_without_enough_to_burn() {
+    new_test_ext(1).execute_with(|| {
+        let old_coldkey = U256::from(3);
+        let new_coldkey = U256::from(4);
+        let hotkey = U256::from(5);
+
+        let burn_cost = SubtensorModule::get_key_swap_cost();
+        assert_noop!(
+            SubtensorModule::schedule_swap_coldkey(
+                <<Test as Config>::RuntimeOrigin>::signed(old_coldkey),
+                new_coldkey
+            ),
+            Error::<Test>::NotEnoughBalanceToPaySwapColdKey
+        );
+    });
+}

runtime/src/lib.rs

@@ -220,7 +220,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: 222,
+    spec_version: 223,
     impl_version: 1,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 1,