Merge pull request #434 from Phoenix-Protocol-Group/kalo/improve-quer…

…ies-in-vesting

Vesting: query_vesting_info now returns the index too

contracts/vesting/src/contract.rs

@@ -3,18 +3,18 @@ use phoenix::{
     utils::{convert_i128_to_u128, convert_u128_to_i128},
 };
 use soroban_sdk::{
-    contract, contractimpl, contractmeta, log, panic_with_error, Address, BytesN, Env, Vec,
+    contract, contractimpl, contractmeta, log, panic_with_error, vec, Address, BytesN, Env, Vec,
 };
 
 #[cfg(feature = "minter")]
 use crate::storage::{get_minter, save_minter, MinterInfo};
 use crate::{
     error::ContractError,
     storage::{
-        get_admin_old, get_all_vestings, get_max_vesting_complexity, get_token_info, get_vesting,
-        is_initialized, save_admin_old, save_max_vesting_complexity, save_token_info, save_vesting,
-        set_initialized, update_vesting, VestingInfo, VestingSchedule, VestingTokenInfo, ADMIN,
-        VESTING_KEY,
+        get_admin_old, get_max_vesting_complexity, get_token_info, get_vesting, is_initialized,
+        save_admin_old, save_max_vesting_complexity, save_token_info, save_vesting,
+        set_initialized, update_vesting, VestingCounterKey, VestingInfo, VestingInfoKey,
+        VestingInfoResponse, VestingSchedule, VestingTokenInfo, ADMIN, VESTING_KEY,
     },
     token_contract,
     utils::{check_duplications, validate_vesting_schedule},
@@ -45,9 +45,9 @@ pub trait VestingTrait {
 
     fn query_balance(env: Env, address: Address) -> i128;
 
-    fn query_vesting_info(env: Env, address: Address, index: u64) -> VestingInfo;
+    fn query_vesting_info(env: Env, address: Address, index: u64) -> VestingInfoResponse;
 
-    fn query_all_vesting_info(env: Env, address: Address) -> Vec<VestingInfo>;
+    fn query_all_vesting_info(env: Env, address: Address) -> Vec<VestingInfoResponse>;
 
     fn query_token_info(env: Env) -> VestingTokenInfo;
 
@@ -189,7 +189,7 @@ impl VestingTrait for Vesting {
             save_vesting(
                 &env,
                 &vesting_schedule.recipient.clone(),
-                &VestingInfo {
+                VestingInfo {
                     balance: vested_amount,
                     recipient: vesting_schedule.recipient,
                     schedule: vesting_schedule.curve.clone(),
@@ -432,18 +432,46 @@ impl VestingTrait for Vesting {
         token_contract::Client::new(&env, &get_token_info(&env).address).balance(&address)
     }
 
-    fn query_vesting_info(env: Env, address: Address, index: u64) -> VestingInfo {
+    fn query_vesting_info(env: Env, address: Address, index: u64) -> VestingInfoResponse {
         env.storage()
             .instance()
             .extend_ttl(INSTANCE_RENEWAL_THRESHOLD, INSTANCE_TARGET_TTL);
-        get_vesting(&env, &address, index)
+        let vesting_info = get_vesting(&env, &address, index);
+        VestingInfoResponse {
+            balance: vesting_info.balance,
+            recipient: vesting_info.recipient,
+            schedule: vesting_info.schedule,
+            index, // use the query parameter index
+        }
     }
 
-    fn query_all_vesting_info(env: Env, address: Address) -> Vec<VestingInfo> {
-        env.storage()
-            .instance()
-            .extend_ttl(INSTANCE_RENEWAL_THRESHOLD, INSTANCE_TARGET_TTL);
-        get_all_vestings(&env, &address)
+    fn query_all_vesting_info(env: Env, address: Address) -> Vec<VestingInfoResponse> {
+        let counter_key = VestingCounterKey {
+            recipient: address.clone(),
+        };
+        let count: u64 = env.storage().persistent().get(&counter_key).unwrap_or(0);
+        let mut responses = vec![&env];
+
+        for i in 0..count {
+            if env
+                .storage()
+                .persistent()
+                .get::<VestingInfoKey, VestingInfo>(&VestingInfoKey {
+                    recipient: address.clone(),
+                    index: i,
+                })
+                .is_some()
+            {
+                let vesting = get_vesting(&env, &address, i);
+                responses.push_back(VestingInfoResponse {
+                    balance: vesting.balance,
+                    recipient: vesting.recipient,
+                    schedule: vesting.schedule,
+                    index: i,
+                });
+            }
+        }
+        responses
     }
 
     fn query_token_info(env: Env) -> VestingTokenInfo {

contracts/vesting/src/storage.rs

@@ -4,8 +4,8 @@ use phoenix::ttl::{
     PERSISTENT_TARGET_TTL,
 };
 use soroban_sdk::{
-    contracttype, log, panic_with_error, symbol_short, vec, Address, ConversionError, Env, String,
-    Symbol, TryFromVal, Val, Vec,
+    contracttype, log, panic_with_error, symbol_short, Address, ConversionError, Env, String,
+    Symbol, TryFromVal, Val,
 };
 
 use crate::error::ContractError;
@@ -60,6 +60,12 @@ pub struct VestingInfo {
     pub schedule: Curve,
 }
 
+#[contracttype]
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct VestingCounterKey {
+    pub recipient: Address,
+}
+
 #[cfg(feature = "minter")]
 #[contracttype]
 #[derive(Clone, Debug, Eq, PartialEq)]
@@ -127,28 +133,37 @@ pub struct VestingInfoKey {
     pub index: u64,
 }
 
-pub fn save_vesting(env: &Env, address: &Address, vesting_info: &VestingInfo) {
-    let mut index = 0u64;
-    let mut vesting_key = VestingInfoKey {
+#[contracttype]
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct VestingInfoResponse {
+    pub balance: u128,
+    pub recipient: Address,
+    pub schedule: Curve,
+    pub index: u64,
+}
+
+pub fn save_vesting(env: &Env, address: &Address, vesting_info: VestingInfo) {
+    let counter_key = VestingCounterKey {
         recipient: address.clone(),
-        index,
     };
 
-    // Find the next available index
-    while env.storage().persistent().has(&vesting_key) {
-        index += 1;
-        vesting_key = VestingInfoKey {
-            recipient: address.clone(),
-            index,
-        };
-    }
+    let next_index: u64 = env.storage().persistent().get(&counter_key).unwrap_or(0);
 
-    env.storage().persistent().set(&vesting_key, vesting_info);
+    let vesting_key = VestingInfoKey {
+        recipient: address.clone(),
+        index: next_index,
+    };
+
+    env.storage().persistent().set(&vesting_key, &vesting_info);
     env.storage().persistent().extend_ttl(
         &vesting_key,
         PERSISTENT_RENEWAL_THRESHOLD,
         PERSISTENT_TARGET_TTL,
     );
+
+    env.storage()
+        .persistent()
+        .set(&counter_key, &(next_index + 1));
 }
 
 pub fn update_vesting(env: &Env, address: &Address, index: u64, vesting_info: &VestingInfo) {
@@ -182,32 +197,6 @@ pub fn get_vesting(env: &Env, recipient: &Address, index: u64) -> VestingInfo {
     vesting_info
 }
 
-pub fn get_all_vestings(env: &Env, address: &Address) -> Vec<VestingInfo> {
-    let mut vestings = vec![&env];
-    let mut index = 0u64;
-
-    loop {
-        let vesting_key = VestingInfoKey {
-            recipient: address.clone(),
-            index,
-        };
-
-        if let Some(vesting_info) = env.storage().persistent().get(&vesting_key) {
-            vestings.push_back(vesting_info);
-            index += 1;
-            env.storage().persistent().extend_ttl(
-                &vesting_key,
-                PERSISTENT_RENEWAL_THRESHOLD,
-                PERSISTENT_TARGET_TTL,
-            );
-        } else {
-            break;
-        }
-    }
-
-    vestings
-}
-
 #[cfg(feature = "minter")]
 pub fn save_minter(env: &Env, minter: &MinterInfo) {
     env.storage().instance().set(&DataKey::Minter, minter);

contracts/vesting/src/tests/claim.rs

@@ -1,5 +1,5 @@
 use crate::{
-    storage::{VestingInfo, VestingSchedule, VestingTokenInfo},
+    storage::{VestingInfoResponse, VestingSchedule, VestingTokenInfo},
     tests::setup::instantiate_vesting_client,
 };
 use curve::{Curve, PiecewiseLinear, SaturatingLinear, Step};
@@ -656,17 +656,18 @@ fn claim_tokens_from_two_distributions() {
         vesting_client.query_all_vesting_info(&vester1),
         vec![
             &env,
-            VestingInfo {
+            VestingInfoResponse {
                 recipient: vester1.clone(),
                 balance: 750, // balance is deducted because it was already once claimed
                 schedule: Curve::SaturatingLinear(SaturatingLinear {
                     min_x: 0,
                     min_y: 1_500,
                     max_x: 100,
                     max_y: 0,
-                })
+                }),
+                index: 0,
             },
-            VestingInfo {
+            VestingInfoResponse {
                 recipient: vester1.clone(),
                 balance: 500,
                 schedule: Curve::PiecewiseLinear(PiecewiseLinear {
@@ -685,7 +686,8 @@ fn claim_tokens_from_two_distributions() {
                             value: 0,
                         },
                     ],
-                })
+                }),
+                index: 1,
             }
         ]
     );
@@ -762,7 +764,7 @@ fn first_mainnet_simulation() {
 
     assert_eq!(
         vesting_client.query_vesting_info(&vester1, &0),
-        VestingInfo {
+        VestingInfoResponse {
             recipient: vester1.clone(),
             balance: 66_666_667,
             schedule: Curve::SaturatingLinear(SaturatingLinear {
@@ -772,6 +774,7 @@ fn first_mainnet_simulation() {
                 max_x: 1716820800,
                 max_y: 0,
             }),
+            index: 0,
         },
     );
 

contracts/vesting/src/tests/instantiate.rs

@@ -1,7 +1,7 @@
 use soroban_sdk::{testutils::Address as _, vec, Address, Env, String};
 
 use crate::{
-    storage::{VestingInfo, VestingSchedule, VestingTokenInfo},
+    storage::{VestingInfoResponse, VestingSchedule, VestingTokenInfo},
     tests::setup::{deploy_token_contract, instantiate_vesting_client},
 };
 use curve::{Curve, SaturatingLinear};
@@ -55,15 +55,16 @@ fn instantiate_contract_successfully() {
         vesting_client.query_all_vesting_info(&vester1),
         vec![
             &env,
-            VestingInfo {
+            VestingInfoResponse {
                 recipient: vester1,
                 balance: 120,
                 schedule: Curve::SaturatingLinear(SaturatingLinear {
                     min_x: 15,
                     min_y: 120,
                     max_x: 60,
                     max_y: 0,
-                })
+                }),
+                index: 0,
             }
         ]
     );