Fix unstaking 99.999% and add tests

Author: gztensor

pallets/subtensor/src/tests/move_stake.rs

@@ -3,7 +3,7 @@ use crate::*;
 use approx::assert_abs_diff_eq;
 use frame_support::{assert_err, assert_noop, assert_ok};
 use sp_core::{Get, U256};
-use substrate_fixed::types::I96F32;
+use substrate_fixed::types::{I96F32, U64F64};
 
 // 1. test_do_move_success
 // Description: Test a successful move of stake between two hotkeys in the same subnet
@@ -1641,3 +1641,112 @@ fn test_stake_transfers_disabled_validate() {
         assert_ok!(result3);
     });
 }
+
+#[test]
+// RUST_LOG=info cargo test --package pallet-subtensor --lib -- tests::staking::test_move_stake_specific_stake_into_subnet_fail --exact --show-output
+fn test_move_stake_specific_stake_into_subnet_fail() {
+    new_test_ext(1).execute_with(|| {
+        let sn_owner_coldkey = U256::from(55453);
+
+        let hotkey_account_id = U256::from(533453);
+        let coldkey_account_id = U256::from(55454);
+        let hotkey_owner_account_id = U256::from(533454);
+
+        let existing_shares: U64F64 =
+            U64F64::from_num(161_986_254).saturating_div(U64F64::from_num(u64::MAX));
+        let existing_stake = 36_711_495_953;
+
+        let tao_in = 2_409_892_148_947;
+        let alpha_in = 15_358_708_513_716;
+
+        let tao_staked = 200_000_000;
+
+        //add network
+        let netuid: u16 = add_dynamic_network(&sn_owner_coldkey, &sn_owner_coldkey);
+
+        let origin_netuid: u16 = add_dynamic_network(&sn_owner_coldkey, &sn_owner_coldkey);
+
+        // Register hotkey on netuid
+        register_ok_neuron(netuid, hotkey_account_id, hotkey_owner_account_id, 0);
+        // Register hotkey on origin netuid
+        register_ok_neuron(origin_netuid, hotkey_account_id, hotkey_owner_account_id, 0);
+
+        // Check we have zero staked
+        assert_eq!(
+            SubtensorModule::get_total_stake_for_hotkey(&hotkey_account_id),
+            0
+        );
+
+        // Set a hotkey pool for the hotkey on destination subnet
+        let mut hotkey_pool = SubtensorModule::get_alpha_share_pool(hotkey_account_id, netuid);
+        hotkey_pool.update_value_for_one(&hotkey_owner_account_id, 1234); // Doesn't matter, will be overridden
+
+        // Adjust the total hotkey stake and shares to match the existing values
+        TotalHotkeyShares::<Test>::insert(hotkey_account_id, netuid, existing_shares);
+        TotalHotkeyAlpha::<Test>::insert(hotkey_account_id, netuid, existing_stake);
+
+        // Make the hotkey a delegate
+        Delegates::<Test>::insert(hotkey_account_id, 0);
+
+        // Setup Subnet pool
+        SubnetAlphaIn::<Test>::insert(netuid, alpha_in);
+        SubnetTAO::<Test>::insert(netuid, tao_in);
+
+        // Give TAO balance to coldkey
+        SubtensorModule::add_balance_to_coldkey_account(
+            &coldkey_account_id,
+            tao_staked + 1_000_000_000,
+        );
+
+        // Setup Subnet pool for origin netuid
+        SubnetAlphaIn::<Test>::insert(origin_netuid, alpha_in + 10_000_000);
+        SubnetTAO::<Test>::insert(origin_netuid, tao_in + 10_000_000);
+
+        // Add stake as new hotkey
+        assert_ok!(SubtensorModule::add_stake(
+            RuntimeOrigin::signed(coldkey_account_id),
+            hotkey_account_id,
+            origin_netuid,
+            tao_staked,
+        ),);
+        let alpha_to_move = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey_account_id,
+            &coldkey_account_id,
+            origin_netuid,
+        );
+
+        // Move stake to destination subnet
+        assert_ok!(SubtensorModule::move_stake(
+            RuntimeOrigin::signed(coldkey_account_id),
+            hotkey_account_id,
+            hotkey_account_id,
+            origin_netuid,
+            netuid,
+            alpha_to_move,
+        ));
+
+        // Check that the stake has been moved
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &hotkey_account_id,
+                &coldkey_account_id,
+                origin_netuid
+            ),
+            0
+        );
+        let fee = DefaultStakingFee::<Test>::get();
+        let alpha_fee: I96F32 = I96F32::from_num(fee) / SubtensorModule::get_alpha_price(netuid);
+        let expected_value = I96F32::from_num(alpha_to_move)
+            * SubtensorModule::get_alpha_price(origin_netuid)
+            / SubtensorModule::get_alpha_price(netuid);
+        assert_abs_diff_eq!(
+            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &hotkey_account_id,
+                &coldkey_account_id,
+                netuid
+            ),
+            (expected_value - alpha_fee).to_num::<u64>(),
+            epsilon = (expected_value / 1000).to_num::<u64>()
+        );
+    });
+}

