[governance] implement locked coin and delegation with it

crates/sui-framework/sources/Coin.move

@@ -212,4 +212,12 @@ module Sui::Coin {
     public fun mint_for_testing<T>(value: u64, ctx: &mut TxContext): Coin<T> {
         Coin { id: TxContext::new_id(ctx), balance: Balance::create_with_value(value) }
     }
+
+    #[test_only]
+    /// Destroy a `Coin` with any value in it for testing purposes.
+    public fun destroy_for_testing<T>(self: Coin<T>): u64 {
+        let Coin { id, balance } = self;
+        ID::delete(id);
+        Balance::destroy_for_testing(balance)
+    }
 }

crates/sui-framework/sources/Governance/Delegation.move

@@ -6,6 +6,7 @@ module Sui::Delegation {
     use Sui::Balance::Balance;
     use Sui::Coin::{Self, Coin};
     use Sui::ID::{Self, VersionedID};
+    use Sui::LockedCoin::{Self, LockedCoin};
     use Sui::SUI::SUI;
     use Sui::Transfer;
     use Sui::TxContext::{Self, TxContext};
@@ -35,6 +36,11 @@ 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
+        /// comes from a Coin<SUI>, this field is None. If it comes from a LockedCoin<SUI>, this
+        /// field not None, and after undelegation the stake will be returned to a LockedCoin<SUI>
+        /// with locked_until_epoch set to this number.
+        coin_locked_until_epoch: Option<u64>,
         /// The delegation target validator.
         validator_address: address,
     }
