Merge pull request #306 from logos-blockchain/marvin/refactor-amm-program

refactor amm program

Cargo.lock

@@ -302,6 +302,25 @@ version = "0.2.21"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923"
 
+[[package]]
+name = "amm_core"
+version = "0.1.0"
+dependencies = [
+ "borsh",
+ "nssa_core",
+ "risc0-zkvm",
+ "serde",
+]
+
+[[package]]
+name = "amm_program"
+version = "0.1.0"
+dependencies = [
+ "amm_core",
+ "nssa_core",
+ "token_core",
+]
+
 [[package]]
 name = "android_system_properties"
 version = "0.1.5"
@@ -4808,8 +4827,8 @@ checksum = "38bf9645c8b145698bb0b18a4637dcacbc421ea49bef2317e4fd8065a387cf21"
 name = "nssa"
 version = "0.1.0"
 dependencies = [
+ "amm_core",
  "borsh",
- "bytemuck",
  "env_logger",
  "hex",
  "hex-literal 1.1.0",
@@ -5339,6 +5358,8 @@ dependencies = [
 name = "programs"
 version = "0.1.0"
 dependencies = [
+ "amm_core",
+ "amm_program",
  "nssa_core",
  "risc0-zkvm",
  "serde",
@@ -7754,6 +7775,7 @@ dependencies = [
 name = "wallet"
 version = "0.1.0"
 dependencies = [
+ "amm_core",
  "anyhow",
  "async-stream",
  "base58",
@@ -7773,7 +7795,6 @@ dependencies = [
  "nssa_core",
  "optfield",
  "rand 0.8.5",
- "risc0-zkvm",
  "serde",
  "serde_json",
  "sha2",

Cargo.toml

@@ -13,6 +13,10 @@ members = [
   "common",
   "nssa",
   "nssa/core",
+  "programs/amm/core",
+  "programs/amm",
+  "programs/token/core",
+  "programs/token",
   "sequencer_core",
   "sequencer_rpc",
   "sequencer_runner",
@@ -50,6 +54,8 @@ wallet = { path = "wallet" }
 wallet-ffi = { path = "wallet-ffi" }
 token_core = { path = "programs/token/core" }
 token_program = { path = "programs/token" }
+amm_core = { path = "programs/amm/core" }
+amm_program = { path = "programs/amm" }
 test_program_methods = { path = "test_program_methods" }
 bedrock_client = { path = "bedrock_client" }
 indexer_core = { path = "indexer_core" }

nssa/Cargo.toml

@@ -16,7 +16,6 @@ borsh.workspace = true
 hex.workspace = true
 secp256k1 = "0.31.1"
 risc0-binfmt = "3.0.2"
-bytemuck = "1.24.0"
 log.workspace = true
 
 [build-dependencies]
@@ -25,6 +24,7 @@ risc0-binfmt = "3.0.2"
 
 [dev-dependencies]
 token_core.workspace = true
+amm_core.workspace = true
 test_program_methods.workspace = true
 
 env_logger.workspace = true

nssa/src/state.rs

@@ -311,6 +311,7 @@ pub mod tests {
 
     use std::collections::HashMap;
 
+    use amm_core::PoolDefinition;
     use nssa_core::{
         Commitment, Nullifier, NullifierPublicKey, NullifierSecretKey, SharedSecretKey,
         account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data},
@@ -2329,137 +2330,6 @@ pub mod tests {
         ));
     }
 
-    // TODO repeated code should ultimately be removed;
-    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),
-        ))
-    }
-
-    fn compute_pool_pda_seed(
-        definition_token_a_id: AccountId,
-        definition_token_b_id: AccountId,
-    ) -> PdaSeed {
-        use risc0_zkvm::sha::{Impl, Sha256};
-
-        let mut i: usize = 0;
-        let (token_1, token_2) = loop {
-            if definition_token_a_id.value()[i] > definition_token_b_id.value()[i] {
-                let token_1 = definition_token_a_id;
-                let token_2 = definition_token_b_id;
-                break (token_1, token_2);
-            } else if definition_token_a_id.value()[i] < definition_token_b_id.value()[i] {
-                let token_1 = definition_token_b_id;
-                let token_2 = definition_token_a_id;
-                break (token_1, token_2);
-            }
-
-            if i == 32 {
-                panic!("Definitions match");
-            } else {
-                i += 1;
-            }
-        };
-
-        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"),
-        )
-    }
-
-    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),
-        ))
-    }
-
-    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"),
-        )
-    }
-
-    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)))
-    }
-
-    fn compute_liquidity_token_pda_seed(pool_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(&[0; 32]);
-
-        PdaSeed::new(
-            Impl::hash_bytes(&bytes)
-                .as_bytes()
-                .try_into()
-                .expect("Hash output must be exactly 32 bytes long"),
-        )
-    }
-
-    const POOL_DEFINITION_DATA_SIZE: usize = 225;
-
-    #[derive(Default)]
-    struct PoolDefinition {
-        definition_token_a_id: AccountId,
-        definition_token_b_id: AccountId,
-        vault_a_id: AccountId,
-        vault_b_id: AccountId,
-        liquidity_pool_id: AccountId,
-        liquidity_pool_supply: u128,
-        reserve_a: u128,
-        reserve_b: u128,
-        fees: u128,
-        active: bool,
-    }
-
-    impl PoolDefinition {
-        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;
-
-            bytes
-                .to_vec()
-                .try_into()
-                .expect("225 bytes should fit into Data")
-        }
-    }
-
     struct PrivateKeysForTests;
 
     impl PrivateKeysForTests {
@@ -2640,7 +2510,7 @@ pub mod tests {
 
     impl IdForTests {
         fn pool_definition_id() -> AccountId {
-            compute_pool_pda(
+            amm_core::compute_pool_pda(
                 Program::amm().id(),
                 IdForTests::token_a_definition_id(),
                 IdForTests::token_b_definition_id(),
@@ -2648,7 +2518,10 @@ pub mod tests {
         }
 
         fn token_lp_definition_id() -> AccountId {
-            compute_liquidity_token_pda(Program::amm().id(), IdForTests::pool_definition_id())
+            amm_core::compute_liquidity_token_pda(
+                Program::amm().id(),
+                IdForTests::pool_definition_id(),
+            )
         }
 
         fn token_a_definition_id() -> AccountId {
@@ -2678,7 +2551,7 @@ pub mod tests {
         }
 
         fn vault_a_id() -> AccountId {
-            compute_vault_pda(
+            amm_core::compute_vault_pda(
                 Program::amm().id(),
                 IdForTests::pool_definition_id(),
                 IdForTests::token_a_definition_id(),
@@ -2686,7 +2559,7 @@ pub mod tests {
         }
 
         fn vault_b_id() -> AccountId {
-            compute_vault_pda(
+            amm_core::compute_vault_pda(
                 Program::amm().id(),
                 IdForTests::pool_definition_id(),
                 IdForTests::token_b_definition_id(),
@@ -2725,7 +2598,7 @@ pub mod tests {
             Account {
                 program_owner: Program::amm().id(),
                 balance: 0u128,
-                data: PoolDefinition::into_data(PoolDefinition {
+                data: Data::from(&PoolDefinition {
                     definition_token_a_id: IdForTests::token_a_definition_id(),
                     definition_token_b_id: IdForTests::token_b_definition_id(),
                     vault_a_id: IdForTests::vault_a_id(),
@@ -2844,7 +2717,7 @@ pub mod tests {
             Account {
                 program_owner: Program::amm().id(),
                 balance: 0u128,
-                data: PoolDefinition::into_data(PoolDefinition {
+                data: Data::from(&PoolDefinition {
                     definition_token_a_id: IdForTests::token_a_definition_id(),
                     definition_token_b_id: IdForTests::token_b_definition_id(),
                     vault_a_id: IdForTests::vault_a_id(),
@@ -2912,7 +2785,7 @@ pub mod tests {
             Account {
                 program_owner: Program::amm().id(),
                 balance: 0u128,
-                data: PoolDefinition::into_data(PoolDefinition {
+                data: Data::from(&PoolDefinition {
                     definition_token_a_id: IdForTests::token_a_definition_id(),
                     definition_token_b_id: IdForTests::token_b_definition_id(),
                     vault_a_id: IdForTests::vault_a_id(),
@@ -2980,7 +2853,7 @@ pub mod tests {
             Account {
                 program_owner: Program::amm().id(),
                 balance: 0u128,
-                data: PoolDefinition::into_data(PoolDefinition {
+                data: Data::from(&PoolDefinition {
                     definition_token_a_id: IdForTests::token_a_definition_id(),
                     definition_token_b_id: IdForTests::token_b_definition_id(),
                     vault_a_id: IdForTests::vault_a_id(),
@@ -3073,7 +2946,7 @@ pub mod tests {
             Account {
                 program_owner: Program::amm().id(),
                 balance: 0u128,
-                data: PoolDefinition::into_data(PoolDefinition {
+                data: Data::from(&PoolDefinition {
                     definition_token_a_id: IdForTests::token_a_definition_id(),
                     definition_token_b_id: IdForTests::token_b_definition_id(),
                     vault_a_id: IdForTests::vault_a_id(),
@@ -3179,7 +3052,7 @@ pub mod tests {
             Account {
                 program_owner: Program::amm().id(),
                 balance: 0u128,
-                data: PoolDefinition::into_data(PoolDefinition {
+                data: Data::from(&PoolDefinition {
                     definition_token_a_id: IdForTests::token_a_definition_id(),
                     definition_token_b_id: IdForTests::token_b_definition_id(),
                     vault_a_id: IdForTests::vault_a_id(),
@@ -3248,7 +3121,7 @@ pub mod tests {
             Account {
                 program_owner: Program::amm().id(),
                 balance: 0u128,
-                data: PoolDefinition::into_data(PoolDefinition {
+                data: Data::from(&PoolDefinition {
                     definition_token_a_id: IdForTests::token_a_definition_id(),
                     definition_token_b_id: IdForTests::token_b_definition_id(),
                     vault_a_id: IdForTests::vault_a_id(),
@@ -3277,11 +3150,6 @@ pub mod tests {
         }
     }
 
-    const AMM_NEW_DEFINITION: u8 = 0;
-    const AMM_SWAP: u8 = 1;
-    const AMM_ADD_LIQUIDITY: u8 = 2;
-    const AMM_REMOVE_LIQUIDITY: u8 = 3;
-
     fn state_for_amm_tests() -> V02State {
         let initial_data = [];
         let mut state =
@@ -3347,11 +3215,11 @@ pub mod tests {
     fn test_simple_amm_remove() {
         let mut state = state_for_amm_tests();
 
-        let mut instruction: Vec<u8> = Vec::new();
+        let instruction = amm_core::Instruction::RemoveLiquidity {
-        instruction.push(AMM_REMOVE_LIQUIDITY);
+            remove_liquidity_amount: BalanceForTests::remove_lp(),
-        instruction.extend_from_slice(&BalanceForTests::remove_lp().to_le_bytes());
+            min_amount_to_remove_token_a: BalanceForTests::remove_min_amount_a(),
-        instruction.extend_from_slice(&BalanceForTests::remove_min_amount_a().to_le_bytes());
+            min_amount_to_remove_token_b: BalanceForTests::remove_min_amount_b(),
-        instruction.extend_from_slice(&BalanceForTests::remove_min_amount_b().to_le_bytes());
+        };
 
         let message = public_transaction::Message::try_new(
             Program::amm().id(),
@@ -3424,12 +3292,11 @@ pub mod tests {
             AccountForTests::token_lp_definition_init_inactive(),
         );
 
-        let mut instruction: Vec<u8> = Vec::new();
+        let instruction = amm_core::Instruction::NewDefinition {
-        instruction.push(AMM_NEW_DEFINITION);
+            token_a_amount: BalanceForTests::vault_a_balance_init(),
-        instruction.extend_from_slice(&BalanceForTests::vault_a_balance_init().to_le_bytes());
+            token_b_amount: BalanceForTests::vault_b_balance_init(),
-        instruction.extend_from_slice(&BalanceForTests::vault_b_balance_init().to_le_bytes());
+            amm_program_id: Program::amm().id(),
-        let amm_program_u8: [u8; 32] = bytemuck::cast(Program::amm().id());
+        };
-        instruction.extend_from_slice(&amm_program_u8);
 
         let message = public_transaction::Message::try_new(
             Program::amm().id(),
@@ -3509,12 +3376,11 @@ pub mod tests {
             AccountForTests::user_token_lp_holding_init_zero(),
         );
 
-        let mut instruction: Vec<u8> = Vec::new();
+        let instruction = amm_core::Instruction::NewDefinition {
-        instruction.push(AMM_NEW_DEFINITION);
+            token_a_amount: BalanceForTests::vault_a_balance_init(),
-        instruction.extend_from_slice(&BalanceForTests::vault_a_balance_init().to_le_bytes());
+            token_b_amount: BalanceForTests::vault_b_balance_init(),
-        instruction.extend_from_slice(&BalanceForTests::vault_b_balance_init().to_le_bytes());
+            amm_program_id: Program::amm().id(),
-        let amm_program_u8: [u8; 32] = bytemuck::cast(Program::amm().id());
+        };
-        instruction.extend_from_slice(&amm_program_u8);
 
         let message = public_transaction::Message::try_new(
             Program::amm().id(),
@@ -3582,12 +3448,11 @@ pub mod tests {
             AccountForTests::vault_b_init_inactive(),
         );
 
-        let mut instruction: Vec<u8> = Vec::new();
+        let instruction = amm_core::Instruction::NewDefinition {
-        instruction.push(AMM_NEW_DEFINITION);
+            token_a_amount: BalanceForTests::vault_a_balance_init(),
-        instruction.extend_from_slice(&BalanceForTests::vault_a_balance_init().to_le_bytes());
+            token_b_amount: BalanceForTests::vault_b_balance_init(),
-        instruction.extend_from_slice(&BalanceForTests::vault_b_balance_init().to_le_bytes());
+            amm_program_id: Program::amm().id(),
-        let amm_program_u8: [u8; 32] = bytemuck::cast(Program::amm().id());
+        };
-        instruction.extend_from_slice(&amm_program_u8);
 
         let message = public_transaction::Message::try_new(
             Program::amm().id(),
@@ -3646,11 +3511,11 @@ pub mod tests {
         env_logger::init();
         let mut state = state_for_amm_tests();
 
-        let mut instruction: Vec<u8> = Vec::new();
+        let instruction = amm_core::Instruction::AddLiquidity {
-        instruction.push(AMM_ADD_LIQUIDITY);
+            min_amount_liquidity: BalanceForTests::add_min_amount_lp(),
-        instruction.extend_from_slice(&BalanceForTests::add_min_amount_lp().to_le_bytes());
+            max_amount_to_add_token_a: BalanceForTests::add_max_amount_a(),
-        instruction.extend_from_slice(&BalanceForTests::add_max_amount_a().to_le_bytes());
+            max_amount_to_add_token_b: BalanceForTests::add_max_amount_b(),
-        instruction.extend_from_slice(&BalanceForTests::add_max_amount_b().to_le_bytes());
+        };
 
         let message = public_transaction::Message::try_new(
             Program::amm().id(),
@@ -3708,11 +3573,11 @@ pub mod tests {
     fn test_simple_amm_swap_1() {
         let mut state = state_for_amm_tests();
 
-        let mut instruction: Vec<u8> = Vec::new();
+        let instruction = amm_core::Instruction::Swap {
-        instruction.push(AMM_SWAP);
+            swap_amount_in: BalanceForTests::swap_amount_in(),
-        instruction.extend_from_slice(&BalanceForTests::swap_amount_in().to_le_bytes());
+            min_amount_out: BalanceForTests::swap_min_amount_out(),
-        instruction.extend_from_slice(&BalanceForTests::swap_min_amount_out().to_le_bytes());
+            token_definition_id_in: IdForTests::token_b_definition_id(),
-        instruction.extend_from_slice(&IdForTests::token_b_definition_id().to_bytes());
+        };
 
         let message = public_transaction::Message::try_new(
             Program::amm().id(),
@@ -3759,12 +3624,11 @@ pub mod tests {
     fn test_simple_amm_swap_2() {
         let mut state = state_for_amm_tests();
 
-        let mut instruction: Vec<u8> = Vec::new();
+        let instruction = amm_core::Instruction::Swap {
-        instruction.push(AMM_SWAP);
+            swap_amount_in: BalanceForTests::swap_amount_in(),
-        instruction.extend_from_slice(&BalanceForTests::swap_amount_in().to_le_bytes());
+            min_amount_out: BalanceForTests::swap_min_amount_out(),
-        instruction.extend_from_slice(&BalanceForTests::swap_min_amount_out().to_le_bytes());
+            token_definition_id_in: IdForTests::token_a_definition_id(),
-        instruction.extend_from_slice(&IdForTests::token_a_definition_id().to_bytes());
+        };
-
         let message = public_transaction::Message::try_new(
             Program::amm().id(),
             vec![

program_methods/guest/Cargo.toml

@@ -8,5 +8,7 @@ license = { workspace = true }
 nssa_core.workspace = true
 token_core.workspace = true
 token_program.workspace = true
+amm_core.workspace = true
+amm_program.workspace = true
 risc0-zkvm.workspace = true
 serde = { workspace = true, default-features = false }

programs/amm/Cargo.toml

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

programs/amm/core/Cargo.toml

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

programs/amm/core/src/lib.rs

@@ -0,0 +1,197 @@
+//! This crate contains core data structures and utilities for the AMM Program.
+
+use borsh::{BorshDeserialize, BorshSerialize};
+use nssa_core::{
+    account::{AccountId, Data},
+    program::{PdaSeed, ProgramId},
+};
+use serde::{Deserialize, Serialize};
+
+/// AMM Program Instruction.
+#[derive(Serialize, Deserialize)]
+pub enum Instruction {
+    /// Initializes a new Pool (or re-initializes an inactive Pool).
+    ///
+    /// Required accounts:
+    /// - 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,
+    },
+
+    /// Adds liquidity to the Pool
+    ///
+    /// Required accounts:
+    /// - 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,
+    },
+
+    /// Removes liquidity from the Pool
+    ///
+    /// Required accounts:
+    /// - 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,
+    },
+
+    /// Swap some quantity of Tokens (either Token A or Token B)
+    /// while maintaining the Pool constant product.
+    ///
+    /// Required accounts:
+    /// - 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,
+    },
+}
+
+#[derive(Clone, Default, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
+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 TryFrom<&Data> for PoolDefinition {
+    type Error = std::io::Error;
+
+    fn try_from(data: &Data) -> Result<Self, Self::Error> {
+        PoolDefinition::try_from_slice(data.as_ref())
+    }
+}
+
+impl From<&PoolDefinition> for Data {
+    fn from(definition: &PoolDefinition) -> Self {
+        // Using size_of_val as size hint for Vec allocation
+        let mut data = Vec::with_capacity(std::mem::size_of_val(definition));
+
+        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")
+    }
+}
+
+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),
+    ))
+}
+
+pub fn compute_pool_pda_seed(
+    definition_token_a_id: AccountId,
+    definition_token_b_id: AccountId,
+) -> PdaSeed {
+    use risc0_zkvm::sha::{Impl, Sha256};
+
+    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"),
+    };
+
+    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"),
+    )
+}
+
+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),
+    ))
+}
+
+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"),
+    )
+}
+
+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)))
+}
+
+pub fn compute_liquidity_token_pda_seed(pool_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(&[0; 32]);
+
+    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,178 @@
+use std::num::NonZeroU128;
+
+use amm_core::{PoolDefinition, compute_liquidity_token_pda_seed};
+use nssa_core::{
+    account::{AccountWithMetadata, Data},
+    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: NonZeroU128,
+    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::try_from(&pool.account.data)
+        .expect("Add liquidity: AMM Program expects valid Pool Definition Account");
+
+    assert_eq!(
+        vault_a.account_id, pool_def_data.vault_a_id,
+        "Vault A was not provided"
+    );
+
+    assert_eq!(
+        pool_def_data.liquidity_pool_id, pool_definition_lp.account_id,
+        "LP definition mismatch"
+    );
+
+    assert_eq!(
+        vault_b.account_id, pool_def_data.vault_b_id,
+        "Vault B was not provided"
+    );
+
+    assert!(
+        max_amount_to_add_token_a != 0 && max_amount_to_add_token_b != 0,
+        "Both max-balances 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"
+        );
+    };
+
+    assert!(pool_def_data.reserve_a != 0, "Reserves must be nonzero");
+    assert!(pool_def_data.reserve_b != 0, "Reserves must be nonzero");
+    assert!(
+        vault_a_balance >= pool_def_data.reserve_a,
+        "Vaults' balances must be at least the reserve amounts"
+    );
+    assert!(
+        vault_b_balance >= pool_def_data.reserve_b,
+        "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
+    assert!(
+        max_amount_to_add_token_a >= actual_amount_a,
+        "Actual trade amounts cannot exceed max_amounts"
+    );
+    assert!(
+        max_amount_to_add_token_b >= actual_amount_b,
+        "Actual trade amounts cannot exceed max_amounts"
+    );
+
+    assert!(actual_amount_a != 0, "A trade amount is 0");
+    assert!(actual_amount_b != 0, "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,
+    );
+
+    assert!(delta_lp != 0, "Payable LP must be nonzero");
+
+    assert!(
+        delta_lp >= min_amount_liquidity.into(),
+        "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 = Data::from(&pool_post_definition);
+    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 amm_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,158 @@
+use std::num::NonZeroU128;
+
+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, Data},
+    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: NonZeroU128,
+    token_b_amount: NonZeroU128,
+    amm_program_id: ProgramId,
+) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
+    // 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;
+
+    assert_eq!(
+        user_holding_b.account.program_owner, token_program,
+        "User Token holdings must use the same Token Program"
+    );
+    assert!(
+        definition_token_a_id != definition_token_b_id,
+        "Cannot set up a swap for a token with itself"
+    );
+    assert_eq!(
+        pool.account_id,
+        compute_pool_pda(amm_program_id, definition_token_a_id, definition_token_b_id),
+        "Pool Definition Account ID does not match PDA"
+    );
+    assert_eq!(
+        vault_a.account_id,
+        compute_vault_pda(amm_program_id, pool.account_id, definition_token_a_id),
+        "Vault ID does not match PDA"
+    );
+    assert_eq!(
+        vault_b.account_id,
+        compute_vault_pda(amm_program_id, pool.account_id, definition_token_b_id),
+        "Vault ID does not match PDA"
+    );
+    assert_eq!(
+        pool_definition_lp.account_id,
+        compute_liquidity_token_pda(amm_program_id, pool.account_id),
+        "Liquidity pool Token Definition Account ID does not match PDA"
+    );
+
+    // TODO: return here
+    // Verify that Pool Account is not active
+    let pool_account_data = if pool.account == Account::default() {
+        PoolDefinition::default()
+    } else {
+        PoolDefinition::try_from(&pool.account.data)
+            .expect("AMM program expects a valid Pool account")
+    };
+
+    assert!(
+        !pool_account_data.active,
+        "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.into(),
+        reserve_a: token_a_amount.into(),
+        reserve_b: token_b_amount.into(),
+        fees: 0u128, // TODO: we assume all fees are 0 for now.
+        active: true,
+    };
+
+    pool_post.data = Data::from(&pool_post_definition);
+    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.into(),
+        },
+    );
+    // 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.into(),
+        },
+    );
+
+    // 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.into(),
+        }
+    } else {
+        token_core::Instruction::Mint {
+            amount_to_mint: token_a_amount.into(),
+        }
+    };
+
+    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

@@ -0,0 +1,166 @@
+use std::num::NonZeroU128;
+
+use amm_core::{PoolDefinition, compute_liquidity_token_pda_seed, compute_vault_pda_seed};
+use nssa_core::{
+    account::{AccountWithMetadata, Data},
+    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: NonZeroU128,
+    min_amount_to_remove_token_a: u128,
+    min_amount_to_remove_token_b: u128,
+) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
+    let remove_liquidity_amount: u128 = remove_liquidity_amount.into();
+
+    // 1. Fetch Pool state
+    let pool_def_data = PoolDefinition::try_from(&pool.account.data)
+        .expect("Remove liquidity: AMM Program expects a valid Pool Definition Account");
+
+    assert!(pool_def_data.active, "Pool is inactive");
+    assert_eq!(
+        pool_def_data.liquidity_pool_id, pool_definition_lp.account_id,
+        "LP definition mismatch"
+    );
+    assert_eq!(
+        vault_a.account_id, pool_def_data.vault_a_id,
+        "Vault A was not provided"
+    );
+    assert_eq!(
+        vault_b.account_id, pool_def_data.vault_b_id,
+        "Vault B was not provided"
+    );
+
+    // Vault addresses do not need to be checked with PDA
+    // calculation for setting authorization since stored
+    // in the Pool Definition.
+    let mut running_vault_a = vault_a.clone();
+    let mut running_vault_b = vault_b.clone();
+    running_vault_a.is_authorized = true;
+    running_vault_b.is_authorized = true;
+
+    assert!(
+        min_amount_to_remove_token_a != 0,
+        "Minimum withdraw amount must be nonzero"
+    );
+    assert!(
+        min_amount_to_remove_token_b != 0,
+        "Minimum withdraw amount must be nonzero"
+    );
+
+    // 2. Compute withdrawal amounts
+    let user_holding_lp_data = token_core::TokenHolding::try_from(&user_holding_lp.account.data)
+        .expect("Remove liquidity: AMM Program expects a valid Token Account for liquidity token");
+    let token_core::TokenHolding::Fungible {
+        definition_id: _,
+        balance: user_lp_balance,
+    } = user_holding_lp_data
+    else {
+        panic!(
+            "Remove liquidity: AMM Program expects a valid Fungible Token Holding Account for liquidity token"
+        );
+    };
+
+    assert!(
+        user_lp_balance <= pool_def_data.liquidity_pool_supply,
+        "Invalid liquidity account provided"
+    );
+    assert_eq!(
+        user_holding_lp_data.definition_id(),
+        pool_def_data.liquidity_pool_id,
+        "Invalid liquidity account provided"
+    );
+
+    let withdraw_amount_a =
+        (pool_def_data.reserve_a * remove_liquidity_amount) / pool_def_data.liquidity_pool_supply;
+    let withdraw_amount_b =
+        (pool_def_data.reserve_b * remove_liquidity_amount) / pool_def_data.liquidity_pool_supply;
+
+    // 3. Validate and slippage check
+    assert!(
+        withdraw_amount_a >= min_amount_to_remove_token_a,
+        "Insufficient minimal withdraw amount (Token A) provided for liquidity amount"
+    );
+    assert!(
+        withdraw_amount_b >= min_amount_to_remove_token_b,
+        "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 * remove_liquidity_amount)
+        / pool_def_data.liquidity_pool_supply;
+
+    let active: bool = pool_def_data.liquidity_pool_supply - delta_lp != 0;
+
+    // 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 - withdraw_amount_a,
+        reserve_b: pool_def_data.reserve_b - withdraw_amount_b,
+        active,
+        ..pool_def_data.clone()
+    };
+
+    pool_post.data = Data::from(&pool_post_definition);
+
+    let token_program_id = user_holding_a.account.program_owner;
+
+    // Chaincall for Token A withdraw
+    let call_token_a = ChainedCall::new(
+        token_program_id,
+        vec![running_vault_a, user_holding_a.clone()],
+        &token_core::Instruction::Transfer {
+            amount_to_transfer: withdraw_amount_a,
+        },
+    )
+    .with_pda_seeds(vec![compute_vault_pda_seed(
+        pool.account_id,
+        pool_def_data.definition_token_a_id,
+    )]);
+    // Chaincall for Token B withdraw
+    let call_token_b = ChainedCall::new(
+        token_program_id,
+        vec![running_vault_b, user_holding_b.clone()],
+        &token_core::Instruction::Transfer {
+            amount_to_transfer: withdraw_amount_b,
+        },
+    )
+    .with_pda_seeds(vec![compute_vault_pda_seed(
+        pool.account_id,
+        pool_def_data.definition_token_b_id,
+    )]);
+    // Chaincall for LP adjustment
+    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, user_holding_lp.clone()],
+        &token_core::Instruction::Burn {
+            amount_to_burn: 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.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,176 @@
+pub use amm_core::{PoolDefinition, compute_liquidity_token_pda_seed, compute_vault_pda_seed};
+use nssa_core::{
+    account::{AccountId, AccountWithMetadata, Data},
+    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::try_from(&pool.account.data)
+        .expect("Swap: AMM Program expects a valid Pool Definition Account");
+
+    assert!(pool_def_data.active, "Pool is inactive");
+    assert_eq!(
+        vault_a.account_id, pool_def_data.vault_a_id,
+        "Vault A was not provided"
+    );
+    assert_eq!(
+        vault_b.account_id, pool_def_data.vault_b_id,
+        "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");
+    };
+
+    assert!(
+        vault_a_balance >= pool_def_data.reserve_a,
+        "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");
+    };
+
+    assert!(
+        vault_b_balance >= pool_def_data.reserve_b,
+        "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 = Data::from(&pool_post_definition);
+
+    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_deposit_vault_amount: u128,
+    reserve_withdraw_vault_amount: u128,
+    pool_id: AccountId,
+) -> (Vec<ChainedCall>, u128, u128) {
+    // 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
+    assert!(
+        min_amount_out <= withdraw_amount,
+        "Withdraw amount is less than minimal amount out"
+    );
+    assert!(withdraw_amount != 0, "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)
+}

wallet/Cargo.toml

@@ -10,6 +10,7 @@ nssa.workspace = true
 common.workspace = true
 key_protocol.workspace = true
 token_core.workspace = true
+amm_core.workspace = true
 
 anyhow.workspace = true
 serde_json.workspace = true
@@ -27,7 +28,6 @@ rand.workspace = true
 itertools.workspace = true
 sha2.workspace = true
 futures.workspace = true
-risc0-zkvm.workspace = true
 async-stream = "0.3.6"
 indicatif = { version = "0.18.3", features = ["improved_unicode"] }
 optfield = "0.4.0"

wallet/src/program_facades/amm.rs

@@ -1,103 +1,9 @@
+use amm_core::{compute_liquidity_token_pda, compute_pool_pda, compute_vault_pda};
 use common::{error::ExecutionFailureKind, rpc_primitives::requests::SendTxResponse};
-use nssa::{AccountId, ProgramId, program::Program};
+use nssa::{AccountId, program::Program};
-use nssa_core::program::PdaSeed;
 use token_core::TokenHolding;
 
 use crate::WalletCore;
-
-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),
-    ))
-}
-
-fn compute_pool_pda_seed(
-    definition_token_a_id: AccountId,
-    definition_token_b_id: AccountId,
-) -> PdaSeed {
-    use risc0_zkvm::sha::{Impl, Sha256};
-
-    let mut i: usize = 0;
-    let (token_1, token_2) = loop {
-        if definition_token_a_id.value()[i] > definition_token_b_id.value()[i] {
-            let token_1 = definition_token_a_id;
-            let token_2 = definition_token_b_id;
-            break (token_1, token_2);
-        } else if definition_token_a_id.value()[i] < definition_token_b_id.value()[i] {
-            let token_1 = definition_token_b_id;
-            let token_2 = definition_token_a_id;
-            break (token_1, token_2);
-        }
-
-        if i == 32 {
-            panic!("Definitions match");
-        } else {
-            i += 1;
-        }
-    };
-
-    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"),
-    )
-}
-
-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),
-    ))
-}
-
-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"),
-    )
-}
-
-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)))
-}
-
-fn compute_liquidity_token_pda_seed(pool_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(&[0; 32]);
-
-    PdaSeed::new(
-        Impl::hash_bytes(&bytes)
-            .as_bytes()
-            .try_into()
-            .expect("Hash output must be exactly 32 bytes long"),
-    )
-}
-
 pub struct Amm<'w>(pub &'w WalletCore);
 
 impl Amm<'_> {
@@ -109,9 +15,13 @@ impl Amm<'_> {
         balance_a: u128,
         balance_b: u128,
     ) -> Result<SendTxResponse, ExecutionFailureKind> {
-        let (instruction, program) = amm_program_preparation_definition(balance_a, balance_b);
+        let program = Program::amm();
-
         let amm_program_id = Program::amm().id();
+        let instruction = amm_core::Instruction::NewDefinition {
+            token_a_amount: balance_a,
+            token_b_amount: balance_b,
+            amm_program_id,
+        };
 
         let user_a_acc = self
             .0
@@ -189,13 +99,16 @@ impl Amm<'_> {
         &self,
         user_holding_a: AccountId,
         user_holding_b: AccountId,
-        amount_in: u128,
+        swap_amount_in: u128,
         min_amount_out: u128,
-        token_definition_id: AccountId,
+        token_definition_id_in: AccountId,
     ) -> Result<SendTxResponse, ExecutionFailureKind> {
-        let (instruction, program) =
+        let instruction = amm_core::Instruction::Swap {
-            amm_program_preparation_swap(amount_in, min_amount_out, token_definition_id);
+            swap_amount_in,
-
+            min_amount_out,
+            token_definition_id_in,
+        };
+        let program = Program::amm();
         let amm_program_id = Program::amm().id();
 
         let user_a_acc = self
@@ -248,12 +161,14 @@ impl Amm<'_> {
         let token_holder_b = TokenHolding::try_from(&token_holder_acc_b.data)
             .map_err(|_| ExecutionFailureKind::AccountDataError(user_holding_b))?;
 
-        if token_holder_a.definition_id() == token_definition_id {
+        if token_holder_a.definition_id() == token_definition_id_in {
             account_id_auth = user_holding_a;
-        } else if token_holder_b.definition_id() == token_definition_id {
+        } else if token_holder_b.definition_id() == token_definition_id_in {
             account_id_auth = user_holding_b;
         } else {
-            return Err(ExecutionFailureKind::AccountDataError(token_definition_id));
+            return Err(ExecutionFailureKind::AccountDataError(
+                token_definition_id_in,
+            ));
         }
 
         let nonces = self
@@ -290,13 +205,16 @@ impl Amm<'_> {
         user_holding_a: AccountId,
         user_holding_b: AccountId,
         user_holding_lp: AccountId,
-        min_amount_lp: u128,
+        min_amount_liquidity: u128,
-        max_amount_a: u128,
+        max_amount_to_add_token_a: u128,
-        max_amount_b: u128,
+        max_amount_to_add_token_b: u128,
     ) -> Result<SendTxResponse, ExecutionFailureKind> {
-        let (instruction, program) =
+        let instruction = amm_core::Instruction::AddLiquidity {
-            amm_program_preparation_add_liq(min_amount_lp, max_amount_a, max_amount_b);
+            min_amount_liquidity,
-
+            max_amount_to_add_token_a,
+            max_amount_to_add_token_b,
+        };
+        let program = Program::amm();
         let amm_program_id = Program::amm().id();
 
         let user_a_acc = self
@@ -376,13 +294,16 @@ impl Amm<'_> {
         user_holding_a: AccountId,
         user_holding_b: AccountId,
         user_holding_lp: AccountId,
-        balance_lp: u128,
+        remove_liquidity_amount: u128,
-        min_amount_a: u128,
+        min_amount_to_remove_token_a: u128,
-        min_amount_b: u128,
+        min_amount_to_remove_token_b: u128,
     ) -> Result<SendTxResponse, ExecutionFailureKind> {
-        let (instruction, program) =
+        let instruction = amm_core::Instruction::RemoveLiquidity {
-            amm_program_preparation_remove_liq(balance_lp, min_amount_a, min_amount_b);
+            remove_liquidity_amount,
-
+            min_amount_to_remove_token_a,
+            min_amount_to_remove_token_b,
+        };
+        let program = Program::amm();
         let amm_program_id = Program::amm().id();
 
         let user_a_acc = self
@@ -448,93 +369,3 @@ impl Amm<'_> {
         Ok(self.0.sequencer_client.send_tx_public(tx).await?)
     }
 }
-
-fn amm_program_preparation_definition(balance_a: u128, balance_b: u128) -> (Vec<u8>, Program) {
-    // 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)]
-    let amm_program_id = Program::amm().id();
-
-    let mut instruction = [0; 65];
-    instruction[1..17].copy_from_slice(&balance_a.to_le_bytes());
-    instruction[17..33].copy_from_slice(&balance_b.to_le_bytes());
-
-    // This can be done less verbose, but it is better to use same way, as in amm program
-    instruction[33..37].copy_from_slice(&amm_program_id[0].to_le_bytes());
-    instruction[37..41].copy_from_slice(&amm_program_id[1].to_le_bytes());
-    instruction[41..45].copy_from_slice(&amm_program_id[2].to_le_bytes());
-    instruction[45..49].copy_from_slice(&amm_program_id[3].to_le_bytes());
-    instruction[49..53].copy_from_slice(&amm_program_id[4].to_le_bytes());
-    instruction[53..57].copy_from_slice(&amm_program_id[5].to_le_bytes());
-    instruction[57..61].copy_from_slice(&amm_program_id[6].to_le_bytes());
-    instruction[61..].copy_from_slice(&amm_program_id[7].to_le_bytes());
-
-    let instruction_data = instruction.to_vec();
-    let program = Program::amm();
-
-    (instruction_data, program)
-}
-
-fn amm_program_preparation_swap(
-    amount_in: u128,
-    min_amount_out: u128,
-    token_definition_id: AccountId,
-) -> (Vec<u8>, Program) {
-    // 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].
-    let mut instruction = [0; 65];
-    instruction[0] = 0x01;
-    instruction[1..17].copy_from_slice(&amount_in.to_le_bytes());
-    instruction[17..33].copy_from_slice(&min_amount_out.to_le_bytes());
-
-    // This can be done less verbose, but it is better to use same way, as in amm program
-    instruction[33..].copy_from_slice(&token_definition_id.to_bytes());
-
-    let instruction_data = instruction.to_vec();
-    let program = Program::amm();
-
-    (instruction_data, program)
-}
-
-fn amm_program_preparation_add_liq(
-    min_amount_lp: u128,
-    max_amount_a: u128,
-    max_amount_b: u128,
-) -> (Vec<u8>, Program) {
-    // 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)].
-    let mut instruction = [0; 49];
-    instruction[0] = 0x02;
-
-    instruction[1..17].copy_from_slice(&min_amount_lp.to_le_bytes());
-    instruction[17..33].copy_from_slice(&max_amount_a.to_le_bytes());
-    instruction[33..49].copy_from_slice(&max_amount_b.to_le_bytes());
-
-    let instruction_data = instruction.to_vec();
-    let program = Program::amm();
-
-    (instruction_data, program)
-}
-
-fn amm_program_preparation_remove_liq(
-    balance_lp: u128,
-    min_amount_a: u128,
-    min_amount_b: u128,
-) -> (Vec<u8>, Program) {
-    // 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)].
-    let mut instruction = [0; 49];
-    instruction[0] = 0x03;
-
-    instruction[1..17].copy_from_slice(&balance_lp.to_le_bytes());
-    instruction[17..33].copy_from_slice(&min_amount_a.to_le_bytes());
-    instruction[33..49].copy_from_slice(&min_amount_b.to_le_bytes());
-
-    let instruction_data = instruction.to_vec();
-    let program = Program::amm();
-
-    (instruction_data, program)
-}