pallets/subtensor/src/tests/staking.rs

@@ -3692,111 +3692,121 @@ fn test_add_stake_specific_stake_into_subnet_fail() {
     });
 }
 
+// cargo test --package pallet-subtensor --lib -- tests::staking::test_remove_99_999_per_cent_stake_removes_all --exact --show-output
 #[test]
-// RUST_LOG=info cargo test --package pallet-subtensor --lib -- tests::staking::test_move_stake_specific_stake_into_subnet_fail --exact --show-output
-fn test_move_stake_specific_stake_into_subnet_fail() {
+fn test_remove_99_9991_per_cent_stake_removes_all() {
+    // When we remove stake, the total issuance of the balances pallet should not change
+    //    this is because the stake should be part of the coldkey account balance (reserved/locked)
+    //    then the removed stake just becomes free balance
     new_test_ext(1).execute_with(|| {
-        let sn_owner_coldkey = U256::from(55453);
-
-        let hotkey_account_id = U256::from(533453);
-        let coldkey_account_id = U256::from(55454);
-        let hotkey_owner_account_id = U256::from(533454);
-
-        let existing_shares: U64F64 =
-            U64F64::from_num(161_986_254).saturating_div(U64F64::from_num(u64::MAX));
-        let existing_stake = 36_711_495_953;
-
-        let tao_in = 2_409_892_148_947;
-        let alpha_in = 15_358_708_513_716;
-
-        let tao_staked = 200_000_000;
+        let subnet_owner_coldkey = U256::from(1);
+        let subnet_owner_hotkey = U256::from(2);
+        let hotkey_account_id = U256::from(581337);
+        let coldkey_account_id = U256::from(81337);
+        let amount = 10_000_000_000;
+        let netuid: u16 = add_dynamic_network(&subnet_owner_hotkey, &subnet_owner_coldkey);
+        let fee = DefaultStakingFee::<Test>::get();
+        register_ok_neuron(netuid, hotkey_account_id, coldkey_account_id, 192213123);
 
-        //add network
-        let netuid: u16 = add_dynamic_network(&sn_owner_coldkey, &sn_owner_coldkey);
+        // Give it some $$$ in his coldkey balance
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey_account_id, amount);
 
-        let origin_netuid: u16 = add_dynamic_network(&sn_owner_coldkey, &sn_owner_coldkey);
+        // Stake to hotkey account, and check if the result is ok
+        assert_ok!(SubtensorModule::add_stake(
+            RuntimeOrigin::signed(coldkey_account_id),
+            hotkey_account_id,
+            netuid,
+            amount
+        ));
 
-        // Register hotkey on netuid
-        register_ok_neuron(netuid, hotkey_account_id, hotkey_owner_account_id, 0);
-        // Register hotkey on origin netuid
-        register_ok_neuron(origin_netuid, hotkey_account_id, hotkey_owner_account_id, 0);
+        // Remove 99.9991% stake
+        let alpha = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey_account_id,
+            &coldkey_account_id,
+            netuid,
+        );
+        assert_ok!(SubtensorModule::remove_stake(
+            RuntimeOrigin::signed(coldkey_account_id),
+            hotkey_account_id,
+            netuid,
+            (U64F64::from_num(alpha) * U64F64::from_num(0.999991)).to_num::<u64>()
+        ));
 
-        // Check we have zero staked
+        // Check that all alpha was unstaked and all TAO balance was returned (less fees)
+        assert_abs_diff_eq!(
+            SubtensorModule::get_coldkey_balance(&coldkey_account_id),
+            amount - fee * 2,
+            epsilon = 10000,
+        );
         assert_eq!(
             SubtensorModule::get_total_stake_for_hotkey(&hotkey_account_id),
             0
         );
-
-        // Set a hotkey pool for the hotkey on destination subnet
-        let mut hotkey_pool = SubtensorModule::get_alpha_share_pool(hotkey_account_id, netuid);
-        hotkey_pool.update_value_for_one(&hotkey_owner_account_id, 1234); // Doesn't matter, will be overridden
-
-        // Adjust the total hotkey stake and shares to match the existing values
-        TotalHotkeyShares::<Test>::insert(hotkey_account_id, netuid, existing_shares);
-        TotalHotkeyAlpha::<Test>::insert(hotkey_account_id, netuid, existing_stake);
-
-        // Make the hotkey a delegate
-        Delegates::<Test>::insert(hotkey_account_id, 0);
-
-        // Setup Subnet pool
-        SubnetAlphaIn::<Test>::insert(netuid, alpha_in);
-        SubnetTAO::<Test>::insert(netuid, tao_in);
-
-        // Give TAO balance to coldkey
-        SubtensorModule::add_balance_to_coldkey_account(
+        let new_alpha = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey_account_id,
             &coldkey_account_id,
-            tao_staked + 1_000_000_000,
+            netuid,
         );
+        assert_eq!(new_alpha, 0);
+    });
+}
+
+// cargo test --package pallet-subtensor --lib -- tests::staking::test_remove_99_9989_per_cent_stake_leaves_a_little --exact --show-output
+#[test]
+fn test_remove_99_9989_per_cent_stake_leaves_a_little() {
+    // When we remove stake, the total issuance of the balances pallet should not change
+    //    this is because the stake should be part of the coldkey account balance (reserved/locked)
+    //    then the removed stake just becomes free balance
+    new_test_ext(1).execute_with(|| {
+        let subnet_owner_coldkey = U256::from(1);
+        let subnet_owner_hotkey = U256::from(2);
+        let hotkey_account_id = U256::from(581337);
+        let coldkey_account_id = U256::from(81337);
+        let amount = 10_000_000_000;
+        let netuid: u16 = add_dynamic_network(&subnet_owner_hotkey, &subnet_owner_coldkey);
+        let fee = DefaultStakingFee::<Test>::get();
+        register_ok_neuron(netuid, hotkey_account_id, coldkey_account_id, 192213123);
 
-        // Setup Subnet pool for origin netuid
-        SubnetAlphaIn::<Test>::insert(origin_netuid, alpha_in + 10_000_000);
-        SubnetTAO::<Test>::insert(origin_netuid, tao_in + 10_000_000);
+        // Give it some $$$ in his coldkey balance
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey_account_id, amount);
 
-        // Add stake as new hotkey
+        // Stake to hotkey account, and check if the result is ok
         assert_ok!(SubtensorModule::add_stake(
             RuntimeOrigin::signed(coldkey_account_id),
             hotkey_account_id,
-            origin_netuid,
-            tao_staked,
-        ),);
-        let alpha_to_move = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+            netuid,
+            amount
+        ));
+
+        // Remove 99.9989% stake
+        let alpha = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
             &hotkey_account_id,
             &coldkey_account_id,
