[governance] add delegation switching and some tests

crates/sui-framework/sources/balance.move

@@ -115,6 +115,12 @@ module sui::balance {
         let Balance { value } = self;
         value
     }
+
+    #[test_only]
+    /// Create a `Supply` of any coin for testing purposes.
+    public fun create_supply_for_testing<T>(value: u64): Supply<T> {
+        Supply { value }
+    }
 }
 
 #[test_only]

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

@@ -3,7 +3,7 @@
 
 module sui::delegation {
     use std::option::{Self, Option};
-    use sui::balance::Balance;
+    use sui::balance::{Self, Balance};
     use sui::coin::{Self, Coin};
     use sui::id::{Self, VersionedID};
     use sui::locked_coin::{Self, LockedCoin};
@@ -107,6 +107,40 @@ module sui::delegation {
         self.ending_epoch = option::some(ending_epoch);
     }
 
+    /// Switch the delegation from the current validator to `new_validator_address`.
+    /// The current `Delegation` object `self` becomes inactive and the balance inside is
+    /// extracted to the new `Delegation` object. 
+    public(friend) fun switch_delegation(
+        self: &mut Delegation,
+        new_validator_address: address,
+        ctx: &mut TxContext,
+    ) {
+        assert!(is_active(self), 0);
+        let current_epoch = tx_context::epoch(ctx);
+        let balance = option::extract(&mut self.active_delegation);
+        let delegate_amount = balance::value(&balance);
+
+        let new_epoch_lock = 
+            if (option::is_some(&self.coin_locked_until_epoch)) {
+                option::some(option::extract(&mut self.coin_locked_until_epoch))
+            } else { 
+                option::none() 
+            };
+
+        self.ending_epoch = option::some(current_epoch);
+
+        let new_delegation = Delegation {
+            id: tx_context::new_id(ctx),
+            active_delegation: option::some(balance),
+            ending_epoch: option::none(),
+            delegate_amount,
+            next_reward_unclaimed_epoch: current_epoch + 1,
+            coin_locked_until_epoch: new_epoch_lock,
+            validator_address: new_validator_address,
+        };
+        transfer::transfer(new_delegation, tx_context::sender(ctx))
+    }
+
     /// Claim delegation reward. Increment next_reward_unclaimed_epoch.
     public(friend) fun claim_reward(
         self: &mut Delegation,
@@ -169,7 +203,7 @@ module sui::delegation {
         self.delegate_amount
     }
 
-    fun is_active(self: &Delegation): bool {
+    public fun is_active(self: &Delegation): bool {
         option::is_some(&self.active_delegation) && option::is_none(&self.ending_epoch)
     }
 }

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

@@ -18,6 +18,9 @@ module sui::sui_system {
 
     friend sui::genesis;
 
+    #[test_only]
+    friend sui::governance_test_utils;
+
     /// A list of system config parameters.
     // TDOO: We will likely add more, a few potential ones:
     // - the change in stake across epochs can be at most +/- x%
@@ -223,6 +226,20 @@ module sui::sui_system {
         delegation::undelegate(delegation, self.epoch, ctx)
     }
 
+    // Switch delegation from the current validator to a new one. 
+    public entry fun request_switch_delegation(
+        self: &mut SuiSystemState,
+        delegation: &mut Delegation,
+        new_validator_address: address,
+        ctx: &mut TxContext,
+    ) {
+        let old_validator_address = delegation::validator(delegation);
+        let amount = delegation::delegate_amount(delegation);
+        validator_set::request_remove_delegation(&mut self.validators, old_validator_address, amount);
+        validator_set::request_add_delegation(&mut self.validators, new_validator_address, amount);
+        delegation::switch_delegation(delegation, new_validator_address, ctx);
+    }
+
     // TODO: Once we support passing vector of object references as arguments,
     // we should support passing a vector of &mut EpochRewardRecord,
     // which will allow delegators to claim all their reward in one transaction.
@@ -262,8 +279,8 @@ module sui::sui_system {
         let storage_reward = balance::increase_supply(&mut self.sui_supply, storage_charge);
         let computation_reward = balance::increase_supply(&mut self.sui_supply, computation_charge);
 
-        let delegation_stake = validator_set::delegation_stake(&self.validators);
-        let validator_stake = validator_set::validator_stake(&self.validators);
+        let delegation_stake = validator_set::total_delegation_stake(&self.validators);
+        let validator_stake = validator_set::total_validator_stake(&self.validators);
         let storage_fund = balance::value(&self.storage_fund);
         let total_stake = delegation_stake + validator_stake + storage_fund;
 
@@ -299,6 +316,18 @@ module sui::sui_system {
         self.epoch
     }
 
+    /// Returns the amount of stake delegated to `validator_addr`. 
+    /// Aborts if `validator_addr` is not an active validator.
+    public fun validator_delegate_amount(self: &SuiSystemState, validator_addr: address): u64 {
+        validator_set::validator_delegate_amount(&self.validators, validator_addr)
+    } 
+
+    /// Returns the amount of delegators who have delegated to `validator_addr`. 
+    /// Aborts if `validator_addr` is not an active validator.
+    public fun validator_delegator_count(self: &SuiSystemState, validator_addr: address): u64 {
+        validator_set::validator_delegator_count(&self.validators, validator_addr)
+    } 
+
     #[test_only]
     public fun set_epoch_for_testing(self: &mut SuiSystemState, epoch_num: u64) {
         self.epoch = epoch_num

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

@@ -21,6 +21,8 @@ module sui::validator {
     friend sui::validator_tests;
     #[test_only]
     friend sui::validator_set_tests;
+    #[test_only]
+    friend sui::governance_test_utils;
 
     struct ValidatorMetadata has store, drop, copy {
         /// The Sui Address of the validator. This is the sender that created the Validator object,

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

@@ -21,10 +21,10 @@ module sui::validator_set {
     struct ValidatorSet has store {
         /// Total amount of stake from all active validators (not including delegation),
         /// at the beginning of the epoch.
-        validator_stake: u64,
+        total_validator_stake: u64,
 
         /// Total amount of stake from delegation, at the beginning of the epoch.
-        delegation_stake: u64,
+        total_delegation_stake: u64,
 
         /// The amount of accumulated stake to reach a quorum among all active validators.
         /// This is always 2/3 of total stake. Keep it here to reduce potential inconsistencies
@@ -48,10 +48,10 @@ module sui::validator_set {
     }
 
     public(friend) fun new(init_active_validators: vector<Validator>): ValidatorSet {
-        let (validator_stake, delegation_stake, quorum_stake_threshold) = calculate_total_stake_and_quorum_threshold(&init_active_validators);
+        let (total_validator_stake, total_delegation_stake, quorum_stake_threshold) = calculate_total_stake_and_quorum_threshold(&init_active_validators);
         let validators = ValidatorSet {
-            validator_stake,
-            delegation_stake,
+            total_validator_stake,
+            total_delegation_stake,
             quorum_stake_threshold,
             active_validators: init_active_validators,
             pending_validators: vector::empty(),
@@ -208,7 +208,7 @@ module sui::validator_set {
         // epoch's stake information to compute reward distribution.
         let rewards = compute_reward_distribution(
             &self.active_validators,
-            self.validator_stake,
+            self.total_validator_stake,
             balance::value(computation_reward),
         );
 
@@ -225,17 +225,32 @@ module sui::validator_set {
         self.next_epoch_validators = derive_next_epoch_validators(self);
 
         let (validator_stake, delegation_stake, quorum_stake_threshold) = calculate_total_stake_and_quorum_threshold(&self.active_validators);
-        self.validator_stake = validator_stake;
-        self.delegation_stake = delegation_stake;
+        self.total_validator_stake = validator_stake;
+        self.total_delegation_stake = delegation_stake;
         self.quorum_stake_threshold = quorum_stake_threshold;
     }
 
-    public fun validator_stake(self: &ValidatorSet): u64 {
-        self.validator_stake
+    public fun total_validator_stake(self: &ValidatorSet): u64 {
+        self.total_validator_stake
     }
 
-    public fun delegation_stake(self: &ValidatorSet): u64 {
-        self.delegation_stake
+    public fun total_delegation_stake(self: &ValidatorSet): u64 {
+        self.total_delegation_stake
+    }
+
+    public fun validator_stake_amount(self: &ValidatorSet, validator_address: address): u64 {
+        let validator = get_validator_ref(&self.active_validators, validator_address);
+        validator::stake_amount(validator)
+    }
+
+    public fun validator_delegate_amount(self: &ValidatorSet, validator_address: address): u64 {
+        let validator = get_validator_ref(&self.active_validators, validator_address);
+        validator::delegate_amount(validator)
+    }
+
+    public fun validator_delegator_count(self: &ValidatorSet, validator_address: address): u64 {
+        let validator = get_validator_ref(&self.active_validators, validator_address);
+        validator::delegator_count(validator)
     }
 
     /// Checks whether a duplicate of `new_validator` is already in `validators`.
@@ -279,6 +294,16 @@ module sui::validator_set {
         vector::borrow_mut(validators, validator_index)
     }
 
+    fun get_validator_ref(
+        validators: &vector<Validator>,
+        validator_address: address,
+    ): &Validator {
+        let validator_index_opt = find_validator(validators, validator_address);
+        assert!(option::is_some(&validator_index_opt), 0);
+        let validator_index = option::extract(&mut validator_index_opt);
+        vector::borrow(validators, validator_index)
+    }
+
     /// Process the pending withdraw requests. For each pending request, the validator
     /// is removed from `validators` and sent back to the address of the validator.
     fun process_pending_removals(
@@ -431,8 +456,8 @@ module sui::validator_set {
         self: ValidatorSet,
     ) {
         let ValidatorSet {
-            validator_stake: _,
-            delegation_stake: _,
+            total_validator_stake: _,
+            total_delegation_stake: _,
             quorum_stake_threshold: _,
             active_validators,
             pending_validators,

crates/sui-framework/tests/delegation_tests.move

@@ -0,0 +1,138 @@
+// Copyright (c) 2022, Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+#[test_only]
+module sui::delegation_tests {
+    use sui::coin;
+    use sui::test_scenario::{Self, Scenario};
+    use sui::sui_system::{Self, SuiSystemState};
+    use sui::delegation::{Self, Delegation};
+    use sui::governance_test_utils::{
+        Self, 
+        create_validator_for_testing, 
+        create_sui_system_state_for_testing
+    };
+
+    const VALIDATOR_ADDR_1: address = @0x1;
+    const VALIDATOR_ADDR_2: address = @0x2;
+
+    const DELEGATOR_ADDR_1: address = @0x42;
+    const DELEGATOR_ADDR_2: address = @0x43;
+
+    #[test]
+    fun test_add_remove_delegation_flow() {
+        let scenario = &mut test_scenario::begin(&VALIDATOR_ADDR_1);
+        set_up_sui_system_state(scenario);
+
+        test_scenario::next_tx(scenario, &DELEGATOR_ADDR_1);
+        {
+            let system_state_wrapper = test_scenario::take_shared<SuiSystemState>(scenario);
+            let system_state_mut_ref = test_scenario::borrow_mut(&mut system_state_wrapper);
+
+            let ctx = test_scenario::ctx(scenario);
+
+            // Create two delegations to VALIDATOR_ADDR_1.
+            sui_system::request_add_delegation(
+                system_state_mut_ref, coin::mint_for_testing(10, ctx), VALIDATOR_ADDR_1, ctx);
+            sui_system::request_add_delegation(
+                system_state_mut_ref, coin::mint_for_testing(60, ctx), VALIDATOR_ADDR_1, ctx);
+
+            // Advance the epoch so that the delegation changes can take into effect.
+            governance_test_utils::advance_epoch(system_state_mut_ref, scenario);
+
+            // Check that the delegation amount and count are changed correctly.
+            assert!(sui_system::validator_delegate_amount(system_state_mut_ref, VALIDATOR_ADDR_1) == 70, 101);
+            assert!(sui_system::validator_delegate_amount(system_state_mut_ref, VALIDATOR_ADDR_2) == 0, 102);
+            assert!(sui_system::validator_delegator_count(system_state_mut_ref, VALIDATOR_ADDR_1) == 2, 103);
+            assert!(sui_system::validator_delegator_count(system_state_mut_ref, VALIDATOR_ADDR_2) == 0, 104);
+            test_scenario::return_shared(scenario, system_state_wrapper);
+        };
+
+
+        test_scenario::next_tx(scenario, &DELEGATOR_ADDR_1);
+        {
+
+            let delegation = test_scenario::take_last_created_owned<Delegation>(scenario);
+            assert!(delegation::delegate_amount(&delegation) == 60, 105);
+
+
+            let system_state_wrapper = test_scenario::take_shared<SuiSystemState>(scenario);
+            let system_state_mut_ref = test_scenario::borrow_mut(&mut system_state_wrapper);
+
+            let ctx = test_scenario::ctx(scenario);
+
+            // Undelegate 60 SUIs from VALIDATOR_ADDR_1
+            sui_system::request_remove_delegation(
+                system_state_mut_ref, &mut delegation, ctx);
+
+            // Check that the delegation object indeed becomes inactive.
+            assert!(!delegation::is_active(&delegation), 106);
+            test_scenario::return_owned(scenario, delegation);
+
+            governance_test_utils::advance_epoch(system_state_mut_ref, scenario);
+
+            assert!(sui_system::validator_delegate_amount(system_state_mut_ref, VALIDATOR_ADDR_1) == 10, 107);
+            assert!(sui_system::validator_delegator_count(system_state_mut_ref, VALIDATOR_ADDR_1) == 1, 108);
+            test_scenario::return_shared(scenario, system_state_wrapper);
+        };
+    }
+
+    #[test]
+    fun test_switch_delegation_flow() {
+        let scenario = &mut test_scenario::begin(&VALIDATOR_ADDR_1);
+        set_up_sui_system_state(scenario);
+
+        test_scenario::next_tx(scenario, &DELEGATOR_ADDR_1);
+        {
+            let system_state_wrapper = test_scenario::take_shared<SuiSystemState>(scenario);
+            let system_state_mut_ref = test_scenario::borrow_mut(&mut system_state_wrapper);
+
+            let ctx = test_scenario::ctx(scenario);
+
+            // Create one delegation to VALIDATOR_ADDR_1, and one to VALIDATOR_ADDR_2.
+            sui_system::request_add_delegation(
+                system_state_mut_ref, coin::mint_for_testing(60, ctx), VALIDATOR_ADDR_1, ctx);
+            sui_system::request_add_delegation(
+                system_state_mut_ref, coin::mint_for_testing(10, ctx), VALIDATOR_ADDR_2, ctx);
+
+            governance_test_utils::advance_epoch(system_state_mut_ref, scenario);
+            test_scenario::return_shared(scenario, system_state_wrapper);
+        };
+
+        test_scenario::next_tx(scenario, &DELEGATOR_ADDR_1);
+        {
+
+            let delegation = test_scenario::take_last_created_owned<Delegation>(scenario);
+
+            let system_state_wrapper = test_scenario::take_shared<SuiSystemState>(scenario);
+            let system_state_mut_ref = test_scenario::borrow_mut(&mut system_state_wrapper);
+
+            let ctx = test_scenario::ctx(scenario);
+
+            // Switch the 10 SUI delegation from VALIDATOR_ADDR_2 to VALIDATOR_ADDR_1
+            sui_system::request_switch_delegation(
+                system_state_mut_ref, &mut delegation, VALIDATOR_ADDR_1, ctx);
+
+            test_scenario::return_owned(scenario, delegation);
+
+            governance_test_utils::advance_epoch(system_state_mut_ref, scenario);
+
+            // Check that now VALIDATOR_ADDR_1 has all the delegations
+            assert!(sui_system::validator_delegate_amount(system_state_mut_ref, VALIDATOR_ADDR_1) == 70, 101);
+            assert!(sui_system::validator_delegate_amount(system_state_mut_ref, VALIDATOR_ADDR_2) == 0, 102);
+            assert!(sui_system::validator_delegator_count(system_state_mut_ref, VALIDATOR_ADDR_1) == 2, 103);
+            assert!(sui_system::validator_delegator_count(system_state_mut_ref, VALIDATOR_ADDR_2) == 0, 104);
+            test_scenario::return_shared(scenario, system_state_wrapper);
+        };
+    }
+
+    fun set_up_sui_system_state(scenario: &mut Scenario) {
+        let ctx = test_scenario::ctx(scenario);
+
+        let validators = vector[
+            create_validator_for_testing(VALIDATOR_ADDR_1, 100, ctx), 
+            create_validator_for_testing(VALIDATOR_ADDR_2, 100, ctx)
+        ];
+        create_sui_system_state_for_testing(validators, 300, 100);
+    }
+}