From 10d42f352b521c0b700c0e31aaf338c2705fdb79 Mon Sep 17 00:00:00 2001 From: jonesmarvin8 <83104039+jonesmarvin8@users.noreply.github.com> Date: Fri, 14 Nov 2025 20:59:42 -0500 Subject: [PATCH] AMM functions written --- nssa/program_methods/guest/src/bin/amm.rs | 681 ++++++++++++++++++ nssa/src/program.rs | 8 +- .../guest/src/bin/pool.rs | 317 -------- 3 files changed, 688 insertions(+), 318 deletions(-) create mode 100644 nssa/program_methods/guest/src/bin/amm.rs delete mode 100644 nssa/test_program_methods/guest/src/bin/pool.rs diff --git a/nssa/program_methods/guest/src/bin/amm.rs b/nssa/program_methods/guest/src/bin/amm.rs new file mode 100644 index 0000000..fc2f260 --- /dev/null +++ b/nssa/program_methods/guest/src/bin/amm.rs @@ -0,0 +1,681 @@ +use nssa_core::{ + account::{Account, AccountId, AccountWithMetadata, Data}, + program::{ProgramId, ProgramInput, ChainedCall, read_nssa_inputs, write_nssa_outputs, write_nssa_outputs_with_chained_call}, +}; + +use bytemuck; + +// The token program has two functions: +// 1. New token definition. +// Arguments to this function are: +// * Two **default** accounts: [definition_account, holding_account]. +// The first default account will be initialized with the token definition account values. The second account will +// be initialized to a token holding account for the new token, holding the entire total supply. +// * An instruction data of 23-bytes, indicating the total supply and the token name, with +// the following layout: +// [0x00 || total_supply (little-endian 16 bytes) || name (6 bytes)] +// The name cannot be equal to [0x00, 0x00, 0x00, 0x00, 0x00, 0x00] +// 2. Token transfer +// Arguments to this function are: +// * Two accounts: [sender_account, recipient_account]. +// * An instruction data byte string of length 23, indicating the total supply with the following layout +// [0x01 || amount (little-endian 16 bytes) || 0x00 || 0x00 || 0x00 || 0x00 || 0x00 || 0x00]. + +const POOL_DEFINITION_DATA_SIZE: usize = 176; + +struct PoolDefinition{ + definition_token_a_id: AccountId, + definition_token_b_id: AccountId, + vault_a_addr: AccountId, + vault_b_addr: AccountId, + liquidity_pool_id: AccountId, + liquidity_pool_cap: u128, + reserve_a: u128, + reserve_b: u128, + token_program_id: ProgramId, +} + + + +impl PoolDefinition { + fn into_data(self) -> Vec { + let u8_token_program_id : [u8;32] = bytemuck::cast(self.token_program_id); + + 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_addr.to_bytes()); + bytes[96..128].copy_from_slice(&self.vault_b_addr.to_bytes()); + bytes[128..160].copy_from_slice(&self.liquidity_pool_id.to_bytes()); + bytes[160..176].copy_from_slice(&self.liquidity_pool_cap.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..].copy_from_slice(&u8_token_program_id); + bytes.into() + } + + fn parse(data: &[u8]) -> Option { + if data.len() != POOL_DEFINITION_DATA_SIZE { + None + } else { + let definition_token_a_id = AccountId::new(data[0..32].try_into().unwrap()); + let definition_token_b_id = AccountId::new(data[32..64].try_into().unwrap()); + let vault_a_addr = AccountId::new(data[64..96].try_into().unwrap()); + let vault_b_addr = AccountId::new(data[96..128].try_into().unwrap()); + let liquidity_pool_id = AccountId::new(data[128..160].try_into().unwrap()); + let liquidity_pool_cap = u128::from_le_bytes(data[160..176].try_into().unwrap()); + let reserve_a = u128::from_le_bytes(data[176..].try_into().unwrap()); + let reserve_b = u128::from_le_bytes(data[192..208].try_into().unwrap()); + + + let token_program_id : &[u32] = bytemuck::cast_slice(&data[208..]); + let token_program_id : ProgramId = token_program_id[0..8].try_into().unwrap(); + Some(Self { + definition_token_a_id, + definition_token_b_id, + vault_a_addr, + vault_b_addr, + liquidity_pool_id, + liquidity_pool_cap, + reserve_a, + reserve_b, + token_program_id, + }) + } + } +} + + +//TODO: remove repeated code for Token_Definition and TokenHoldling +const TOKEN_DEFINITION_TYPE: u8 = 0; +const TOKEN_DEFINITION_DATA_SIZE: usize = 23; +const TOKEN_HOLDING_TYPE: u8 = 1; +const TOKEN_HOLDING_DATA_SIZE: usize = 49; + +struct TokenHolding { + account_type: u8, + definition_id: AccountId, + balance: u128, +} + +impl TokenHolding { + fn new(definition_id: &AccountId) -> Self { + Self { + account_type: TOKEN_HOLDING_TYPE, + definition_id: definition_id.clone(), + balance: 0, + } + } + + fn parse(data: &[u8]) -> Option { + if data.len() != TOKEN_HOLDING_DATA_SIZE || data[0] != TOKEN_HOLDING_TYPE { + None + } else { + let account_type = data[0]; + let definition_id = AccountId::new(data[1..33].try_into().unwrap()); + let balance = u128::from_le_bytes(data[33..].try_into().unwrap()); + Some(Self { + definition_id, + balance, + account_type, + }) + } + } + + fn into_data(self) -> Data { + let mut bytes = [0; TOKEN_HOLDING_DATA_SIZE]; + bytes[0] = self.account_type; + bytes[1..33].copy_from_slice(&self.definition_id.to_bytes()); + bytes[33..].copy_from_slice(&self.balance.to_le_bytes()); + bytes.into() + } +} + + +fn new_definition( + pre_states: &[AccountWithMetadata], + balance_in: &[u128], + token_program: ProgramId, + ) -> (Vec, Vec) { + + //Pool accounts: pool itself, and its 2 vaults and LP token + //2 accounts for funding tokens + //initial funder's LP account + if pre_states.len() != 7 { + panic!("Invalid number of input account") + } + + /* + if pool_definitions.len() != 2 { + panic!("Invalid number of token definitions") + }*/ + + if balance_in.len() != 2 { + panic!("Invalid number of balance") + } + + let pool = pre_states[0].clone(); + let vault_a = pre_states[1].clone(); + let vault_b = pre_states[2].clone(); + let pool_lp = pre_states[3].clone(); + let user1 = pre_states[4].clone(); + let user2 = pre_states[5].clone(); + let user_lp = pre_states[6].clone(); + + if pool.account != Account::default() || !pool.is_authorized { + panic!("TODO-1"); + } + + // TODO: temporary band-aid to prevent vault's from being + // owned by the amm program. + if vault_a.account == Account::default() || vault_b.account == Account::default() { + panic!("Vault accounts must be initialized first; issue to be fixed") + } + if pool_lp.account == Account::default() { + panic!("Pool LP must be initialized first; issue to be fixed") + } + + let amount_a = balance_in[0]; + let amount_b = balance_in[1]; + + // Prevents pool constant coefficient (k) from being 0. + assert!(amount_a > 0); + assert!(amount_b > 0); + + // Verify token_a and token_b are different + //TODO: crucial fix. + let definition_token_a_id = TokenHolding::parse(&vault_a.account.data).unwrap().definition_id; + let definition_token_b_id = TokenHolding::parse(&vault_b.account.data).unwrap().definition_id; + let user1_id = TokenHolding::parse(&vault_a.account.data).unwrap().definition_id; + + // 5. Update pool account + let mut pool_post = Account::default(); + let pool_post_definition = PoolDefinition { + definition_token_a_id, + definition_token_b_id, + vault_a_addr: vault_a.account_id, + vault_b_addr: vault_b.account_id, + liquidity_pool_id: pool_lp.account_id, + liquidity_pool_cap: amount_a, + reserve_a: amount_a, + reserve_b: amount_b, + token_program_id: token_program, + }; + + pool_post.data = pool_post_definition.into_data(); + + let mut chained_call = Vec::new(); + + let mut instruction_data = [0; 23]; + instruction_data[0] = 1; + + instruction_data[1..17].copy_from_slice(&amount_a.to_le_bytes()); + let call_token_a = ChainedCall{ + program_id: token_program, + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + account_indices: vec![1] + }; + + + instruction_data[1..17].copy_from_slice(&amount_b.to_le_bytes()); + let call_token_b = ChainedCall{ + program_id: token_program, + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + account_indices: vec![1] + }; + + instruction_data[1..17].copy_from_slice(&amount_a.to_le_bytes()); + let call_token_lp = ChainedCall{ + program_id: token_program, + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + account_indices: vec![1] + }; + + chained_call.push(call_token_lp); + chained_call.push(call_token_b); + chained_call.push(call_token_a); + + + let post_states = vec![pool_post.clone(), + pre_states[1].account.clone(), + pre_states[2].account.clone(), + pre_states[3].account.clone(), + pre_states[4].account.clone(), + pre_states[5].account.clone(), + pre_states[6].account.clone()]; + + (post_states.clone(), chained_call) +} + + + +type Instruction = Vec; +//deserialize is not implemented for 33??? +fn main() { + let ProgramInput { + pre_states, + instruction, + } = read_nssa_inputs::(); + + match instruction[0] { + 0 => { + /* + program_id: pool_def_data.token_program_id.clone(), + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + pre_states: &[AccountWithMetadata], + token_program: ProgramId, + balance_in: &[u128] + + */ + + // let token_program_id : &[u32] = bytemuck::cast_slice(&instruction[33..55]); + + /* + let (post_states, chained_call) = new_definition(&pre_states, + &[u128::from_le_bytes(instruction[1..17].try_into().unwrap()), + u128::from_le_bytes(instruction[16..33].try_into().unwrap()),], + ); + + write_nssa_outputs_with_chained_call(pre_states, post_states, chained_call);*/ + } + 1 => { + let intent = SwapIntent { + token_id: AccountId::new(instruction[1..33].try_into().unwrap()), + amount: u128::from_le_bytes(instruction[33..49].try_into().unwrap()), + }; + + let (post_states, chained_call) = swap(&pre_states, &intent); + + write_nssa_outputs_with_chained_call(pre_states, post_states, chained_call); + } + 2 => { + let (post_states, chained_call) = add_liquidity(&pre_states, + &[u128::from_le_bytes(instruction[1..17].try_into().unwrap()), + u128::from_le_bytes(instruction[16..33].try_into().unwrap()),], + AccountId::new(instruction[33..65].try_into().unwrap())); + + write_nssa_outputs_with_chained_call(pre_states, post_states, chained_call); + } + 3 => { + + let (post_states, chained_call) = remove_liquidity(&pre_states); + + write_nssa_outputs_with_chained_call(pre_states, post_states, chained_call); + } + _ => panic!("Invalid instruction"), + }; +} + +struct SwapIntent { + token_id: AccountId, + amount: u128, +} + +fn swap( + pre_states: &[AccountWithMetadata], + intent: &SwapIntent, + ) -> (Vec, Vec) { + + if pre_states.len() != 5 { + panic!("Invalid number of input accounts"); + } + + let pool = &pre_states[0]; + let vault1 = &pre_states[1]; + let vault2 = &pre_states[2]; + + // Verify vaults are in fact vaults + let mut pool_def_data = PoolDefinition::parse(&pool.account.data).unwrap(); + + let mut vault_a = Account::default(); + let mut vault_b = Account::default(); + + if vault1.account_id == pool_def_data.definition_token_a_id { + vault_a = vault1.account.clone(); + } else if vault2.account_id == pool_def_data.definition_token_a_id { + vault_a = vault2.account.clone(); + } else { + panic!("Vault A was no provided"); + } + + if vault1.account_id == pool_def_data.definition_token_b_id { + vault_b = vault1.account.clone(); + } else if vault2.account_id == pool_def_data.definition_token_b_id { + vault_b = vault2.account.clone(); + } else { + panic!("Vault B was no provided"); + } + + // 1. Identify swap direction (a -> b or b -> a) + let mut deposit_a = 0; + let mut deposit_b = 0; + let a_to_b; + if intent.token_id == pool_def_data.definition_token_a_id { + deposit_a = intent.amount; + a_to_b = true; + } else if intent.token_id == pool_def_data.definition_token_b_id { + deposit_b = intent.amount; + a_to_b = false; + } else { + panic!("Intent address is not a token type for the pool"); + } + + // 2. fetch pool reserves + //validates reserves is at least the vaults' balances + assert!(vault_a.balance >= pool_def_data.reserve_a); + assert!(vault_b.balance >= pool_def_data.reserve_b); + //Cannot swap if a reserve is 0 + assert!(pool_def_data.reserve_a > 0); + assert!(pool_def_data.reserve_b > 0); + + // 3. Compute output amount + // Note: no fees + // Compute pool's exchange constant + // let k = pool_def_data.reserve_a * pool_def_data.reserve_b; + let withdraw_a = if a_to_b { 0 } + else { (pool_def_data.reserve_a * deposit_b)/(pool_def_data.reserve_b + deposit_b) }; + let withdraw_b = if a_to_b { (pool_def_data.reserve_b * deposit_a)/(pool_def_data.reserve_a + deposit_a)} + else { 0 }; + + // 4. Slippage check + if a_to_b { + assert!(withdraw_b == 0); } + else{ + assert!(withdraw_a == 0); } + + // 5. Update pool account + let mut pool_post = pool.account.clone(); + let pool_post_definition = PoolDefinition { + definition_token_a_id: pool_def_data.definition_token_a_id.clone(), + definition_token_b_id: pool_def_data.definition_token_b_id.clone(), + vault_a_addr: pool_def_data.vault_a_addr.clone(), + vault_b_addr: pool_def_data.vault_b_addr.clone(), + liquidity_pool_id: pool_def_data.liquidity_pool_id.clone(), + liquidity_pool_cap: pool_def_data.liquidity_pool_cap.clone(), + reserve_a: pool_def_data.reserve_a + deposit_a - withdraw_a, + reserve_b: pool_def_data.reserve_b + deposit_b - withdraw_b, + token_program_id: pool_def_data.token_program_id.clone(), + }; + + pool_post.data = pool_post_definition.into_data(); + + let mut chained_call = Vec::new(); + + let mut instruction_data = [0; 23]; + instruction_data[0] = 1; + + let call_token_a = if a_to_b { + instruction_data[1..17].copy_from_slice(&withdraw_a.to_le_bytes()); + ChainedCall{ + program_id: pool_def_data.token_program_id.clone(), + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + account_indices: vec![0] + } + } else { + instruction_data[1..17].copy_from_slice(&deposit_a.to_le_bytes()); + ChainedCall{ + program_id: pool_def_data.token_program_id.clone(), + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + account_indices: vec![0] + } + }; + + let call_token_b = if a_to_b { + instruction_data[1..17].copy_from_slice(&deposit_b.to_le_bytes()); + ChainedCall{ + program_id: pool_def_data.token_program_id.clone(), + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + account_indices: vec![1] + } + } else { + instruction_data[1..17].copy_from_slice(&withdraw_b.to_le_bytes()); + ChainedCall{ + program_id: pool_def_data.token_program_id.clone(), + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + account_indices: vec![1] + } + }; + + chained_call.push(call_token_a); + chained_call.push(call_token_b); + + let post_states = vec![pool_post.clone(), + pre_states[1].account.clone(), + pre_states[2].account.clone(), + pre_states[3].account.clone(), + pre_states[4].account.clone()]; + + (post_states.clone(), chained_call) +} + + +fn add_liquidity(pre_states: &[AccountWithMetadata], + max_balance_in: &[u128], + main_token: AccountId) -> (Vec, Vec) { + + if pre_states.len() != 7 { + panic!("Invalid number of input accounts"); + } + + let pool = &pre_states[0]; + let vault1 = &pre_states[1]; + let vault2 = &pre_states[2]; + let pool_lp = &pre_states[3]; + let user_a = &pre_states[4]; + let user_b = &pre_states[5]; + let user_lp = &pre_states[6]; + + let mut vault_a = Account::default(); + let mut vault_b = Account::default(); + + let pool_def_data = PoolDefinition::parse(&pool.account.data).unwrap(); + + if vault1.account_id == pool_def_data.definition_token_a_id { + vault_a = vault1.account.clone(); + } else if vault2.account_id == pool_def_data.definition_token_a_id { + vault_a = vault2.account.clone(); + } else { + panic!("Vault A was no provided"); + } + + if vault1.account_id == pool_def_data.definition_token_b_id { + vault_b = vault1.account.clone(); + } else if vault2.account_id == pool_def_data.definition_token_b_id { + vault_b = vault2.account.clone(); + } else { + panic!("Vault B was no provided"); + } + + + if max_balance_in.len() != 2 { + panic!("Invalid number of input balances"); + } + + let max_amount_a = max_balance_in[0]; + let max_amount_b = max_balance_in[1]; + + // 2. Determine deposit amounts + let mut actual_amount_a = 0; + let mut actual_amount_b = 0; + + if main_token == pool_def_data.definition_token_a_id { + actual_amount_a = max_amount_a; + actual_amount_b = (vault_b.balance/vault_a.balance)*actual_amount_a; + } else if main_token == pool_def_data.definition_token_b_id { + actual_amount_b = max_amount_b; + actual_amount_a = (vault_a.balance/vault_b.balance)*actual_amount_b; + } else { + panic!("Mismatch of token types"); //main token does not match with vaults. + } + + + // 3. Validate amounts + assert!(user_a.account.balance >= actual_amount_a && actual_amount_a > 0); + assert!(user_b.account.balance >= actual_amount_b && actual_amount_b > 0); + + // 4. Calculate LP to mint + let delta_lp : u128 = pool_def_data.liquidity_pool_cap * (actual_amount_b/pool_def_data.reserve_b); + + // 5. Update pool account + let mut pool_post = pool.account.clone(); + let pool_post_definition = PoolDefinition { + definition_token_a_id: pool_def_data.definition_token_a_id.clone(), + definition_token_b_id: pool_def_data.definition_token_b_id.clone(), + vault_a_addr: pool_def_data.vault_a_addr.clone(), + vault_b_addr: pool_def_data.vault_b_addr.clone(), + liquidity_pool_id: pool_def_data.liquidity_pool_id.clone(), + liquidity_pool_cap: pool_def_data.liquidity_pool_cap + delta_lp, + reserve_a: pool_def_data.reserve_a + actual_amount_a, + reserve_b: pool_def_data.reserve_b + actual_amount_b, + token_program_id: pool_def_data.token_program_id.clone(), + }; + + pool_post.data = pool_post_definition.into_data(); + + let mut chained_call = Vec::new(); + + let mut instruction_data = [0; 23]; + instruction_data[0] = 1; + + instruction_data[1..17].copy_from_slice(&actual_amount_a.to_le_bytes()); + let call_token_a = ChainedCall{ + program_id: pool_def_data.token_program_id.clone(), + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + account_indices: vec![1] + }; + + + instruction_data[1..17].copy_from_slice(&actual_amount_b.to_le_bytes()); + let call_token_b = ChainedCall{ + program_id: pool_def_data.token_program_id.clone(), + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + account_indices: vec![1] + }; + + instruction_data[1..17].copy_from_slice(&delta_lp.to_le_bytes()); + let call_token_lp = ChainedCall{ + program_id: pool_def_data.token_program_id.clone(), + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + account_indices: vec![1] + }; + + chained_call.push(call_token_lp); + chained_call.push(call_token_b); + chained_call.push(call_token_a); + + + let post_states = vec![pool_post.clone(), + pre_states[1].account.clone(), + pre_states[2].account.clone(), + pre_states[3].account.clone(), + pre_states[4].account.clone(), + pre_states[5].account.clone(), + pre_states[6].account.clone(),]; + + (post_states.clone(), chained_call) + +} + +fn remove_liquidity(pre_states: &[AccountWithMetadata]) -> (Vec, Vec) { + + if pre_states.len() != 7 { + panic!("Invalid number of input accounts"); + } + + let pool = &pre_states[0]; + let vault1 = &pre_states[1]; + let vault2 = &pre_states[2]; + let pool_lp = &pre_states[3]; + let user_a = &pre_states[4]; + let user_b = &pre_states[5]; + let user_lp = &pre_states[6]; + + let mut vault_a = Account::default(); + let mut vault_b = Account::default(); + + let pool_def_data = PoolDefinition::parse(&pool.account.data).unwrap(); + + if vault1.account_id == pool_def_data.definition_token_a_id { + vault_a = vault1.account.clone(); + } else if vault2.account_id == pool_def_data.definition_token_a_id { + vault_a = vault2.account.clone(); + } else { + panic!("Vault A was no provided"); + } + + if vault1.account_id == pool_def_data.definition_token_b_id { + vault_b = vault1.account.clone(); + } else if vault2.account_id == pool_def_data.definition_token_b_id { + vault_b = vault2.account.clone(); + } else { + panic!("Vault B was no provided"); + } + + // 2. Determine deposit amounts + let withdraw_amount_a = pool_def_data.reserve_a * (user_lp.account.balance/pool_def_data.liquidity_pool_cap); + let withdraw_amount_b = pool_def_data.reserve_b * (user_lp.account.balance/pool_def_data.liquidity_pool_cap); + + //3. Validate amounts handled by token programs + + // 4. Calculate LP to reduce cap by + let delta_lp : u128 = (pool_def_data.liquidity_pool_cap*pool_def_data.liquidity_pool_cap - user_lp.account.balance)/pool_def_data.liquidity_pool_cap; + + // 5. Update pool account + let mut pool_post = pool.account.clone(); + let pool_post_definition = PoolDefinition { + definition_token_a_id: pool_def_data.definition_token_a_id.clone(), + definition_token_b_id: pool_def_data.definition_token_b_id.clone(), + vault_a_addr: pool_def_data.vault_a_addr.clone(), + vault_b_addr: pool_def_data.vault_b_addr.clone(), + liquidity_pool_id: pool_def_data.liquidity_pool_id.clone(), + liquidity_pool_cap: pool_def_data.liquidity_pool_cap - delta_lp, + reserve_a: pool_def_data.reserve_a - withdraw_amount_a, + reserve_b: pool_def_data.reserve_b - withdraw_amount_b, + token_program_id: pool_def_data.token_program_id.clone(), + }; + + pool_post.data = pool_post_definition.into_data(); + + let mut chained_call = Vec::new(); + + let mut instruction_data = [0; 23]; + instruction_data[0] = 1; + + instruction_data[1..17].copy_from_slice(&withdraw_amount_a.to_le_bytes()); + let call_token_a = ChainedCall{ + program_id: pool_def_data.token_program_id.clone(), + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + account_indices: vec![1] + }; + + + instruction_data[1..17].copy_from_slice(&withdraw_amount_b.to_le_bytes()); + let call_token_b = ChainedCall{ + program_id: pool_def_data.token_program_id.clone(), + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + account_indices: vec![1] + }; + + instruction_data[1..17].copy_from_slice(&delta_lp.to_le_bytes()); + let call_token_lp = ChainedCall{ + program_id: pool_def_data.token_program_id.clone(), + instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + account_indices: vec![1] + }; + + chained_call.push(call_token_lp); + chained_call.push(call_token_b); + chained_call.push(call_token_a); + + + let post_states = vec![pool_post.clone(), + pre_states[1].account.clone(), + pre_states[2].account.clone(), + pre_states[3].account.clone(), + pre_states[4].account.clone(), + pre_states[5].account.clone(), + pre_states[6].account.clone()]; + + (post_states.clone(), chained_call) + +} \ No newline at end of file diff --git a/nssa/src/program.rs b/nssa/src/program.rs index d3f28b5..c524a17 100644 --- a/nssa/src/program.rs +++ b/nssa/src/program.rs @@ -7,7 +7,7 @@ use serde::Serialize; use crate::{ error::NssaError, - program_methods::{AUTHENTICATED_TRANSFER_ELF, PINATA_ELF, TOKEN_ELF}, + program_methods::{AUTHENTICATED_TRANSFER_ELF, PINATA_ELF, TOKEN_ELF, AMM_ELF}, }; /// Maximum number of cycles for a public execution. @@ -95,6 +95,12 @@ impl Program { // `program_methods` Self::new(TOKEN_ELF.to_vec()).unwrap() } + + pub fn amm() -> Self { + // This unwrap wont panic since the `AMM_ELF` comes from risc0 build of + // `program_methods` + Self::new(AMM_ELF.to_vec()).unwrap() + } } // TODO: Testnet only. Refactor to prevent compilation on mainnet. diff --git a/nssa/test_program_methods/guest/src/bin/pool.rs b/nssa/test_program_methods/guest/src/bin/pool.rs deleted file mode 100644 index d720fbf..0000000 --- a/nssa/test_program_methods/guest/src/bin/pool.rs +++ /dev/null @@ -1,317 +0,0 @@ -use nssa_core::{ - account::{Account, AccountId, AccountWithMetadata, Data}, - program::{ProgramInput, read_nssa_inputs, write_nssa_outputs}, -}; - -// The token program has two functions: -// 1. New token definition. -// Arguments to this function are: -// * Two **default** accounts: [definition_account, holding_account]. -// The first default account will be initialized with the token definition account values. The second account will -// be initialized to a token holding account for the new token, holding the entire total supply. -// * An instruction data of 23-bytes, indicating the total supply and the token name, with -// the following layout: -// [0x00 || total_supply (little-endian 16 bytes) || name (6 bytes)] -// The name cannot be equal to [0x00, 0x00, 0x00, 0x00, 0x00, 0x00] -// 2. Token transfer -// Arguments to this function are: -// * Two accounts: [sender_account, recipient_account]. -// * An instruction data byte string of length 23, indicating the total supply with the following layout -// [0x01 || amount (little-endian 16 bytes) || 0x00 || 0x00 || 0x00 || 0x00 || 0x00 || 0x00]. - - -//TODO: pool should have 2 tokens - - -//TODO: correct these values -const TOKEN_DEFINITION_TYPE: u8 = 0; -const POOL_DEFINITION_DATA_SIZE: usize = 19; - -const TOKEN_HOLDING_TYPE: u8 = 1; -const TOKEN_HOLDING_DATA_SIZE: usize = 49; - -struct PoolDefinition{ - account_type: u8, - name_pool: [u8; 6], //TODO: unsure - name_token_a: [u8; 6], //TODO: specifies token A - name_token_b: [u8; 6], //TODO: specifies token B -} - -struct PoolHolding { - account_type: u8, - definition_pool_id: AccountId, - definition_token_a_id: AccountId, - definition_token_b_id: AccountId, - definition_token_lp_id: AccountId, -} - - -impl PoolDefinition { - fn into_data(self) -> Vec { - let mut bytes = [0; POOL_DEFINITION_DATA_SIZE]; - bytes[0] = self.account_type; - bytes[1..7].copy_from_slice(&self.name_pool); - bytes[7..13].copy_from_slice(&self.name_token_a); - bytes[13..].copy_from_slice(&self.name_token_b); - bytes.into(); - } -} - -impl PoolHolding { - fn new(definition_pool_id: &AccountId, - definition_token_a_id: &AccountId, - definition_token_b_id: &AccountId, - definition_token_lp_id: &AccountId, - ) -> Self { - Self { - account_type: TOKEN_HOLDING_TYPE, //TODO - definition_pool_id: definition_pool_id.clone(), - definition_token_a_id: definition_token_a_id.clone(), - definition_token_b_id: definition_token_b_id.clone(), - definition_token_lp_id: definition_token_lp_id.clone(), - } - } - - fn parse(data: &[u8]) -> Option { - if data.len() != TOKEN_HOLDING_DATA_SIZE || data[0] != TOKEN_HOLDING_TYPE { - None - } else { - let account_type = data[0]; - let definition_pool_id = AccountId::new(data[1..33].try_into().unwrap()); - let definition_token_a_id = AccountId::new(data[33..65].try_into().unwrap()); - let definition_token_b_id = AccountId::new(data[65..97].try_into().unwrap()); - let definition_token_lp_id = AccountId::new(data[97..129]); - Some(Self { - definition_pool_id, - definition_token_a, - definition_token_b, - definition_token_lp_id, - }) - } - } - - fn into_data(self) -> Data { - let mut bytes = [0; TOKEN_HOLDING_DATA_SIZE]; - bytes[0] = self.account_type; - bytes[1..33].copy_from_slice(&self.definition_pool_id.to_bytes()); - bytes[33..65].copy_from_slice(&self.definition_token_a_id.to_bytes()); - bytes[65..97].copy_from_slice(&self.definition_token_b_id.to_bytes()); - bytes[97..].copy_from_slice(&self.definition_token_lp_id.to_bytes()); - bytes.into() - } -} - - -fn initialize_pool(pre_state: &[AccountWithMetadata], balance_in: [u128]) { - //Pool accounts: pool itself, and its 2 vaults and LP token - //2 accounts for funding tokens - //initial funder's LP account - if pre_states.len() != 7 { - panic!("Invalid number of input account") - } - - if balance_in.len() != 2 { - panic!("Invalid number of balance") - } - - let mut pool = pre_state[0]; - let mut vault_a = pre_state[1]; - let mut vault_b = pre_state[2]; - let mut pool_lp = pre_state[3]; - let mut fund_a = pre_state[4]; - let mut fund_b = pre_state[5]; - let mut user_lp = pre_state[6]; - - if pool.account != Account::default() || !pool.is_authorized { - return; - } - - if vault_a.account != Account::default() || !vault_a.is_authorized { - return; - } - - if pool_b.account != Account::default() || !vault_b.is_authorized { - return; - } - - if pool_lp.account != Account::default() || !pool_lp.account.is_authorized { - return; - } - - if !fund_a.is_authorized || !fund_b.is_authorized { - return; - } - - if user_lp.account != Account::default() || !user_lp.account.is_authorized { - return; - } - - - let balance_a = balance_in[0]; - let balance_b = balance_in[1]; - - // Prevents pool constant coefficient (k) from being 0. - assert!(balance_a > 0); - assert!(balance_b > 0); - - // Verify token_a and token_b are different - token_a_id = fund_a.account.data.parse().definition_id; - token_b_id = fund_b.account.data.parse().definition_id; - assert!(token_a_id != token_b_id); - - // 1. Account verification - //TODO: check a pool for (tokenA, tokenB) does not already exist? - - - // 2. Initialize stake - let pool_data = PoolDefinition::new(pool_id, - token_a_id, - token_b_id).into_data(); - - - // 3. LP token minting calculations - //TODO - - // 4. Cross program calls - //TODO -} - -fn swap(pre_states: &[AccountWithMetadata], balance_in: [u128], min_amount_out: u128) { - //Does not require pool's LP account - if pre_states.len() != 5 { - panic!("Invalid number of input accounts"); - } - let pool = &pre_states[0]; - let vault_a = &pre_states[1]; - let vault_b = &pre_states[2]; - let user_a = &pre_states[3]; - let user_b = &pre_states[4]; - - if balance_in.len() != 2 { - panic!("Invalid number of input balances"); - } - - //TODO: return here - let mut pool_holding = - PoolHolding::parse(&pool.account.data).expect("Invalid pool data"); - - //TODO: return here - //TODO: a new account must be minted for the recipient regardless. - //So, we should receive 3 accounts for pre_state. - //TODO: fix sender_holding - let mut user_holding = if recipient.account == Account::default() { - TokenHolding::new(&sender_holding.definition_id); - }; - - - // 1. Identify swap direction (a -> b or b -> a) - // Exactly one should be 0. - let in_a = balance_in[0]; - let in_b = balance_in[1]; - assert!( in_a == 0 || in_b == 0); - assert!( in_a > 0 || in_b > 0); - let a_to_b: bool = if in_a > 0 { true } else { false }; - - // 2. fetch pool reserves - assert!(vault_a.account.balance > 0); - assert!(vault_b.account.balance > 0); - - // 3. Compute output amount - // Note: no fees - // Compute pool's exchange constant - let k = vault_a.account.balance * vault_b.account.balance; - let net_in_a = in_a; - let net_in_b = in_b; - let amount_out_a = if a_to_b { (vault_b.balance * net_in_b)/(vault_a.account.balance + net_in_a)} - else { 0 }; - let amount_out_b = if a_to_b { 0 } - else { - (vault_a.account.balance * net_in_a)/(vault_b.account.balance + net_in_b) }; - - // 4. Slippage check - if a_to_b { - assert!(amount_out_a > min_amount_out); } - else{ - assert!(amount_out_b > min_amount_out); } - - //TODO Note to self: step 5 unnecessary (update reserves) - - // 6. Transfer tokens (Cross call) - //TODO - - // 7. Result - //TODO - -} - - - -fn add_liquidity(pre_state: &[AccountWithMetadata], max_balance_in: [u128], main_token: AccountId) { - 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_lp = &pre_states[3]; - let user_a = &pre_states[4]; - let user_b = &pre_states[5]; - let user_lp = &pre_state[6]; - - if balance_in.len() != 2 { - panic!("Invalid number of input balances"); - } - - //TODO: add authorization checks if need be; - //might be redundant - - max_amount_a = balance_in[0]; - max_amount_b = balance_in[1]; - - // 2. Determine deposit amounts - pool_data = pool.account.data.parse(); - let mut actual_amount_a = 0; - let mut actual_amount_b = 0; - - if main_token == pool_data.definition_token_a { - actual_amount_a = max_amount_a; - actual_amount_b = (vault_b.account.balance/vault_a.account.balance)*actual_amount_a; - } else if main_token == pool_data.definition_token_b { - actual_amount_b = max_amount_b; - actual_amount_a = (vault_a.account.balance/vault_b.account.balance)*actual_amount_b; - } else { - return; //main token does not match with vaults. - } - - // 3. Validate amounts - assert!(user_a.account.balance >= actual_amount_a && actual_amount_a > 0); - assert!(user_b.account.balance >= actual_amount_b && actual_amount_b > 0) - - // 4. Calculate LP to mint - //TODO -} - - - - -fn remove_liquidity(pre_state: &[AccountWithMetadata], max_balance_in: [u128], main_token: AccountId) { - 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_lp = &pre_states[3]; - let user_a = &pre_states[4]; - let user_b = &pre_states[5]; - let user_lp = &pre_states[6]; - - if balance_in.len() != 2 { - panic!("Invalid number of input balances"); - } - - assert!(user_lp.account.balance) - //TODO -} \ No newline at end of file