From a9613c3624bff25d4dd965593affbecca42191a2 Mon Sep 17 00:00:00 2001 From: teolhyn Date: Wed, 15 Jan 2025 14:10:12 +0200 Subject: [PATCH] fix: fixed withdraw calculations to include interest --- contracts/loan_pool/src/contract.rs | 73 +++++++++++++++++------- contracts/loan_pool/src/dto.rs | 7 ++- contracts/loan_pool/src/pool.rs | 60 ++++++++++--------- contracts/loan_pool/src/positions.rs | 8 +-- contracts/loan_pool/src/storage_types.rs | 8 +-- src/contexts/pool-context.tsx | 6 +- 6 files changed, 95 insertions(+), 67 deletions(-) diff --git a/contracts/loan_pool/src/contract.rs b/contracts/loan_pool/src/contract.rs index c304ea12..080a050f 100644 --- a/contracts/loan_pool/src/contract.rs +++ b/contracts/loan_pool/src/contract.rs @@ -75,39 +75,64 @@ impl LoanPoolContract { } /// Transfers share tokens back, burns them and gives corresponding amount of tokens back to user. Returns amount of tokens withdrawn - pub fn withdraw(e: Env, user: Address, amount: i128) -> Result<(i128, i128), Error> { + pub fn withdraw(e: Env, user: Address, amount: i128) -> Result { user.require_auth(); Self::add_interest_to_accrual(e.clone())?; // Get users receivables - let Positions { receivables, .. } = positions::read_positions(&e, &user); + let Positions { + receivable_shares, .. + } = positions::read_positions(&e, &user); // Check that user is not trying to move more than receivables (TODO: also include collateral?) - assert!( - amount <= receivables, - "Amount can not be greater than receivables!" - ); - let available_balance = Self::get_available_balance(e.clone())?; - assert!(amount <= available_balance); + if amount > receivable_shares { + return Err(Error::WithdrawIsNegative); + } - // TODO: Decrease AvailableBalance - pool::change_available_balance(&e, amount.checked_neg().ok_or(Error::OverOrUnderFlow)?)?; - // TODO: Decrease TotalShares - Positions should have shares if we use them - // TODO: Decrease TotalBalance - pool::change_total_balance(&e, amount.checked_neg().ok_or(Error::OverOrUnderFlow)?)?; + let available_balance_tokens = Self::get_available_balance(e.clone())?; + if amount > available_balance_tokens { + return Err(Error::WithdrawOverBalance); + } + let total_balance_shares = Self::get_total_balance_shares(e.clone())?; + let total_balance_tokens = Self::get_contract_balance(e.clone())?; + let shares_to_decrease = amount + .checked_mul(total_balance_shares) + .ok_or(Error::OverOrUnderFlow)? + .checked_div(total_balance_tokens) + .ok_or(Error::OverOrUnderFlow)?; - // Decrease users position in pool as they withdraw + let new_available_balance_tokens = pool::change_available_balance( + &e, + amount.checked_neg().ok_or(Error::OverOrUnderFlow)?, + )?; + let new_total_balance_tokens = + pool::change_total_balance(&e, amount.checked_neg().ok_or(Error::OverOrUnderFlow)?)?; + let new_total_balance_shares = pool::change_total_shares(&e, shares_to_decrease)?; let liabilities: i128 = 0; let collateral: i128 = 0; - positions::decrease_positions(&e, user.clone(), amount, liabilities, collateral)?; + positions::decrease_positions( + &e, + user.clone(), + shares_to_decrease, + liabilities, + collateral, + )?; // Transfer tokens from pool to user let token_address = &pool::read_currency(&e)?.token_address; let client = token::Client::new(&e, token_address); client.transfer(&e.current_contract_address(), &user, &amount); - Ok((amount, amount)) + let new_annual_interest_rate = Self::get_interest(e.clone())?; + + let pool_state = PoolState { + total_balance_tokens: new_total_balance_tokens, + available_balance_tokens: new_available_balance_tokens, + total_balance_shares: new_total_balance_shares, + annual_interest_rate: new_annual_interest_rate, + }; + Ok(pool_state) } /// Borrow tokens from the pool @@ -247,6 +272,10 @@ impl LoanPoolContract { pool::read_total_balance(&e) } + pub fn get_total_balance_shares(e: Env) -> Result { + pool::read_total_shares(&e) + } + pub fn get_available_balance(e: Env) -> Result { pool::read_available_balance(&e) } @@ -261,9 +290,9 @@ impl LoanPoolContract { pub fn get_pool_state(e: Env) -> Result { Ok(PoolState { - total_balance: pool::read_total_balance(&e)?, - available_balance: pool::read_available_balance(&e)?, - total_shares: pool::read_total_shares(&e)?, + total_balance_tokens: pool::read_total_balance(&e)?, + available_balance_tokens: pool::read_available_balance(&e)?, + total_balance_shares: pool::read_total_shares(&e)?, annual_interest_rate: interest::get_interest(e)?, }) } @@ -529,7 +558,7 @@ mod test { } #[test] - #[should_panic(expected = "Amount can not be greater than receivables!")] + #[should_panic(expected = "Error(Contract, #12)")] fn withdraw_more_than_balance() { let e = Env::default(); e.mock_all_auths(); @@ -597,9 +626,9 @@ mod test { contract_client.borrow(&user2, &500); - let withdraw_result: (i128, i128) = contract_client.withdraw(&user, &amount); + let withdraw_result = contract_client.withdraw(&user, &amount); - assert_eq!(withdraw_result, (amount, amount)); + assert_eq!(withdraw_result, contract_client.get_pool_state()); } #[test] fn add_accrual_full_usage() { diff --git a/contracts/loan_pool/src/dto.rs b/contracts/loan_pool/src/dto.rs index b06118d7..ebb11836 100644 --- a/contracts/loan_pool/src/dto.rs +++ b/contracts/loan_pool/src/dto.rs @@ -1,9 +1,10 @@ use soroban_sdk::contracttype; #[contracttype] +#[derive(Debug, PartialEq)] pub struct PoolState { - pub total_balance: i128, - pub available_balance: i128, - pub total_shares: i128, + pub total_balance_tokens: i128, + pub available_balance_tokens: i128, + pub total_balance_shares: i128, pub annual_interest_rate: i128, } diff --git a/contracts/loan_pool/src/pool.rs b/contracts/loan_pool/src/pool.rs index bb675600..0f87711a 100644 --- a/contracts/loan_pool/src/pool.rs +++ b/contracts/loan_pool/src/pool.rs @@ -21,6 +21,8 @@ pub enum Error { AccrualLastUpdated = 8, OverOrUnderFlow = 9, NegativeDeposit = 10, + WithdrawOverBalance = 11, + WithdrawIsNegative = 12, } pub fn write_loan_manager_addr(e: &Env, loan_manager_addr: Address) { @@ -65,14 +67,14 @@ pub fn write_liquidation_threshold(e: &Env, threshold: i128) { } pub fn write_total_shares(e: &Env, amount: i128) { - let key: PoolDataKey = PoolDataKey::TotalShares; + let key: PoolDataKey = PoolDataKey::TotalBalanceShares; e.storage().persistent().set(&key, &amount); extend_persistent(e.clone(), &key); } pub fn read_total_shares(e: &Env) -> Result { - let key: PoolDataKey = PoolDataKey::TotalBalance; + let key: PoolDataKey = PoolDataKey::TotalBalanceTokens; if let Some(total_shares) = e.storage().persistent().get(&key) { total_shares @@ -81,27 +83,25 @@ pub fn read_total_shares(e: &Env) -> Result { } } -pub fn change_total_shares(e: &Env, amount: i128) -> Result<(), Error> { +pub fn change_total_shares(e: &Env, amount: i128) -> Result { let current_balance = read_total_shares(e)?; - write_total_shares( - e, - amount - .checked_add(current_balance) - .ok_or(Error::OverOrUnderFlow)?, - ); - Ok(()) + let new_amount = amount + .checked_add(current_balance) + .ok_or(Error::OverOrUnderFlow)?; + write_total_shares(e, new_amount); + Ok(new_amount) } pub fn write_total_balance(e: &Env, amount: i128) { - let key: PoolDataKey = PoolDataKey::TotalBalance; + let key: PoolDataKey = PoolDataKey::TotalBalanceTokens; e.storage().persistent().set(&key, &amount); extend_persistent(e.clone(), &key); } pub fn read_total_balance(e: &Env) -> Result { - let key: PoolDataKey = PoolDataKey::TotalBalance; + let key: PoolDataKey = PoolDataKey::TotalBalanceTokens; if let Some(total_balance) = e.storage().persistent().get(&key) { total_balance @@ -110,27 +110,25 @@ pub fn read_total_balance(e: &Env) -> Result { } } -pub fn change_total_balance(e: &Env, amount: i128) -> Result<(), Error> { +pub fn change_total_balance(e: &Env, amount: i128) -> Result { let current_balance = read_total_balance(e)?; - write_total_balance( - e, - amount - .checked_add(current_balance) - .ok_or(Error::OverOrUnderFlow)?, - ); - Ok(()) + let new_amount = amount + .checked_add(current_balance) + .ok_or(Error::OverOrUnderFlow)?; + write_total_balance(e, new_amount); + Ok(new_amount) } pub fn write_available_balance(e: &Env, amount: i128) { - let key: PoolDataKey = PoolDataKey::AvailableBalance; + let key: PoolDataKey = PoolDataKey::AvailableBalanceTokens; e.storage().persistent().set(&key, &amount); extend_persistent(e.clone(), &key); } pub fn read_available_balance(e: &Env) -> Result { - let key: PoolDataKey = PoolDataKey::AvailableBalance; + let key: PoolDataKey = PoolDataKey::AvailableBalanceTokens; if let Some(available_balance) = e.storage().persistent().get(&key) { available_balance @@ -139,16 +137,14 @@ pub fn read_available_balance(e: &Env) -> Result { } } -pub fn change_available_balance(e: &Env, amount: i128) -> Result<(), Error> { +pub fn change_available_balance(e: &Env, amount: i128) -> Result { let current_balance = read_available_balance(e)?; - write_available_balance( - e, - amount - .checked_add(current_balance) - .ok_or(Error::OverOrUnderFlow)?, - ); - Ok(()) + let new_amount = amount + .checked_add(current_balance) + .ok_or(Error::OverOrUnderFlow)?; + write_available_balance(e, new_amount); + Ok(new_amount) } pub fn write_accrual(e: &Env, accrual: i128) { @@ -168,11 +164,13 @@ pub fn read_accrual(e: &Env) -> Result { } } -pub fn write_accrual_last_updated(e: &Env, sequence: u64) { +pub fn write_accrual_last_updated(e: &Env, sequence: u64) -> u64 { let key = PoolDataKey::AccrualLastUpdate; e.storage().persistent().set(&key, &sequence); extend_persistent(e.clone(), &key); + + sequence } pub fn read_accrual_last_updated(e: &Env) -> Result { diff --git a/contracts/loan_pool/src/positions.rs b/contracts/loan_pool/src/positions.rs index 3a2437f1..5d01bc79 100644 --- a/contracts/loan_pool/src/positions.rs +++ b/contracts/loan_pool/src/positions.rs @@ -10,7 +10,7 @@ pub fn read_positions(e: &Env, addr: &Address) -> Positions { positions } else { Positions { - receivables: 0, + receivable_shares: 0, liabilities: 0, collateral: 0, } @@ -21,7 +21,7 @@ fn write_positions(e: &Env, addr: Address, receivables: i128, liabilities: i128, let key: PoolDataKey = PoolDataKey::Positions(addr); let positions: Positions = Positions { - receivables, + receivable_shares: receivables, liabilities, collateral, }; @@ -42,7 +42,7 @@ pub fn increase_positions( ) -> Result<(), Error> { let positions = read_positions(e, &addr); - let receivables_now: i128 = positions.receivables; + let receivables_now: i128 = positions.receivable_shares; let liabilities_now: i128 = positions.liabilities; let collateral_now = positions.collateral; write_positions( @@ -71,7 +71,7 @@ pub fn decrease_positions( let positions = read_positions(e, &addr); // TODO: Might need to use get rather than get_unchecked and convert from Option to V - let receivables_now = positions.receivables; + let receivables_now = positions.receivable_shares; let liabilities_now = positions.liabilities; let collateral_now = positions.collateral; diff --git a/contracts/loan_pool/src/storage_types.rs b/contracts/loan_pool/src/storage_types.rs index cdc16d46..20551e1f 100644 --- a/contracts/loan_pool/src/storage_types.rs +++ b/contracts/loan_pool/src/storage_types.rs @@ -21,7 +21,7 @@ pub struct PoolConfig { #[contracttype] pub struct Positions { // struct names under 9 characters are marginally more efficient. Need to think if we value marginal efficiency over readibility - pub receivables: i128, + pub receivable_shares: i128, pub liabilities: i128, pub collateral: i128, } @@ -38,11 +38,11 @@ pub enum PoolDataKey { // Users positions in the pool Positions(Address), // Total amount of shares in circulation - TotalShares, + TotalBalanceShares, // Total balance of pool - TotalBalance, + TotalBalanceTokens, // Available balance of pool - AvailableBalance, + AvailableBalanceTokens, // Pool interest accrual index Accrual, // Last update ledger of accrual diff --git a/src/contexts/pool-context.tsx b/src/contexts/pool-context.tsx index e9a186d9..142ce138 100644 --- a/src/contexts/pool-context.tsx +++ b/src/contexts/pool-context.tsx @@ -58,9 +58,9 @@ const fetchPoolState = async (ticker: SupportedCurrency): Promise => if (result.isOk()) { const value = result.unwrap(); return { - totalBalance: value.total_balance, - availableBalance: value.available_balance, - totalShares: value.total_shares, + totalBalance: value.total_balance_tokens, + availableBalance: value.available_balance_tokens, + totalShares: value.total_balance_shares, annualInterestRate: value.annual_interest_rate, }; }