fix: fixed withdraw calculations to include interest

Author: Teolhyn

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<PoolState, Error> {
         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<i128, Error> {
+        pool::read_total_shares(&e)
+    }
+
     pub fn get_available_balance(e: Env) -> Result<i128, Error> {
         pool::read_available_balance(&e)
     }
@@ -261,9 +290,9 @@ impl LoanPoolContract {
 
     pub fn get_pool_state(e: Env) -> Result<PoolState, Error> {
         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() {

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,
 }

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<i128, Error> {
-    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<i128, Error> {
     }
 }
 
-pub fn change_total_shares(e: &Env, amount: i128) -> Result<(), Error> {
+pub fn change_total_shares(e: &Env, amount: i128) -> Result<i128, Error> {
     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<i128, Error> {
-    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<i128, Error> {
     }
 }
 
-pub fn change_total_balance(e: &Env, amount: i128) -> Result<(), Error> {
+pub fn change_total_balance(e: &Env, amount: i128) -> Result<i128, Error> {
     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<i128, Error> {
-    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<i128, Error> {
     }
 }
 
-pub fn change_available_balance(e: &Env, amount: i128) -> Result<(), Error> {
+pub fn change_available_balance(e: &Env, amount: i128) -> Result<i128, Error> {
     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<i128, Error> {
     }
 }
 
-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<u64, Error> {

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<V> to V
-    let receivables_now = positions.receivables;
+    let receivables_now = positions.receivable_shares;
     let liabilities_now = positions.liabilities;
     let collateral_now = positions.collateral;
 

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

src/contexts/pool-context.tsx

@@ -58,9 +58,9 @@ const fetchPoolState = async (ticker: SupportedCurrency): Promise<PoolState> =>
   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,
     };
   }