Merge pull request #1 from solendprotocol/hook

Weight-based delegation hook, change decrease_validator_stake to accept address, not validator index

Author: 0xripleys

contracts/sources/hooks/weight.move

@@ -0,0 +1,144 @@
+/// This module allows for the dynamic rebalancing of validator stakes based 
+/// on a set of validator addresses and weights. Rebalance can be called permissionlessly
+module liquid_staking::weight {
+    use sui_system::sui_system::{SuiSystemState};
+    use liquid_staking::liquid_staking::{LiquidStakingInfo, AdminCap};
+    use sui::vec_map::{Self, VecMap};
+    use sui::bag::{Self, Bag};
+    use liquid_staking::version::{Self, Version};
+
+    /* Constants */
+    const CURRENT_VERSION: u16 = 1;
+
+    public struct WeightHook<phantom P> has key, store {
+        id: UID,
+        validator_addresses_and_weights: VecMap<address, u64>,
+        total_weight: u64,
+        admin_cap: AdminCap<P>,
+        version: Version,
+        extra_fields: Bag
+    }
+
+    public struct WeightHookAdminCap<phantom P> has key, store {
+        id: UID
+    }
+
+    public fun new<P>(
+        admin_cap: AdminCap<P>,
+        ctx: &mut TxContext
+    ): (WeightHook<P>, WeightHookAdminCap<P>) {
+        (
+            WeightHook {
+                id: object::new(ctx),
+                validator_addresses_and_weights: vec_map::empty(),
+                total_weight: 0,
+                admin_cap,
+                version: version::new(CURRENT_VERSION),
+                extra_fields: bag::new(ctx)
+            },
+            WeightHookAdminCap { id: object::new(ctx) }
+        )
+    }
+
+    public fun set_validator_addresses_and_weights<P>(
+        self: &mut WeightHook<P>,
+        _: &WeightHookAdminCap<P>,
+        validator_addresses_and_weights: VecMap<address, u64>,
+    ) {
+        self.version.assert_version_and_upgrade(CURRENT_VERSION);
+        self.validator_addresses_and_weights = validator_addresses_and_weights;
+
+        let mut total_weight = 0;
+        self.validator_addresses_and_weights.keys().length().do!(|i| {
+            let (_, weight) = self.validator_addresses_and_weights.get_entry_by_idx(i);
+            total_weight = total_weight + *weight;
+        });
+
+        self.total_weight = total_weight;
+    }
+
+    public fun rebalance<P>(
+        self: &mut WeightHook<P>,
+        system_state: &mut SuiSystemState,
+        liquid_staking_info: &mut LiquidStakingInfo<P>,
+        ctx: &mut TxContext
+    ) {
+        self.version.assert_version_and_upgrade(CURRENT_VERSION);
+        if (self.total_weight == 0) {
+            return
+        };
+
+        liquid_staking_info.refresh(system_state, ctx);
+
+        let mut validator_addresses_and_weights = self.validator_addresses_and_weights;
+
+        // 1. add all validators that exist in lst_info.validators() to the validator_address_to_weight map if they don't already exist
+        liquid_staking_info.storage().validators().do_ref!(|validator| {
+            let validator_address = validator.validator_address();
+            if (!validator_addresses_and_weights.contains(&validator_address)) {
+                validator_addresses_and_weights.insert(validator_address, 0);
+            };
+        });
+
+        // 2. calculate current and target amounts of sui for each validator
+        let (validator_addresses, validator_weights) = validator_addresses_and_weights.into_keys_values();
+
+        let total_sui_supply = liquid_staking_info.storage().total_sui_supply(); // we want to allocate the unaccrued spread fees as well
+
+        let validator_target_amounts  = validator_weights.map!(|weight| {
+            ((total_sui_supply as u128) * (weight as u128) / (self.total_weight as u128)) as u64
+        });
+
+        let validator_current_amounts = validator_addresses.map_ref!(|validator_address| {
+            let validator_index = liquid_staking_info.storage().find_validator_index_by_address(*validator_address);
+            if (validator_index >= liquid_staking_info.storage().validators().length()) {
+                return 0
+            };
+
+            let validator = liquid_staking_info.storage().validators().borrow(validator_index);
+            validator.total_sui_amount()
+        });
+
+        // 3. decrease the stake for validators that have more stake than the target amount
+        validator_addresses.length().do!(|i| {
+            if (validator_current_amounts[i] > validator_target_amounts[i]) {
+                liquid_staking_info.decrease_validator_stake(
+                    &self.admin_cap,
+                    system_state,
+                    validator_addresses[i],
+                    validator_current_amounts[i] - validator_target_amounts[i],
+                    ctx
+                );
+            };
+        });
+
+        // 4. increase the stake for validators that have less stake than the target amount
+        validator_addresses.length().do!(|i| {
+            if (validator_current_amounts[i] < validator_target_amounts[i]) {
+                liquid_staking_info.increase_validator_stake(
+                    &self.admin_cap,
+                    system_state,
+                    validator_addresses[i],
+                    validator_target_amounts[i] - validator_current_amounts[i],
+                    ctx
+                );
+            };
+        });
+    }
+
+    public fun eject<P>(
+        mut self: WeightHook<P>,
+        admin_cap: WeightHookAdminCap<P>,
+    ): AdminCap<P> {
+        self.version.assert_version_and_upgrade(CURRENT_VERSION);
+
+        let WeightHookAdminCap { id } = admin_cap;
+        object::delete(id);
+
+        let WeightHook { id, admin_cap, extra_fields, .. } = self;
+        extra_fields.destroy_empty();
+        object::delete(id);
+
+        admin_cap
+    }
+}

