From 68ebff4a01cb0a07a19bb4a285225a83d596c731 Mon Sep 17 00:00:00 2001
From: Serge Farny <serge.farny@gmail.com>
Date: Mon, 24 Jun 2024 16:54:41 +0200
Subject: [PATCH] Program: charge collateral fee directly on borrowed tokens

---
 .../token_charge_collateral_fees.rs           | 37 +++++++++++++++----
 .../tests/cases/test_collateral_fees.rs       | 21 +++++++++--
 2 files changed, 47 insertions(+), 11 deletions(-)

diff --git a/programs/mango-v4/src/instructions/token_charge_collateral_fees.rs b/programs/mango-v4/src/instructions/token_charge_collateral_fees.rs
index 7a06245e1..3484bf446 100644
--- a/programs/mango-v4/src/instructions/token_charge_collateral_fees.rs
+++ b/programs/mango-v4/src/instructions/token_charge_collateral_fees.rs
@@ -87,6 +87,8 @@ pub fn token_charge_collateral_fees(ctx: Context<TokenChargeCollateralFees>) ->
 
     let scaling = asset_usage_scaling * time_scaling;
 
+    let mut total_collateral_fees_in_usd = I80F48::ZERO;
+
     let token_position_count = account.active_token_positions().count();
     for bank_ai in &ctx.remaining_accounts[0..token_position_count] {
         let mut bank = bank_ai.load_mut::<Bank>()?;
@@ -103,16 +105,11 @@ pub fn token_charge_collateral_fees(ctx: Context<TokenChargeCollateralFees>) ->
         let fee = token_balance * scaling * I80F48::from_num(bank.collateral_fee_per_day);
         assert!(fee <= token_balance);
 
-        let is_active = bank.withdraw_without_fee(token_position, fee, now_ts)?;
-        if !is_active {
-            account.deactivate_token_position_and_log(raw_token_index, ctx.accounts.account.key());
-        }
+        let token_info = health_cache.token_info(bank.token_index)?;
 
-        bank.collected_fees_native += fee;
-        bank.collected_collateral_fees += fee;
+        total_collateral_fees_in_usd += fee * token_info.prices.oracle;
 
-        let token_info = health_cache.token_info(bank.token_index)?;
-        let token_position = account.token_position(bank.token_index)?;
+        bank.collected_collateral_fees += fee;
 
         emit_stack(TokenCollateralFeeLog {
             mango_group: ctx.accounts.group.key(),
@@ -122,6 +119,30 @@ pub fn token_charge_collateral_fees(ctx: Context<TokenChargeCollateralFees>) ->
             asset_usage_fraction: asset_usage_scaling.to_bits(),
             price: token_info.prices.oracle.to_bits(),
         });
+    }
+
+    for bank_ai in &ctx.remaining_accounts[0..token_position_count] {
+        let mut bank = bank_ai.load_mut::<Bank>()?;
+
+        let (token_position, raw_token_index) = account.token_position_mut(bank.token_index)?;
+        let token_balance = token_position.native(&bank);
+        let token_info = health_cache.token_info(bank.token_index)?;
+        let health = token_info.health_contribution(HealthType::Maint, token_balance);
+
+        if health >= 0 {
+            continue;
+        }
+
+        let borrow_scaling = (health / total_liab_health).abs();
+        let fee = borrow_scaling * total_collateral_fees_in_usd / token_info.prices.oracle;
+
+        let is_active = bank.withdraw_without_fee(token_position, fee, now_ts)?;
+        if !is_active {
+            account.deactivate_token_position_and_log(raw_token_index, ctx.accounts.account.key());
+        }
+
+        bank.collected_fees_native += fee;
+        let token_position = account.token_position(bank.token_index)?;
 
         emit_stack(TokenBalanceLog {
             mango_group: ctx.accounts.group.key(),
diff --git a/programs/mango-v4/tests/cases/test_collateral_fees.rs b/programs/mango-v4/tests/cases/test_collateral_fees.rs
index c4d852f7f..d67ae698e 100644
--- a/programs/mango-v4/tests/cases/test_collateral_fees.rs
+++ b/programs/mango-v4/tests/cases/test_collateral_fees.rs
@@ -186,12 +186,20 @@ async fn test_collateral_fees() -> Result<(), TransportError> {
         .await
         .unwrap();
     last_time = solana.clock_timestamp().await;
+
+    let fee = 1500.0 * (0.1 * (9.0 / 24.0) * (600.0 / 1200.0));
+    println!("fee -> {}", fee);
     assert_eq_f64!(
         account_position_f64(solana, account, tokens[0].bank).await,
-        1500.0 * (1.0 - 0.1 * (9.0 / 24.0) * (600.0 / 1200.0)),
+        1500.0,
+        0.01
+    );
+    assert_eq_f64!(
+        account_position_f64(solana, account, tokens[1].bank).await,
+        -500.0 - fee,
         0.01
     );
-    let last_balance = account_position_f64(solana, account, tokens[0].bank).await;
+    let last_balance = account_position_f64(solana, account, tokens[1].bank).await;
 
     //
     // TEST: More borrows
@@ -217,9 +225,16 @@ async fn test_collateral_fees() -> Result<(), TransportError> {
         .await
         .unwrap();
     //last_time = solana.clock_timestamp().await;
+    let fee = 1500.0 * 0.1 * (7.0 / 24.0) * ((last_balance.abs() + 100.0) * 1.2 / (1500.0 * 0.8));
+    println!("fee -> {}", fee);
     assert_eq_f64!(
         account_position_f64(solana, account, tokens[0].bank).await,
-        last_balance * (1.0 - 0.1 * (7.0 / 24.0) * (720.0 / (last_balance * 0.8))),
+        1500.0,
+        0.01
+    );
+    assert_eq_f64!(
+        account_position_f64(solana, account, tokens[1].bank).await,
+        -(last_balance.abs() + 100.0) - fee,
         0.01
     );