bigz/expand-rust-test-suite

.devcontainer/Dockerfile

@@ -38,7 +38,8 @@ RUN . $NVM_DIR/nvm.sh && \
     nvm install ${NODE_VERSION} && \
     nvm use ${NODE_VERSION} && \
     nvm alias default node && \
-    npm install -g yarn
+    npm install -g yarn && \
+    yarn add ts-mocha
 
 # Install Solana tools.
 RUN sh -c "$(curl -sSfL https://release.solana.com/v${SOLANA_CLI}/install)"
@@ -47,6 +48,7 @@ RUN sh -c "$(curl -sSfL https://release.solana.com/v${SOLANA_CLI}/install)"
 RUN cargo install --git https://github.com/coral-xyz/anchor avm --locked --force
 RUN avm install ${ANCHOR_CLI} && avm use ${ANCHOR_CLI}
 
+RUN solana-keygen new --no-bip39-passphrase
 
 WORKDIR /workdir
 #be sure to add `/root/.avm/bin` to your PATH to be able to run the installed binaries

programs/amm/src/error.rs

@@ -6,6 +6,8 @@ pub enum ErrorCode {
     AddLiquidityCalculationError,
     #[msg("Error in decimal scale conversion")]
     DecimalScaleError,
+    #[msg("Swap invariant error")]
+    SwapInvariantError,
 }
 
 #[macro_export]
@@ -18,3 +20,26 @@ macro_rules! print_error {
         }
     }};
 }
+
+#[macro_export]
+macro_rules! validate {
+        ($assert:expr, $err:expr) => {{
+            if ($assert) {
+                Ok(())
+            } else {
+                let error_code: ErrorCode = $err;
+                msg!("Error {} thrown at {}:{}", error_code, file!(), line!());
+                Err(error_code)
+            }
+        }};
+        ($assert:expr, $err:expr, $($arg:tt)+) => {{
+        if ($assert) {
+            Ok(())
+        } else {
+            let error_code: ErrorCode = $err;
+            msg!("Error {} thrown at {}:{}", error_code, file!(), line!());
+            msg!($($arg)*);
+            Err(error_code)
+        }
+    }};
+}

programs/amm/src/instructions/swap.rs

