[governance] implement custodial staking and staking with locked coins

crates/sui-framework/sources/epoch_time_lock.move

@@ -26,4 +26,9 @@ module sui::epoch_time_lock {
         let EpochTimeLock { epoch } = lock;
         assert!(tx_context::epoch(ctx) >= epoch, EEPOCH_STILL_LOCKED);
     }
+
+    /// Getter for the epoch number.
+    public fun epoch(lock: &EpochTimeLock): u64 {
+        lock.epoch
+    }
 }

crates/sui-framework/sources/governance/delegation.move

@@ -37,7 +37,7 @@ module sui::delegation {
         /// is the next epoch that the delegator can claim epoch. Whenever the delegator
         /// claims reward for an epoch, this value increments by one.
         next_reward_unclaimed_epoch: u64,
-        /// The epoch until which the delegated coin is locked until. If the delegated stake
+        /// The epoch until which the delegated coin is locked. If the delegated stake
         /// comes from a Coin<SUI>, this field is None. If it comes from a LockedCoin<SUI>, this
         /// field is not None, and after undelegation the stake will be returned to a LockedCoin<SUI>
         /// with locked_until_epoch set to this epoch.

crates/sui-framework/sources/governance/genesis.move

@@ -9,6 +9,7 @@ module sui::genesis {
     use sui::sui_system;
     use sui::tx_context::TxContext;
     use sui::validator;
+    use std::option;
 
     /// The initial amount of SUI locked in the storage fund.
     /// 10^14, an arbitrary number.
@@ -32,7 +33,7 @@ module sui::genesis {
         validator_names: vector<vector<u8>>,
         validator_net_addresses: vector<vector<u8>>,
         validator_stakes: vector<u64>,
-        _ctx: &mut TxContext,
+        ctx: &mut TxContext,
     ) {
         let sui_supply = sui::new();
         let storage_fund = balance::increase_supply(&mut sui_supply, INIT_STORAGE_FUND);
@@ -58,6 +59,8 @@ module sui::genesis {
                 name,
                 net_address,
                 balance::increase_supply(&mut sui_supply, stake),
+                option::none(),
+                ctx
             ));
             i = i + 1;
         };

crates/sui-framework/sources/governance/stake.move

@@ -0,0 +1,77 @@
+// Copyright (c) 2022, Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+module sui::stake {
+    use std::option::{Self, Option};
+    use sui::balance::Balance;
+    use sui::id::VersionedID;
+    use sui::locked_coin;
+    use sui::sui::SUI;
+    use sui::transfer;
+    use sui::tx_context::{Self, TxContext};
+    use sui::epoch_time_lock::EpochTimeLock;
+    use sui::epoch_time_lock;
+    use sui::balance;
+    use sui::math;
+
+    friend sui::sui_system;
+    friend sui::validator;
+
+    /// A custodial stake object holding the staked SUI coin.
+    struct Stake has key {
+        id: VersionedID,
+        /// The staked SUI tokens.
+        balance: Balance<SUI>,
+        /// The epoch until which the staked coin is locked. If the stake
+        /// comes from a Coin<SUI>, this field is None. If it comes from a LockedCoin<SUI>, this
+        /// field will record the original lock expiration epoch, to be used when unstaking.
+        coin_locked_until_epoch: Option<EpochTimeLock>,
+    }
+
+    /// The number of epochs the withdrawn stake is locked for.
+    /// TODO: this is a placehodler number and may be changed.
+    const BONDING_PERIOD: u64 = 1;
+
+    /// Create a stake object from a SUI balance. If the balance comes from a
+    /// `LockedCoin`, an EpochTimeLock is passed in to keep track of locking period.
+    public(friend) fun create(
+        balance: Balance<SUI>,
+        recipient: address,
+        coin_locked_until_epoch: Option<EpochTimeLock>,
+        ctx: &mut TxContext,
+    ) {
+        let stake = Stake {
+            id: tx_context::new_id(ctx),
+            balance,
+            coin_locked_until_epoch,
+        };
+        transfer::transfer(stake, recipient)
+    }
+
+    /// Withdraw `amount` from the balance of `stake`.
+    public(friend) fun withdraw_stake(
+        self: &mut Stake,
+        amount: u64,
+        ctx: &mut TxContext,
+    ) {
+        let sender = tx_context::sender(ctx);
+        let unlock_epoch = tx_context::epoch(ctx) + BONDING_PERIOD;
+        let balance = balance::split(&mut self.balance, amount);
+
+        if (option::is_none(&self.coin_locked_until_epoch)) {
+            // If the stake didn't come from a locked coin, we give back the stake and
+            // lock the coin for `BONDING_PERIOD`.
+            locked_coin::new_from_balance(balance, epoch_time_lock::new(unlock_epoch, ctx), sender, ctx);
+        } else {
+            // If the stake did come from a locked coin, we lock the coin for
+            // max(BONDING_PERIOD, remaining_lock_time).
+            let original_unlock_epoch = epoch_time_lock::epoch(option::borrow(&self.coin_locked_until_epoch));
+            let unlock_epoch = math::max(original_unlock_epoch, unlock_epoch);
+            locked_coin::new_from_balance(balance, epoch_time_lock::new(unlock_epoch, ctx), sender, ctx);
+        };
+    }
+
+    public fun value(self: &Stake): u64 {
+        balance::value(&self.balance)
+    }
+}

