[sui-system] Allow inactive StakedSui objects to be withdrawn immediately

crates/sui-framework/packages/sui-system/sources/staking_pool.move

@@ -130,6 +130,14 @@ module sui_system::staking_pool {
         staked_sui: StakedSui,
         ctx: &TxContext
     ) : Balance<SUI> {
+        // stake is inactive
+        if (staked_sui.stake_activation_epoch > ctx.epoch()) {
+            let principal = unwrap_staked_sui(staked_sui);
+            pool.pending_stake = pool.pending_stake - principal.value();
+
+            return principal
+        };
+
         let (pool_token_withdraw_amount, mut principal_withdraw) =
             withdraw_from_principal(pool, staked_sui);
         let principal_withdraw_amount = principal_withdraw.value();

crates/sui-framework/packages/sui-system/sources/sui_system_state_inner.move

@@ -4,7 +4,7 @@
 module sui_system::sui_system_state_inner {
     use sui::balance::{Self, Balance};
     use sui::coin::Coin;
-    use sui_system::staking_pool::{stake_activation_epoch, StakedSui};
+    use sui_system::staking_pool::StakedSui;
     use sui::sui::SUI;
     use sui_system::validator::{Self, Validator};
     use sui_system::validator_set::{Self, ValidatorSet};
@@ -214,7 +214,6 @@ module sui_system::sui_system_state_inner {
     const ECannotReportOneself: u64 = 3;
     const EReportRecordNotFound: u64 = 4;
     const EBpsTooLarge: u64 = 5;
-    const EStakeWithdrawBeforeActivation: u64 = 6;
     const ESafeModeGasNotProcessed: u64 = 7;
     const EAdvancedToWrongEpoch: u64 = 8;
 
@@ -516,10 +515,6 @@ module sui_system::sui_system_state_inner {
         staked_sui: StakedSui,
         ctx: &TxContext,
     ) : Balance<SUI> {
-        assert!(
-            stake_activation_epoch(&staked_sui) <= ctx.epoch(),
-            EStakeWithdrawBeforeActivation
-        );
         self.validators.request_withdraw_stake(staked_sui, ctx)
     }
 

crates/sui-framework/packages/sui-system/tests/sui_system_tests.move

@@ -9,7 +9,7 @@
 module sui_system::sui_system_tests {
     use sui::test_scenario::{Self, Scenario};
     use sui::sui::SUI;
-    use sui_system::governance_test_utils::{add_validator_full_flow, advance_epoch, remove_validator, set_up_sui_system_state, create_sui_system_state_for_testing};
+    use sui_system::governance_test_utils::{add_validator_full_flow, advance_epoch, remove_validator, set_up_sui_system_state, create_sui_system_state_for_testing, stake_with, unstake};
     use sui_system::sui_system::SuiSystemState;
     use sui_system::sui_system_state_inner;
     use sui_system::validator::{Self, Validator};
@@ -1002,4 +1002,54 @@ module sui_system::sui_system_tests {
         test_scenario::return_shared(system_state);
         scenario.next_epoch(@0x0);
     }
+
+    #[test]
+    fun test_withdraw_inactive_stake() {
+        let mut scenario_val = test_scenario::begin(@0x0);
+        let scenario = &mut scenario_val;
+        // Epoch duration is set to be 42 here.
+        set_up_sui_system_state(vector[@0x1, @0x2]);
+
+        {
+            scenario.next_tx(@0x0);
+            let mut system_state = scenario.take_shared<SuiSystemState>();
+            let staking_pool = system_state.active_validator_by_address(@0x1).get_staking_pool_ref();
+
+            assert!(staking_pool.pending_stake_amount() == 0, 0);
+            assert!(staking_pool.pending_stake_withdraw_amount() == 0, 0);
+            assert!(staking_pool.sui_balance() == 100 * 1_000_000_000, 0);
+
+            test_scenario::return_shared(system_state);
+        };
+
+        stake_with(@0x0, @0x1, 1, scenario);
+
+        {
+            scenario.next_tx(@0x0);
+            let mut system_state = scenario.take_shared<SuiSystemState>();
+            let staking_pool = system_state.active_validator_by_address(@0x1).get_staking_pool_ref();
+
+            assert!(staking_pool.pending_stake_amount() == 1_000_000_000, 0);
+            assert!(staking_pool.pending_stake_withdraw_amount() == 0, 0);
+            assert!(staking_pool.sui_balance() == 100 * 1_000_000_000, 0);
+
+            test_scenario::return_shared(system_state);
+        };
+
+        unstake(@0x0, 0, scenario);
+
+        {
+            scenario.next_tx(@0x0);
+            let mut system_state = scenario.take_shared<SuiSystemState>();
+            let staking_pool = system_state.active_validator_by_address(@0x1).get_staking_pool_ref();
+
+            assert!(staking_pool.pending_stake_amount() == 0, 0);
+            assert!(staking_pool.pending_stake_withdraw_amount() == 0, 0);
+            assert!(staking_pool.sui_balance() == 100 * 1_000_000_000, 0);
+
+            test_scenario::return_shared(system_state);
+        };
+
+        scenario_val.end();
+    }
 }