From 9351d0652df70bc5f8753f32a3d5d4666339c7f9 Mon Sep 17 00:00:00 2001
From: Christian Kamm <mail@ckamm.de>
Date: Thu, 17 Mar 2022 12:05:07 +0100
Subject: [PATCH] Serum: health based on worst case order outcome

---
 programs/mango-v4/src/state/health.rs | 75 +++++++++++++++++++--------
 1 file changed, 54 insertions(+), 21 deletions(-)

diff --git a/programs/mango-v4/src/state/health.rs b/programs/mango-v4/src/state/health.rs
index 4f04db135..40448c75d 100644
--- a/programs/mango-v4/src/state/health.rs
+++ b/programs/mango-v4/src/state/health.rs
@@ -29,6 +29,25 @@ struct TokenInfo<'a> {
     balance: I80F48,
 }
 
+fn health_contribution(bank: &Bank, price: I80F48, balance: I80F48) -> Result<I80F48> {
+    Ok(if balance.is_negative() {
+        cm!(balance * price * bank.init_liab_weight)
+    } else {
+        cm!(balance * price * bank.init_asset_weight)
+    })
+}
+
+fn pair_health(
+    info1: &TokenInfo,
+    balance1: I80F48,
+    info2: &TokenInfo,
+    balance2: I80F48,
+) -> Result<I80F48> {
+    let health1 = health_contribution(&info1.bank, info1.oracle_price, balance1)?;
+    let health2 = health_contribution(&info2.bank, info2.oracle_price, balance2)?;
+    Ok(cm!(health1 + health2))
+}
+
 fn strip_dex_padding<'a>(acc: &'a AccountInfo) -> Result<Ref<'a, [u8]>> {
     require!(acc.data_len() >= 12, MangoError::SomeError);
     let unpadded_data: Ref<[u8]> = Ref::map(acc.try_borrow_data()?, |data| {
@@ -103,38 +122,52 @@ fn compute_health_detail(
             .position(|ti| ti.bank.token_index == serum_account.quote_token_index)
             .ok_or(error!(MangoError::SomeError))?;
 
+        let base_info = &token_infos[base_index];
+        let quote_info = &token_infos[quote_index];
+        let mut base = base_info.balance;
+        let mut quote = quote_info.balance;
+
         let oo = load_open_orders(oo_ai)?;
 
         // add the amounts that are freely settleable
-        token_infos[base_index].balance += I80F48::from_num(oo.native_coin_free);
-        token_infos[quote_index].balance +=
-            I80F48::from_num(oo.native_pc_free + oo.referrer_rebates_accrued);
+        let base_free = I80F48::from_num(oo.native_coin_free);
+        let quote_free = I80F48::from_num(cm!(oo.native_pc_free + oo.referrer_rebates_accrued));
+        base = cm!(base + base_free);
+        quote = cm!(quote + quote_free);
 
         // for the amounts that are reserved for orders, compute the worst case for health
         // by checking if everything-is-base or everything-is-quote produces worse
         // outcomes
-        // TODO: that kind of approach may no longer be possible with each
-        // market potentially having two different tokens involved?
-        let reserved_base = oo.native_coin_total - oo.native_coin_free;
-        let reserved_quote = oo.native_pc_total - oo.native_pc_free;
-        // TODO: do it, this is just a stub
-        token_infos[base_index].balance += I80F48::from_num(reserved_base);
-        token_infos[quote_index].balance += I80F48::from_num(reserved_quote);
+        let reserved_base = I80F48::from_num(cm!(oo.native_coin_total - oo.native_coin_free));
+        let reserved_quote = I80F48::from_num(cm!(oo.native_pc_total - oo.native_pc_free));
+        let all_in_base = cm!(base
+            + reserved_base
+            + reserved_quote * quote_info.oracle_price / base_info.oracle_price);
+        let all_in_quote = cm!(quote
+            + reserved_quote
+            + reserved_base * base_info.oracle_price / quote_info.oracle_price);
+        if pair_health(base_info, all_in_base, quote_info, quote)?
+            < pair_health(base_info, base, quote_info, all_in_quote)?
+        {
+            base = all_in_base;
+        } else {
+            quote = all_in_quote;
+        }
+
+        token_infos[base_index].balance = base;
+        token_infos[quote_index].balance = quote;
     }
 
     // convert the token balance to health
-    let mut asset_health = I80F48::ZERO;
-    let mut liability_health = I80F48::ZERO; // positive
+    let mut health = I80F48::ZERO;
     for token_info in token_infos.iter() {
-        let bank = &token_info.bank;
-        if token_info.balance.is_negative() {
-            liability_health = cm!(liability_health
-                - bank.init_liab_weight * token_info.balance * token_info.oracle_price);
-        } else {
-            asset_health = cm!(asset_health
-                + bank.init_asset_weight * token_info.balance * token_info.oracle_price);
-        }
+        let contrib = health_contribution(
+            &token_info.bank,
+            token_info.oracle_price,
+            token_info.balance,
+        )?;
+        health = cm!(health + contrib);
     }
 
-    Ok(cm!(asset_health - liability_health))
+    Ok(health)
 }