crates/sui-framework/sources/governance/sui_system.move

@@ -13,6 +13,8 @@ module sui::sui_system {
     use sui::tx_context::{Self, TxContext};
     use sui::validator::{Self, Validator};
     use sui::validator_set::{Self, ValidatorSet};
+    use sui::stake::Stake;
+    use std::option;
 
     friend sui::genesis;
 
@@ -108,7 +110,9 @@ module sui::sui_system {
             pubkey_bytes,
             name,
             net_address,
-            coin::into_balance(stake)
+            coin::into_balance(stake),
+            option::none(),
+            ctx
         );
 
         validator_set::request_add_validator(&mut self.validators, validator);
@@ -138,6 +142,22 @@ module sui::sui_system {
         validator_set::request_add_stake(
             &mut self.validators,
             coin::into_balance(new_stake),
+            option::none(),
+            ctx,
+        )
+    }
+
+    /// A validator can request adding more stake using a locked coin. This will be processed at the end of epoch.
+    public entry fun request_add_stake_with_locked_coin(
+        self: &mut SuiSystemState,
+        new_stake: LockedCoin<SUI>,
+        ctx: &mut TxContext,
+    ) {
+        let (balance, epoch_time_lock) = locked_coin::into_balance(new_stake);
+        validator_set::request_add_stake(
+            &mut self.validators,
+            balance,
+            option::some(epoch_time_lock),
             ctx,
         )
     }
@@ -149,11 +169,13 @@ module sui::sui_system {
     /// If the sender represents an active validator, the request will be processed at the end of epoch.
     public entry fun request_withdraw_stake(
         self: &mut SuiSystemState,
+        stake: &mut Stake,
         withdraw_amount: u64,
         ctx: &mut TxContext,
     ) {
         validator_set::request_withdraw_stake(
             &mut self.validators,
+            stake,
             withdraw_amount,
             self.parameters.min_validator_stake,
             ctx,

crates/sui-framework/sources/governance/validator.move

@@ -3,14 +3,15 @@
 
 module sui::validator {
     use std::ascii;
-    use std::option::{Self, Option};
     use std::vector;
 
     use sui::balance::{Self, Balance};
-    use sui::coin;
     use sui::sui::SUI;
-    use sui::transfer;
     use sui::tx_context::TxContext;
+    use sui::stake;
+    use sui::stake::Stake;
+    use sui::epoch_time_lock::EpochTimeLock;
+    use std::option::Option;
 
     friend sui::genesis;
     friend sui::sui_system;
@@ -41,13 +42,13 @@ module sui::validator {
     struct Validator has store {
         /// Summary of the validator.
         metadata: ValidatorMetadata,
-        /// The current active stake. This will not change during an epoch. It can only
+        /// The current active stake amount. This will not change during an epoch. It can only
         /// be updated at the end of epoch.
-        stake: Balance<SUI>,
+        stake_amount: u64,
         /// Amount of delegated stake from token holders.
         delegation: u64,
-        /// Pending stake deposits. It will be put into `stake` at the end of epoch.
-        pending_stake: Option<Balance<SUI>>,
+        /// Pending stake deposit amount, processed at end of epoch.
+        pending_stake: u64,
         /// Pending withdraw amount, processed at end of epoch.
         pending_withdraw: u64,
         /// Pending delegation deposits.
@@ -70,6 +71,8 @@ module sui::validator {
         name: vector<u8>,
         net_address: vector<u8>,
         stake: Balance<SUI>,
+        coin_locked_until_epoch: Option<EpochTimeLock>,
+        ctx: &mut TxContext
     ): Validator {
         assert!(
             // TODO: These constants are arbitrary, will adjust once we know more.
@@ -78,17 +81,19 @@ module sui::validator {
         );
         // Check that the name is human-readable.
         ascii::string(copy name);
+        let stake_amount = balance::value(&stake);
+        stake::create(stake, sui_address, coin_locked_until_epoch, ctx);
         Validator {
             metadata: ValidatorMetadata {
                 sui_address,
                 pubkey_bytes,
                 name,
                 net_address,
-                next_epoch_stake: balance::value(&stake),
+                next_epoch_stake: stake_amount,
             },
-            stake,
+            stake_amount,
             delegation: 0,
-            pending_stake: option::none(),
+            pending_stake: 0,
             pending_withdraw: 0,
             pending_delegation: 0,
             pending_delegation_withdraw: 0,
@@ -98,46 +103,33 @@ module sui::validator {
         }
     }
 
-    public(friend) fun destroy(self: Validator, ctx: &mut TxContext) {
+    public(friend) fun destroy(self: Validator) {
         let Validator {
-            metadata,
-            stake,
+            metadata: _,
+            stake_amount: _,
             delegation: _,
-            pending_stake,
-            pending_withdraw,
+            pending_stake: _,
+            pending_withdraw: _,
             pending_delegation: _,
             pending_delegation_withdraw: _,
             delegator_count: _,
             pending_delegator_count: _,
             pending_delegator_withdraw_count: _,
         } = self;
-
-        assert!(pending_withdraw == 0, 0);
-        if (option::is_some(&pending_stake)) {
-            // pending_stake can be non-empty as it can contain the gas reward from the last epoch.
-            let pending_stake_balance = option::extract(&mut pending_stake);
-            balance::join(&mut stake, pending_stake_balance);
-        };
-        option::destroy_none(pending_stake);
-        transfer::transfer(coin::from_balance(stake, ctx), metadata.sui_address);
     }
 
     /// Add stake to an active validator. The new stake is added to the pending_stake field,
     /// which will be processed at the end of epoch.
     public(friend) fun request_add_stake(
         self: &mut Validator,
         new_stake: Balance<SUI>,
+        coin_locked_until_epoch: Option<EpochTimeLock>,
+        ctx: &mut TxContext,
     ) {
         let new_stake_value = balance::value(&new_stake);
-        let pending_stake = if (option::is_some(&self.pending_stake)) {
-            let pending_stake = option::extract(&mut self.pending_stake);
-            balance::join(&mut pending_stake, new_stake);
-            pending_stake
-        } else {
-            new_stake
-        };
-        option::fill(&mut self.pending_stake, pending_stake);
+        self.pending_stake = self.pending_stake + new_stake_value;
         self.metadata.next_epoch_stake = self.metadata.next_epoch_stake + new_stake_value;
+        stake::create(new_stake, self.metadata.sui_address, coin_locked_until_epoch, ctx);
     }
 
     /// Withdraw stake from an active validator. Since it's active, we need
@@ -146,26 +138,23 @@ module sui::validator {
     /// stake still satisfy the minimum requirement.
     public(friend) fun request_withdraw_stake(
         self: &mut Validator,
+        stake: &mut Stake,
         withdraw_amount: u64,
         min_validator_stake: u64,
+        ctx: &mut TxContext,
     ) {
         assert!(self.metadata.next_epoch_stake >= withdraw_amount + min_validator_stake, 0);
         self.pending_withdraw = self.pending_withdraw + withdraw_amount;
         self.metadata.next_epoch_stake = self.metadata.next_epoch_stake - withdraw_amount;
+        stake::withdraw_stake(stake, withdraw_amount, ctx);
     }
 
     /// Process pending stake and pending withdraws.
-    public(friend) fun adjust_stake(self: &mut Validator, ctx: &mut TxContext) {
-        if (option::is_some(&self.pending_stake)) {
-            let pending_stake = option::extract(&mut self.pending_stake);
-            balance::join(&mut self.stake, pending_stake);
-        };
-        if (self.pending_withdraw > 0) {
-            let coin = coin::take(&mut self.stake, self.pending_withdraw, ctx);
-            coin::transfer(coin, self.metadata.sui_address);
-            self.pending_withdraw = 0;
-        };
-        assert!(balance::value(&self.stake) == self.metadata.next_epoch_stake, 0);
+    public(friend) fun adjust_stake(self: &mut Validator) {
+        self.stake_amount = self.stake_amount + self.pending_stake - self.pending_withdraw;
+        self.pending_stake = 0;
+        self.pending_withdraw = 0;
+        assert!(self.stake_amount == self.metadata.next_epoch_stake, 0);
 
         self.delegation = self.delegation + self.pending_delegation - self.pending_delegation_withdraw;
         self.pending_delegation = 0;
@@ -196,7 +185,7 @@ module sui::validator {
     }
 
     public fun stake_amount(self: &Validator): u64 {
-        balance::value(&self.stake)
+        self.stake_amount
     }
 
     public fun delegate_amount(self: &Validator): u64 {
@@ -208,11 +197,7 @@ module sui::validator {
     }
 
     public fun pending_stake_amount(self: &Validator): u64 {
-        if (option::is_some(&self.pending_stake)) {
-            balance::value(option::borrow(&self.pending_stake))
-        } else {
-            0
-        }
+        self.pending_stake
     }
 
     public fun pending_withdraw(self: &Validator): u64 {