refactor AMM

programs/amm/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "amm_program"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+nssa_core.workspace = true
+token_core.workspace = true
+amm_core.workspace = true

programs/amm/core/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "amm_core"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+nssa_core.workspace = true
+serde.workspace = true
+borsh.workspace = true
+risc0-zkvm.workspace = true

programs/amm/core/src/lib.rs

@@ -1,330 +1,248 @@
-//! This crate contains core data structures and utilities for the Token Program.
+//! This crate contains core data structures and utilities for the AMM Program.
 
-use borsh::{BorshDeserialize, BorshSerialize};
-use nssa_core::account::{AccountId, Data};
+use nssa_core::{
+    account::{AccountId, Data},
+    program::{PdaSeed, ProgramId},
+};
 use serde::{Deserialize, Serialize};
 
-pub const CURRENT_VERSION: u8 = 1;
-
-
-// The AMM program has five functions (four directly accessible via instructions):
-// 1. New AMM definition. Arguments to this function are:
-//      * Seven accounts: [amm_pool, vault_holding_a, vault_holding_b, pool_lp, user_holding_a,
-//        user_holding_b, user_holding_lp]. For new AMM Pool: amm_pool, vault_holding_a,
-//        vault_holding_b, pool_lp and user_holding_lp are default accounts. amm_pool is a default
-//        account that will initiate the amm definition account values vault_holding_a is a token
-//        holding account for token a vault_holding_b is a token holding account for token b pool_lp
-//        is a token holding account for the pool's lp token user_holding_a is a token holding
-//        account for token a user_holding_b is a token holding account for token b user_holding_lp
-//        is a token holding account for lp token
-//      * PDA remark: Accounts amm_pool, vault_holding_a, vault_holding_b and pool_lp are PDA. The
-//        AccountId for these accounts must be computed using: amm_pool AccountId <-
-//        compute_pool_pda vault_holding_a, vault_holding_b <- compute_vault_pda pool_lp
-//        <-compute_liquidity_token_pda
-//      * Requires authorization: user_holding_a, user_holding_b
-//      * An instruction data of 65-bytes, indicating the initial amm reserves' balances and
-//        token_program_id with the following layout: [0x00 || array of balances (little-endian 16
-//        bytes) || AMM_PROGRAM_ID)]
-//      * Internally, calls compute_liquidity_token_pda_seed, compute_vault_pda_seed to authorize
-//        transfers.
-//      * Internally, calls compute_pool_da, compute_vault_pda and compute_vault_pda to check
-//        various AccountIds are correct.
-// 3. Add liquidity Arguments to this function are:
-//      * Seven accounts: [amm_pool, vault_holding_a, vault_holding_b, pool_lp, user_holding_a,
-//        user_holding_a, user_holding_lp].
-//      * Requires authorization: user_holding_a, user_holding_b
-//      * An instruction data byte string of length 49, amounts for minimum amount of liquidity from
-//        add (min_amount_lp),
-//      * max amount added for each token (max_amount_a and max_amount_b); indicate [0x02 || array
-//        of of balances (little-endian 16 bytes)].
-//      * Internally, calls compute_liquidity_token_pda_seed to compute liquidity pool PDA seed.
-// 4. Remove liquidity
-//      * Seven accounts: [amm_pool, vault_holding_a, vault_holding_b, pool_lp, user_holding_a,
-//        user_holding_a, user_holding_lp].
-//      * Requires authorization: user_holding_lp
-//      * An instruction data byte string of length 49, amounts for minimum amount of liquidity to
-//        redeem (balance_lp),
-//      * minimum balance of each token to remove (min_amount_a and min_amount_b); indicate [0x03 ||
-//        array of balances (little-endian 16 bytes)].
-//      * Internally, calls compute_vault_pda_seed to compute vault_a and vault_b's PDA seed.
-
-
-
 /// AMM Program Instruction.
 #[derive(Serialize, Deserialize)]
 pub enum Instruction {
-
-    /// Create a new fungible token definition without metadata.
+    /// Initializes a new Pool (or re-initializes an inactive Pool).
     ///
     /// Required accounts:
-    /// - Token Definition account (uninitialized),
-    /// - Token Holding account (uninitialized).
-    NewDefinition { name: String, total_supply: u128 },
+    /// - AMM Pool
+    /// - Vault Holding Account for Token A
+    /// - Vault Holding Account for Token B
+    /// - Pool Liquidity Token Definition
+    /// - User Holding Account for Token A (authorized)
+    /// - User Holding Account for Token B (authorized)
+    /// - User Holding Account for Pool Liquidity
+    NewDefinition {
+        token_a_amount: u128,
+        token_b_amount: u128,
+        amm_program_id: ProgramId,
+    },
 
-    /// Create a new fungible or non-fungible token definition with metadata.
+    /// Adds liquidity to the Pool
     ///
     /// Required accounts:
-    /// - Token Definition account (uninitialized),
-    /// - Token Holding account (uninitialized),
-    /// - Token Metadata account (uninitialized).
-    NewDefinitionWithMetadata {
-        new_definition: NewTokenDefinition,
-        /// Boxed to avoid large enum variant size
-        metadata: Box<NewTokenMetadata>,
+    /// - AMM Pool (initialized)
+    /// - Vault Holding Account for Token A (initialized)
+    /// - Vault Holding Account for Token B (initialized)
+    /// - Pool Liquidity Token Definition (initialized)
+    /// - User Holding Account for Token A (authorized)
+    /// - User Holding Account for Token B (authorized)
+    /// - User Holding Account for Pool Liquidity
+    AddLiquidity {
+        min_amount_liquidity: u128,
+        max_amount_to_add_token_a: u128,
+        max_amount_to_add_token_b: u128,
     },
 
-
-    
-    /// Initialize a token holding account for a given token definition.
+    /// Removes liquidity from the Pool
     ///
     /// Required accounts:
-    /// - Token Definition account (initialized),
-    /// - Token Holding account (uninitialized),
-    InitializeAccount,
+    /// - AMM Pool (initialized)
+    /// - Vault Holding Account for Token A (initialized)
+    /// - Vault Holding Account for Token B (initialized)
+    /// - Pool Liquidity Token Definition (initialized)
+    /// - User Holding Account for Token A (initialized)
+    /// - User Holding Account for Token B (initialized)
+    /// - User Holding Account for Pool Liquidity (authorized)
+    RemoveLiquidity {
+        remove_liquidity_amount: u128,
+        min_amount_to_remove_token_a: u128,
+        min_amount_to_remove_token_b: u128,
+    },
 
-
-// 2. Swap assets Arguments to this function are:
-//      * Five accounts: [amm_pool, vault_holding_a, vault_holding_b, user_holding_a,
-//        user_holding_b].
-//      * Requires authorization: user holding account associated to TOKEN_DEFINITION_ID (either
-//        user_holding_a or user_holding_b)
-//      * An instruction data byte string of length 65, indicating which token type to swap,
-//        quantity of tokens put into the swap (of type TOKEN_DEFINITION_ID) and min_amount_out.
-//        [0x01 || amount (little-endian 16 bytes) || TOKEN_DEFINITION_ID].
-//      * Internally, calls swap logic.
-//              * Four accounts: [user_deposit, vault_deposit, vault_withdraw, user_withdraw].
-//                user_deposit and vault_deposit define deposit transaction. vault_withdraw and
-//                user_withdraw define withdraw transaction.
-//              * deposit_amount is the amount for user_deposit -> vault_deposit transfer.
-//              * reserve_amounts is the pool's reserves; used to compute the withdraw amount.
-//              * Outputs the token transfers as a Vec<ChainedCall> and the withdraw amount.
-
-
-    ///TODO update description
-    /// Burn tokens from the holder's account.
+    /// Swap some quantity of Tokens (either Token A or Token B)
+    /// while maintaining the Pool constant product.
     ///
     /// Required accounts:
-    /// - Token Definition account (initialized),
-    /// - Token Holding account (authorized).
-    Burn { amount_to_burn: u128 },
-
-
-/*
-fn add_liquidity(
-    pre_states: &[AccountWithMetadata],
-    balances: &[u128],
-) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
-    if pre_states.len() != 7 {
-        panic!("Invalid number of input accounts");
-    }
-
-    let pool = &pre_states[0];
-    let vault_a = &pre_states[1];
-    let vault_b = &pre_states[2];
-    let pool_definition_lp = &pre_states[3];
-    let user_holding_a = &pre_states[4];
-    let user_holding_b = &pre_states[5];
-    let user_holding_lp = &pre_states[6];
-
-        let min_amount_lp = balances[0];
-    let max_amount_a = balances[1];
-    let max_amount_b = balances[2];
-
-*/
-
-    ///TODO: update for add
-    /// Mint new tokens to the holder's account.
-    ///
-    /// Required accounts:
-    /// - Token Definition account (authorized),
-    /// - Token Holding account.
-    AddLiquidity { min_amount_liquidity: u128, max_amount_to_add_token_a: u128, max_amount_to_add_token_b: u128 },
-
-
-    /*
-    fn remove_liquidity(
-    pre_states: &[AccountWithMetadata],
-    amounts: &[u128],
-) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
-    if pre_states.len() != 7 {
-        panic!("Invalid number of input accounts");
-    }
-
-    let pool = &pre_states[0];
-    let vault_a = &pre_states[1];
-    let vault_b = &pre_states[2];
-    let pool_definition_lp = &pre_states[3];
-    let user_holding_a = &pre_states[4];
-    let user_holding_b = &pre_states[5];
-    let user_holding_lp = &pre_states[6];
-
-    if amounts.len() != 3 {
-        panic!("Invalid number of balances");
-    }
-
-    let amount_lp = amounts[0];
-    let amount_min_a = amounts[1];
-    let amount_min_b = amounts[2];
-    
-     */
-    //TODO types
-    RemoveLiquidity { remove_liquidity_amount: u128, min_amount_to_remove_token_a: u128, min_amount_to_remove_token_b: u128 }
-}
-
-/*
-
-#[derive(Serialize, Deserialize)]
-pub enum NewTokenDefinition {
-    Fungible { name: String, total_supply: u128 },
-    NonFungible { name: String, print_balance: u128 },
-}
-
-#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
-pub enum TokenDefinition {
-    Fungible {
-        name: String,
-        total_supply: u128,
-        metadata_id: Option<AccountId>,
-    },
-    NonFungible {
-        name: String,
-        metadata_id: AccountId,
+    /// - AMM Pool (initialized)
+    /// - Vault Holding Account for Token A (initialized)
+    /// - Vault Holding Account for Token B (initialized)
+    /// - User Holding Account for Token A
+    /// - User Holding Account for Token B
+    /// Either User Holding Account for Token A or Token B is authorized.
+    Swap {
+        swap_amount_in: u128,
+        min_amount_out: u128,
+        token_definition_id_in: AccountId,
     },
 }
 
-impl TryFrom<&Data> for TokenDefinition {
-    type Error = std::io::Error;
+const POOL_DEFINITION_DATA_SIZE: usize = 225;
 
-    fn try_from(data: &Data) -> Result<Self, Self::Error> {
-        TokenDefinition::try_from_slice(data.as_ref())
-    }
+#[derive(Clone, Default)]
+pub struct PoolDefinition {
+    pub definition_token_a_id: AccountId,
+    pub definition_token_b_id: AccountId,
+    pub vault_a_id: AccountId,
+    pub vault_b_id: AccountId,
+    pub liquidity_pool_id: AccountId,
+    pub liquidity_pool_supply: u128,
+    pub reserve_a: u128,
+    pub reserve_b: u128,
+    /// Fees are currently not used
+    pub fees: u128,
+    /// A pool becomes inactive (active = false)
+    /// once all of its liquidity has been removed (e.g., reserves are emptied and
+    /// liquidity_pool_supply = 0)
+    pub active: bool,
 }
 
-impl From<&TokenDefinition> for Data {
-    fn from(definition: &TokenDefinition) -> Self {
-        // Using size_of_val as size hint for Vec allocation
-        let mut data = Vec::with_capacity(std::mem::size_of_val(definition));
+impl PoolDefinition {
+    pub fn into_data(self) -> Data {
+        let mut bytes = [0; POOL_DEFINITION_DATA_SIZE];
+        bytes[0..32].copy_from_slice(&self.definition_token_a_id.to_bytes());
+        bytes[32..64].copy_from_slice(&self.definition_token_b_id.to_bytes());
+        bytes[64..96].copy_from_slice(&self.vault_a_id.to_bytes());
+        bytes[96..128].copy_from_slice(&self.vault_b_id.to_bytes());
+        bytes[128..160].copy_from_slice(&self.liquidity_pool_id.to_bytes());
+        bytes[160..176].copy_from_slice(&self.liquidity_pool_supply.to_le_bytes());
+        bytes[176..192].copy_from_slice(&self.reserve_a.to_le_bytes());
+        bytes[192..208].copy_from_slice(&self.reserve_b.to_le_bytes());
+        bytes[208..224].copy_from_slice(&self.fees.to_le_bytes());
+        bytes[224] = self.active as u8;
 
-        BorshSerialize::serialize(definition, &mut data)
-            .expect("Serialization to Vec should not fail");
-
-        Data::try_from(data).expect("Token definition encoded data should fit into Data")
-    }
-}
-
-#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
-pub enum TokenHolding {
-    Fungible {
-        definition_id: AccountId,
-        balance: u128,
-    },
-    NftMaster {
-        definition_id: AccountId,
-        /// The amount of printed copies left - 1 (1 reserved for master copy itself).
-        print_balance: u128,
-    },
-    NftPrintedCopy {
-        definition_id: AccountId,
-        /// Whether nft is owned by the holder.
-        owned: bool,
-    },
-}
-
-impl TokenHolding {
-    pub fn zeroized_clone_from(other: &Self) -> Self {
-        match other {
-            TokenHolding::Fungible { definition_id, .. } => TokenHolding::Fungible {
-                definition_id: *definition_id,
-                balance: 0,
-            },
-            TokenHolding::NftMaster { definition_id, .. } => TokenHolding::NftMaster {
-                definition_id: *definition_id,
-                print_balance: 0,
-            },
-            TokenHolding::NftPrintedCopy { definition_id, .. } => TokenHolding::NftPrintedCopy {
-                definition_id: *definition_id,
-                owned: false,
-            },
-        }
+        bytes
+            .to_vec()
+            .try_into()
+            .expect("225 bytes should fit into Data")
     }
 
-    pub fn zeroized_from_definition(
-        definition_id: AccountId,
-        definition: &TokenDefinition,
-    ) -> Self {
-        match definition {
-            TokenDefinition::Fungible { .. } => TokenHolding::Fungible {
-                definition_id,
-                balance: 0,
-            },
-            TokenDefinition::NonFungible { .. } => TokenHolding::NftPrintedCopy {
-                definition_id,
-                owned: false,
-            },
-        }
-    }
+    pub fn parse(data: &[u8]) -> Option<Self> {
+        if data.len() != POOL_DEFINITION_DATA_SIZE {
+            None
+        } else {
+            let definition_token_a_id = AccountId::new(data[0..32].try_into().expect("Parse data: The AMM program must be provided a valid AccountId for Token A definition"));
+            let definition_token_b_id = AccountId::new(data[32..64].try_into().expect("Parse data: The AMM program must be provided a valid AccountId for Vault B definition"));
+            let vault_a_id = AccountId::new(data[64..96].try_into().expect(
+                "Parse data: The AMM program must be provided a valid AccountId for Vault A",
+            ));
+            let vault_b_id = AccountId::new(data[96..128].try_into().expect(
+                "Parse data: The AMM program must be provided a valid AccountId for Vault B",
+            ));
+            let liquidity_pool_id = AccountId::new(data[128..160].try_into().expect("Parse data: The AMM program must be provided a valid AccountId for Token liquidity pool definition"));
+            let liquidity_pool_supply = u128::from_le_bytes(data[160..176].try_into().expect(
+                "Parse data: The AMM program must be provided a valid u128 for liquidity cap",
+            ));
+            let reserve_a = u128::from_le_bytes(data[176..192].try_into().expect(
+                "Parse data: The AMM program must be provided a valid u128 for reserve A balance",
+            ));
+            let reserve_b = u128::from_le_bytes(data[192..208].try_into().expect(
+                "Parse data: The AMM program must be provided a valid u128 for reserve B balance",
+            ));
+            let fees = u128::from_le_bytes(
+                data[208..224]
+                    .try_into()
+                    .expect("Parse data: The AMM program must be provided a valid u128 for fees"),
+            );
 
-    pub fn definition_id(&self) -> AccountId {
-        match self {
-            TokenHolding::Fungible { definition_id, .. } => *definition_id,
-            TokenHolding::NftMaster { definition_id, .. } => *definition_id,
-            TokenHolding::NftPrintedCopy { definition_id, .. } => *definition_id,
+            let active = match data[224] {
+                0 => false,
+                1 => true,
+                _ => panic!("Parse data: The AMM program must be provided a valid bool for active"),
+            };
+
+            Some(Self {
+                definition_token_a_id,
+                definition_token_b_id,
+                vault_a_id,
+                vault_b_id,
+                liquidity_pool_id,
+                liquidity_pool_supply,
+                reserve_a,
+                reserve_b,
+                fees,
+                active,
+            })
         }
     }
 }
 
-impl TryFrom<&Data> for TokenHolding {
-    type Error = std::io::Error;
-
-    fn try_from(data: &Data) -> Result<Self, Self::Error> {
-        TokenHolding::try_from_slice(data.as_ref())
-    }
+pub fn compute_pool_pda(
+    amm_program_id: ProgramId,
+    definition_token_a_id: AccountId,
+    definition_token_b_id: AccountId,
+) -> AccountId {
+    AccountId::from((
+        &amm_program_id,
+        &compute_pool_pda_seed(definition_token_a_id, definition_token_b_id),
+    ))
 }
 
-impl From<&TokenHolding> for Data {
-    fn from(holding: &TokenHolding) -> Self {
-        // Using size_of_val as size hint for Vec allocation
-        let mut data = Vec::with_capacity(std::mem::size_of_val(holding));
+pub fn compute_pool_pda_seed(
+    definition_token_a_id: AccountId,
+    definition_token_b_id: AccountId,
+) -> PdaSeed {
+    use risc0_zkvm::sha::{Impl, Sha256};
 
-        BorshSerialize::serialize(holding, &mut data)
-            .expect("Serialization to Vec should not fail");
+    let (token_1, token_2) = match definition_token_a_id
+        .value()
+        .cmp(definition_token_b_id.value())
+    {
+        std::cmp::Ordering::Less => (definition_token_b_id, definition_token_a_id),
+        std::cmp::Ordering::Greater => (definition_token_a_id, definition_token_b_id),
+        std::cmp::Ordering::Equal => panic!("Definitions match"),
+    };
 
-        Data::try_from(data).expect("Token holding encoded data should fit into Data")
-    }
+    let mut bytes = [0; 64];
+    bytes[0..32].copy_from_slice(&token_1.to_bytes());
+    bytes[32..].copy_from_slice(&token_2.to_bytes());
+
+    PdaSeed::new(
+        Impl::hash_bytes(&bytes)
+            .as_bytes()
+            .try_into()
+            .expect("Hash output must be exactly 32 bytes long"),
+    )
 }
 
-#[derive(Serialize, Deserialize)]
-pub struct NewTokenMetadata {
-    pub uri: String,
-    pub creators: String,
+pub fn compute_vault_pda(
+    amm_program_id: ProgramId,
+    pool_id: AccountId,
+    definition_token_id: AccountId,
+) -> AccountId {
+    AccountId::from((
+        &amm_program_id,
+        &compute_vault_pda_seed(pool_id, definition_token_id),
+    ))
 }
 
-#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
-pub struct TokenMetadata {
-    pub version: u8,
-    pub definition_id: AccountId,
-    pub uri: String,
-    pub creators: String,
-    /// Block id
-    pub primary_sale_date: u64,
+pub fn compute_vault_pda_seed(pool_id: AccountId, definition_token_id: AccountId) -> PdaSeed {
+    use risc0_zkvm::sha::{Impl, Sha256};
+
+    let mut bytes = [0; 64];
+    bytes[0..32].copy_from_slice(&pool_id.to_bytes());
+    bytes[32..].copy_from_slice(&definition_token_id.to_bytes());
+
+    PdaSeed::new(
+        Impl::hash_bytes(&bytes)
+            .as_bytes()
+            .try_into()
+            .expect("Hash output must be exactly 32 bytes long"),
+    )
 }
 
-impl TryFrom<&Data> for TokenMetadata {
-    type Error = std::io::Error;
-
-    fn try_from(data: &Data) -> Result<Self, Self::Error> {
-        TokenMetadata::try_from_slice(data.as_ref())
-    }
+pub fn compute_liquidity_token_pda(amm_program_id: ProgramId, pool_id: AccountId) -> AccountId {
+    AccountId::from((&amm_program_id, &compute_liquidity_token_pda_seed(pool_id)))
 }
 
-impl From<&TokenMetadata> for Data {
-    fn from(metadata: &TokenMetadata) -> Self {
-        // Using size_of_val as size hint for Vec allocation
-        let mut data = Vec::with_capacity(std::mem::size_of_val(metadata));
+pub fn compute_liquidity_token_pda_seed(pool_id: AccountId) -> PdaSeed {
+    use risc0_zkvm::sha::{Impl, Sha256};
 
-        BorshSerialize::serialize(metadata, &mut data)
-            .expect("Serialization to Vec should not fail");
+    let mut bytes = [0; 64];
+    bytes[0..32].copy_from_slice(&pool_id.to_bytes());
+    bytes[32..].copy_from_slice(&[0; 32]);
 
-        Data::try_from(data).expect("Token metadata encoded data should fit into Data")
-    }
-}*/
+    PdaSeed::new(
+        Impl::hash_bytes(&bytes)
+            .as_bytes()
+            .try_into()
+            .expect("Hash output must be exactly 32 bytes long"),
+    )
+}

