From dafaf5a2b0cc26ccb35138f7553563477a0bfec1 Mon Sep 17 00:00:00 2001
From: metaproph3t <metaproph3t@protonmail.com>
Date: Sat, 24 Aug 2024 00:00:00 +0000
Subject: [PATCH] Add an assets >= liabilities invariant

---
 .../src/instructions/merge_tokens.rs          |  6 +++
 .../src/instructions/redeem_tokens.rs         |  6 +++
 .../src/instructions/split_tokens.rs          |  6 +++
 .../src/state/conditional_vault.rs            | 41 ++++++++++++++++++-
 4 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/programs/conditional_vault/src/instructions/merge_tokens.rs b/programs/conditional_vault/src/instructions/merge_tokens.rs
index 4d4ef645..14b920e5 100644
--- a/programs/conditional_vault/src/instructions/merge_tokens.rs
+++ b/programs/conditional_vault/src/instructions/merge_tokens.rs
@@ -90,6 +90,12 @@ impl<'info, 'c: 'info> InteractWithNewVault<'info> {
             assert!(account.amount == *expected_balance);
         }
 
+        ctx.accounts.vault.invariant(
+            &ctx.accounts.question,
+            conditional_token_mints.iter().map(|mint| mint.supply).collect::<Vec<u64>>(),
+            ctx.accounts.vault_underlying_token_account.amount,
+        )?;
+
         Ok(())
     }
 }
diff --git a/programs/conditional_vault/src/instructions/redeem_tokens.rs b/programs/conditional_vault/src/instructions/redeem_tokens.rs
index cea47180..96b71240 100644
--- a/programs/conditional_vault/src/instructions/redeem_tokens.rs
+++ b/programs/conditional_vault/src/instructions/redeem_tokens.rs
@@ -108,6 +108,12 @@ impl<'info, 'c: 'info> InteractWithNewVault<'info> {
             assert!(mint.supply == *expected_supply);
         }
 
+        ctx.accounts.vault.invariant(
+            &ctx.accounts.question,
+            conditional_token_mints.iter().map(|mint| mint.supply).collect::<Vec<u64>>(),
+            ctx.accounts.vault_underlying_token_account.amount,
+        )?;
+
         Ok(())
     }
 }
diff --git a/programs/conditional_vault/src/instructions/split_tokens.rs b/programs/conditional_vault/src/instructions/split_tokens.rs
index 031fb56b..3f39a1d6 100644
--- a/programs/conditional_vault/src/instructions/split_tokens.rs
+++ b/programs/conditional_vault/src/instructions/split_tokens.rs
@@ -73,6 +73,12 @@ impl<'info, 'c: 'info> InteractWithNewVault<'info> {
             assert!(acc.amount == pre_conditional_user_balances[i] + amount);
         }
 
+        ctx.accounts.vault.invariant(
+            &ctx.accounts.question,
+            conditional_token_mints.iter().map(|mint| mint.supply).collect::<Vec<u64>>(),
+            ctx.accounts.vault_underlying_token_account.amount,
+        )?;
+
         Ok(())
     }
 }
diff --git a/programs/conditional_vault/src/state/conditional_vault.rs b/programs/conditional_vault/src/state/conditional_vault.rs
index 34fe9872..171cd07a 100644
--- a/programs/conditional_vault/src/state/conditional_vault.rs
+++ b/programs/conditional_vault/src/state/conditional_vault.rs
@@ -17,6 +17,45 @@ pub struct NewConditionalVault {
     pub decimals: u8,
 }
 
+impl NewConditionalVault {
+    /// Checks that the vault's assets are always greater than its potential
+    /// liabilities. Should be called anytime you mint or burn conditional
+    /// tokens.
+    ///
+    /// `conditional_token_supplies` should be in the same order as
+    /// `vault.conditional_token_mints`.
+    pub fn invariant(
+        &self,
+        question: &Question,
+        conditional_token_supplies: Vec<u64>,
+        vault_underlying_balance: u64,
+    ) -> Result<()> {
+        // if the question isn't resolved, the vault should have more underlying
+        // tokens than ANY conditional token mint's supply
+
+        // if the question is resolved, the vault should have more underlying
+        // tokens than the sum of the conditional token mint's supplies multiplied
+        // by their respective payouts
+
+        let max_possible_liability = if !question.is_resolved() {
+            // safe because conditional_token_supplies is non-empty
+            *conditional_token_supplies.iter().max().unwrap()
+        } else {
+            conditional_token_supplies
+                .iter()
+                .enumerate()
+                .map(|(i, supply)| {
+                    *supply * question.payout_numerators[i] as u64 / question.payout_denominator as u64
+                })
+                .sum::<u64>()
+        };
+
+        assert!(vault_underlying_balance >= max_possible_liability);
+
+        Ok(())
+    }
+}
+
 #[macro_export]
 macro_rules! generate_new_vault_seeds {
     ($vault:expr) => {{
@@ -56,4 +95,4 @@ macro_rules! generate_vault_seeds {
             &[$vault.pda_bump],
         ]
     }};
-}
+}
\ No newline at end of file