diff --git a/nssa/core/src/program.rs b/nssa/core/src/program.rs index 054f993..881cc00 100644 --- a/nssa/core/src/program.rs +++ b/nssa/core/src/program.rs @@ -12,7 +12,7 @@ pub struct ProgramInput { pub instruction: T, } -#[derive(Serialize, Deserialize, Clone)] +#[derive(Serialize, Deserialize, Clone, Default)] #[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))] pub struct ChainedCall { pub program_id: ProgramId, diff --git a/nssa/program_methods/guest/src/bin/amm.rs b/nssa/program_methods/guest/src/bin/amm.rs index 3a25310..ef1a505 100644 --- a/nssa/program_methods/guest/src/bin/amm.rs +++ b/nssa/program_methods/guest/src/bin/amm.rs @@ -1,4 +1,5 @@ use nssa_core::{ + address::Address, account::{Account, AccountId, AccountWithMetadata, Data}, program::{ProgramId, ProgramInput, ChainedCall, read_nssa_inputs, write_nssa_outputs_with_chained_call}, }; @@ -36,8 +37,8 @@ use bytemuck; const POOL_DEFINITION_DATA_SIZE: usize = 240; struct PoolDefinition{ - definition_token_a_id: AccountId, - definition_token_b_id: AccountId, + definition_token_a_id: Address, + definition_token_b_id: Address, vault_a_addr: AccountId, vault_b_addr: AccountId, liquidity_pool_id: AccountId, @@ -50,9 +51,7 @@ struct PoolDefinition{ impl PoolDefinition { fn into_data(self) -> Vec { - //let u8_token_program_id : [u8;32] = bytemuck::cast(self.token_program_id); - //TODO: fix to include token Program - + 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()); @@ -63,7 +62,7 @@ impl PoolDefinition { 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[208..].copy_from_slice(&u8_token_program_id); bytes.into() } @@ -133,148 +132,6 @@ impl TokenHolding { } } -/* - //TODO-useful -fn convert_u8_vec_to_u32_vec() { - -} -*/ - -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 accounts") - } - - if balance_in.len() != 2 { - panic!("Invalid number of balance") - } - - 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]; - - /* - //TODO: fix? - if pool.account == Account::default() || !pool.is_authorized { - panic!("Pool account is uninitiated or not authorized"); - } - */ - - // 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 uninitialized") - } - if pool_lp.account == Account::default() { - panic!("Pool LP must be initialized first") - } - - let amount_a = balance_in[0]; - let amount_b = balance_in[1]; - - // Prevents pool constant coefficient (k) from being 0. - if amount_a == 0 || amount_b == 0 { - panic!("Balances must be nonzero") - } - - // Verify token_a and token_b are different - 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; - - if definition_token_a_id == definition_token_b_id { - panic!("Vaults are for the same token") - } - - // 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.clone(), - vault_b_addr: vault_b.account_id.clone(), - liquidity_pool_id: pool_lp.account_id.clone(), - 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: [u8;32] = [0; 32]; - instruction[0] = 1; - instruction[1..17].copy_from_slice(&amount_a.to_le_bytes()); - - //works? - let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap(); - /* - //works? - let mut instruction_data: [u32;8] = [0; 8]; - instruction_data[0] = u32::from_le_bytes(instruction[0..4].try_into().unwrap()); - instruction_data[1] = u32::from_le_bytes(instruction[4..8].try_into().unwrap()); - instruction_data[2] = u32::from_le_bytes(instruction[8..12].try_into().unwrap()); - instruction_data[3] = u32::from_le_bytes(instruction[12..16].try_into().unwrap()); - instruction_data[4] = u32::from_le_bytes(instruction[16..20].try_into().unwrap()); - instruction_data[5] = u32::from_le_bytes(instruction[20..24].try_into().unwrap()); - instruction_data[6] = u32::from_le_bytes(instruction[24..28].try_into().unwrap()); - instruction_data[7] = u32::from_le_bytes(instruction[28..32].try_into().unwrap()); -*/ - //TODO bytemuck seems to cause issues (non-deterministic) -// instruction_data = bytemuck::try_cast_slice(&instruction); - - //TODO: current testing location issue - let call_token_a = ChainedCall{ - program_id: token_program, // This is likely the issue - instruction_data: &instruction_data,//bytemuck::try_cast_slice(&instruction),//.to_vec(), - pre_states: vec![user_a.clone(), vault_a.clone()] - }; - - /* - 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(), - pre_states: vec![user_b.clone(), vault_b.clone()] - }; - - 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(), - pre_states: vec![pool_lp.clone(), user_lp.clone()] - }; - - 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; fn main() { @@ -306,18 +163,28 @@ fn main() { write_nssa_outputs_with_chained_call(pre_states, post_states, chained_call); } 1 => { - let token_id = AccountId::new(instruction[1..33].try_into().unwrap()); - let amount = u128::from_le_bytes(instruction[33..49].try_into().unwrap()); + let mut token_addr: [u8;32] = [0;32]; + token_addr[0..].copy_from_slice(&instruction[17..49]); + + let token_addr = Address::new(token_addr); + + let amount = u128::from_le_bytes(instruction[1..17].try_into().unwrap()); - let (post_states, chained_call) = swap(&pre_states, amount, token_id); + let (post_states, chained_call) = swap(&pre_states, amount, token_addr); write_nssa_outputs_with_chained_call(pre_states, post_states, chained_call); } 2 => { + + let balance_a = u128::from_le_bytes(instruction[1..17].try_into().unwrap()); + let balance_b = u128::from_le_bytes(instruction[17..33].try_into().unwrap()); + + let mut token_addr: [u8;32] = [0;32]; + token_addr[0..].copy_from_slice(&instruction[33..65]); + let token_addr = Address::new(token_addr); + 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())); + &[balance_a, balance_b], token_addr.clone()); write_nssa_outputs_with_chained_call(pre_states, post_states, chained_call); } 3 => { @@ -330,10 +197,127 @@ fn main() { }; } +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 accounts") + } + + if balance_in.len() != 2 { + panic!("Invalid number of balance") + } + + 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 pool.account != Account::default() || !pool.is_authorized { + panic!("Pool account is initiated or not authorized"); + } + + // 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 uninitialized") + } + if pool_lp.account == Account::default() { + panic!("Pool LP must be initialized first") + } + + let amount_a = balance_in[0]; + let amount_b = balance_in[1]; + + // Prevents pool constant coefficient (k) from being 0. + if amount_a == 0 || amount_b == 0 { + panic!("Balances must be nonzero") + } + + // Verify token_a and token_b are different + let definition_token_a_id = TokenHolding::parse(&user_a.account.data).unwrap().definition_id; + let definition_token_b_id = TokenHolding::parse(&user_b.account.data).unwrap().definition_id; + + if definition_token_a_id == definition_token_b_id { + panic!("Cannot set up a swap for a token with itself.") + } + + // 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.clone(), + vault_b_addr: vault_b.account_id.clone(), + liquidity_pool_id: pool_lp.account_id.clone(), + 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: [u8;32] = [0; 32]; + instruction[0] = 1; + instruction[1..17].copy_from_slice(&amount_a.to_le_bytes()); + let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap(); + + let call_token_a = ChainedCall{ + program_id: token_program, + instruction_data: instruction_data, + pre_states: vec![user_a.clone(), vault_a.clone()] + }; + + instruction[1..17].copy_from_slice(&amount_b.to_le_bytes()); + let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap(); + + let call_token_b = ChainedCall{ + program_id: token_program, + instruction_data: instruction_data, + pre_states: vec![user_b.clone(), vault_b.clone()] + }; + + instruction[1..17].copy_from_slice(&amount_a.to_le_bytes()); + let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap(); + + let call_token_lp = ChainedCall{ + program_id: token_program, + instruction_data: instruction_data, + pre_states: vec![pool_lp.clone(), user_lp.clone()] + }; + + 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 swap( pre_states: &[AccountWithMetadata], amount: u128, - token_id: AccountId, + token_id: Address, ) -> (Vec, Vec) { if pre_states.len() != 5 { @@ -351,7 +335,6 @@ fn swap( let vault1_data = TokenHolding::parse(&vault1.account.data).unwrap(); let vault2_data = TokenHolding::parse(&vault2.account.data).unwrap(); - let mut vault_a = AccountWithMetadata::default(); let mut vault_b = AccountWithMetadata::default(); @@ -387,8 +370,8 @@ fn swap( // 2. fetch pool reserves //validates reserves is at least the vaults' balances - assert!(vault_a.account.balance >= pool_def_data.reserve_a); - assert!(vault_b.account.balance >= pool_def_data.reserve_b); + 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); //Cannot swap if a reserve is 0 assert!(pool_def_data.reserve_a > 0); assert!(pool_def_data.reserve_b > 0); @@ -404,9 +387,9 @@ fn swap( // 4. Slippage check if a_to_b { - assert!(withdraw_b == 0); } - else{ assert!(withdraw_a == 0); } + else { + assert!(withdraw_b == 0); } // 5. Update pool account let mut pool_post = pool.account.clone(); @@ -425,47 +408,63 @@ fn swap( pool_post.data = pool_post_definition.into_data(); let mut chained_call = Vec::new(); + + let call_token_a = ChainedCall::default(); + let call_token_b = ChainedCall::default(); - let mut instruction_data = [0; 23]; - instruction_data[0] = 1; - - let call_token_a = if a_to_b { + 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()); - ChainedCall{ - program_id: pool_def_data.token_program_id.clone(), - instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), - pre_states: vec![user_a.clone(), vault_a.clone()] - } - } else { - 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(), - pre_states: vec![vault_a.clone(), user_a.clone()] - } - }; + let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); - 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(), - pre_states: vec![user_b.clone(), vault_b.clone()] - } - } else { + let call_token_a = ChainedCall{ + program_id: pool_def_data.token_program_id, + instruction_data: instruction_data, + pre_states: vec![user_a.clone(), vault_a.clone()] + }; + + let mut instruction_data = [0; 23]; + instruction_data[0] = 1; 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(), - pre_states: vec![vault_b.clone(), user_b.clone()] - } - }; + let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); + let call_token_b = ChainedCall{ + program_id: pool_def_data.token_program_id, + instruction_data: instruction_data, + pre_states: vec![user_b.clone(), vault_b.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(); + + let call_token_a = ChainedCall{ + program_id: pool_def_data.token_program_id, + instruction_data: instruction_data, + pre_states: vec![vault_a.clone(), user_a.clone()] + }; + + 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(); + + let call_token_b = ChainedCall{ + program_id: pool_def_data.token_program_id, + instruction_data: instruction_data, + pre_states: vec![vault_b.clone(), user_b.clone()] + }; + } + + chained_call.push(call_token_a); chained_call.push(call_token_b); - + let post_states = vec![ - pool_post.clone(), + pool_post.clone(), pre_states[1].account.clone(), pre_states[2].account.clone(), pre_states[3].account.clone(), @@ -474,10 +473,9 @@ fn swap( (post_states.clone(), chained_call) } - fn add_liquidity(pre_states: &[AccountWithMetadata], max_balance_in: &[u128], - main_token: AccountId) -> (Vec, Vec) { + main_token: Address) -> (Vec, Vec) { if pre_states.len() != 7 { panic!("Invalid number of input accounts"); @@ -528,16 +526,17 @@ fn add_liquidity(pre_states: &[AccountWithMetadata], // 2. Determine deposit amounts let mut actual_amount_a = 0; let mut actual_amount_b = 0; - - if vault_b.account.balance == 0 || vault_a.account.balance == 0 { + + let vault_b_balance = TokenHolding::parse(&vault_b.account.data).unwrap().balance; + let vault_a_balance = TokenHolding::parse(&vault_a.account.data).unwrap().balance; + if vault_a_balance == 0 || vault_b_balance == 0 { panic!("Vaults must have nonzero balances"); } - + if pool_def_data.reserve_a == 0 || pool_def_data.reserve_b == 0 { panic!("Reserves must be nonzero"); } - if main_token == pool_def_data.definition_token_a_id { actual_amount_a += max_amount_a; actual_amount_b += (pool_def_data.reserve_b*actual_amount_a)/pool_def_data.reserve_a; @@ -548,14 +547,15 @@ fn add_liquidity(pre_states: &[AccountWithMetadata], panic!("Mismatch of token types"); //main token does not match with vaults. } - // 3. Validate amounts + let user_a_balance = TokenHolding::parse(&user_a.account.data).unwrap().balance; + let user_b_balance = TokenHolding::parse(&user_b.account.data).unwrap().balance; assert!(max_amount_a >= actual_amount_a && max_amount_b >= actual_amount_b); - if user_a.account.balance < actual_amount_a { + if user_a_balance < actual_amount_a { panic!("Insufficient balance"); } - if user_b.account.balance < actual_amount_b { + if user_b_balance < actual_amount_b { panic!("Insufficient balance"); } @@ -563,9 +563,8 @@ fn add_liquidity(pre_states: &[AccountWithMetadata], panic!("A trade amount is 0"); } - // 4. Calculate LP to mint - let delta_lp = (pool_def_data.liquidity_pool_cap *actual_amount_b)/pool_def_data.reserve_b; + let delta_lp = (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(); @@ -580,41 +579,45 @@ fn add_liquidity(pre_states: &[AccountWithMetadata], 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 instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); let call_token_a = ChainedCall{ - program_id: pool_def_data.token_program_id.clone(), - instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + program_id: pool_def_data.token_program_id, + instruction_data: instruction_data, pre_states: vec![user_a.clone(), vault_a] }; - + let mut instruction_data = [0; 23]; + instruction_data[0] = 1; instruction_data[1..17].copy_from_slice(&actual_amount_b.to_le_bytes()); + let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); let call_token_b = ChainedCall{ - program_id: pool_def_data.token_program_id.clone(), - instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + program_id: pool_def_data.token_program_id, + instruction_data: instruction_data, pre_states: vec![user_b.clone(), vault_b] }; - + + let mut instruction_data = [0; 23]; + instruction_data[0] = 1; instruction_data[1..17].copy_from_slice(&delta_lp.to_le_bytes()); + let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap(); let call_token_lp = ChainedCall{ - program_id: pool_def_data.token_program_id.clone(), - instruction_data: bytemuck::cast_slice(&instruction_data).to_vec(), + program_id: pool_def_data.token_program_id, + instruction_data: instruction_data, pre_states: vec![pool_lp.clone(), user_lp.clone()] }; + 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(), @@ -628,7 +631,6 @@ fn add_liquidity(pre_states: &[AccountWithMetadata], } - fn remove_liquidity(pre_states: &[AccountWithMetadata]) -> (Vec, Vec) { if pre_states.len() != 7 { @@ -667,8 +669,9 @@ fn remove_liquidity(pre_states: &[AccountWithMetadata]) -> (Vec, Vec (Vec, Vec (Vec, Vec (Vec, Vec Vec { + /* + //TODO: delete? + let mut u8_token_program_id: [u8;32] = [0;32]; + u8_token_program_id[0..4].clone_from_slice(&self.token_program_id[0].to_le_bytes()); + u8_token_program_id[4..8].clone_from_slice(&self.token_program_id[1].to_le_bytes()); + u8_token_program_id[8..12].clone_from_slice(&self.token_program_id[2].to_le_bytes()); + u8_token_program_id[12..16].clone_from_slice(&self.token_program_id[3].to_le_bytes()); + u8_token_program_id[16..20].clone_from_slice(&self.token_program_id[4].to_le_bytes()); + u8_token_program_id[20..24].clone_from_slice(&self.token_program_id[5].to_le_bytes()); + u8_token_program_id[24..28].clone_from_slice(&self.token_program_id[6].to_le_bytes()); + u8_token_program_id[28..32].clone_from_slice(&self.token_program_id[7].to_le_bytes()); +*/ + + 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..192].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, + }) + } + } +} + + + + + #[test] - fn test_simple_amm_initialized_attempt2() { - //let token_program_id = program_methods::TOKEN_ID; - let program_amm = Program::amm(); + fn test_simple_amm_initialized_attempt_initialize() { let program_token = Program::token().id(); let initial_data = []; @@ -2443,8 +2513,10 @@ pub mod tests { let tx = PublicTransaction::new(message, witness_set); state.transition_from_public_transaction(&tx).unwrap(); + let amt: u128 = token_lp_supply - 1; + //Pool LP should be initialized at full amount (token_lp_supply) instruction[0] = 1; //transfer - instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes()); + instruction[1..17].copy_from_slice(&amt.to_le_bytes()); let message = public_transaction::Message::try_new( Program::token().id(), @@ -2465,25 +2537,12 @@ pub mod tests { let init_balance_a: u128 = 1000; let init_balance_b: u128 = 1000; let token_program_u8: [u8;32] = bytemuck::cast(program_token); - let mut instruction: Vec = Vec::new(); instruction.push(0); instruction.extend_from_slice(&init_balance_a.to_le_bytes()); instruction.extend_from_slice(&init_balance_b.to_le_bytes()); instruction.extend_from_slice(&token_program_u8); - - assert!(token_program_u8.len() == 32); - - - let mut instruction: Vec = Vec::new(); - instruction.push(0); - instruction.extend_from_slice(&init_balance_a.to_le_bytes()); - instruction.extend_from_slice(&init_balance_b.to_le_bytes()); - instruction.extend_from_slice(&token_program_u8); - - assert!(token_program_u8.len() == 32); - let message = public_transaction::Message::try_new( Program::amm().id(), @@ -2497,7 +2556,7 @@ pub mod tests { user_lp_holding_address ], vec![0,0,0,0], - instruction, //TODO: bah + instruction, ) .unwrap(); @@ -2515,37 +2574,1066 @@ pub mod tests { state.transition_from_public_transaction(&tx).unwrap(); + let pool_post = state.get_account_by_address(&pool_address); + let vault_a_post = state.get_account_by_address(&vault_a_address); + let vault_b_post = state.get_account_by_address(&vault_b_address); + let pool_lp_post = state.get_account_by_address(&pool_lp_holding_address); + let user_a_post = state.get_account_by_address(&user_a_holding_address); + let user_b_post = state.get_account_by_address(&user_b_holding_address); + let user_lp_post = state.get_account_by_address(&user_lp_holding_address); + + //TODO: token accounts initialized with temp_amt balance. + // this ensures that amm program does not own any of the token accounts + assert!(TokenHolding::parse(&vault_a_post.data).unwrap().balance == init_balance_a+temp_amt); + assert!(TokenHolding::parse(&vault_b_post.data).unwrap().balance == init_balance_b+temp_amt); + assert!(TokenHolding::parse(&user_lp_post.data).unwrap().balance == init_balance_a+temp_amt); + assert!(PoolDefinition::parse(&pool_post.data).unwrap().reserve_a == init_balance_a); + assert!(PoolDefinition::parse(&pool_post.data).unwrap().reserve_b == init_balance_b); + } + + +#[test] + fn test_simple_amm_initialized_attempt_remove() { + let program_token = Program::token().id(); + + let initial_data = []; + let mut state = + V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs(); + + + let token_a_holding_key = PrivateKey::try_new([1;32]).unwrap(); + let token_a_definition_key = PrivateKey::try_new([2;32]).unwrap(); + let token_a_holding_address = Address::from(&PublicKey::new_from_private_key(&token_a_holding_key)); + let token_a_definition_address = Address::from(&PublicKey::new_from_private_key(&token_a_definition_key)); + let token_a_supply: u128 = 30000; + + let token_b_holding_key = PrivateKey::try_new([3;32]).unwrap(); + let token_b_definition_key = PrivateKey::try_new([4;32]).unwrap(); + let token_b_holding_address = Address::from(&PublicKey::new_from_private_key(&token_b_holding_key)); + let token_b_definition_address = Address::from(&PublicKey::new_from_private_key(&token_b_definition_key)); + let token_b_supply: u128 = 50000; + + let token_lp_holding_key = PrivateKey::try_new([5;32]).unwrap(); + let token_lp_definition_key = PrivateKey::try_new([6;32]).unwrap(); + let token_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&token_lp_holding_key)); + let token_lp_definition_address = Address::from(&PublicKey::new_from_private_key(&token_lp_definition_key)); + let token_lp_supply: u128 = 300000; + + let user_a_holding_key = PrivateKey::try_new([7;32]).unwrap(); + let user_a_holding_address = Address::from(&PublicKey::new_from_private_key(&user_a_holding_key)); + let user_a_amount: u128 = 10000; + + let user_b_holding_key = PrivateKey::try_new([8;32]).unwrap(); + let user_b_holding_address = Address::from(&PublicKey::new_from_private_key(&user_b_holding_key)); + let user_b_amount: u128 = 10000; + let vault_a_key = PrivateKey::try_new([9;32]).unwrap(); + let vault_a_address = Address::from(&PublicKey::new_from_private_key(&vault_a_key)); + let vault_b_key = PrivateKey::try_new([10;32]).unwrap(); + let vault_b_address = Address::from(&PublicKey::new_from_private_key(&vault_b_key)); - /* + let user_lp_holding_key = PrivateKey::try_new([11;32]).unwrap(); + let user_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&user_lp_holding_key)); - assert_eq!(state.get_account_by_address(&to), Account::default()); + let pool_lp_holding_key = PrivateKey::try_new([12;32]).unwrap(); + let pool_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&pool_lp_holding_key)); - let expected_recipient_post = Account { - program_owner: program.id(), - balance: amount, - ..Account::default() - }; + let pool_key = PrivateKey::try_new([13;32]).unwrap(); + let pool_address = Address::from(&PublicKey::new_from_private_key(&pool_key)); - let message = - public_transaction::Message::try_new(program.id(), vec![from, to], vec![0], amount) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key]); + //initialize Tokens: Token A, Token B and Took LP + let mut instruction: [u8;23] = [0;23]; + instruction[1..17].copy_from_slice(&token_a_supply.to_le_bytes()); + instruction[18] = 0x01; //name is not default. + instruction[19] = 0x02; + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_a_definition_address, token_a_holding_address], + vec![], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[], + ); let tx = PublicTransaction::new(message, witness_set); - state.transition_from_public_transaction(&tx).unwrap(); - let recipient_post = state.get_account_by_address(&to); + instruction[1..17].copy_from_slice(&token_b_supply.to_le_bytes()); + instruction[18] = 0x03; //name is not default. + instruction[19] = 0x02; - assert_eq!(recipient_post, expected_recipient_post); + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_b_definition_address, token_b_holding_address], + vec![], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + instruction[1..17].copy_from_slice(&token_lp_supply.to_le_bytes()); + instruction[18] = 0x03; //name is not default. + instruction[19] = 0x04; + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_lp_definition_address, token_lp_holding_address], + vec![], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + + // Initialize User accounts for Token A and B + let mut instruction: [u8;23] = [0;23]; + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&user_a_amount.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_a_holding_address, user_a_holding_address], + vec![0], + 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(); + + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&user_b_amount.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_b_holding_address, user_b_holding_address], + vec![0], + 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(); + + + //TODO: initialize vaults - ideally, we won't need to do this. + let temp_amt = 1u128; + let mut instruction: [u8;23] = [0;23]; + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_a_holding_address, vault_a_address], + vec![1], + 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(); + + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_b_holding_address, vault_b_address], + vec![1], + 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(); + + //need LP accounts for both AMM and user + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_lp_holding_address, user_lp_holding_address], + vec![0], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[&token_lp_holding_key], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + let amt: u128 = token_lp_supply - 1; + //Pool LP should be initialized at full amount (token_lp_supply) + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&amt.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_lp_holding_address, pool_lp_holding_address], + vec![1], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[&token_lp_holding_key], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + //Set up instruction to initialize AMM for (Token-A, Token-B) + let init_balance_a: u128 = 1000; + let init_balance_b: u128 = 1000; + let token_program_u8: [u8;32] = bytemuck::cast(program_token); - */ + let mut instruction: Vec = Vec::new(); + instruction.push(0); + instruction.extend_from_slice(&init_balance_a.to_le_bytes()); + instruction.extend_from_slice(&init_balance_b.to_le_bytes()); + instruction.extend_from_slice(&token_program_u8); + + let message = public_transaction::Message::try_new( + Program::amm().id(), + vec![ + pool_address, + vault_a_address, + vault_b_address, + pool_lp_holding_address, + user_a_holding_address, + user_b_holding_address, + user_lp_holding_address + ], + vec![0,0,0,0], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[ + &pool_key, + &pool_lp_holding_key, + &user_a_holding_key, + &user_b_holding_key, + ], + ); + + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + let mut instruction: Vec = Vec::new(); + instruction.push(3); + + let message = public_transaction::Message::try_new( + Program::amm().id(), + vec![ + pool_address, + vault_a_address, + vault_b_address, + pool_lp_holding_address, + user_a_holding_address, + user_b_holding_address, + user_lp_holding_address + ], + vec![ + state.get_account_by_address(&pool_address).nonce, + state.get_account_by_address(&user_lp_holding_address).nonce, + state.get_account_by_address(&vault_a_address).nonce, + state.get_account_by_address(&vault_b_address).nonce, + ], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[ + &pool_key, + &user_lp_holding_key, + &vault_a_key, + &vault_b_key, + ], + ); + + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + //TODO: + //user holdings + let pool_post = state.get_account_by_address(&pool_address); + let vault_a_post = state.get_account_by_address(&vault_a_address); + let vault_b_post = state.get_account_by_address(&vault_b_address); + let pool_lp_post = state.get_account_by_address(&pool_lp_holding_address); + let user_a_post = state.get_account_by_address(&user_a_holding_address); + let user_b_post = state.get_account_by_address(&user_b_holding_address); + let user_lp_post = state.get_account_by_address(&user_lp_holding_address); + + //TODO: token accounts initialized with temp_amt balance. + // this ensures that amm program does not own any of the token accounts + assert!(TokenHolding::parse(&vault_a_post.data).unwrap().balance == temp_amt); + assert!(TokenHolding::parse(&vault_b_post.data).unwrap().balance == temp_amt); + + assert!(PoolDefinition::parse(&pool_post.data).unwrap().reserve_a == 0); + assert!(PoolDefinition::parse(&pool_post.data).unwrap().reserve_b == 0); + assert!(TokenHolding::parse(&user_lp_post.data).unwrap().balance == 0); + + + + } + +#[test] + fn test_simple_amm_initialized_attempt_add() { + let program_token = Program::token().id(); + + let initial_data = []; + let mut state = + V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs(); + + + let token_a_holding_key = PrivateKey::try_new([1;32]).unwrap(); + let token_a_definition_key = PrivateKey::try_new([2;32]).unwrap(); + let token_a_holding_address = Address::from(&PublicKey::new_from_private_key(&token_a_holding_key)); + let token_a_definition_address = Address::from(&PublicKey::new_from_private_key(&token_a_definition_key)); + let token_a_supply: u128 = 30000; + + let token_b_holding_key = PrivateKey::try_new([3;32]).unwrap(); + let token_b_definition_key = PrivateKey::try_new([4;32]).unwrap(); + let token_b_holding_address = Address::from(&PublicKey::new_from_private_key(&token_b_holding_key)); + let token_b_definition_address = Address::from(&PublicKey::new_from_private_key(&token_b_definition_key)); + let token_b_supply: u128 = 50000; + + let token_lp_holding_key = PrivateKey::try_new([5;32]).unwrap(); + let token_lp_definition_key = PrivateKey::try_new([6;32]).unwrap(); + let token_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&token_lp_holding_key)); + let token_lp_definition_address = Address::from(&PublicKey::new_from_private_key(&token_lp_definition_key)); + let token_lp_supply: u128 = 300000; + + let user_a_holding_key = PrivateKey::try_new([7;32]).unwrap(); + let user_a_holding_address = Address::from(&PublicKey::new_from_private_key(&user_a_holding_key)); + let user_a_amount: u128 = 10000; + + let user_b_holding_key = PrivateKey::try_new([8;32]).unwrap(); + let user_b_holding_address = Address::from(&PublicKey::new_from_private_key(&user_b_holding_key)); + let user_b_amount: u128 = 10000; + + + + let vault_a_key = PrivateKey::try_new([9;32]).unwrap(); + let vault_a_address = Address::from(&PublicKey::new_from_private_key(&vault_a_key)); + + let vault_b_key = PrivateKey::try_new([10;32]).unwrap(); + let vault_b_address = Address::from(&PublicKey::new_from_private_key(&vault_b_key)); + + let user_lp_holding_key = PrivateKey::try_new([11;32]).unwrap(); + let user_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&user_lp_holding_key)); + + let pool_lp_holding_key = PrivateKey::try_new([12;32]).unwrap(); + let pool_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&pool_lp_holding_key)); + + let pool_key = PrivateKey::try_new([13;32]).unwrap(); + let pool_address = Address::from(&PublicKey::new_from_private_key(&pool_key)); + + //initialize Tokens: Token A, Token B and Took LP + let mut instruction: [u8;23] = [0;23]; + instruction[1..17].copy_from_slice(&token_a_supply.to_le_bytes()); + instruction[18] = 0x01; //name is not default. + instruction[19] = 0x02; + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_a_definition_address, token_a_holding_address], + vec![], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + instruction[1..17].copy_from_slice(&token_b_supply.to_le_bytes()); + instruction[18] = 0x03; //name is not default. + instruction[19] = 0x02; + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_b_definition_address, token_b_holding_address], + vec![], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + instruction[1..17].copy_from_slice(&token_lp_supply.to_le_bytes()); + instruction[18] = 0x03; //name is not default. + instruction[19] = 0x04; + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_lp_definition_address, token_lp_holding_address], + vec![], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + + // Initialize User accounts for Token A and B + let mut instruction: [u8;23] = [0;23]; + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&user_a_amount.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_a_holding_address, user_a_holding_address], + vec![0], + 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(); + + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&user_b_amount.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_b_holding_address, user_b_holding_address], + vec![0], + 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(); + + + //TODO: initialize vaults - ideally, we won't need to do this. + let temp_amt = 1u128; + let mut instruction: [u8;23] = [0;23]; + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_a_holding_address, vault_a_address], + vec![1], + 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(); + + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_b_holding_address, vault_b_address], + vec![1], + 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(); + + //need LP accounts for both AMM and user + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_lp_holding_address, user_lp_holding_address], + vec![0], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[&token_lp_holding_key], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + let amt: u128 = token_lp_supply - 1; + //Pool LP should be initialized at full amount (token_lp_supply) + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&amt.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_lp_holding_address, pool_lp_holding_address], + vec![1], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[&token_lp_holding_key], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + //Set up instruction to initialize AMM for (Token-A, Token-B) + let init_balance_a: u128 = 1000; + let init_balance_b: u128 = 1000; + let token_program_u8: [u8;32] = bytemuck::cast(program_token); + + let mut instruction: Vec = Vec::new(); + instruction.push(0); + instruction.extend_from_slice(&init_balance_a.to_le_bytes()); + instruction.extend_from_slice(&init_balance_b.to_le_bytes()); + instruction.extend_from_slice(&token_program_u8); + + let message = public_transaction::Message::try_new( + Program::amm().id(), + vec![ + pool_address, + vault_a_address, + vault_b_address, + pool_lp_holding_address, + user_a_holding_address, + user_b_holding_address, + user_lp_holding_address + ], + vec![0,0,0,0], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[ + &pool_key, + &pool_lp_holding_key, + &user_a_holding_key, + &user_b_holding_key, + ], + ); + + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + let add_a: u128 = 500; + let add_b: u128 = 500; + let main_addr = token_a_definition_address; + let mut instruction: Vec = Vec::new(); + instruction.push(2); + instruction.extend_from_slice(&add_a.to_le_bytes()); + instruction.extend_from_slice(&add_b.to_le_bytes()); + instruction.extend_from_slice(main_addr.value()); + + let message = public_transaction::Message::try_new( + Program::amm().id(), + vec![ + pool_address, + vault_a_address, + vault_b_address, + pool_lp_holding_address, + user_a_holding_address, + user_b_holding_address, + user_lp_holding_address + ], + vec![ + state.get_account_by_address(&pool_address).nonce, + state.get_account_by_address(&pool_lp_holding_address).nonce, + state.get_account_by_address(&user_a_holding_address).nonce, + state.get_account_by_address(&user_b_holding_address).nonce, + ], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[ + &pool_key, + &pool_lp_holding_key, + &user_a_holding_key, + &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_address(&pool_address); + let vault_a_post = state.get_account_by_address(&vault_a_address); + let vault_b_post = state.get_account_by_address(&vault_b_address); + let pool_lp_post = state.get_account_by_address(&pool_lp_holding_address); + let user_a_post = state.get_account_by_address(&user_a_holding_address); + let user_b_post = state.get_account_by_address(&user_b_holding_address); + let user_lp_post = state.get_account_by_address(&user_lp_holding_address); + + //TODO: token accounts initialized with temp_amt balance. + // this ensures that amm program does not own any of the token accounts + assert!(TokenHolding::parse(&vault_a_post.data).unwrap().balance == init_balance_a + add_a + temp_amt); + assert!(TokenHolding::parse(&vault_b_post.data).unwrap().balance == init_balance_b + add_b + temp_amt); + + assert!(PoolDefinition::parse(&pool_post.data).unwrap().reserve_a == init_balance_a + add_a); + assert!(PoolDefinition::parse(&pool_post.data).unwrap().reserve_b == init_balance_b + add_b); + assert!(TokenHolding::parse(&user_lp_post.data).unwrap().balance == init_balance_a + add_a + temp_amt); + + } + + +#[test] + fn test_simple_amm_initialized_attempt_swap() { + let program_token = Program::token().id(); + + let initial_data = []; + let mut state = + V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs(); + + + let token_a_holding_key = PrivateKey::try_new([1;32]).unwrap(); + let token_a_definition_key = PrivateKey::try_new([2;32]).unwrap(); + let token_a_holding_address = Address::from(&PublicKey::new_from_private_key(&token_a_holding_key)); + let token_a_definition_address = Address::from(&PublicKey::new_from_private_key(&token_a_definition_key)); + let token_a_supply: u128 = 30000; + + let token_b_holding_key = PrivateKey::try_new([3;32]).unwrap(); + let token_b_definition_key = PrivateKey::try_new([4;32]).unwrap(); + let token_b_holding_address = Address::from(&PublicKey::new_from_private_key(&token_b_holding_key)); + let token_b_definition_address = Address::from(&PublicKey::new_from_private_key(&token_b_definition_key)); + let token_b_supply: u128 = 50000; + + let token_lp_holding_key = PrivateKey::try_new([5;32]).unwrap(); + let token_lp_definition_key = PrivateKey::try_new([6;32]).unwrap(); + let token_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&token_lp_holding_key)); + let token_lp_definition_address = Address::from(&PublicKey::new_from_private_key(&token_lp_definition_key)); + let token_lp_supply: u128 = 300000; + + let user_a_holding_key = PrivateKey::try_new([7;32]).unwrap(); + let user_a_holding_address = Address::from(&PublicKey::new_from_private_key(&user_a_holding_key)); + let user_a_amount: u128 = 10000; + + let user_b_holding_key = PrivateKey::try_new([8;32]).unwrap(); + let user_b_holding_address = Address::from(&PublicKey::new_from_private_key(&user_b_holding_key)); + let user_b_amount: u128 = 10000; + + + + let vault_a_key = PrivateKey::try_new([9;32]).unwrap(); + let vault_a_address = Address::from(&PublicKey::new_from_private_key(&vault_a_key)); + + let vault_b_key = PrivateKey::try_new([10;32]).unwrap(); + let vault_b_address = Address::from(&PublicKey::new_from_private_key(&vault_b_key)); + + let user_lp_holding_key = PrivateKey::try_new([11;32]).unwrap(); + let user_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&user_lp_holding_key)); + + let pool_lp_holding_key = PrivateKey::try_new([12;32]).unwrap(); + let pool_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&pool_lp_holding_key)); + + let pool_key = PrivateKey::try_new([13;32]).unwrap(); + let pool_address = Address::from(&PublicKey::new_from_private_key(&pool_key)); + + //initialize Tokens: Token A, Token B and Took LP + let mut instruction: [u8;23] = [0;23]; + instruction[1..17].copy_from_slice(&token_a_supply.to_le_bytes()); + instruction[18] = 0x01; //name is not default. + instruction[19] = 0x02; + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_a_definition_address, token_a_holding_address], + vec![], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + instruction[1..17].copy_from_slice(&token_b_supply.to_le_bytes()); + instruction[18] = 0x03; //name is not default. + instruction[19] = 0x02; + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_b_definition_address, token_b_holding_address], + vec![], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + instruction[1..17].copy_from_slice(&token_lp_supply.to_le_bytes()); + instruction[18] = 0x03; //name is not default. + instruction[19] = 0x04; + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_lp_definition_address, token_lp_holding_address], + vec![], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + + // Initialize User accounts for Token A and B + let mut instruction: [u8;23] = [0;23]; + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&user_a_amount.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_a_holding_address, user_a_holding_address], + vec![0], + 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(); + + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&user_b_amount.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_b_holding_address, user_b_holding_address], + vec![0], + 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(); + + + //TODO: initialize vaults - ideally, we won't need to do this. + let temp_amt = 1u128; + let mut instruction: [u8;23] = [0;23]; + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_a_holding_address, vault_a_address], + vec![1], + 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(); + + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_b_holding_address, vault_b_address], + vec![1], + 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(); + + //need LP accounts for both AMM and user + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_lp_holding_address, user_lp_holding_address], + vec![0], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[&token_lp_holding_key], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + let amt: u128 = token_lp_supply - 1; + //Pool LP should be initialized at full amount (token_lp_supply) + instruction[0] = 1; //transfer + instruction[1..17].copy_from_slice(&amt.to_le_bytes()); + + let message = public_transaction::Message::try_new( + Program::token().id(), + vec![token_lp_holding_address, pool_lp_holding_address], + vec![1], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[&token_lp_holding_key], + ); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + //Set up instruction to initialize AMM for (Token-A, Token-B) + let init_balance_a: u128 = 1000; + let init_balance_b: u128 = 1000; + let token_program_u8: [u8;32] = bytemuck::cast(program_token); + + let mut instruction: Vec = Vec::new(); + instruction.push(0); + instruction.extend_from_slice(&init_balance_a.to_le_bytes()); + instruction.extend_from_slice(&init_balance_b.to_le_bytes()); + instruction.extend_from_slice(&token_program_u8); + + let message = public_transaction::Message::try_new( + Program::amm().id(), + vec![ + pool_address, + vault_a_address, + vault_b_address, + pool_lp_holding_address, + user_a_holding_address, + user_b_holding_address, + user_lp_holding_address + ], + vec![0,0,0,0], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message( + &message, + &[ + &pool_key, + &pool_lp_holding_key, + &user_a_holding_key, + &user_b_holding_key, + ], + ); + + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx).unwrap(); + + + //Initialize swap user accounts + let swap_user_a_holding_key = PrivateKey::try_new([7;32]).unwrap(); + let swap_user_a_holding_address = Address::from(&PublicKey::new_from_private_key(&user_a_holding_key)); + let swap_user_a_amount: u128 = 10000; + + let swap_user_b_holding_key = PrivateKey::try_new([8;32]).unwrap(); + let swap_user_b_holding_address = Address::from(&PublicKey::new_from_private_key(&user_b_holding_key)); + let swap_user_b_amount: u128 = 10000; + + // Initialize Swap User accounts for Token A and B + 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_address, swap_user_a_holding_address], + vec![state.get_account_by_address(&token_a_holding_address).nonce], + 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(); + + + 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_address, swap_user_b_holding_address], + vec![state.get_account_by_address(&token_b_holding_address).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(); + + + let main_addr = token_a_definition_address; + let swap_a: u128 = 500; + + let mut instruction: Vec = Vec::new(); + instruction.push(1); + instruction.extend_from_slice(&swap_a.to_le_bytes()); + instruction.extend_from_slice(main_addr.value()); + + + let message = public_transaction::Message::try_new( + Program::amm().id(), + vec![ + pool_address, + vault_a_address, + vault_b_address, + swap_user_a_holding_address, + swap_user_b_holding_address + ], + vec![ + state.get_account_by_address(&pool_address).nonce, + state.get_account_by_address(&vault_a_address).nonce, + state.get_account_by_address(&vault_b_address).nonce, + state.get_account_by_address(&swap_user_a_holding_address).nonce, + state.get_account_by_address(&swap_user_b_holding_address).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(); + + //TODO: + //Check AMM pairing: + // reserve_a = vault_a + // reserve_b = vault_b + let pool_post = state.get_account_by_address(&pool_address); + let vault_a_post = state.get_account_by_address(&vault_a_address); + let vault_b_post = state.get_account_by_address(&vault_b_address); + let pool_lp_post = state.get_account_by_address(&pool_lp_holding_address); + let swap_a_post = state.get_account_by_address(&swap_user_a_holding_address); + let swap_b_post = state.get_account_by_address(&swap_user_b_holding_address); + let user_lp_post = state.get_account_by_address(&user_lp_holding_address); + + assert!(TokenHolding::parse(&swap_a_post.data).unwrap().balance == swap_user_a_amount - swap_a); + assert!(TokenHolding::parse(&swap_b_post.data).unwrap().balance == swap_user_b_amount + swap_a); + //todo: reserves reduced }