programs/amm/src/add.rs

@@ -0,0 +1,169 @@
+use amm_core::{PoolDefinition, compute_liquidity_token_pda_seed};
+use nssa_core::{
+    account::AccountWithMetadata,
+    program::{AccountPostState, ChainedCall},
+};
+
+#[expect(clippy::too_many_arguments, reason = "TODO: Fix later")]
+pub fn add_liquidity(
+    pool: AccountWithMetadata,
+    vault_a: AccountWithMetadata,
+    vault_b: AccountWithMetadata,
+    pool_definition_lp: AccountWithMetadata,
+    user_holding_a: AccountWithMetadata,
+    user_holding_b: AccountWithMetadata,
+    user_holding_lp: AccountWithMetadata,
+    min_amount_liquidity: u128,
+    max_amount_to_add_token_a: u128,
+    max_amount_to_add_token_b: u128,
+) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
+    // 1. Fetch Pool state
+    let pool_def_data = PoolDefinition::parse(&pool.account.data)
+        .expect("Add liquidity: AMM Program expects valid Pool Definition Account");
+    if vault_a.account_id != pool_def_data.vault_a_id {
+        panic!("Vault A was not provided");
+    }
+
+    if pool_def_data.liquidity_pool_id != pool_definition_lp.account_id {
+        panic!("LP definition mismatch");
+    }
+
+    if vault_b.account_id != pool_def_data.vault_b_id {
+        panic!("Vault B was not provided");
+    }
+
+    if max_amount_to_add_token_a == 0 || max_amount_to_add_token_b == 0 {
+        panic!("Both max-balances must be nonzero");
+    }
+
+    if min_amount_liquidity == 0 {
+        panic!("Min-lp must be nonzero");
+    }
+
+    // 2. Determine deposit amount
+    let vault_b_token_holding = token_core::TokenHolding::try_from(&vault_b.account.data)
+        .expect("Add liquidity: AMM Program expects valid Token Holding Account for Vault B");
+    let token_core::TokenHolding::Fungible {
+        definition_id: _,
+        balance: vault_b_balance,
+    } = vault_b_token_holding
+    else {
+        panic!(
+            "Add liquidity: AMM Program expects valid Fungible Token Holding Account for Vault B"
+        );
+    };
+
+    let vault_a_token_holding = token_core::TokenHolding::try_from(&vault_a.account.data)
+        .expect("Add liquidity: AMM Program expects valid Token Holding Account for Vault A");
+    let token_core::TokenHolding::Fungible {
+        definition_id: _,
+        balance: vault_a_balance,
+    } = vault_a_token_holding
+    else {
+        panic!(
+            "Add liquidity: AMM Program expects valid Fungible Token Holding Account for Vault A"
+        );
+    };
+
+    if pool_def_data.reserve_a == 0 || pool_def_data.reserve_b == 0 {
+        panic!("Reserves must be nonzero");
+    }
+
+    if vault_a_balance < pool_def_data.reserve_a || vault_b_balance < pool_def_data.reserve_b {
+        panic!("Vaults' balances must be at least the reserve amounts");
+    }
+
+    // Calculate actual_amounts
+    let ideal_a: u128 =
+        (pool_def_data.reserve_a * max_amount_to_add_token_b) / pool_def_data.reserve_b;
+    let ideal_b: u128 =
+        (pool_def_data.reserve_b * max_amount_to_add_token_a) / pool_def_data.reserve_a;
+
+    let actual_amount_a = if ideal_a > max_amount_to_add_token_a {
+        max_amount_to_add_token_a
+    } else {
+        ideal_a
+    };
+    let actual_amount_b = if ideal_b > max_amount_to_add_token_b {
+        max_amount_to_add_token_b
+    } else {
+        ideal_b
+    };
+
+    // 3. Validate amounts
+    if max_amount_to_add_token_a < actual_amount_a || max_amount_to_add_token_b < actual_amount_b {
+        panic!("Actual trade amounts cannot exceed max_amounts");
+    }
+
+    if actual_amount_a == 0 || actual_amount_b == 0 {
+        panic!("A trade amount is 0");
+    }
+
+    // 4. Calculate LP to mint
+    let delta_lp = std::cmp::min(
+        pool_def_data.liquidity_pool_supply * actual_amount_a / pool_def_data.reserve_a,
+        pool_def_data.liquidity_pool_supply * actual_amount_b / pool_def_data.reserve_b,
+    );
+
+    if delta_lp == 0 {
+        panic!("Payable LP must be nonzero");
+    }
+
+    if delta_lp < min_amount_liquidity {
+        panic!("Payable LP is less than provided minimum LP amount");
+    }
+
+    // 5. Update pool account
+    let mut pool_post = pool.account.clone();
+    let pool_post_definition = PoolDefinition {
+        liquidity_pool_supply: pool_def_data.liquidity_pool_supply + delta_lp,
+        reserve_a: pool_def_data.reserve_a + actual_amount_a,
+        reserve_b: pool_def_data.reserve_b + actual_amount_b,
+        ..pool_def_data
+    };
+
+    pool_post.data = pool_post_definition.into_data();
+    let token_program_id = user_holding_a.account.program_owner;
+
+    // Chain call for Token A (UserHoldingA -> Vault_A)
+    let call_token_a = ChainedCall::new(
+        token_program_id,
+        vec![user_holding_a.clone(), vault_a.clone()],
+        &token_core::Instruction::Transfer {
+            amount_to_transfer: actual_amount_a,
+        },
+    );
+    // Chain call for Token B (UserHoldingB -> Vault_B)
+    let call_token_b = ChainedCall::new(
+        token_program_id,
+        vec![user_holding_b.clone(), vault_b.clone()],
+        &token_core::Instruction::Transfer {
+            amount_to_transfer: actual_amount_b,
+        },
+    );
+    // Chain call for LP (mint new tokens for user_holding_lp)
+    let mut pool_definition_lp_auth = pool_definition_lp.clone();
+    pool_definition_lp_auth.is_authorized = true;
+    let call_token_lp = ChainedCall::new(
+        token_program_id,
+        vec![pool_definition_lp_auth.clone(), user_holding_lp.clone()],
+        &token_core::Instruction::Mint {
+            amount_to_mint: delta_lp,
+        },
+    )
+    .with_pda_seeds(vec![compute_liquidity_token_pda_seed(pool.account_id)]);
+
+    let chained_calls = vec![call_token_lp, call_token_b, call_token_a];
+
+    let post_states = vec![
+        AccountPostState::new(pool_post),
+        AccountPostState::new(vault_a.account.clone()),
+        AccountPostState::new(vault_b.account.clone()),
+        AccountPostState::new(pool_definition_lp.account.clone()),
+        AccountPostState::new(user_holding_a.account.clone()),
+        AccountPostState::new(user_holding_b.account.clone()),
+        AccountPostState::new(user_holding_lp.account.clone()),
+    ];
+
+    (post_states, chained_calls)
+}