-            origin_netuid,
+            netuid,
         );
-
-        // Move stake to destination subnet
-        assert_ok!(SubtensorModule::move_stake(
+        assert_ok!(SubtensorModule::remove_stake(
             RuntimeOrigin::signed(coldkey_account_id),
             hotkey_account_id,
-            hotkey_account_id,
-            origin_netuid,
             netuid,
-            alpha_to_move,
+            (U64F64::from_num(alpha) * U64F64::from_num(0.99)).to_num::<u64>()
         ));
 
-        // Check that the stake has been moved
-        assert_eq!(
-            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
-                &hotkey_account_id,
-                &coldkey_account_id,
-                origin_netuid
-            ),
-            0
+        // Check that all alpha was unstaked and 99% TAO balance was returned (less fees)
+        assert_abs_diff_eq!(
+            SubtensorModule::get_coldkey_balance(&coldkey_account_id),
+            (amount as f64 * 0.99) as u64 - fee * 2,
+            epsilon = amount / 1000,
         );
-        let fee = DefaultStakingFee::<Test>::get();
-        let alpha_fee: I96F32 = I96F32::from_num(fee) / SubtensorModule::get_alpha_price(netuid);
-        let expected_value = I96F32::from_num(alpha_to_move)
-            * SubtensorModule::get_alpha_price(origin_netuid)
-            / SubtensorModule::get_alpha_price(netuid);
         assert_abs_diff_eq!(
-            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
-                &hotkey_account_id,
-                &coldkey_account_id,
-                netuid
-            ),
-            (expected_value - alpha_fee).to_num::<u64>(),
-            epsilon = (expected_value / 1000).to_num::<u64>()
+            SubtensorModule::get_total_stake_for_hotkey(&hotkey_account_id),
+            (amount as f64 * 0.01) as u64,
+            epsilon = amount / 1000,
+        );
+        let new_alpha = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey_account_id,
+            &coldkey_account_id,
+            netuid,
         );