contracts/sources/liquid_staking.move

@@ -19,6 +19,7 @@ module liquid_staking::liquid_staking {
     const EInvalidLstCreation: u64 = 0;
     const EMintInvariantViolated: u64 = 1;
     const ERedeemInvariantViolated: u64 = 2;
+    const EValidatorNotFound: u64 = 3;
 
     /* Constants */
     const CURRENT_VERSION: u16 = 1;
@@ -322,12 +323,15 @@ module liquid_staking::liquid_staking {
         self: &mut LiquidStakingInfo<P>,
         _: &AdminCap<P>,
         system_state: &mut SuiSystemState,
-        validator_index: u64,
+        validator_address: address,
         max_sui_amount: u64,
         ctx: &mut TxContext
     ): u64 {
         self.refresh(system_state, ctx);
 
+        let validator_index = self.storage.find_validator_index_by_address(validator_address);
+        assert!(validator_index < self.storage.validators().length(), EValidatorNotFound);
+
         let sui_amount = self.storage.unstake_approx_n_sui_from_validator(
             system_state,
             validator_index,

contracts/sources/storage.move

@@ -95,6 +95,19 @@ module liquid_staking::storage {
         self.total_sui_amount
     }
 
+    public(package) fun find_validator_index_by_address(self: &Storage, validator_address: address): u64 {
+        let mut i = 0;
+        while (i < self.validator_infos.length()) {
+            if (self.validator_infos[i].validator_address == validator_address) {
+                return i
+            };
+
+            i = i + 1;
+        };
+
+        i
+    }
+
     fun is_empty(self: &ValidatorInfo): bool {
         self.active_stake.is_none() && self.inactive_stake.is_none() && self.total_sui_amount == 0
     }

contracts/sources/version.move

@@ -6,7 +6,7 @@ module liquid_staking::version {
     const EIncorrectVersion: u64 = 0;
 
     /// Capability object given to the pool creator
-    public struct Version has store(u16)
+    public struct Version has store, drop (u16)
 
     public(package) fun new(
         version: u16,

contracts/tests/liquid_staking_tests.move

@@ -364,7 +364,7 @@ module liquid_staking::liquid_staking_tests {
         lst_info.decrease_validator_stake(
             &admin_cap, 
             &mut system_state, 
-            1,
+            @0x1,
             40 * MIST_PER_SUI, 
             scenario.ctx()
         );
@@ -692,7 +692,7 @@ module liquid_staking::liquid_staking_tests {
         let unstaked_amount = lst_info.decrease_validator_stake(
             &admin_cap,
             &mut system_state,
-            0,
+            @0x0,
             unstake_amount,
             scenario.ctx()
         );

contracts/tests/weight_tests.move

@@ -0,0 +1,121 @@
+#[test_only]
+module liquid_staking::weight_tests {
+    /* Tests */
+    use sui::vec_map::{Self};
+    use sui::test_scenario::{Self, Scenario};
+    use sui::coin::{Self};
+    use sui::address;
+    use sui_system::governance_test_utils::{
+        create_validators_with_stakes,
+        create_sui_system_state_for_testing,
+        advance_epoch_with_reward_amounts,
+    };
+    use sui_system::sui_system::{SuiSystemState};
+    use sui_system::staking_pool::StakedSui;
+    const MIST_PER_SUI: u64 = 1_000_000_000;
+    use liquid_staking::fees::{Self};
+    use liquid_staking::liquid_staking::{create_lst};
+    use liquid_staking::weight::{Self, WeightHook};
+
+    public struct TEST has drop {}
+
+    #[test_only]
+    public fun stake_with(validator_index: u64, amount: u64, scenario: &mut Scenario): StakedSui {
+        scenario.next_tx(@0x0);
+
+        let mut system_state = scenario.take_shared<SuiSystemState>();
+
+        let ctx = scenario.ctx();
+
+        let staked_sui = system_state.request_add_stake_non_entry(
+            coin::mint_for_testing(amount * MIST_PER_SUI, ctx), 
+            address::from_u256(validator_index as u256), 
+            ctx
+        );
+
+        test_scenario::return_shared(system_state);
+        scenario.next_tx(@0x0);
+
+        staked_sui
+    }
+
+    #[test_only]
+    fun setup_sui_system(scenario: &mut Scenario, stakes: vector<u64>) {
+        let validators = create_validators_with_stakes(stakes, scenario.ctx());
+        create_sui_system_state_for_testing(validators, 0, 0, scenario.ctx());
+
+        advance_epoch_with_reward_amounts(0, 0, scenario);
+    }
+
+
+     #[test]
+     fun test_rebalance() {
+        let mut scenario = test_scenario::begin(@0x0);
+
+        setup_sui_system(&mut scenario, vector[100, 100, 100]);
+        scenario.next_tx(@0x0);
+
+        let (admin_cap, mut lst_info) = create_lst<TEST>(
+            fees::new_builder(scenario.ctx()).to_fee_config(),
+            coin::create_treasury_cap_for_testing(scenario.ctx()),
+            scenario.ctx()
+        );
+
+        let mut system_state = scenario.take_shared<SuiSystemState>();
+        let sui = coin::mint_for_testing(100 * MIST_PER_SUI, scenario.ctx());
+        let lst = lst_info.mint(&mut system_state, sui, scenario.ctx());
+
+        assert!(lst_info.total_lst_supply() == 100 * MIST_PER_SUI, 0);
+        assert!(lst_info.storage().total_sui_supply() == 100 * MIST_PER_SUI, 0);
+
+        let (mut weight_hook, weight_hook_admin_cap) = weight::new(admin_cap, scenario.ctx());
+
+        weight_hook.set_validator_addresses_and_weights(
+            &weight_hook_admin_cap, 
+            {
+                let mut map = vec_map::empty();
+                map.insert(address::from_u256(0), 100);
+                map.insert(address::from_u256(1), 300);
+
+                map
+            }
+        );
+
+        weight_hook.rebalance(&mut system_state, &mut lst_info, scenario.ctx());
+
+        std::debug::print(&lst_info);
+
+        assert!(lst_info.storage().validators().borrow(0).total_sui_amount() == 25 * MIST_PER_SUI, 0);
+        assert!(lst_info.storage().validators().borrow(1).total_sui_amount() == 75 * MIST_PER_SUI, 0);
+
+        weight_hook.set_validator_addresses_and_weights(
+            &weight_hook_admin_cap, 
+            {
+                let mut map = vec_map::empty();
+                map.insert(address::from_u256(2), 100);
+
+                map
+            }
+        );
+        weight_hook.rebalance(&mut system_state, &mut lst_info, scenario.ctx());
+
+        assert!(lst_info.storage().validators().borrow(0).total_sui_amount() == 0, 0);
+        assert!(lst_info.storage().validators().borrow(1).total_sui_amount() == 0, 0);
+        assert!(lst_info.storage().validators().borrow(2).total_sui_amount() == 100 * MIST_PER_SUI, 0);
+
+        // sharing to make sure shared object deletion actually works lol
+        transfer::public_share_object(weight_hook);
+        scenario.next_tx(@0x0);
+
+        let weight_hook = scenario.take_shared<WeightHook<TEST>>();
+        let admin_cap = weight_hook.eject(weight_hook_admin_cap);
+
+        test_scenario::return_shared(system_state);
+
+        sui::test_utils::destroy(admin_cap);
+        sui::test_utils::destroy(lst_info);
+        sui::test_utils::destroy(lst);
+
+        scenario.end(); 
+     }
+}