Merge pull request #181 from paltalabs/feat/refactorModels

Feat/refactor models

apps/contracts/factory/src/events.rs

@@ -1,6 +1,6 @@
 //! Definition of the Events used in the contract
 use soroban_sdk::{contracttype, symbol_short, Address, Env, Vec};
-use crate::defindex::AssetAllocation;
+use crate::defindex::AssetStrategySet;
 
 // INITIALIZED
 #[contracttype]
@@ -29,7 +29,7 @@ pub struct CreateDeFindexEvent {
     pub fee_receiver: Address, 
     pub manager: Address,
     pub vault_fee: u32,
-    pub assets: Vec<AssetAllocation>
+    pub assets: Vec<AssetStrategySet>
 }
 
 /// Publishes an `CreateDeFindexEvent` to the event stream.
@@ -39,7 +39,7 @@ pub(crate) fn emit_create_defindex_vault(
     fee_receiver: Address, 
     manager: Address,
     vault_fee: u32,
-    assets: Vec<AssetAllocation>,
+    assets: Vec<AssetStrategySet>,
 ) {
     let event = CreateDeFindexEvent { 
       emergency_manager,

apps/contracts/factory/src/lib.rs

@@ -9,7 +9,7 @@ use soroban_sdk::{
     contract, contractimpl, Address, BytesN, Env, Map, String, Vec
 };
 use error::FactoryError;
-use defindex::{create_contract, AssetAllocation};
+use defindex::{create_contract, AssetStrategySet};
 use storage::{ add_new_defindex, extend_instance_ttl, get_admin, get_defi_wasm_hash, get_defindex_receiver, get_deployed_defindexes, get_fee_rate, has_admin, put_admin, put_defi_wasm_hash, put_defindex_receiver, put_defindex_fee };
 
 fn check_initialized(e: &Env) -> Result<(), FactoryError> {
@@ -49,7 +49,7 @@ pub trait FactoryTrait {
     /// * `vault_name` - The name of the vault.
     /// * `vault_symbol` - The symbol of the vault.
     /// * `manager` - The address assigned as the vault manager.
-    /// * `assets` - A vector of `AssetAllocation` structs that define the assets managed by the vault.
+    /// * `assets` - A vector of `AssetStrategySet` structs that define the assets managed by the vault.
     /// * `salt` - A salt used for ensuring unique addresses for each deployed vault.
     ///
     /// # Returns
@@ -62,7 +62,7 @@ pub trait FactoryTrait {
         vault_name: String,
         vault_symbol: String,
         manager: Address,
-        assets: Vec<AssetAllocation>,
+        assets: Vec<AssetStrategySet>,
         salt: BytesN<32>
     ) -> Result<Address, FactoryError>;
 
@@ -76,7 +76,7 @@ pub trait FactoryTrait {
     /// * `vault_name` - The name of the vault.
     /// * `vault_symbol` - The symbol of the vault.
     /// * `manager` - The address assigned as the vault manager.
-    /// * `assets` - A vector of `AssetAllocation` structs that define the assets managed by the vault.
+    /// * `assets` - A vector of `AssetStrategySet` structs that define the assets managed by the vault.
     /// * `amounts` - A vector of `AssetAmounts` structs that define the initial deposit amounts.
     /// * `salt` - A salt used for ensuring unique addresses for each deployed vault.
     ///
@@ -91,7 +91,7 @@ pub trait FactoryTrait {
         vault_name: String,
         vault_symbol: String,
         manager: Address,
-        assets: Vec<AssetAllocation>,
+        assets: Vec<AssetStrategySet>,
         amounts: Vec<i128>,
         salt: BytesN<32>
     ) -> Result<Address, FactoryError>;
@@ -213,7 +213,7 @@ impl FactoryTrait for DeFindexFactory {
     /// * `fee_receiver` - The address designated to receive fees from the vault.
     /// * `vault_fee` - The percentage share of fees allocated to the vault's fee receiver.
     /// * `manager` - The address assigned as the vault manager.
-    /// * `assets` - A vector of `AssetAllocation` structs that define the assets managed by the vault.
+    /// * `assets` - A vector of `AssetStrategySet` structs that define the assets managed by the vault.
     /// * `salt` - A salt used for ensuring unique addresses for each deployed vault.
     ///
     /// # Returns
@@ -226,7 +226,7 @@ impl FactoryTrait for DeFindexFactory {
         vault_name: String,
         vault_symbol: String,
         manager: Address,
-        assets: Vec<AssetAllocation>,
+        assets: Vec<AssetStrategySet>,
         salt: BytesN<32>
     ) -> Result<Address, FactoryError> {
         extend_instance_ttl(&e);
@@ -265,7 +265,7 @@ impl FactoryTrait for DeFindexFactory {
     /// * `vault_name` - The name of the vault.
     /// * `vault_symbol` - The symbol of the vault.
     /// * `manager` - The address assigned as the vault manager.
-    /// * `assets` - A vector of `AssetAllocation` structs that define the assets managed by the vault.
+    /// * `assets` - A vector of `AssetStrategySet` structs that define the assets managed by the vault.
     /// * `amounts` - A vector of `AssetAmounts` structs that define the initial deposit amounts.
     /// * `salt` - A salt used for ensuring unique addresses for each deployed vault.
     ///
@@ -280,7 +280,7 @@ impl FactoryTrait for DeFindexFactory {
         vault_name: String,
         vault_symbol: String,
         manager: Address,
-        assets: Vec<AssetAllocation>,
+        assets: Vec<AssetStrategySet>,
         amounts: Vec<i128>,
         salt: BytesN<32>
     ) -> Result<Address, FactoryError> {

apps/contracts/factory/src/test.rs

@@ -1,6 +1,6 @@
 #![cfg(test)]
 extern crate std;
-use crate::defindex::{AssetAllocation, Strategy};
+use crate::defindex::{AssetStrategySet, Strategy};
 use crate::{DeFindexFactory, DeFindexFactoryClient};
 use soroban_sdk::token::{
     StellarAssetClient as SorobanTokenAdminClient, TokenClient as SorobanTokenClient,
@@ -53,10 +53,10 @@ pub(crate) fn get_token_admin_client<'a>(
     SorobanTokenAdminClient::new(e, address)
 }
 
-pub(crate) fn create_asset_params(test: &DeFindexFactoryTest) -> Vec<AssetAllocation> {
+pub(crate) fn create_asset_params(test: &DeFindexFactoryTest) -> Vec<AssetStrategySet> {
     sorobanvec![
         &test.env,
-        AssetAllocation {
+        AssetStrategySet {
             address: test.token0.address.clone(),
             strategies: sorobanvec![
                 &test.env,
@@ -67,7 +67,7 @@ pub(crate) fn create_asset_params(test: &DeFindexFactoryTest) -> Vec<AssetAlloca
                 }
             ],
         },
-        AssetAllocation {
+        AssetStrategySet {
             address: test.token1.address.clone(),
             strategies: sorobanvec![
                 &test.env,

apps/contracts/factory/src/test/all_flow.rs

@@ -1,6 +1,12 @@
 use soroban_sdk::{vec, BytesN, String};
 
-use crate::test::{create_asset_params, defindex_vault_contract::{self, Investment}, DeFindexFactoryTest};
+use crate::test::{
+  create_asset_params, 
+  defindex_vault_contract::{
+    self, 
+    AssetInvestmentAllocation,  
+    StrategyInvestment, 
+  }, DeFindexFactoryTest};
 
 #[test]
 fn test_deposit_success() {
@@ -49,7 +55,7 @@ fn test_deposit_success() {
   defindex_contract.deposit(&vec![&test.env, amount_token0, amount_token1], &vec![&test.env, 0, 0], &users[0]);
 
   let df_balance = defindex_contract.balance(&users[0]);
-  assert_eq!(df_balance, amount_token0 + amount_token1); // TODO: The amount of dfTokens minted is the sum of both asset deposited?
+  assert_eq!(df_balance, amount_token0 + amount_token1 - 1000); // TODO: The amount of dfTokens minted is the sum of both asset deposited?
 
 
   // defindex_contract.withdraw(&df_balance, &users[0]);
@@ -109,7 +115,7 @@ fn test_withdraw_success() {
   defindex_contract.deposit(&vec![&test.env, amount_token0, amount_token1], &vec![&test.env, 0, 0], &users[0]);
 
   let df_balance = defindex_contract.balance(&users[0]);
-  assert_eq!(df_balance.clone(), amount_token0 + amount_token1); // TODO: The amount of dfTokens minted is the sum of both asset deposited?
+  assert_eq!(df_balance.clone(), amount_token0 + amount_token1 - 1000); // TODO: The amount of dfTokens minted is the sum of both asset deposited?
 
   let vault_token0_balance = test.token0.balance(&defindex_contract.address);
   assert_eq!(vault_token0_balance, amount_token0);
@@ -118,15 +124,28 @@ fn test_withdraw_success() {
   assert_eq!(vault_token1_balance, amount_token1);
 
   let investments = vec![
-    &test.env, 
-    Investment {
-      amount: amount_token0, 
-      strategy: test.strategy_contract_token0.address.clone()
-    }, 
-    Investment {
-      amount: amount_token1, 
-      strategy: test.strategy_contract_token1.address.clone()
-    }];
+        &test.env,
+        Some(AssetInvestmentAllocation {
+            asset: test.token0.address.clone(),
+            strategy_investments: vec![
+                &test.env,
+                Some(StrategyInvestment {
+                    strategy: test.strategy_contract_token0.address.clone(),
+                    amount: amount_token0,
+                }),
+            ],
+        }),
+        Some(AssetInvestmentAllocation {
+          asset: test.token1.address.clone(),
+          strategy_investments: vec![
+              &test.env,
+              Some(StrategyInvestment {
+                  strategy: test.strategy_contract_token1.address.clone(),
+                  amount: amount_token1,
+              }),
+          ],
+      })
+    ];
 
 
   defindex_contract.invest(&investments);

apps/contracts/vault/src/error.rs

@@ -21,6 +21,10 @@ pub enum ContractError {
     UnsupportedAsset = 116,
     InsufficientAmount = 117,
     NoOptimalAmounts = 118, //this should not happen
+    WrongInvestmentLength = 119,
+    WrongAssetAddress = 122,
+    WrongStrategiesLength = 123,
+
 
     // Arithmetic Errors (12x)
     ArithmeticError = 120,
@@ -35,6 +39,7 @@ pub enum ContractError {
     StrategyPausedOrNotFound = 141,
     StrategyWithdrawError = 142,
     StrategyInvestError = 143,
+    StrategyPaused = 144,
 
     // Asset Errors (15x)
     AssetNotFound = 150,

apps/contracts/vault/src/events.rs

@@ -1,5 +1,5 @@
 //! Definition of the Events used in the DeFindex Vault contract
-use crate::models::AssetAllocation;
+use crate::models::AssetStrategySet;
 use soroban_sdk::{contracttype, symbol_short, Address, Env, Vec};
 
 // INITIALIZED VAULT EVENT
@@ -10,7 +10,7 @@ pub struct InitializedVaultEvent {
     pub vault_fee_receiver: Address,
     pub manager: Address,
     pub defindex_protocol_receiver: Address,
-    pub assets: Vec<AssetAllocation>,
+    pub assets: Vec<AssetStrategySet>,
 }
 
 /// Publishes an `InitializedVaultEvent` to the event stream.
@@ -20,7 +20,7 @@ pub(crate) fn emit_initialized_vault(
     vault_fee_receiver: Address,
     manager: Address,
     defindex_protocol_receiver: Address,
-    assets: Vec<AssetAllocation>,
+    assets: Vec<AssetStrategySet>,
 ) {
     let event = InitializedVaultEvent {
         emergency_manager,

apps/contracts/vault/src/funds.rs

@@ -1,11 +1,11 @@
 use soroban_sdk::token::TokenClient;
 use soroban_sdk::{Address, Env, Map};
 
-use crate::models::AssetAllocation;
+use crate::models::AssetStrategySet;
 use crate::storage::get_assets;
 use crate::strategies::get_strategy_client;
 
-// Funds for AssetAllocation
+// Funds for AssetStrategySet
 
 /// Fetches the idle funds for a given asset. Idle funds refer to the balance of the asset
 /// that is currently not invested in any strategies.
@@ -16,8 +16,8 @@ use crate::strategies::get_strategy_client;
 ///
 /// # Returns
 /// * The idle balance (i128) of the asset in the current contract address.
-fn fetch_idle_funds_for_asset(e: &Env, asset: &AssetAllocation) -> i128 {
-    TokenClient::new(e, &asset.address).balance(&e.current_contract_address())
+fn fetch_idle_funds_for_asset(e: &Env, asset: &Address) -> i128 {
+    TokenClient::new(e, &asset).balance(&e.current_contract_address())
 }
 
 /// Fetches the total funds that are invested for a given asset.
@@ -43,7 +43,7 @@ pub fn fetch_invested_funds_for_strategy(e: &Env, strategy_address: &Address) ->
 ///
 /// # Returns
 /// * The total invested balance (i128) of the asset across all strategies.
-pub fn fetch_invested_funds_for_asset(e: &Env, asset: &AssetAllocation) -> i128 {
+pub fn fetch_invested_funds_for_asset(e: &Env, asset: &AssetStrategySet) -> i128 {
     let mut invested_funds = 0;
     for strategy in asset.strategies.iter() {
         invested_funds += fetch_invested_funds_for_strategy(e, &strategy.address);
@@ -65,7 +65,7 @@ pub fn fetch_current_idle_funds(e: &Env) -> Map<Address, i128> {
     let assets = get_assets(e);
     let mut map: Map<Address, i128> = Map::new(e);
     for asset in assets {
-        map.set(asset.address.clone(), fetch_idle_funds_for_asset(e, &asset));
+        map.set(asset.address.clone(), fetch_idle_funds_for_asset(e, &asset.address));
     }
     map
 }
@@ -103,7 +103,7 @@ pub fn fetch_total_managed_funds(e: &Env) -> Map<Address, i128> {
     let assets = get_assets(e);
     let mut map: Map<Address, i128> = Map::new(e);
     for asset in assets {
-        let idle_funds = fetch_idle_funds_for_asset(e, &asset);
+        let idle_funds = fetch_idle_funds_for_asset(e, &asset.address);
         let invested_funds = fetch_invested_funds_for_asset(e, &asset);
         map.set(asset.address.clone(), idle_funds + invested_funds);
     }

apps/contracts/vault/src/interface.rs

@@ -1,7 +1,7 @@
 use soroban_sdk::{Address, Env, Map, String, Vec};
 
 use crate::{
-    models::{AssetAllocation, Instruction, Investment},
+    models::{AssetStrategySet, Instruction, AssetInvestmentAllocation},
     ContractError,
 };
 
@@ -32,7 +32,7 @@ pub trait VaultTrait {
     ///
     fn initialize(
         e: Env,
-        assets: Vec<AssetAllocation>,
+        assets: Vec<AssetStrategySet>,
         manager: Address,
         emergency_manager: Address,
         vault_fee_receiver: Address,
@@ -156,8 +156,8 @@ pub trait VaultTrait {
     /// * `e` - The environment.
     ///
     /// # Returns:
-    /// * `Vec<AssetAllocation>` - A vector of `AssetAllocation` structs representing the assets managed by the vault.
-    fn get_assets(e: Env) -> Vec<AssetAllocation>;
+    /// * `Vec<AssetStrategySet>` - A vector of `AssetStrategySet` structs representing the assets managed by the vault.
+    fn get_assets(e: Env) -> Vec<AssetStrategySet>;
 
     /// Returns the total managed funds of the vault, including both invested and idle funds.
     ///
@@ -267,16 +267,37 @@ pub trait AdminInterfaceTrait {
 }
 
 pub trait VaultManagementTrait {
-    /// Invests the vault's idle funds into the specified strategies.
+    /// Executes the investment of the vault's idle funds based on the specified asset allocations.
+    /// This function allows partial investments by providing an optional allocation for each asset,
+    /// and it ensures proper authorization and validation checks before proceeding with investments.
     ///
-    /// # Arguments:
-    /// * `e` - The environment.
-    /// * `investment` - A vector of `Investment` structs representing the amount to invest in each strategy.
-    /// * `caller` - The address of the caller.
+    /// # Arguments
+    /// * `e` - The current environment reference.
+    /// * `asset_investments` - A vector of optional `AssetInvestmentAllocation` structures, where each element 
+    ///   represents an allocation for a specific asset. The vector must match the number of vault assets in length.
     ///
-    /// # Returns:
-    /// * `Result<(), ContractError>` - Ok if successful, otherwise returns a ContractError.
-    fn invest(e: Env, investment: Vec<Investment>) -> Result<(), ContractError>;
+    /// # Returns
+    /// * `Result<(), ContractError>` - Returns `Ok(())` if the investments are successful or a `ContractError`
+    ///   if any issue occurs during validation or execution.
+    ///
+    /// # Function Flow
+    /// 1. **Extend Instance TTL**: Extends the contract instance's time-to-live to keep the instance active.
+    /// 2. **Check Initialization**: Verifies that the vault is properly initialized before proceeding.
+    /// 3. **Access Control**: Ensures the caller has the `Manager` role required to initiate investments.
+    /// 4. **Asset Count Validation**: Verifies that the length of the `asset_investments` vector matches
+    ///    the number of assets managed by the vault. If they don't match, a `WrongInvestmentLength` error is returned.
+    /// 5. **Investment Execution**: Calls the `check_and_execute_investments` function to perform the investment
+    ///    after validating the inputs and ensuring correct execution flows for each asset allocation.
+    ///
+    /// # Errors
+    /// * Returns `ContractError::WrongInvestmentLength` if the length of `asset_investments` does not match the vault assets.
+    /// * Returns `ContractError` if access control validation fails or if investment execution encounters an issue.
+    ///
+    /// # Security
+    /// - Only addresses with the `Manager` role can call this function, ensuring restricted access to managing investments.
+    fn invest(e: Env, 
+        asset_investments: Vec<Option<AssetInvestmentAllocation>>
+    ) -> Result<(), ContractError>;
 
     /// Rebalances the vault by executing a series of instructions.
     ///