+        assert_abs_diff_eq!(new_alpha, (alpha as f64 * 0.01) as u64, epsilon = 10);
     });
 }

primitives/share-pool/src/lib.rs

@@ -124,16 +124,14 @@ where
         let current_share: U64F64 = self.state_ops.get_share(key);
         let denominator: U64F64 = self.state_ops.get_denominator();
         let initial_value: i64 = self.get_value(key) as i64;
-
-        // First, update shared value
-        self.update_value_for_all(update);
-        let new_shared_value: U64F64 = self.state_ops.get_shared_value();
+        let mut actual_update: i64 = update;
 
         // Then, update this key's share
         if denominator == 0 {
             // Initialize the pool. The first key gets all.
-            self.state_ops.set_denominator(new_shared_value);
-            self.state_ops.set_share(key, new_shared_value);
+            let update_fixed: U64F64 = U64F64::saturating_from_num(update);
+            self.state_ops.set_denominator(update_fixed);
+            self.state_ops.set_share(key, update_fixed);
         } else {
             let shares_per_update: I64F64 =
                 self.get_shares_per_update(update, &shared_value, &denominator);
@@ -168,15 +166,19 @@ where
                     // yes, precision is low, just remove all
                     new_share = U64F64::saturating_from_num(0);
                     new_denominator = denominator.saturating_sub(current_share);
+                    actual_update = initial_value.neg();
                 }
 
                 self.state_ops.set_denominator(new_denominator);
                 self.state_ops.set_share(key, new_share);
             }
         }
 
-        let final_value: i64 = self.get_value(key) as i64;
-        final_value.saturating_sub(initial_value)
+        // Update shared value
+        self.update_value_for_all(actual_update);
+
+        // Return actual udate
+        actual_update
     }
 }
 
@@ -335,8 +337,8 @@ mod tests {
 
         // First to stake gets all accumulated emission if there are no other stakers
         // (which is artificial situation because there will be no emissions if there is no stake)
-        assert!((value1 - 1_001_000_000_000_000).abs() < 10);
-        assert!((value2 - 1_000_000_000_000).abs() < 10);
+        assert!((value1 - 1_001_000_000_000_000).abs() < 100);
+        assert!((value2 - 1_000_000_000_000).abs() < 100);
     }
 
     // cargo test --package share-pool --lib -- tests::test_denom_high_precision_many_small_unstakes --exact --show-output