@@ -52,6 +58,27 @@ module Sui::Delegation {
             ending_epoch: option::none(),
             delegate_amount,
             next_reward_unclaimed_epoch: starting_epoch,
+            coin_locked_until_epoch: Option::none(),
+            validator_address,
+        };
+        Transfer::transfer(delegation, TxContext::sender(ctx))
+    }
+
+    public(friend) fun create_from_locked_coin(
+        starting_epoch: u64,
+        validator_address: address,
+        stake: LockedCoin<SUI>,
+        ctx: &mut TxContext,
+    ) {
+        let delegate_amount = LockedCoin::value(&stake);
+        let coin_locked_until_epoch = Option::some(LockedCoin::locked_until_epoch(&stake));
+        let delegation = Delegation {
+            id: TxContext::new_id(ctx),
+            active_delegation: Option::some(LockedCoin::into_balance(stake)),
+            ending_epoch: Option::none(),
+            delegate_amount,
+            next_reward_unclaimed_epoch: starting_epoch,
+            coin_locked_until_epoch,
             validator_address,
         };
         Transfer::transfer(delegation, TxContext::sender(ctx))
@@ -68,7 +95,13 @@ module Sui::Delegation {
 
         let stake = option::extract(&mut self.active_delegation);
         let sender = TxContext::sender(ctx);
-        Transfer::transfer(Coin::from_balance(stake, ctx), sender);
+
+        if (Option::is_none(&self.coin_locked_until_epoch)) {
+            Transfer::transfer(Coin::from_balance(stake, ctx), sender);
+        } else {
+            let locked_until_epoch = *Option::borrow(&self.coin_locked_until_epoch);
+            Transfer::transfer(LockedCoin::from_balance(stake, locked_until_epoch, ctx), sender);
+        };
 
         self.ending_epoch = option::some(ending_epoch);
     }
@@ -96,6 +129,7 @@ module Sui::Delegation {
             ending_epoch,
             delegate_amount: _,
             next_reward_unclaimed_epoch,
+            coin_locked_until_epoch: _,
             validator_address: _,
         } = self;
         ID::delete(id);

crates/sui-framework/sources/Governance/SuiSystem.move

@@ -7,6 +7,7 @@ module Sui::SuiSystem {
     use Sui::Delegation::{Self, Delegation};
     use Sui::EpochRewardRecord::{Self, EpochRewardRecord};
     use Sui::ID::{Self, VersionedID};
+    use Sui::LockedCoin::{Self, LockedCoin};
     use Sui::SUI::SUI;
     use Sui::Transfer;
     use Sui::TxContext::{Self, TxContext};
@@ -169,6 +170,20 @@ module Sui::SuiSystem {
         Delegation::create(starting_epoch, validator_address, delegate_stake, ctx);
     }
 
+    public entry fun request_add_delegation_with_locked_coin(
+        self: &mut SuiSystemState,
+        delegate_stake: LockedCoin<SUI>,
+        validator_address: address,
+        ctx: &mut TxContext,
+    ) {
+        let amount = LockedCoin::value(&delegate_stake);
+        ValidatorSet::request_add_delegation(&mut self.validators, validator_address, amount);
+
+        // Delegation starts from the next epoch.
+        let starting_epoch = self.epoch + 1;
+        Delegation::create_from_locked_coin(starting_epoch, validator_address, delegate_stake, ctx);
+    }
+
     public entry fun request_remove_delegation(
         self: &mut SuiSystemState,
         delegation: &mut Delegation,

crates/sui-framework/sources/LockedCoin.move

@@ -0,0 +1,68 @@
+// Copyright (c) 2022, Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+module Sui::LockedCoin {
+    use Sui::Balance::{Self, Balance};
+    use Sui::Coin::{Self, Coin};
+    use Std::Errors;
+    use Sui::ID::{Self, VersionedID};
+    use Sui::Transfer;
+    use Sui::TxContext::{Self, TxContext};
+
+    /// The locked_until time passed into the creation of a locked coin is invalid.
+    const EINVALID_LOCK_UNTIL: u64 = 0;
+    /// Attempt is made to unlock a locked coin that has not been unlocked yet.
+    const ECOIN_STILL_LOCKED: u64 = 1;
+
+    /// A coin of type `T` locked until `locked_until_epoch`.
+    struct LockedCoin<phantom T> has key, store {
+        id: VersionedID,
+        balance: Balance<T>,
+        locked_until_epoch: u64
+    }
+
+    /// Returns the epoch until which the coin is locked.
+    public fun locked_until_epoch<T>(locked_coin: &LockedCoin<T>) : u64 {
+        locked_coin.locked_until_epoch
+    }
+
+    /// Wrap a balance into a LockedCoin.
+    public fun from_balance<T>(balance: Balance<T>, locked_until_epoch: u64, ctx: &mut TxContext): LockedCoin<T> {
+        LockedCoin { id: TxContext::new_id(ctx), balance, locked_until_epoch }
+    }
+
+    /// Destruct a LockedCoin wrapper and keep the balance.
+    public fun into_balance<T>(coin: LockedCoin<T>): Balance<T> {
+        let LockedCoin { id, locked_until_epoch: _, balance } = coin;
+        ID::delete(id);
+        balance
+    }
+
+    /// Public getter for the locked coin's value
+    public fun value<T>(self: &LockedCoin<T>): u64 {
+        Balance::value(&self.balance)
+    }
+
+    /// Lock a coin up until `locked_until_epoch`. The input Coin<T> is deleted and a LockedCoin<T>
+    /// is transferred to the `recipient`. This function aborts if the `locked_until_epoch` is less than
+    /// or equal to the current epoch.
+    public entry fun lock_coin<T>(
+        coin: Coin<T>, recipient: address, locked_until_epoch: u64, ctx: &mut TxContext
+    ) {
+        assert!(TxContext::epoch(ctx) < locked_until_epoch, Errors::invalid_argument(EINVALID_LOCK_UNTIL));
+        let balance = Coin::into_balance(coin);
+        let locked_coin = LockedCoin { id: TxContext::new_id(ctx), balance, locked_until_epoch };
+        Transfer::transfer(locked_coin, recipient);
+    }
+
+    /// Unlock a locked coin. The function aborts if the current epoch is less than the `locked_until_epoch`
+    /// of the coin. If the check is successful, the locked coin is deleted and a Coin<T> is transferred back
+    /// to the sender.
+    public entry fun unlock_coin<T>(locked_coin: LockedCoin<T>, ctx: &mut TxContext) {
+        let LockedCoin { id, balance, locked_until_epoch } = locked_coin;
+        assert!(TxContext::epoch(ctx) >= locked_until_epoch, Errors::invalid_argument(ECOIN_STILL_LOCKED));
+        ID::delete(id);
+        let coin = Coin::from_balance(balance, ctx);
+        Transfer::transfer(coin, TxContext::sender(ctx));
+    }
+}

crates/sui-framework/sources/TestScenario.move

@@ -110,7 +110,13 @@ module Sui::TestScenario {
         // create a seed for new transaction digest to ensure that this tx has a different
         // digest (and consequently, different object ID's) than the previous tx
         let new_tx_digest_seed = (vector::length(&scenario.event_start_indexes) as u8);
-        scenario.ctx = TxContext::new_from_address(*sender, new_tx_digest_seed);
+        let epoch = TxContext::epoch(&scenario.ctx);
+        scenario.ctx = TxContext::new_with_epoch(*sender, epoch, new_tx_digest_seed);
+    }
+
+    /// Advance the scenario to a new epoch.
+    public fun next_epoch(scenario: &mut Scenario) {
+        TxContext::advance_epoch(&mut scenario.ctx);
     }
 
     /// Remove the object of type `T` from the inventory of the current tx sender in `scenario`.

crates/sui-framework/sources/TxContext.move

@@ -73,35 +73,31 @@ module Sui::TxContext {
 
     #[test_only]
     /// Create a `TxContext` for testing
-    public fun new(signer: signer, tx_hash: vector<u8>, ids_created: u64): TxContext {
+    public fun new(signer: signer, tx_hash: vector<u8>, epoch: u64, ids_created: u64): TxContext {
         assert!(
             vector::length(&tx_hash) == TX_HASH_LENGTH,
             errors::invalid_argument(EBadTxHashLength)
         );
-        TxContext { signer, tx_hash, epoch: 0, ids_created }
+        TxContext { signer, tx_hash, epoch, ids_created }
     }
 
     #[test_only]
     /// Create a `TxContext` for testing, with a potentially non-zero epoch number.
-    public fun new_with_epoch(signer: signer, tx_hash: vector<u8>, epoch: u64, ids_created: u64): TxContext {
-        assert!(
-            vector::length(&tx_hash) == TX_HASH_LENGTH,
-            errors::invalid_argument(EBadTxHashLength)
-        );
-        TxContext { signer, tx_hash, epoch, ids_created }
+    public fun new_with_epoch(a: address, epoch: u64, hint: u8): TxContext {
+        new(new_signer_from_address(a), dummy_tx_hash_with_hint(hint), epoch, 0)
     }
 
     #[test_only]
     /// Create a `TxContext` with sender `a` for testing, and a tx hash derived from `hint`
     public fun new_from_address(a: address, hint: u8): TxContext {
-        new(new_signer_from_address(a), dummy_tx_hash_with_hint(hint), 0)
+        new(new_signer_from_address(a), dummy_tx_hash_with_hint(hint), 0, 0)
     }
 
     #[test_only]
     /// Create a dummy `TxContext` for testing
     public fun dummy(): TxContext {
         let tx_hash = x"3a985da74fe225b2045c172d6bd390bd855f086e3e9d525b46bfe24511431532";
-        new(new_signer_from_address(@0x0), tx_hash, 0)
+        new(new_signer_from_address(@0x0), tx_hash, 0, 0)
     }
 
     #[test_only]
@@ -130,6 +126,11 @@ module Sui::TxContext {
         ID::new(derive_id(*&self.tx_hash, ids_created - 1))
     }
 
+    #[test_only]
+    public fun advance_epoch(self: &mut TxContext) {
+        self.epoch = self.epoch + 1
+    }
+
     #[test_only]
     /// Test-only function for creating a new signer from `signer_address`.
     native fun new_signer_from_address(signer_address: address): signer;

crates/sui-framework/tests/CoinBalanceTests.move

@@ -7,6 +7,11 @@ module Sui::TestCoin {
     use Sui::Coin;
     use Sui::Balance;
     use Sui::SUI::SUI;
+    use Sui::LockedCoin::LockedCoin;
+    use Std::Errors;
+    use Sui::TxContext;
+    use Sui::LockedCoin;
+    use Sui::Coin::Coin;
 
     #[test]
     fun type_morphing() {
@@ -33,4 +38,54 @@ module Sui::TestCoin {
         let coin = Coin::from_balance(balance, ctx(test));
         Coin::keep(coin, ctx(test));
     }
+
+    const TEST_SENDER_ADDR: address = @0xA11CE;
+    const TEST_RECIPIENT_ADDR: address = @0xB0B;
+
+    #[test]
+    public(script) fun test_locked_coin_valid() {
+        let scenario = &mut TestScenario::begin(&TEST_SENDER_ADDR);
+        let ctx = TestScenario::ctx(scenario);
+        let coin = Coin::mint_for_testing<SUI>(42, ctx);
+
+        TestScenario::next_tx(scenario, &TEST_SENDER_ADDR);
+        // Lock up the coin until epoch 2.
+        LockedCoin::lock_coin(coin, TEST_RECIPIENT_ADDR, 2, TestScenario::ctx(scenario));
+
+        // Advance the epoch by 2.
+        TestScenario::next_epoch(scenario);
+        TestScenario::next_epoch(scenario);
+        assert!(TxContext::epoch(TestScenario::ctx(scenario)) == 2, Errors::invalid_state(1));
+
+        TestScenario::next_tx(scenario, &TEST_RECIPIENT_ADDR);
+        let locked_coin = TestScenario::take_owned<LockedCoin<SUI>>(scenario);
+        // The unlock should go through since epoch requirement is met.
+        LockedCoin::unlock_coin(locked_coin, TestScenario::ctx(scenario));
+
+        TestScenario::next_tx(scenario, &TEST_RECIPIENT_ADDR);
+        let unlocked_coin = TestScenario::take_owned<Coin<SUI>>(scenario);
+        assert!(Coin::value(&unlocked_coin) == 42, Errors::invalid_state(2));
+        Coin::destroy_for_testing(unlocked_coin);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = 263)]
+    public(script) fun test_locked_coin_invalid() {
+        let scenario = &mut TestScenario::begin(&TEST_SENDER_ADDR);
+        let ctx = TestScenario::ctx(scenario);
+        let coin = Coin::mint_for_testing<SUI>(42, ctx);
+
+        TestScenario::next_tx(scenario, &TEST_SENDER_ADDR);
+        // Lock up the coin until epoch 2.
+        LockedCoin::lock_coin(coin, TEST_RECIPIENT_ADDR, 2, TestScenario::ctx(scenario));
+
+        // Advance the epoch by 1.
+        TestScenario::next_epoch(scenario);
+        assert!(TxContext::epoch(TestScenario::ctx(scenario)) == 1, Errors::invalid_state(1));
+
+        TestScenario::next_tx(scenario, &TEST_RECIPIENT_ADDR);
+        let locked_coin = TestScenario::take_owned<LockedCoin<SUI>>(scenario);
+        // The unlock should fail.
+        LockedCoin::unlock_coin(locked_coin, TestScenario::ctx(scenario));
+    }
 }