programs/amm/src/lib.rs

@@ -0,0 +1,10 @@
+//! The AMM Program implementation.
+
+pub use token_core as core;
+
+pub mod add;
+pub mod new_definition;
+pub mod remove;
+pub mod swap;
+
+mod tests;

programs/amm/src/new_definition.rs

@@ -0,0 +1,156 @@
+use amm_core::{
+    PoolDefinition, compute_liquidity_token_pda, compute_liquidity_token_pda_seed,
+    compute_pool_pda, compute_vault_pda,
+};
+use nssa_core::{
+    account::{Account, AccountWithMetadata},
+    program::{AccountPostState, ChainedCall, ProgramId},
+};
+
+#[expect(clippy::too_many_arguments, reason = "TODO: Fix later")]
+pub fn new_definition(
+    pool: AccountWithMetadata,
+    vault_a: AccountWithMetadata,
+    vault_b: AccountWithMetadata,
+    pool_definition_lp: AccountWithMetadata,
+    user_holding_a: AccountWithMetadata,
+    user_holding_b: AccountWithMetadata,
+    user_holding_lp: AccountWithMetadata,
+    token_a_amount: u128,
+    token_b_amount: u128,
+    amm_program_id: ProgramId,
+) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
+    // Prevents pool constant coefficient (k) from being 0.
+    if token_a_amount == 0 || token_b_amount == 0 {
+        panic!("Balances must be nonzero")
+    }
+
+    // Verify token_a and token_b are different
+    let definition_token_a_id = token_core::TokenHolding::try_from(&user_holding_a.account.data)
+        .expect("New definition: AMM Program expects valid Token Holding account for Token A")
+        .definition_id();
+    let definition_token_b_id = token_core::TokenHolding::try_from(&user_holding_b.account.data)
+        .expect("New definition: AMM Program expects valid Token Holding account for Token B")
+        .definition_id();
+
+    // both instances of the same token program
+    let token_program = user_holding_a.account.program_owner;
+
+    if user_holding_b.account.program_owner != token_program {
+        panic!("User Token holdings must use the same Token Program");
+    }
+
+    if definition_token_a_id == definition_token_b_id {
+        panic!("Cannot set up a swap for a token with itself")
+    }
+
+    if pool.account_id
+        != compute_pool_pda(amm_program_id, definition_token_a_id, definition_token_b_id)
+    {
+        panic!("Pool Definition Account ID does not match PDA");
+    }
+
+    if vault_a.account_id
+        != compute_vault_pda(amm_program_id, pool.account_id, definition_token_a_id)
+        || vault_b.account_id
+            != compute_vault_pda(amm_program_id, pool.account_id, definition_token_b_id)
+    {
+        panic!("Vault ID does not match PDA");
+    }
+
+    if pool_definition_lp.account_id != compute_liquidity_token_pda(amm_program_id, pool.account_id)
+    {
+        panic!("Liquidity pool Token Definition Account ID does not match PDA");
+    }
+
+    // Verify that Pool Account is not active
+    let pool_account_data = if pool.account == Account::default() {
+        PoolDefinition::default()
+    } else {
+        PoolDefinition::parse(&pool.account.data).expect("AMM program expects a valid Pool account")
+    };
+
+    if pool_account_data.active {
+        panic!("Cannot initialize an active Pool Definition")
+    }
+
+    // LP Token minting calculation
+    // We assume LP is based on the initial deposit amount for Token_A.
+
+    // Update pool account
+    let mut pool_post = pool.account.clone();
+    let pool_post_definition = PoolDefinition {
+        definition_token_a_id,
+        definition_token_b_id,
+        vault_a_id: vault_a.account_id,
+        vault_b_id: vault_b.account_id,
+        liquidity_pool_id: pool_definition_lp.account_id,
+        liquidity_pool_supply: token_a_amount,
+        reserve_a: token_a_amount,
+        reserve_b: token_b_amount,
+        fees: 0u128, // TODO: we assume all fees are 0 for now.
+        active: true,
+    };
+
+    pool_post.data = pool_post_definition.into_data();
+    let pool_post: AccountPostState = if pool.account == Account::default() {
+        AccountPostState::new_claimed(pool_post.clone())
+    } else {
+        AccountPostState::new(pool_post.clone())
+    };
+
+    let token_program_id = user_holding_a.account.program_owner;
+
+    // Chain call for Token A (user_holding_a -> Vault_A)
+    let call_token_a = ChainedCall::new(
+        token_program_id,
+        vec![user_holding_a.clone(), vault_a.clone()],
+        &token_core::Instruction::Transfer {
+            amount_to_transfer: token_a_amount,
+        },
+    );
+    // Chain call for Token B (user_holding_b -> Vault_B)
+    let call_token_b = ChainedCall::new(
+        token_program_id,
+        vec![user_holding_b.clone(), vault_b.clone()],
+        &token_core::Instruction::Transfer {
+            amount_to_transfer: token_b_amount,
+        },
+    );
+
+    // Chain call for liquidity token (TokenLP definition -> User LP Holding)
+    let instruction = if pool.account == Account::default() {
+        token_core::Instruction::NewFungibleDefinition {
+            name: String::from("LP Token"),
+            total_supply: token_a_amount,
+        }
+    } else {
+        token_core::Instruction::Mint {
+            amount_to_mint: token_a_amount,
+        }
+    };
+
+    let mut pool_lp_auth = pool_definition_lp.clone();
+    pool_lp_auth.is_authorized = true;
+
+    let call_token_lp = ChainedCall::new(
+        token_program_id,
+        vec![pool_lp_auth.clone(), user_holding_lp.clone()],
+        &instruction,
+    )
+    .with_pda_seeds(vec![compute_liquidity_token_pda_seed(pool.account_id)]);
+
+    let chained_calls = vec![call_token_lp, call_token_b, call_token_a];
+
+    let post_states = vec![
+        pool_post.clone(),
+        AccountPostState::new(vault_a.account.clone()),
+        AccountPostState::new(vault_b.account.clone()),
+        AccountPostState::new(pool_definition_lp.account.clone()),
+        AccountPostState::new(user_holding_a.account.clone()),
+        AccountPostState::new(user_holding_b.account.clone()),
+        AccountPostState::new(user_holding_lp.account.clone()),
+    ];
+
+    (post_states.clone(), chained_calls)
+}

