From 90d65803773444d08e0b129350638c2f1c8e3318 Mon Sep 17 00:00:00 2001 From: jonesmarvin8 <83104039+jonesmarvin8@users.noreply.github.com> Date: Tue, 2 Dec 2025 17:21:49 -0500 Subject: [PATCH] simplified swap logic removed branching logic from swap logic, and adjusted tests accordingly. --- nssa/program_methods/guest/src/bin/amm.rs | 516 ++++++---------------- nssa/src/state.rs | 215 --------- 2 files changed, 137 insertions(+), 594 deletions(-) diff --git a/nssa/program_methods/guest/src/bin/amm.rs b/nssa/program_methods/guest/src/bin/amm.rs index f4ed7a3..168996b 100644 --- a/nssa/program_methods/guest/src/bin/amm.rs +++ b/nssa/program_methods/guest/src/bin/amm.rs @@ -5,7 +5,7 @@ use nssa_core::{ use bytemuck; -// The AMM program has four functions: +// The AMM program has five functions (four directly accessible via instructions): // 1. New AMM definition. // Arguments to this function are: // * Seven **default** accounts: [amm_pool, vault_holding_a, vault_holding_b, pool_lp, user_holding_a, user_holding_b, user_holding_lp]. @@ -32,6 +32,14 @@ use bytemuck; // [0x02 || array of max amounts (little-endian 16 bytes) || TOKEN_DEFINITION_ID (for primary)]. // 4. Remove liquidity // * Input instruction set [0x03]. +// - Swap logic +// Arguments of this function are: +// * Four accounts: [user_deposit_tx, vault_deposit_tx, vault_withdraw_tx, user_withdraw_tx]. +// user_deposit_tx and vault_deposit_tx define deposit transaction. +// vault_withdraw_tx and user_withdraw_tx define withdraw transaction. +// * deposit_amount is the amount for user_deposit_tx -> vault_deposit_tx transfer. +// * reserve_amounts is the pool's reserves; used to compute the withdraw amount. +// * Outputs the token transfers as a Vec and the withdraw amount. const POOL_DEFINITION_DATA_SIZE: usize = 240; @@ -325,46 +333,25 @@ fn swap( if pre_states.len() != 5 { panic!("Invalid number of input accounts"); } - //TODO: get rid of - let pool = &pre_states[0]; - let vault1 = &pre_states[1]; - let vault2 = &pre_states[2]; + let vault_a = &pre_states[1]; + let vault_b = &pre_states[2]; let user_a = &pre_states[3]; let user_b = &pre_states[4]; // Verify vaults are in fact vaults let pool_def_data = PoolDefinition::parse(&pool.account.data).unwrap(); - let vault_a = if vault1.account_id == pool_def_data.vault_a_addr { - vault1.clone() - } else if vault2.account_id == pool_def_data.vault_a_addr { - vault2.clone() - } else { + if vault_a.account_id != pool_def_data.vault_a_addr { panic!("Vault A was not provided"); - }; + } - let vault_b = if vault1.account_id == pool_def_data.vault_b_addr { - vault1.clone() - } else if vault2.account_id == pool_def_data.vault_b_addr { - vault2.clone() - } else { + if vault_b.account_id != pool_def_data.vault_b_addr { panic!("Vault B was not provided"); - }; + } - // 1. Identify swap direction (a -> b or b -> a) - let a_to_b = if token_id == pool_def_data.definition_token_a_id { - true - } else if token_id == pool_def_data.definition_token_b_id { false } - else { - panic!("AccountId is not a token type for the pool"); - }; - - let deposit_a = if a_to_b { amount } else { 0 }; - let deposit_b = if a_to_b { 0 } else { amount }; - - // 2. fetch pool reserves + // fetch pool reserves //validates reserves is at least the vaults' balances assert!(TokenHolding::parse(&vault_a.account.data).unwrap().balance >= pool_def_data.reserve_a); assert!(TokenHolding::parse(&vault_b.account.data).unwrap().balance >= pool_def_data.reserve_b); @@ -372,24 +359,24 @@ fn swap( 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 }; + let (chained_call, [deposit_a, withdraw_a], [deposit_b, withdraw_b]) + = if token_id == pool_def_data.definition_token_a_id { + let (chained_call, withdraw_b) = swap_logic(&[user_a.clone(), vault_a.clone(), vault_b.clone(), user_b.clone()], + amount, + &[pool_def_data.reserve_a, pool_def_data.reserve_b]); + + (chained_call, [amount, 0], [0, withdraw_b]) + } else if token_id == pool_def_data.definition_token_b_id { + let (chained_call, withdraw_a) = swap_logic(&[user_b.clone(), vault_b.clone(), vault_a.clone(), user_a.clone()], + amount, + &[pool_def_data.reserve_b, pool_def_data.reserve_a]); - // 4. Slippage check - if a_to_b { - assert!(withdraw_b != 0); - assert!(withdraw_a == 0); } - else { - assert!(withdraw_a != 0); - assert!(withdraw_b == 0); } + (chained_call, [0, withdraw_a], [amount, 0]) + } else { + panic!("AccountId is not a token type for the pool"); + }; - // 5. Update pool account + // 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(), @@ -404,59 +391,6 @@ fn swap( }; pool_post.data = pool_post_definition.into_data(); - - let mut chained_call = Vec::new(); - - let call_token_a = if a_to_b { - let mut instruction_data = [0; 23]; - instruction_data[0] = 1; - instruction_data[1..17].copy_from_slice(&deposit_a.to_le_bytes()); - let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); - - ChainedCall{ - program_id: pool_def_data.token_program_id, - instruction_data: instruction_data, - pre_states: vec![user_a.clone(), vault_a.clone()] - } - } else { - let mut instruction_data = [0; 23]; - instruction_data[0] = 1; - instruction_data[1..17].copy_from_slice(&withdraw_a.to_le_bytes()); - let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); - - ChainedCall{ - program_id: pool_def_data.token_program_id, - instruction_data: instruction_data, - pre_states: vec![vault_a.clone(), user_a.clone()] - } - }; - - let call_token_b = if a_to_b { - let mut instruction_data = [0; 23]; - instruction_data[0] = 1; - instruction_data[1..17].copy_from_slice(&withdraw_b.to_le_bytes()); - let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); - ChainedCall{ - program_id: pool_def_data.token_program_id, - instruction_data: instruction_data, - pre_states: vec![vault_b.clone(), user_b.clone()] - } - } else { - let mut instruction_data = [0; 23]; - instruction_data[0] = 1; - instruction_data[1..17].copy_from_slice(&deposit_b.to_le_bytes()); - let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); - - ChainedCall{ - program_id: pool_def_data.token_program_id, - instruction_data: instruction_data, - pre_states: vec![user_b.clone(), vault_b.clone()] - } - - }; - - chained_call.push(call_token_a); - chained_call.push(call_token_b); let post_states = vec![ pool_post.clone(), @@ -468,6 +402,58 @@ fn swap( (post_states.clone(), chained_call) } +fn swap_logic( + pre_states: &[AccountWithMetadata], + deposit_amount: u128, + reserve_amounts: &[u128], +) -> (Vec, u128) +{ + + let user_deposit_tx = pre_states[0].clone(); + let vault_deposit_tx = pre_states[1].clone(); + let vault_withdraw_tx = pre_states[2].clone(); + let user_withdraw_tx = pre_states[3].clone(); + + let reserve_deposit_vault_amount = reserve_amounts[0]; + let reserve_withdraw_vault_amount = reserve_amounts[1]; + + // Compute withdraw amount + // Compute pool's exchange constant + // let k = pool_def_data.reserve_a * pool_def_data.reserve_b; + let withdraw_amount = (reserve_withdraw_vault_amount * deposit_amount)/(reserve_deposit_vault_amount + deposit_amount); + + //Slippage check + assert!(withdraw_amount != 0); + + let mut chained_call = Vec::new(); + + let mut instruction_data = [0;23]; + instruction_data[0] = 1; + instruction_data[1..17].copy_from_slice(&deposit_amount.to_le_bytes()); + let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); + chained_call.push( + ChainedCall{ + program_id: vault_deposit_tx.account.program_owner, + instruction_data: instruction_data, + pre_states: vec![user_deposit_tx.clone(), vault_deposit_tx.clone()] + } + ); + + let mut instruction_data = [0;23]; + instruction_data[0] = 1; + instruction_data[1..17].copy_from_slice(&withdraw_amount.to_le_bytes()); + let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); + chained_call.push( + ChainedCall{ + program_id: vault_deposit_tx.account.program_owner, + instruction_data: instruction_data, + pre_states: vec![vault_withdraw_tx.clone(), user_withdraw_tx.clone()] + } + ); + + (chained_call, withdraw_amount) +} + fn add_liquidity(pre_states: &[AccountWithMetadata], max_balance_in: &[u128], main_token: AccountId) -> (Vec, Vec) { @@ -3579,30 +3565,39 @@ mod tests { #[test] fn test_call_swap_incorrect_token_type() { let mut pool = Account::default(); - let mut vault1 = Account::default(); - let mut vault2 = Account::default(); + let mut vault_a = Account::default(); + let mut vault_b = Account::default(); let mut pool_lp = Account::default(); - + let mut user_a = Account::default(); + let mut user_b = Account::default(); let definition_token_a_id = AccountId::new([1;32]); let definition_token_b_id = AccountId::new([2;32]); - vault1.data = TokenHolding::into_data( + vault_a.data = TokenHolding::into_data( TokenHolding { account_type: TOKEN_HOLDING_TYPE, definition_id:definition_token_a_id.clone(), balance: 15u128 } ); - vault2.data = TokenHolding::into_data( + vault_b.data = TokenHolding::into_data( TokenHolding { account_type: TOKEN_HOLDING_TYPE, definition_id:definition_token_b_id.clone(), - balance: 15u128 } + balance: 20u128 } ); pool_lp.data = vec![ 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]; + user_a.data = vec![ + 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 + ]; + user_b.data = vec![ + 1, 1, 1, 1, 1, 1, 1, 12, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 + ]; let definition_token_a_id = AccountId::new([1;32]); let definition_token_b_id = AccountId::new([2;32]); @@ -3610,46 +3605,48 @@ mod tests { let vault_b_addr = AccountId::new([6;32]); let liquidity_pool_id = AccountId::new([7;32]); let liquidity_pool_cap: u128 = 30u128; - let reserve_a: u128 = 10; + let reserve_a: u128 = 15; let reserve_b: u128 = 20; - let token_program_id: [u32;8] = [0; 8]; + let token_program_id: [u32;8] = [5; 8]; - pool.data = PoolDefinition::into_data( PoolDefinition { - definition_token_a_id, - definition_token_b_id, - vault_a_addr: vault_a_addr.clone(), - vault_b_addr: vault_b_addr.clone(), - liquidity_pool_id, - liquidity_pool_cap, - reserve_a, - reserve_b, - token_program_id, - }); + pool.data = PoolDefinition::into_data( + PoolDefinition { + definition_token_a_id: definition_token_a_id.clone(), + definition_token_b_id: definition_token_b_id.clone(), + vault_a_addr: vault_a_addr.clone(), + vault_b_addr: vault_b_addr.clone(), + liquidity_pool_id, + liquidity_pool_cap, + reserve_a, + reserve_b, + token_program_id, + } + ); let pre_states = vec![AccountWithMetadata { - account: pool, + account: pool.clone(), is_authorized: true, account_id: AccountId::new([0; 32])}, AccountWithMetadata { - account: vault1, + account: vault_a.clone(), is_authorized: true, account_id: vault_a_addr.clone()}, AccountWithMetadata { - account: vault2, + account: vault_b.clone(), is_authorized: true, account_id: vault_b_addr.clone()}, AccountWithMetadata { - account: pool_lp, + account: user_a.clone(), is_authorized: true, account_id: AccountId::new([4; 32])}, AccountWithMetadata { - account: Account::default(), + account: user_b.clone(), is_authorized: true, account_id: AccountId::new([5; 32])} ]; let amount = 15u128; - let main_token = AccountId::new([9;32]); - let _post_states = swap(&pre_states, amount, main_token); + let token_addr = AccountId::new([42;32]); + let (post_accounts, chain_calls) = swap(&pre_states, amount, token_addr); } #[should_panic(expected = "Vault A was not provided")] @@ -3852,6 +3849,11 @@ mod tests { let reserve_b: u128 = 20; let token_program_id: [u32;8] = [5; 8]; + user_a.program_owner = token_program_id; + user_b.program_owner = token_program_id; + vault_a.program_owner = token_program_id; + vault_b.program_owner = token_program_id; + pool.data = PoolDefinition::into_data( PoolDefinition { definition_token_a_id: definition_token_a_id.clone(), @@ -3888,6 +3890,7 @@ mod tests { account_id: AccountId::new([5; 32])} ]; let amount = 15u128; + let token_addr = definition_token_a_id; let (post_accounts, chain_calls) = swap(&pre_states, amount, token_addr); @@ -3895,16 +3898,16 @@ mod tests { let pool_pre_data = PoolDefinition::parse(&pool.data).unwrap(); let pool_post_data = PoolDefinition::parse(&pool_post.data).unwrap(); assert!(pool_post_data.reserve_a == pool_pre_data.reserve_a + amount); - + let expected_withdraw = (pool_pre_data.reserve_b * amount)/(pool_pre_data.reserve_a + amount); assert!(pool_post_data.reserve_b == pool_pre_data.reserve_b - expected_withdraw); - + let chain_call_a = chain_calls[0].clone(); let chain_call_b = chain_calls[1].clone(); - + assert!(chain_call_b.program_id == token_program_id); assert!(chain_call_a.program_id == token_program_id); - + let mut instruction_data = [0; 23]; instruction_data[0] = 1; instruction_data[1..17].copy_from_slice(&amount.to_le_bytes()); @@ -3926,7 +3929,6 @@ mod tests { assert!(chain_call_b.instruction_data == expected_instruction_data_1); assert!(chain_call_b_account0 == vault_b); assert!(chain_call_b_account1 == user_b); - } #[test] @@ -3976,129 +3978,10 @@ mod tests { let reserve_b: u128 = 20; let token_program_id: [u32;8] = [5; 8]; - pool.data = PoolDefinition::into_data( - PoolDefinition { - definition_token_a_id: definition_token_a_id.clone(), - definition_token_b_id: definition_token_b_id.clone(), - vault_a_addr: vault_a_addr.clone(), - vault_b_addr: vault_b_addr.clone(), - liquidity_pool_id, - liquidity_pool_cap, - reserve_a, - reserve_b, - token_program_id, - } - ); - //swapped order of vault_a and vault_b - let pre_states = vec![AccountWithMetadata { - account: pool.clone(), - is_authorized: true, - account_id: AccountId::new([0; 32])}, - AccountWithMetadata { - account: vault_b.clone(), - is_authorized: true, - account_id: vault_b_addr.clone()}, - AccountWithMetadata { - account: vault_a.clone(), - is_authorized: true, - account_id: vault_a_addr.clone()}, - AccountWithMetadata { - account: user_a.clone(), - is_authorized: true, - account_id: AccountId::new([4; 32])}, - AccountWithMetadata { - account: user_b.clone(), - is_authorized: true, - account_id: AccountId::new([5; 32])} - ]; - let amount = 15u128; - let token_addr = definition_token_a_id; - let (post_accounts, chain_calls) = swap(&pre_states, amount, token_addr); - - let pool_post = post_accounts[0].clone(); - let pool_pre_data = PoolDefinition::parse(&pool.data).unwrap(); - let pool_post_data = PoolDefinition::parse(&pool_post.data).unwrap(); - assert!(pool_post_data.reserve_a == pool_pre_data.reserve_a + amount); - - let expected_withdraw = (pool_pre_data.reserve_b * amount)/(pool_pre_data.reserve_a + amount); - assert!(pool_post_data.reserve_b == pool_pre_data.reserve_b - expected_withdraw); - - let chain_call_a = chain_calls[0].clone(); - let chain_call_b = chain_calls[1].clone(); - - assert!(chain_call_b.program_id == token_program_id); - assert!(chain_call_a.program_id == token_program_id); - - let mut instruction_data = [0; 23]; - instruction_data[0] = 1; - instruction_data[1..17].copy_from_slice(&amount.to_le_bytes()); - let expected_instruction_data_0 = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); - let mut instruction_data = [0; 23]; - instruction_data[0] = 1; - instruction_data[1..17].copy_from_slice(&expected_withdraw.to_le_bytes()); - let expected_instruction_data_1 = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); - - let chain_call_a_account0 = chain_call_a.pre_states[0].account.clone(); - let chain_call_a_account1 = chain_call_a.pre_states[1].account.clone(); - - let chain_call_b_account0 = chain_call_b.pre_states[0].account.clone(); - let chain_call_b_account1 = chain_call_b.pre_states[1].account.clone(); - - assert!(chain_call_a.instruction_data == expected_instruction_data_0); - assert!(chain_call_a_account0 == user_a); - assert!(chain_call_a_account1 == vault_a); - assert!(chain_call_b.instruction_data == expected_instruction_data_1); - assert!(chain_call_b_account0 == vault_b); - assert!(chain_call_b_account1 == user_b); - - } - - #[test] - fn test_call_swap_successful_chain_call_3() { - let mut pool = Account::default(); - let mut vault_a = Account::default(); - let mut vault_b = Account::default(); - let mut pool_lp = Account::default(); - let mut user_a = Account::default(); - let mut user_b = Account::default(); - - let definition_token_a_id = AccountId::new([1;32]); - let definition_token_b_id = AccountId::new([2;32]); - - vault_a.data = TokenHolding::into_data( - TokenHolding { account_type: TOKEN_HOLDING_TYPE, - definition_id:definition_token_a_id.clone(), - balance: 15u128 } - ); - - vault_b.data = TokenHolding::into_data( - TokenHolding { account_type: TOKEN_HOLDING_TYPE, - definition_id:definition_token_b_id.clone(), - balance: 20u128 } - ); - - pool_lp.data = vec![ - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 - ]; - user_a.data = vec![ - 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 - ]; - user_b.data = vec![ - 1, 1, 1, 1, 1, 1, 1, 12, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 - ]; - - let definition_token_a_id = AccountId::new([1;32]); - let definition_token_b_id = AccountId::new([2;32]); - let vault_a_addr = AccountId::new([5;32]); - let vault_b_addr = AccountId::new([6;32]); - let liquidity_pool_id = AccountId::new([7;32]); - let liquidity_pool_cap: u128 = 30u128; - let reserve_a: u128 = 15; - let reserve_b: u128 = 20; - let token_program_id: [u32;8] = [5; 8]; + user_a.program_owner = token_program_id; + user_b.program_owner = token_program_id; + vault_a.program_owner = token_program_id; + vault_b.program_owner = token_program_id; pool.data = PoolDefinition::into_data( PoolDefinition { @@ -4147,133 +4030,8 @@ mod tests { let expected_withdraw = (pool_pre_data.reserve_a * amount)/(pool_pre_data.reserve_b + amount); assert!(pool_post_data.reserve_a == pool_pre_data.reserve_a - expected_withdraw); - let chain_call_a = chain_calls[0].clone(); - let chain_call_b = chain_calls[1].clone(); - - assert!(chain_call_b.program_id == token_program_id); - assert!(chain_call_a.program_id == token_program_id); - - let mut instruction_data = [0; 23]; - instruction_data[0] = 1; - instruction_data[1..17].copy_from_slice(&expected_withdraw.to_le_bytes()); - let expected_instruction_data_0 = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); - let mut instruction_data = [0; 23]; - instruction_data[0] = 1; - instruction_data[1..17].copy_from_slice(&amount.to_le_bytes()); - let expected_instruction_data_1 = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); - - let chain_call_a_account0 = chain_call_a.pre_states[0].account.clone(); - let chain_call_a_account1 = chain_call_a.pre_states[1].account.clone(); - - let chain_call_b_account0 = chain_call_b.pre_states[0].account.clone(); - let chain_call_b_account1 = chain_call_b.pre_states[1].account.clone(); - - assert!(chain_call_a.instruction_data == expected_instruction_data_0); - assert!(chain_call_a_account0 == vault_a); - assert!(chain_call_a_account1 == user_a); - assert!(chain_call_b.instruction_data == expected_instruction_data_1); - assert!(chain_call_b_account0 == user_b); - assert!(chain_call_b_account1 == vault_b); - - } - - #[test] - fn test_call_swap_successful_chain_call_4() { - let mut pool = Account::default(); - let mut vault_a = Account::default(); - let mut vault_b = Account::default(); - let mut pool_lp = Account::default(); - let mut user_a = Account::default(); - let mut user_b = Account::default(); - - let definition_token_a_id = AccountId::new([1;32]); - let definition_token_b_id = AccountId::new([2;32]); - - vault_a.data = TokenHolding::into_data( - TokenHolding { account_type: TOKEN_HOLDING_TYPE, - definition_id:definition_token_a_id.clone(), - balance: 15u128 } - ); - - vault_b.data = TokenHolding::into_data( - TokenHolding { account_type: TOKEN_HOLDING_TYPE, - definition_id:definition_token_b_id.clone(), - balance: 20u128 } - ); - - pool_lp.data = vec![ - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 - ]; - user_a.data = vec![ - 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 - ]; - user_b.data = vec![ - 1, 1, 1, 1, 1, 1, 1, 12, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 - ]; - - let definition_token_a_id = AccountId::new([1;32]); - let definition_token_b_id = AccountId::new([2;32]); - let vault_a_addr = AccountId::new([5;32]); - let vault_b_addr = AccountId::new([6;32]); - let liquidity_pool_id = AccountId::new([7;32]); - let liquidity_pool_cap: u128 = 30u128; - let reserve_a: u128 = 15; - let reserve_b: u128 = 20; - let token_program_id: [u32;8] = [5; 8]; - - pool.data = PoolDefinition::into_data( - PoolDefinition { - definition_token_a_id: definition_token_a_id.clone(), - definition_token_b_id: definition_token_b_id.clone(), - vault_a_addr: vault_a_addr.clone(), - vault_b_addr: vault_b_addr.clone(), - liquidity_pool_id, - liquidity_pool_cap, - reserve_a, - reserve_b, - token_program_id, - } - ); - - //swapped order of vaults - let pre_states = vec![AccountWithMetadata { - account: pool.clone(), - is_authorized: true, - account_id: AccountId::new([0; 32])}, - AccountWithMetadata { - account: vault_a.clone(), - is_authorized: true, - account_id: vault_a_addr.clone()}, - AccountWithMetadata { - account: vault_b.clone(), - is_authorized: true, - account_id: vault_b_addr.clone()}, - AccountWithMetadata { - account: user_a.clone(), - is_authorized: true, - account_id: AccountId::new([4; 32])}, - AccountWithMetadata { - account: user_b.clone(), - is_authorized: true, - account_id: AccountId::new([5; 32])} - ]; - let amount = 15u128; - let token_addr = definition_token_b_id; - let (post_accounts, chain_calls) = swap(&pre_states, amount, token_addr); - - let pool_post = post_accounts[0].clone(); - let pool_pre_data = PoolDefinition::parse(&pool.data).unwrap(); - let pool_post_data = PoolDefinition::parse(&pool_post.data).unwrap(); - assert!(pool_post_data.reserve_b == pool_pre_data.reserve_b + amount); - - let expected_withdraw = (pool_pre_data.reserve_a * amount)/(pool_pre_data.reserve_b + amount); - assert!(pool_post_data.reserve_a == pool_pre_data.reserve_a - expected_withdraw); - - let chain_call_a = chain_calls[0].clone(); - let chain_call_b = chain_calls[1].clone(); + let chain_call_b = chain_calls[0].clone(); + let chain_call_a = chain_calls[1].clone(); assert!(chain_call_b.program_id == token_program_id); assert!(chain_call_a.program_id == token_program_id); diff --git a/nssa/src/state.rs b/nssa/src/state.rs index 9c373ed..0052799 100644 --- a/nssa/src/state.rs +++ b/nssa/src/state.rs @@ -3445,219 +3445,4 @@ pub mod tests { assert!(pool_post == expected_pool); } - - #[test] - fn test_simple_amm_swap_3() { - let (state, vec_private_keys, vec_id, vec_amounts) = initialize_amm(); - let mut state: V02State = state; - - let temp_amt = vec_amounts[0]; - let init_balance_a = vec_amounts[1]; - let init_balance_b = vec_amounts[2]; - let user_a_amount = vec_amounts[3]; - let user_b_amount = vec_amounts[4]; - - let token_a_holding_key = &vec_private_keys[0]; - let token_a_definition_key = &vec_private_keys[1]; - let token_b_holding_key = &vec_private_keys[2]; - let token_b_definition_key = &vec_private_keys[3]; - let pool_lp_definition_key = &vec_private_keys[4]; - let user_a_holding_key = &vec_private_keys[5]; - let user_b_holding_key = &vec_private_keys[6]; - let vault_a_key = &vec_private_keys[7]; - let vault_b_key = &vec_private_keys[8]; - let user_lp_holding_key = &vec_private_keys[9]; - let pool_key = &vec_private_keys[10]; - let pool_lp_holding_key = &vec_private_keys[11]; - - let token_a_holding_id = vec_id[0]; - let token_a_definition_id = vec_id[1]; - let token_b_holding_id = vec_id[2]; - let token_b_definition_id = vec_id[3]; - let token_lp_definition_id = vec_id[4]; - let user_a_holding_id = vec_id[5]; - let user_b_holding_id = vec_id[6]; - let vault_a_id = vec_id[7]; - let vault_b_id = vec_id[8]; - let user_lp_holding_id = vec_id[9]; - let pool_id = vec_id[10]; - let pool_lp_holding_id = vec_id[11]; - - //Initialize swap user accounts - let swap_user_a_holding_key = PrivateKey::try_new([21; 32]).unwrap(); - let swap_user_a_holding_id = - AccountId::from(&PublicKey::new_from_private_key(&swap_user_a_holding_key)); - let swap_user_a_amount: u128 = 5000; - - let swap_user_b_holding_key = PrivateKey::try_new([22; 32]).unwrap(); - let swap_user_b_holding_id = - AccountId::from(&PublicKey::new_from_private_key(&swap_user_b_holding_key)); - let swap_user_b_amount: u128 = 5000; - - // Initialize Swap User account for Token A - let mut instruction: [u8; 23] = [0; 23]; - instruction[0] = 1; //transfer - instruction[1..17].copy_from_slice(&swap_user_a_amount.to_le_bytes()); - - let message = public_transaction::Message::try_new( - Program::token().id(), - vec![token_a_holding_id, swap_user_a_holding_id], - vec![2], - instruction, - ) - .unwrap(); - - let witness_set = - public_transaction::WitnessSet::for_message(&message, &[&token_a_holding_key]); - let tx = PublicTransaction::new(message, witness_set); - state.transition_from_public_transaction(&tx).unwrap(); - - // Initialize Swap User account for Token B - let mut instruction: [u8; 23] = [0; 23]; - instruction[0] = 1; //transfer - instruction[1..17].copy_from_slice(&swap_user_b_amount.to_le_bytes()); - - let message = public_transaction::Message::try_new( - Program::token().id(), - vec![token_b_holding_id, swap_user_b_holding_id], - vec![state.get_account_by_id(&token_b_holding_id).nonce], - instruction, - ) - .unwrap(); - - let witness_set = - public_transaction::WitnessSet::for_message(&message, &[&token_b_holding_key]); - let tx = PublicTransaction::new(message, witness_set); - state.transition_from_public_transaction(&tx).unwrap(); - - // Swap - let main_addr = token_b_definition_id; - let swap_b: u128 = 500; - - let mut instruction: Vec = Vec::new(); - instruction.push(1); - instruction.extend_from_slice(&swap_b.to_le_bytes()); - instruction.extend_from_slice(main_addr.value()); - - //order of vault_a and vault_b are swapped. - let message = public_transaction::Message::try_new( - Program::amm().id(), - vec![ - pool_id, - vault_b_id, - vault_a_id, - swap_user_a_holding_id, - swap_user_b_holding_id, - ], - vec![ - state.get_account_by_id(&pool_id).nonce, - state.get_account_by_id(&vault_b_id).nonce, - state.get_account_by_id(&vault_a_id).nonce, - state - .get_account_by_id(&swap_user_a_holding_id) - .nonce, - state - .get_account_by_id(&swap_user_b_holding_id) - .nonce, - ], - instruction, - ) - .unwrap(); - - let witness_set = public_transaction::WitnessSet::for_message( - &message, - &[ - &pool_key, - &vault_a_key, - &vault_b_key, - &swap_user_a_holding_key, - &swap_user_b_holding_key, - ], - ); - - let tx = PublicTransaction::new(message, witness_set); - state.transition_from_public_transaction(&tx).unwrap(); - - let pool_post = state.get_account_by_id(&pool_id); - let vault_a_post = state.get_account_by_id(&vault_a_id); - let vault_b_post = state.get_account_by_id(&vault_b_id); - let swap_user_a_post = state.get_account_by_id(&swap_user_a_holding_id); - let swap_user_b_post = state.get_account_by_id(&swap_user_b_holding_id); - - let withdraw_a = (init_balance_a * swap_b)/(init_balance_b + swap_b); - - let expected_pool = Account { - program_owner: Program::amm().id(), - balance: 0u128, - data: PoolDefinition::into_data( - PoolDefinition { - definition_token_a_id: token_a_definition_id, - definition_token_b_id: token_b_definition_id, - vault_a_addr: vault_a_id, - vault_b_addr: vault_b_id, - liquidity_pool_id: token_lp_definition_id, - liquidity_pool_cap: init_balance_a, - reserve_a: init_balance_a - withdraw_a, - reserve_b: init_balance_b + swap_b, - token_program_id: Program::token().id(), - }), - nonce: 2, - }; - - let expected_vault_a = Account { - program_owner: Program::token().id(), - balance: 0u128, - data: TokenHolding::into_data( - TokenHolding{ - account_type: TOKEN_HOLDING_TYPE, - definition_id: token_a_definition_id, - balance: init_balance_a + temp_amt - withdraw_a, - }), - nonce: 1 - }; - - let expected_vault_b = Account { - program_owner: Program::token().id(), - balance: 0u128, - data: TokenHolding::into_data( - TokenHolding{ - account_type: TOKEN_HOLDING_TYPE, - definition_id: token_b_definition_id, - balance: init_balance_b + temp_amt + swap_b, - }), - nonce: 1 - }; - - let expected_swap_user_a = Account { - program_owner: Program::token().id(), - balance: 0u128, - data: TokenHolding::into_data( - TokenHolding{ - account_type: TOKEN_HOLDING_TYPE, - definition_id: token_a_definition_id, - balance: swap_user_a_amount + withdraw_a, - }), - nonce: 1 - }; - - let expected_swap_user_b = Account { - program_owner: Program::token().id(), - balance: 0u128, - data: TokenHolding::into_data( - TokenHolding{ - account_type: TOKEN_HOLDING_TYPE, - definition_id: token_b_definition_id, - balance: swap_user_b_amount - swap_b, - }), - nonce: 1 - }; - - assert!(vault_a_post == expected_vault_a); - assert!(vault_b_post == expected_vault_b); - assert!(swap_user_a_post == expected_swap_user_a); - assert!(swap_user_b_post == expected_swap_user_b); - assert!(pool_post == expected_pool); - - } - }