diff --git a/nssa/program_methods/guest/src/bin/amm.rs b/nssa/program_methods/guest/src/bin/amm.rs index 8625429..418f7f6 100644 --- a/nssa/program_methods/guest/src/bin/amm.rs +++ b/nssa/program_methods/guest/src/bin/amm.rs @@ -5,21 +5,33 @@ use nssa_core::{ use bytemuck; -// The token program has two functions: -// 1. New token definition. +// The AMM program has four functions: +// 1. New AMM 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 +// * Seven **default** accounts: [amm_pool, vault_holding_a, vault_holding_b, pool_lp, user_holding_a, user_holding_b, user_holding_lp]. +// amm_pool is a default account that will initiate the amm definition account values +// vault_holding_a is a token holding account for token a +// vault_holding_b is a token holding account for token b +// pool_lp is a token holding account for the pool's lp token +// user_holding_a is a token holding account for token a +// user_holding_b is a token holding account for token b +// user_holding_lp is a token holding account for lp token +// TODO: ideally, vault_holding_a, vault_holding_b, pool_lp and user_holding_lp are uninitated. +// * An instruction data of 55-bytes, indicating the initial amm reserves' balances and token_program_id 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 +// [0x00 || array of balances (little-endian 16 bytes) || TOKEN_PROGRAM_ID)] +// 2. Swap assets // 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]. +// * Two accounts: [amm_pool, vault_holding_1, vault_holding_2, user_holding_a, user_holding_b]. +// * An instruction data byte string of length 49, indicating which token type to swap and maximum amount with the following layout +// [0x01 || amount (little-endian 16 bytes) || TOKEN_DEFINITION_ID]. +// 3. Add liquidity +// Arguments to this function are: +// * Two accounts: [amm_pool, vault_holding_a, vault_holding_b, pool_lp, user_holding_a, user_holding_b, user_holding_lp]. +// * An instruction data byte string of length 65, amounts to add +// [0x02 || array of max amounts (little-endian 16 bytes) || TOKEN_DEFINITION_ID (for primary)]. +// 4. Remove liquidity +// * Input instruction set [0x03]. const POOL_DEFINITION_DATA_SIZE: usize = 240; @@ -36,7 +48,6 @@ struct PoolDefinition{ } - impl PoolDefinition { fn into_data(self) -> Vec { let u8_token_program_id : [u8;32] = bytemuck::cast(self.token_program_id); @@ -259,7 +270,7 @@ fn main() { let balance_b: u128 = u128::from_le_bytes(instruction[17..33].try_into().unwrap()); - let token_program_id : &[u32] = bytemuck::cast_slice(&instruction[33..55]); + let token_program_id : &[u32] = bytemuck::cast_slice(&instruction[33..65]); let token_program_id : [u32;8] = token_program_id.try_into().unwrap(); let (post_states, chained_call) = new_definition(&pre_states, @@ -312,21 +323,24 @@ fn swap( // Verify vaults are in fact vaults let pool_def_data = PoolDefinition::parse(&pool.account.data).unwrap(); + 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(); - if vault1.account_id == pool_def_data.definition_token_a_id { + if vault1_data.definition_id == pool_def_data.definition_token_a_id { vault_a = vault1.clone(); - } else if vault2.account_id == pool_def_data.definition_token_a_id { + } else if vault2_data.definition_id == pool_def_data.definition_token_a_id { vault_a = vault2.clone(); } else { panic!("Vault A was not provided"); } - if vault1.account_id == pool_def_data.definition_token_b_id { + if vault1_data.definition_id == pool_def_data.definition_token_b_id { vault_b = vault1.clone(); - } else if vault2.account_id == pool_def_data.definition_token_b_id { + } else if vault2_data.definition_id == pool_def_data.definition_token_b_id { vault_b = vault2.clone(); } else { panic!("Vault B was not provided"); @@ -456,39 +470,54 @@ fn add_liquidity(pre_states: &[AccountWithMetadata], let pool_def_data = PoolDefinition::parse(&pool.account.data).unwrap(); + let vault1_data = TokenHolding::parse(&vault1.account.data).unwrap(); + let vault2_data = TokenHolding::parse(&vault2.account.data).unwrap(); + + if vault1_data.definition_id == pool_def_data.definition_token_a_id { + vault_a = vault1.clone(); + } else if vault2_data.definition_id == pool_def_data.definition_token_a_id { + vault_a = vault2.clone(); + } else { + panic!("Vault A was not provided"); + } + + if vault1_data.definition_id == pool_def_data.definition_token_b_id { + vault_b = vault1.clone(); + } else if vault2_data.definition_id == pool_def_data.definition_token_b_id { + vault_b = vault2.clone(); + } else { + panic!("Vault B was not 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]; - - if vault1.account_id == pool_def_data.definition_token_a_id { - vault_a = vault1.clone(); - } else if vault2.account_id == pool_def_data.definition_token_a_id { - vault_a = vault2.clone(); - } else { - panic!("Vault A was not provided"); - } - - if vault1.account_id == pool_def_data.definition_token_b_id { - vault_b = vault1.clone(); - } else if vault2.account_id == pool_def_data.definition_token_b_id { - vault_b = vault2.clone(); - } else { - panic!("Vault B was not provided"); + if max_amount_a == 0 || max_amount_b == 0 { + panic!("Both max-balances must be nonzero"); } - + // 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 { + 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 = (vault_b.account.balance/vault_a.account.balance)*actual_amount_a; + actual_amount_b = (pool_def_data.reserve_b*actual_amount_a)/pool_def_data.reserve_a; } else if main_token == pool_def_data.definition_token_b_id { actual_amount_b = max_amount_b; - actual_amount_a = (vault_a.account.balance/vault_b.account.balance)*actual_amount_b; + actual_amount_a = (pool_def_data.reserve_a*actual_amount_b)/pool_def_data.reserve_b; } else { panic!("Mismatch of token types"); //main token does not match with vaults. } @@ -496,11 +525,22 @@ fn add_liquidity(pre_states: &[AccountWithMetadata], // 3. Validate amounts assert!(max_amount_a >= actual_amount_a && max_amount_b >= actual_amount_b); - assert!(user_a.account.balance >= actual_amount_a && actual_amount_a > 0); - assert!(user_b.account.balance >= actual_amount_b && actual_amount_b > 0); + if user_a.account.balance < actual_amount_a { + panic!("Insufficient balance"); + } + + if user_b.account.balance < actual_amount_b { + panic!("Insufficient balance"); + } + + if actual_amount_a == 0 || actual_amount_b == 0 { + panic!("A trade amount is 0"); + } + // 4. Calculate LP to mint - let delta_lp : u128 = pool_def_data.liquidity_pool_cap * (actual_amount_b/pool_def_data.reserve_b); + let mut delta_lp: u128 = 0; + 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(); @@ -581,18 +621,20 @@ fn remove_liquidity(pre_states: &[AccountWithMetadata]) -> (Vec, Vec (Vec, Vec (Vec, Vec Vec { + let mut bytes = [0; TOKEN_DEFINITION_DATA_SIZE]; + bytes[0] = self.account_type; + bytes[1..7].copy_from_slice(&self.name); + bytes[7..].copy_from_slice(&self.total_supply.to_le_bytes()); + bytes.into() + } +} + +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() + } +} + +#[derive(Serialize)] +struct NewAMMInstructions { + option: u8, + balance_a: u128, + balance_b: u128, + token_id: ProgramId, +} + #[test] + fn test_simple_amm_initialized() { + let token_program_id = program_methods::TOKEN_ID; + let program = Program::amm(); + + //initialize AMM accounts + let mut token_a_definition_data = TokenDefinition::into_data(TokenDefinition{ + account_type: TOKEN_DEFINITION_TYPE, + name: [1u8;6], + total_supply: 1000u128 + }); + + let mut token_b_definition_data = TokenDefinition::into_data(TokenDefinition{ + account_type: TOKEN_DEFINITION_TYPE, + name: [2u8;6], + total_supply: 1000u128 + }); + + let mut pool_lp_definition_data = TokenDefinition::into_data(TokenDefinition{ + account_type: TOKEN_DEFINITION_TYPE, + name: [2u8;6], + total_supply: u128::MAX + }); + + let amm_key = PrivateKey::try_new([1; 32]).unwrap(); + let vault_a_key = PrivateKey::try_new([2; 32]).unwrap(); + let vault_b_key = PrivateKey::try_new([3; 32]).unwrap(); + let user_a_key = PrivateKey::try_new([4; 32]).unwrap(); + let user_b_key = PrivateKey::try_new([5; 32]).unwrap(); + let user_lp_key = PrivateKey::try_new([6; 32]).unwrap(); + let pool_lp_key = PrivateKey::try_new([7; 32]).unwrap(); + + let mut definition_a_account = Account::default(); + let mut definition_b_account = Account::default(); + let mut pool_lp_account = Account::default(); + + let definition_a_address = Address::new([1u8;32]); + let definition_b_address= Address::new([2u8;32]); + let definition_lp_address = Address::new([3u8;32]); + let vault_a_address = Address::from(&PublicKey::new_from_private_key(&vault_a_key)); + let vault_b_address = Address::from(&PublicKey::new_from_private_key(&vault_b_key)); + let user_a_address = Address::from(&PublicKey::new_from_private_key(&user_a_key)); + let user_b_address = Address::from(&PublicKey::new_from_private_key(&user_b_key)); + let user_lp_address= Address::from(&PublicKey::new_from_private_key(&pool_lp_key)); + let pool_lp_address = Address::from(&PublicKey::new_from_private_key(&user_lp_key)); + let amm_pool_address= Address::from(&PublicKey::new_from_private_key(&amm_key)); + + definition_a_account.data = token_a_definition_data; + definition_b_account.data = token_b_definition_data; + pool_lp_account.data = pool_lp_definition_data; + + let mut vault_a_account = Account::default(); + let mut vault_b_account = Account::default(); + let mut user_a_account = Account::default(); + let mut user_b_account = Account::default(); + let mut pool_lp_account = Account::default(); + let mut user_lp_account = Account::default(); + + vault_a_account.program_owner = program_methods::TOKEN_ID; + vault_b_account.program_owner = program_methods::TOKEN_ID; + user_a_account.program_owner = program_methods::TOKEN_ID; + user_b_account.program_owner = program_methods::TOKEN_ID; + pool_lp_account.program_owner = program_methods::TOKEN_ID; + user_lp_account.program_owner = program_methods::TOKEN_ID; + + user_lp_account.data = TokenHolding::into_data( + TokenHolding{ + account_type: TOKEN_HOLDING_TYPE, + definition_id: definition_lp_address.clone(), + balance: 0u128 + + } + ); + + vault_a_account.data = TokenHolding::into_data( + TokenHolding{ + account_type: TOKEN_HOLDING_TYPE, + definition_id: definition_a_address.clone(), + balance: 0u128 + + } + ); + + vault_b_account.data = TokenHolding::into_data( + TokenHolding{ + account_type: TOKEN_HOLDING_TYPE, + definition_id: definition_b_address.clone(), + balance: 0u128 + + } + ); + + user_a_account.data = TokenHolding::into_data( + TokenHolding{ + account_type: TOKEN_HOLDING_TYPE, + definition_id: definition_a_address.clone(), + balance: 30u128 + + } + ); + + user_b_account.data = TokenHolding::into_data( + TokenHolding{ + account_type: TOKEN_HOLDING_TYPE, + definition_id: definition_b_address.clone(), //TODO + balance: 30u128 + + } + ); + + let amm_pool = AccountWithMetadata::default(); + let mut user_a_meta = AccountWithMetadata::default(); + let mut user_b_meta = AccountWithMetadata::default(); + let mut vault_a_meta = AccountWithMetadata::default(); + let mut vault_b_meta = AccountWithMetadata::default(); + let mut pool_lp_meta = AccountWithMetadata::default(); + let mut user_lp_meta = AccountWithMetadata::default(); + + user_a_meta.account = user_a_account.clone(); + user_a_meta.account_id = user_a_address; + + user_b_meta.account = user_b_account.clone(); + user_a_meta.account_id = user_b_address; + + vault_a_meta.account = vault_a_account.clone(); + vault_a_meta.account_id = vault_a_address; + + vault_b_meta.account = vault_b_account.clone(); + vault_b_meta.account_id = vault_b_address; + + pool_lp_meta.account = pool_lp_account.clone(); + pool_lp_meta.account_id = pool_lp_address; + + user_lp_meta.account = user_lp_account.clone(); + user_lp_meta.account_id = user_lp_address; + + + let key = PrivateKey::try_new([1; 32]).unwrap(); + let from_address = Address::from(&PublicKey::new_from_private_key(&key)); + let to_address = Address::new([2; 32]); + let initial_balance = 100; + let initial_data = [(from_address, initial_balance), (to_address, 0)]; + let mut state = + V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs(); + + + state.force_insert_account(user_a_address, user_a_account.clone()); + state.force_insert_account(user_b_address, user_b_account.clone()); + state.force_insert_account(vault_a_address, vault_a_account.clone()); + state.force_insert_account(vault_b_address, vault_b_account.clone()); + state.force_insert_account(pool_lp_address, pool_lp_account.clone()); + state.force_insert_account(user_lp_address, user_lp_account.clone()); + + let amount_a: u128 = 10; + let amount_b: u128 = 10; + + let u8_token_program_id : [u8;32] = bytemuck::cast(token_program_id); + let mut instruction= NewAMMInstructions { + option: 0u8, + balance_a: amount_a, + balance_b: amount_b, + token_id: token_program_id, + }; + + let message = public_transaction::Message::try_new( + program.id(), + vec![amm_pool_address, vault_a_address, vault_b_address, pool_lp_address, user_a_address, user_b_address, user_lp_address], + vec![0], + instruction, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&amm_key, &vault_a_key, &vault_b_key, &pool_lp_key, &user_a_key, &user_b_key, &user_lp_key]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx); + /* assert!(matches!( + result, + Err(NssaError::MaxChainedCallsDepthExceeded) + ));*/ + } + }