@@ -89,17 +89,6 @@ pub fn handler(
 
     amm.update_ltwap()?;
 
-    let base_amount_start = amm.base_amount as u128;
-    let quote_amount_start = amm.quote_amount as u128;
-
-    let k = base_amount_start.checked_mul(quote_amount_start).unwrap();
-
-    let input_amount_minus_fee = input_amount
-        .checked_mul(BPS_SCALE.checked_sub(amm.swap_fee_bps).unwrap())
-        .unwrap()
-        .checked_div(BPS_SCALE)
-        .unwrap() as u128;
-
     let base_mint_key = base_mint.key();
     let quote_mint_key = quote_mint.key();
     let swap_fee_bps_bytes = amm.swap_fee_bps.to_le_bytes();
@@ -113,21 +102,9 @@ pub fn handler(
         amm.bump
     );
 
-    let output_amount = if is_quote_to_base {
-        let temp_quote_amount = quote_amount_start
-            .checked_add(input_amount_minus_fee)
-            .unwrap();
-        let temp_base_amount = k.checked_div(temp_quote_amount).unwrap();
-
-        let output_amount_base = base_amount_start
-            .checked_sub(temp_base_amount)
-            .unwrap()
-            .to_u64()
-            .unwrap();
-
-        amm.quote_amount = amm.quote_amount.checked_add(input_amount).unwrap();
-        amm.base_amount = amm.base_amount.checked_sub(output_amount_base).unwrap();
+    let output_amount = amm.swap(input_amount, is_quote_to_base)?;
 
+    if is_quote_to_base {
         // send user quote tokens to vault
         token_transfer(
             input_amount,
@@ -139,30 +116,14 @@ pub fn handler(
 
         // send vault base tokens to user
         token_transfer_signed(
-            output_amount_base,
+            output_amount,
             token_program,
             vault_ata_base,
             user_ata_base,
             amm,
             seeds,
         )?;
-
-        output_amount_base
     } else {
-        let temp_base_amount = base_amount_start
-            .checked_add(input_amount_minus_fee)
-            .unwrap();
-        let temp_quote_amount = k.checked_div(temp_base_amount).unwrap();
-
-        let output_amount_quote = quote_amount_start
-            .checked_sub(temp_quote_amount)
-            .unwrap()
-            .to_u64()
-            .unwrap();
-
-        amm.base_amount = amm.base_amount.checked_add(input_amount).unwrap();
-        amm.quote_amount = amm.quote_amount.checked_sub(output_amount_quote).unwrap();
-
         // send user base tokens to vault
         token_transfer(
             input_amount,
@@ -174,22 +135,14 @@ pub fn handler(
 
         // send vault quote tokens to user
         token_transfer_signed(
-            output_amount_quote,
+            output_amount,
             token_program,
             vault_ata_quote,
             user_ata_quote,
             amm,
             seeds,
         )?;
-
-        output_amount_quote
-    };
-
-    let new_k = (amm.base_amount as u128)
-        .checked_mul(amm.quote_amount as u128)
-        .unwrap();
-
-    assert!(new_k >= k); // with non-zero fees, k should always increase
+    }
     assert!(output_amount >= output_amount_min);
 
     Ok(())

programs/amm/src/lib.rs

@@ -24,6 +24,10 @@ use crate::error::*;
 use crate::instructions::*;
 use crate::state::*;
 use crate::utils::*;
+use crate::BPS_SCALE;
+
+#[cfg(test)]
+mod tests;
 
 declare_id!("Ens7Gx99whnA8zZm6ZiFnWgGq3x76nXbSmh5gaaJqpAz");
 #[program]

programs/amm/src/state/amm.rs

@@ -2,11 +2,12 @@ use anchor_lang::prelude::*;
 use num_traits::{FromPrimitive, ToPrimitive};
 use rust_decimal::Decimal;
 
-use crate::error::ErrorCode;
 use crate::utils::anchor_decimal::*;
 use crate::utils::*;
-
+use crate::BPS_SCALE;
+use crate::{error::ErrorCode, validate};
 #[account]
+#[derive(Default, Eq, PartialEq, Debug)]
 pub struct Amm {
     pub bump: u8,
 
@@ -122,4 +123,69 @@ impl Amm {
 
         Ok(quote_amount_d / quote_decimal_scale_d)
     }
+    pub fn k(&self) -> Result<u128> {
+        Ok((self.base_amount as u128)
+            .checked_mul(self.quote_amount as u128)
+            .unwrap())
+    }
+
+    pub fn swap(&mut self, input_amount: u64, is_quote_to_base: bool) -> Result<u64> {
+        let base_amount_start = self.base_amount as u128;
+        let quote_amount_start = self.quote_amount as u128;
+
+        let k = base_amount_start.checked_mul(quote_amount_start).unwrap();
+
+        let input_amount_minus_fee = input_amount
+            .checked_mul(BPS_SCALE.checked_sub(self.swap_fee_bps).unwrap())
+            .unwrap()
+            .checked_div(BPS_SCALE)
+            .unwrap() as u128;
+
+        let output_amount = if is_quote_to_base {
+            let temp_quote_amount = quote_amount_start
+                .checked_add(input_amount_minus_fee)
+                .unwrap();
+            let temp_base_amount = k.checked_div(temp_quote_amount).unwrap();
+
+            let output_amount_base = base_amount_start
+                .checked_sub(temp_base_amount)
+                .unwrap()
+                .to_u64()
+                .unwrap();
+
+            self.quote_amount = self.quote_amount.checked_add(input_amount).unwrap();
+            self.base_amount = self.base_amount.checked_sub(output_amount_base).unwrap();
+            output_amount_base
+        } else {
+            let temp_base_amount = base_amount_start
+                .checked_add(input_amount_minus_fee)
+                .unwrap();
+            let temp_quote_amount = k.checked_div(temp_base_amount).unwrap();
+
+            let output_amount_quote = quote_amount_start
+                .checked_sub(temp_quote_amount)
+                .unwrap()
+                .to_u64()
+                .unwrap();
+
+            self.base_amount = self.base_amount.checked_add(input_amount).unwrap();
+            self.quote_amount = self.quote_amount.checked_sub(output_amount_quote).unwrap();
+            output_amount_quote
+        };
+
+        let new_k = (self.base_amount as u128)
+            .checked_mul(self.quote_amount as u128)
+            .unwrap();
+
+        // with non-zero fees, k should always increase
+        validate!(
+            new_k >= k,
+            ErrorCode::SwapInvariantError,
+            "new_k={} is smaller than original k={}",
+            new_k,
+            k
+        )?;
+
+        Ok(output_amount)
+    }
 }

programs/amm/src/tests.rs

@@ -0,0 +1,51 @@
+#[cfg(test)]
+mod simple_amm_tests {
+    use crate::state::*;
+    use crate::utils::*;
+
+    #[test]
+    pub fn base_case_amm() {
+        let mut amm = Amm { ..Amm::default() };
+        assert_eq!(amm.get_ltwap().unwrap(), 0);
+        assert_eq!(amm.swap(1, true).unwrap(), 0);
+        assert_eq!(amm.swap(1, false).unwrap(), 1);
+        assert_eq!(amm.k().unwrap(), 0);
+    }
+
+    #[test]
+    pub fn medium_amm() {
+        let mut amm = Amm {
+            base_amount: 10000,
+            quote_amount: 10000,
+            swap_fee_bps: 1,
+            ..Amm::default()
+        };
+
+        assert_eq!(amm.get_ltwap().unwrap(), 0);
+        assert_eq!(amm.swap(1, true).unwrap(), 0);
+        assert_eq!(amm.swap(1, false).unwrap(), 0);
+        assert_eq!(amm.k().unwrap(), 100020001);
+
+        assert_eq!(amm.swap(100, true).unwrap(), 99);
+        assert_eq!(amm.swap(100, false).unwrap(), 100);
+        assert_eq!(amm.k().unwrap(), 100030002);
+
+        assert_eq!(amm.swap(1000, true).unwrap(), 909);
+        assert_eq!(amm.swap(1000, false).unwrap(), 1089);
+        assert_eq!(amm.k().unwrap(), 100041816);
+    }
+
+    #[test]
+    pub fn medium_amm_with_swap_err() {
+        let mut amm = Amm {
+            base_amount: 10000,
+            quote_amount: 10000,
+            swap_fee_bps: 1,
+            ..Amm::default()
+        };
+
+        // todo?
+        assert!(amm.swap(amm.quote_amount - 1, true).is_err());
+    }
+
+}

programs/amm/src/utils/anchor_decimal.rs

@@ -1,7 +1,7 @@
 use anchor_lang::prelude::*;
 use rust_decimal::Decimal;
 
-#[derive(Debug, Clone, Copy, AnchorSerialize, AnchorDeserialize, PartialEq, Eq)]
+#[derive(Default, Debug, Clone, Copy, AnchorSerialize, AnchorDeserialize, PartialEq, Eq)]
 pub struct AnchorDecimal {
     data: [u8; 16],
 }

programs/amm/src/utils/token.rs

@@ -1,7 +1,13 @@
 use anchor_lang::prelude::*;
 use anchor_spl::token;
 
-pub fn token_mint_signed<'info, P: ToAccountInfo<'info>, M: ToAccountInfo<'info>, T: ToAccountInfo<'info>, A: ToAccountInfo<'info>>(
+pub fn token_mint_signed<
+    'info,
+    P: ToAccountInfo<'info>,
+    M: ToAccountInfo<'info>,
+    T: ToAccountInfo<'info>,
+    A: ToAccountInfo<'info>,
+>(
     amount: u64,
     token_program: &P,
     mint: &M,
@@ -27,7 +33,13 @@ pub fn token_mint_signed<'info, P: ToAccountInfo<'info>, M: ToAccountInfo<'info>
     Ok(())
 }
 
-pub fn token_burn<'info, P: ToAccountInfo<'info>, M: ToAccountInfo<'info>, F: ToAccountInfo<'info>, A: ToAccountInfo<'info>>(
+pub fn token_burn<
+    'info,
+    P: ToAccountInfo<'info>,
+    M: ToAccountInfo<'info>,
+    F: ToAccountInfo<'info>,
+    A: ToAccountInfo<'info>,
+>(
     amount: u64,
     token_program: &P,
     mint: &M,
@@ -51,7 +63,13 @@ pub fn token_burn<'info, P: ToAccountInfo<'info>, M: ToAccountInfo<'info>, F: To
     Ok(())
 }
 
-pub fn token_transfer<'info, P: ToAccountInfo<'info>, F: ToAccountInfo<'info>, T: ToAccountInfo<'info>, A: ToAccountInfo<'info>>(
+pub fn token_transfer<
+    'info,
+    P: ToAccountInfo<'info>,
+    F: ToAccountInfo<'info>,
+    T: ToAccountInfo<'info>,
+    A: ToAccountInfo<'info>,
+>(
     amount: u64,
     token_program: &P,
     from: &F,
@@ -75,7 +93,13 @@ pub fn token_transfer<'info, P: ToAccountInfo<'info>, F: ToAccountInfo<'info>, T
     Ok(())
 }
 
-pub fn token_transfer_signed<'info, P: ToAccountInfo<'info>, F: ToAccountInfo<'info>, T: ToAccountInfo<'info>, A: ToAccountInfo<'info>>(
+pub fn token_transfer_signed<
+    'info,
+    P: ToAccountInfo<'info>,
+    F: ToAccountInfo<'info>,
+    T: ToAccountInfo<'info>,
+    A: ToAccountInfo<'info>,
+>(
     amount: u64,
     token_program: &P,
     from: &F,

programs/autocrat/src/utils/mod.rs

@@ -1,11 +1,11 @@
-use anchor_lang::prelude::*;
 use crate::error::ErrorCode;
+use anchor_lang::prelude::*;
 
-pub use token::*;
 pub use seeds::*;
+pub use token::*;
 
-pub mod token;
 pub mod seeds;
+pub mod token;
 
 use crate::state::*;
 
@@ -36,7 +36,7 @@ pub fn get_decimal_scale_u64(decimals: u8) -> Result<u64> {
 
 pub fn get_instructions_size(instructions: &Vec<ProposalInstruction>) -> usize {
     instructions.iter().fold(0, |accumulator, ix| {
-        accumulator + 
+        accumulator +
         32 + // program id
         4 + // accounts vec prefix
         ix.accounts.len() * (32 + 1 + 1) + // pubkey + 2 bools per account