programs/amm/src/remove.rs

@@ -1,27 +1,22 @@
-fn remove_liquidity(
-    pre_states: &[AccountWithMetadata],
-    amounts: &[u128],
+use amm_core::{PoolDefinition, compute_liquidity_token_pda_seed, compute_vault_pda_seed};
+use nssa_core::{
+    account::AccountWithMetadata,
+    program::{AccountPostState, ChainedCall},
+};
+
+#[expect(clippy::too_many_arguments, reason = "TODO: Fix later")]
+pub fn remove_liquidity(
+    pool: AccountWithMetadata,
+    vault_a: AccountWithMetadata,
+    vault_b: AccountWithMetadata,
+    pool_definition_lp: AccountWithMetadata,
+    user_holding_a: AccountWithMetadata,
+    user_holding_b: AccountWithMetadata,
+    user_holding_lp: AccountWithMetadata,
+    remove_liquidity_amount: u128,
+    min_amount_to_remove_token_a: u128,
+    min_amount_to_remove_token_b: u128,
 ) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
-    if pre_states.len() != 7 {
-        panic!("Invalid number of input accounts");
-    }
-
-    let pool = &pre_states[0];
-    let vault_a = &pre_states[1];
-    let vault_b = &pre_states[2];
-    let pool_definition_lp = &pre_states[3];
-    let user_holding_a = &pre_states[4];
-    let user_holding_b = &pre_states[5];
-    let user_holding_lp = &pre_states[6];
-
-    if amounts.len() != 3 {
-        panic!("Invalid number of balances");
-    }
-
-    let amount_lp = amounts[0];
-    let amount_min_a = amounts[1];
-    let amount_min_b = amounts[2];
-
     // 1. Fetch Pool state
     let pool_def_data = PoolDefinition::parse(&pool.account.data)
         .expect("Remove liquidity: AMM Program expects a valid Pool Definition Account");
@@ -50,11 +45,11 @@ fn remove_liquidity(
     running_vault_a.is_authorized = true;
     running_vault_b.is_authorized = true;
 
-    if amount_min_a == 0 || amount_min_b == 0 {
+    if min_amount_to_remove_token_a == 0 || min_amount_to_remove_token_b == 0 {
         panic!("Minimum withdraw amount must be nonzero");
     }
 
-    if amount_lp == 0 {
+    if remove_liquidity_amount == 0 {
         panic!("Liquidity amount must be nonzero");
     }
 
@@ -78,21 +73,21 @@ fn remove_liquidity(
     }
 
     let withdraw_amount_a =
-        (pool_def_data.reserve_a * amount_lp) / pool_def_data.liquidity_pool_supply;
+        (pool_def_data.reserve_a * remove_liquidity_amount) / pool_def_data.liquidity_pool_supply;
     let withdraw_amount_b =
-        (pool_def_data.reserve_b * amount_lp) / pool_def_data.liquidity_pool_supply;
+        (pool_def_data.reserve_b * remove_liquidity_amount) / pool_def_data.liquidity_pool_supply;
 
     // 3. Validate and slippage check
-    if withdraw_amount_a < amount_min_a {
+    if withdraw_amount_a < min_amount_to_remove_token_a {
         panic!("Insufficient minimal withdraw amount (Token A) provided for liquidity amount");
     }
-    if withdraw_amount_b < amount_min_b {
+    if withdraw_amount_b < min_amount_to_remove_token_b {
         panic!("Insufficient minimal withdraw amount (Token B) provided for liquidity amount");
     }
 
     // 4. Calculate LP to reduce cap by
-    let delta_lp: u128 =
-        (pool_def_data.liquidity_pool_supply * amount_lp) / pool_def_data.liquidity_pool_supply;
+    let delta_lp: u128 = (pool_def_data.liquidity_pool_supply * remove_liquidity_amount)
+        / pool_def_data.liquidity_pool_supply;
 
     let active: bool = pool_def_data.liquidity_pool_supply - delta_lp != 0;
 
@@ -150,13 +145,13 @@ fn remove_liquidity(
 
     let post_states = vec![
         AccountPostState::new(pool_post.clone()),
-        AccountPostState::new(pre_states[1].account.clone()),
-        AccountPostState::new(pre_states[2].account.clone()),
-        AccountPostState::new(pre_states[3].account.clone()),
-        AccountPostState::new(pre_states[4].account.clone()),
-        AccountPostState::new(pre_states[5].account.clone()),
-        AccountPostState::new(pre_states[6].account.clone()),
+        AccountPostState::new(vault_a.account.clone()),
+        AccountPostState::new(vault_b.account.clone()),
+        AccountPostState::new(pool_definition_lp.account.clone()),
+        AccountPostState::new(user_holding_a.account.clone()),
+        AccountPostState::new(user_holding_b.account.clone()),
+        AccountPostState::new(user_holding_lp.account.clone()),
     ];
 
     (post_states, chained_calls)
-}
+}

programs/amm/src/swap.rs

@@ -0,0 +1,177 @@
+pub use amm_core::{PoolDefinition, compute_liquidity_token_pda_seed, compute_vault_pda_seed};
+use nssa_core::{
+    account::{AccountId, AccountWithMetadata},
+    program::{AccountPostState, ChainedCall},
+};
+
+#[expect(clippy::too_many_arguments, reason = "TODO: Fix later")]
+pub fn swap(
+    pool: AccountWithMetadata,
+    vault_a: AccountWithMetadata,
+    vault_b: AccountWithMetadata,
+    user_holding_a: AccountWithMetadata,
+    user_holding_b: AccountWithMetadata,
+    swap_amount_in: u128,
+    min_amount_out: u128,
+    token_in_id: AccountId,
+) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
+    // Verify vaults are in fact vaults
+    let pool_def_data = PoolDefinition::parse(&pool.account.data)
+        .expect("Swap: AMM Program expects a valid Pool Definition Account");
+
+    if !pool_def_data.active {
+        panic!("Pool is inactive");
+    }
+
+    if vault_a.account_id != pool_def_data.vault_a_id {
+        panic!("Vault A was not provided");
+    }
+
+    if vault_b.account_id != pool_def_data.vault_b_id {
+        panic!("Vault B was not provided");
+    }
+
+    // fetch pool reserves
+    // validates reserves is at least the vaults' balances
+    let vault_a_token_holding = token_core::TokenHolding::try_from(&vault_a.account.data)
+        .expect("Swap: AMM Program expects a valid Token Holding Account for Vault A");
+    let token_core::TokenHolding::Fungible {
+        definition_id: _,
+        balance: vault_a_balance,
+    } = vault_a_token_holding
+    else {
+        panic!("Swap: AMM Program expects a valid Fungible Token Holding Account for Vault A");
+    };
+    if vault_a_balance < pool_def_data.reserve_a {
+        panic!("Reserve for Token A exceeds vault balance");
+    }
+
+    let vault_b_token_holding = token_core::TokenHolding::try_from(&vault_b.account.data)
+        .expect("Swap: AMM Program expects a valid Token Holding Account for Vault B");
+    let token_core::TokenHolding::Fungible {
+        definition_id: _,
+        balance: vault_b_balance,
+    } = vault_b_token_holding
+    else {
+        panic!("Swap: AMM Program expects a valid Fungible Token Holding Account for Vault B");
+    };
+
+    if vault_b_balance < pool_def_data.reserve_b {
+        panic!("Reserve for Token B exceeds vault balance");
+    }
+
+    let (chained_calls, [deposit_a, withdraw_a], [deposit_b, withdraw_b]) =
+        if token_in_id == pool_def_data.definition_token_a_id {
+            let (chained_calls, deposit_a, withdraw_b) = swap_logic(
+                user_holding_a.clone(),
+                vault_a.clone(),
+                vault_b.clone(),
+                user_holding_b.clone(),
+                swap_amount_in,
+                min_amount_out,
+                &[pool_def_data.reserve_a, pool_def_data.reserve_b],
+                pool.account_id,
+            );
+
+            (chained_calls, [deposit_a, 0], [0, withdraw_b])
+        } else if token_in_id == pool_def_data.definition_token_b_id {
+            let (chained_calls, deposit_b, withdraw_a) = swap_logic(
+                user_holding_b.clone(),
+                vault_b.clone(),
+                vault_a.clone(),
+                user_holding_a.clone(),
+                swap_amount_in,
+                min_amount_out,
+                &[pool_def_data.reserve_b, pool_def_data.reserve_a],
+                pool.account_id,
+            );
+
+            (chained_calls, [0, withdraw_a], [deposit_b, 0])
+        } else {
+            panic!("AccountId is not a token type for the pool");
+        };
+
+    // Update pool account
+    let mut pool_post = pool.account.clone();
+    let pool_post_definition = PoolDefinition {
+        reserve_a: pool_def_data.reserve_a + deposit_a - withdraw_a,
+        reserve_b: pool_def_data.reserve_b + deposit_b - withdraw_b,
+        ..pool_def_data
+    };
+
+    pool_post.data = pool_post_definition.into_data();
+
+    let post_states = vec![
+        AccountPostState::new(pool_post.clone()),
+        AccountPostState::new(vault_a.account.clone()),
+        AccountPostState::new(vault_b.account.clone()),
+        AccountPostState::new(user_holding_a.account.clone()),
+        AccountPostState::new(user_holding_b.account.clone()),
+    ];
+
+    (post_states, chained_calls)
+}
+
+#[expect(clippy::too_many_arguments, reason = "TODO: Fix later")]
+fn swap_logic(
+    user_deposit: AccountWithMetadata,
+    vault_deposit: AccountWithMetadata,
+    vault_withdraw: AccountWithMetadata,
+    user_withdraw: AccountWithMetadata,
+    swap_amount_in: u128,
+    min_amount_out: u128,
+    reserve_amounts: &[u128],
+    pool_id: AccountId,
+) -> (Vec<ChainedCall>, u128, u128) {
+    let reserve_deposit_vault_amount = reserve_amounts[0];
+    let reserve_withdraw_vault_amount = reserve_amounts[1];
+
+    // Compute withdraw amount
+    // Maintains pool constant product
+    // k = pool_def_data.reserve_a * pool_def_data.reserve_b;
+    let withdraw_amount = (reserve_withdraw_vault_amount * swap_amount_in)
+        / (reserve_deposit_vault_amount + swap_amount_in);
+
+    // Slippage check
+    if min_amount_out > withdraw_amount {
+        panic!("Withdraw amount is less than minimal amount out");
+    }
+
+    if withdraw_amount == 0 {
+        panic!("Withdraw amount should be nonzero");
+    }
+
+    let token_program_id = user_deposit.account.program_owner;
+
+    let mut chained_calls = Vec::new();
+    chained_calls.push(ChainedCall::new(
+        token_program_id,
+        vec![user_deposit, vault_deposit],
+        &token_core::Instruction::Transfer {
+            amount_to_transfer: swap_amount_in,
+        },
+    ));
+
+    let mut vault_withdraw = vault_withdraw.clone();
+    vault_withdraw.is_authorized = true;
+
+    let pda_seed = compute_vault_pda_seed(
+        pool_id,
+        token_core::TokenHolding::try_from(&vault_withdraw.account.data)
+            .expect("Swap Logic: AMM Program expects valid token data")
+            .definition_id(),
+    );
+
+    chained_calls.push(
+        ChainedCall::new(
+            token_program_id,
+            vec![vault_withdraw, user_withdraw],
+            &token_core::Instruction::Transfer {
+                amount_to_transfer: withdraw_amount,
+            },
+        )
+        .with_pda_seeds(vec![pda_seed]),
+    );
+
+    (chained_calls, swap_amount_in, withdraw_amount)
+}