logos-blockchain/lssa
1
0
Fork 0
mirror of https://github.com/logos-blockchain/lssa.git synced 2026-02-17 20:03:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor AMM

Browse Source
This commit is contained in:
jonesmarvin8 2026-01-23 16:30:54 -05:00
parent e4e476fde9
commit 2367bf343b
35 changed files with 2728 additions and 732 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

BIN
artifacts/program_methods/amm.bin
View File

Binary file not shown.

BIN
artifacts/program_methods/authenticated_transfer.bin
View File

Binary file not shown.

BIN
artifacts/program_methods/pinata.bin
View File

Binary file not shown.

BIN
artifacts/program_methods/pinata_token.bin
View File

Binary file not shown.

BIN
artifacts/program_methods/privacy_preserving_circuit.bin
View File

Binary file not shown.

BIN
artifacts/program_methods/token.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/burner.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/chain_caller.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/changer_claimer.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/claimer.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/data_changer.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/extra_output.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/malicious_authorization_changer.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/minter.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/missing_output.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/modified_transfer.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/nonce_changer.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/noop.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/program_owner_changer.bin
View File

Binary file not shown.

BIN
artifacts/test_program_methods/simple_balance_transfer.bin
View File

Binary file not shown.

1
nssa/Cargo.toml
View File

@ -24,6 +24,7 @@ risc0-binfmt = "3.0.2"
[dev-dependencies]
token_core.workspace = true
amm_core.workspace = true
test_program_methods.workspace = true
env_logger.workspace = true

222
nssa/src/state.rs
View File

@ -267,6 +267,7 @@ pub mod tests {
use std::collections::HashMap;
use amm_core::PoolDefinition;
use nssa_core::{
Commitment, Nullifier, NullifierPublicKey, NullifierSecretKey, SharedSecretKey,
account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data},
@ -2285,137 +2286,6 @@ pub mod tests {
));
}
// TODO repeated code should ultimately be removed;
fn compute_pool_pda(
amm_program_id: ProgramId,
definition_token_a_id: AccountId,
definition_token_b_id: AccountId,
) -> AccountId {
AccountId::from((
&amm_program_id,
&compute_pool_pda_seed(definition_token_a_id, definition_token_b_id),
))
}
fn compute_pool_pda_seed(
definition_token_a_id: AccountId,
definition_token_b_id: AccountId,
) -> PdaSeed {
use risc0_zkvm::sha::{Impl, Sha256};
let mut i: usize = 0;
let (token_1, token_2) = loop {
if definition_token_a_id.value()[i] > definition_token_b_id.value()[i] {
let token_1 = definition_token_a_id;
let token_2 = definition_token_b_id;
break (token_1, token_2);
} else if definition_token_a_id.value()[i] < definition_token_b_id.value()[i] {
let token_1 = definition_token_b_id;
let token_2 = definition_token_a_id;
break (token_1, token_2);
}
if i == 32 {
panic!("Definitions match");
} else {
i += 1;
}
};
let mut bytes = [0; 64];
bytes[0..32].copy_from_slice(&token_1.to_bytes());
bytes[32..].copy_from_slice(&token_2.to_bytes());
PdaSeed::new(
Impl::hash_bytes(&bytes)
.as_bytes()
.try_into()
.expect("Hash output must be exactly 32 bytes long"),
)
}
fn compute_vault_pda(
amm_program_id: ProgramId,
pool_id: AccountId,
definition_token_id: AccountId,
) -> AccountId {
AccountId::from((
&amm_program_id,
&compute_vault_pda_seed(pool_id, definition_token_id),
))
}
fn compute_vault_pda_seed(pool_id: AccountId, definition_token_id: AccountId) -> PdaSeed {
use risc0_zkvm::sha::{Impl, Sha256};
let mut bytes = [0; 64];
bytes[0..32].copy_from_slice(&pool_id.to_bytes());
bytes[32..].copy_from_slice(&definition_token_id.to_bytes());
PdaSeed::new(
Impl::hash_bytes(&bytes)
.as_bytes()
.try_into()
.expect("Hash output must be exactly 32 bytes long"),
)
}
fn compute_liquidity_token_pda(amm_program_id: ProgramId, pool_id: AccountId) -> AccountId {
AccountId::from((&amm_program_id, &compute_liquidity_token_pda_seed(pool_id)))
}
fn compute_liquidity_token_pda_seed(pool_id: AccountId) -> PdaSeed {
use risc0_zkvm::sha::{Impl, Sha256};
let mut bytes = [0; 64];
bytes[0..32].copy_from_slice(&pool_id.to_bytes());
bytes[32..].copy_from_slice(&[0; 32]);
PdaSeed::new(
Impl::hash_bytes(&bytes)
.as_bytes()
.try_into()
.expect("Hash output must be exactly 32 bytes long"),
)
}
const POOL_DEFINITION_DATA_SIZE: usize = 225;
#[derive(Default)]
struct PoolDefinition {
definition_token_a_id: AccountId,
definition_token_b_id: AccountId,
vault_a_id: AccountId,
vault_b_id: AccountId,
liquidity_pool_id: AccountId,
liquidity_pool_supply: u128,
reserve_a: u128,
reserve_b: u128,
fees: u128,
active: bool,
}
impl PoolDefinition {
fn into_data(self) -> Data {
let mut bytes = [0; POOL_DEFINITION_DATA_SIZE];
bytes[0..32].copy_from_slice(&self.definition_token_a_id.to_bytes());
bytes[32..64].copy_from_slice(&self.definition_token_b_id.to_bytes());
bytes[64..96].copy_from_slice(&self.vault_a_id.to_bytes());
bytes[96..128].copy_from_slice(&self.vault_b_id.to_bytes());
bytes[128..160].copy_from_slice(&self.liquidity_pool_id.to_bytes());
bytes[160..176].copy_from_slice(&self.liquidity_pool_supply.to_le_bytes());
bytes[176..192].copy_from_slice(&self.reserve_a.to_le_bytes());
bytes[192..208].copy_from_slice(&self.reserve_b.to_le_bytes());
bytes[208..224].copy_from_slice(&self.fees.to_le_bytes());
bytes[224] = self.active as u8;
bytes
.to_vec()
.try_into()
.expect("225 bytes should fit into Data")
}
}
struct PrivateKeysForTests;
impl PrivateKeysForTests {
@ -2596,7 +2466,7 @@ pub mod tests {
impl IdForTests {
fn pool_definition_id() -> AccountId {
compute_pool_pda(
amm_core::compute_pool_pda(
Program::amm().id(),
IdForTests::token_a_definition_id(),
IdForTests::token_b_definition_id(),
@ -2604,7 +2474,10 @@ pub mod tests {
}
fn token_lp_definition_id() -> AccountId {
compute_liquidity_token_pda(Program::amm().id(), IdForTests::pool_definition_id())
amm_core::compute_liquidity_token_pda(
Program::amm().id(),
IdForTests::pool_definition_id(),
)
}
fn token_a_definition_id() -> AccountId {
@ -2634,7 +2507,7 @@ pub mod tests {
}
fn vault_a_id() -> AccountId {
compute_vault_pda(
amm_core::compute_vault_pda(
Program::amm().id(),
IdForTests::pool_definition_id(),
IdForTests::token_a_definition_id(),
@ -2642,7 +2515,7 @@ pub mod tests {
}
fn vault_b_id() -> AccountId {
compute_vault_pda(
amm_core::compute_vault_pda(
Program::amm().id(),
IdForTests::pool_definition_id(),
IdForTests::token_b_definition_id(),
@ -3233,11 +3106,6 @@ pub mod tests {
}
}
const AMM_NEW_DEFINITION: u8 = 0;
const AMM_SWAP: u8 = 1;
const AMM_ADD_LIQUIDITY: u8 = 2;
const AMM_REMOVE_LIQUIDITY: u8 = 3;
fn state_for_amm_tests() -> V02State {
let initial_data = [];
let mut state =
@ -3303,11 +3171,11 @@ pub mod tests {
fn test_simple_amm_remove() {
let mut state = state_for_amm_tests();
let mut instruction: Vec<u8> = Vec::new();
instruction.push(AMM_REMOVE_LIQUIDITY);
instruction.extend_from_slice(&BalanceForTests::remove_lp().to_le_bytes());
instruction.extend_from_slice(&BalanceForTests::remove_min_amount_a().to_le_bytes());
instruction.extend_from_slice(&BalanceForTests::remove_min_amount_b().to_le_bytes());
let instruction = amm_core::Instruction::RemoveLiquidity {
remove_liquidity_amount: BalanceForTests::remove_lp(),
min_amount_to_remove_token_a: BalanceForTests::remove_min_amount_a(),
min_amount_to_remove_token_b: BalanceForTests::remove_min_amount_b(),
};
let message = public_transaction::Message::try_new(
Program::amm().id(),
@ -3380,12 +3248,11 @@ pub mod tests {
AccountForTests::token_lp_definition_init_inactive(),
);
let mut instruction: Vec<u8> = Vec::new();
instruction.push(AMM_NEW_DEFINITION);
instruction.extend_from_slice(&BalanceForTests::vault_a_balance_init().to_le_bytes());
instruction.extend_from_slice(&BalanceForTests::vault_b_balance_init().to_le_bytes());
let amm_program_u8: [u8; 32] = bytemuck::cast(Program::amm().id());
instruction.extend_from_slice(&amm_program_u8);
let instruction = amm_core::Instruction::NewDefinition {
token_a_amount: BalanceForTests::vault_a_balance_init(),
token_b_amount: BalanceForTests::vault_b_balance_init(),
amm_program_id: Program::amm().id(),
};
let message = public_transaction::Message::try_new(
Program::amm().id(),
@ -3465,12 +3332,11 @@ pub mod tests {
AccountForTests::user_token_lp_holding_init_zero(),
);
let mut instruction: Vec<u8> = Vec::new();
instruction.push(AMM_NEW_DEFINITION);
instruction.extend_from_slice(&BalanceForTests::vault_a_balance_init().to_le_bytes());
instruction.extend_from_slice(&BalanceForTests::vault_b_balance_init().to_le_bytes());
let amm_program_u8: [u8; 32] = bytemuck::cast(Program::amm().id());
instruction.extend_from_slice(&amm_program_u8);
let instruction = amm_core::Instruction::NewDefinition {
token_a_amount: BalanceForTests::vault_a_balance_init(),
token_b_amount: BalanceForTests::vault_b_balance_init(),
amm_program_id: Program::amm().id(),
};
let message = public_transaction::Message::try_new(
Program::amm().id(),
@ -3538,12 +3404,11 @@ pub mod tests {
AccountForTests::vault_b_init_inactive(),
);
let mut instruction: Vec<u8> = Vec::new();
instruction.push(AMM_NEW_DEFINITION);
instruction.extend_from_slice(&BalanceForTests::vault_a_balance_init().to_le_bytes());
instruction.extend_from_slice(&BalanceForTests::vault_b_balance_init().to_le_bytes());
let amm_program_u8: [u8; 32] = bytemuck::cast(Program::amm().id());
instruction.extend_from_slice(&amm_program_u8);
let instruction = amm_core::Instruction::NewDefinition {
token_a_amount: BalanceForTests::vault_a_balance_init(),
token_b_amount: BalanceForTests::vault_b_balance_init(),
amm_program_id: Program::amm().id(),
};
let message = public_transaction::Message::try_new(
Program::amm().id(),
@ -3602,11 +3467,11 @@ pub mod tests {
env_logger::init();
let mut state = state_for_amm_tests();
let mut instruction: Vec<u8> = Vec::new();
instruction.push(AMM_ADD_LIQUIDITY);
instruction.extend_from_slice(&BalanceForTests::add_min_amount_lp().to_le_bytes());
instruction.extend_from_slice(&BalanceForTests::add_max_amount_a().to_le_bytes());
instruction.extend_from_slice(&BalanceForTests::add_max_amount_b().to_le_bytes());
let instruction = amm_core::Instruction::AddLiquidity {
min_amount_liquidity: BalanceForTests::add_min_amount_lp(),
max_amount_to_add_token_a: BalanceForTests::add_max_amount_a(),
max_amount_to_add_token_b: BalanceForTests::add_max_amount_b(),
};
let message = public_transaction::Message::try_new(
Program::amm().id(),
@ -3664,11 +3529,11 @@ pub mod tests {
fn test_simple_amm_swap_1() {
let mut state = state_for_amm_tests();
let mut instruction: Vec<u8> = Vec::new();
instruction.push(AMM_SWAP);
instruction.extend_from_slice(&BalanceForTests::swap_amount_in().to_le_bytes());
instruction.extend_from_slice(&BalanceForTests::swap_min_amount_out().to_le_bytes());
instruction.extend_from_slice(&IdForTests::token_b_definition_id().to_bytes());
let instruction = amm_core::Instruction::Swap {
swap_amount_in: BalanceForTests::swap_amount_in(),
min_amount_out: BalanceForTests::swap_min_amount_out(),
token_definition_id_in: IdForTests::token_b_definition_id(),
};
let message = public_transaction::Message::try_new(
Program::amm().id(),
@ -3715,12 +3580,11 @@ pub mod tests {
fn test_simple_amm_swap_2() {
let mut state = state_for_amm_tests();
let mut instruction: Vec<u8> = Vec::new();
instruction.push(AMM_SWAP);
instruction.extend_from_slice(&BalanceForTests::swap_amount_in().to_le_bytes());
instruction.extend_from_slice(&BalanceForTests::swap_min_amount_out().to_le_bytes());
instruction.extend_from_slice(&IdForTests::token_a_definition_id().to_bytes());
let instruction = amm_core::Instruction::Swap {
swap_amount_in: BalanceForTests::swap_amount_in(),
min_amount_out: BalanceForTests::swap_min_amount_out(),
token_definition_id_in: IdForTests::token_a_definition_id(),
};
let message = public_transaction::Message::try_new(
Program::amm().id(),
vec![

2
program_methods/guest/Cargo.toml
View File

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

148
program_methods/guest/src/bin/amm.rs
View File

@ -1,11 +1,17 @@
use nssa_core::{
account::{Account, AccountId, AccountWithMetadata, Data},
program::{
AccountPostState, ChainedCall, PdaSeed, ProgramId, ProgramInput, read_nssa_inputs,
write_nssa_outputs_with_chained_call,
},
//! The AMM Program.
//!
//! This program implements a simple AMM that supports multiple AMM pools (a single pool per
//! token pair).
//!
//! AMM program accepts [`Instruction`] as input, refer to the corresponding documentation
//! for more details.
use amm_core::Instruction;
use amm_program;
use nssa_core::program::{
AccountPostState, ChainedCall, ProgramInput, read_nssa_inputs,
write_nssa_outputs_with_chained_call,
};
use amm_program::core::Instruction;
fn main() {
let (
@ -18,24 +24,120 @@ fn main() {
let pre_states_clone = pre_states.clone();
let (post_states, chained_calls) = match instruction {
Intruction::NewAMM => {
}
Instruction::RemoveLiquidity{ remove_liquidity_amount, min_amount_to_remove_token_a, min_amount_to_remove_token_b } =>{
let [pool, vault_a, vault_b, pool_definition_lp, user_holding_a, user_holding_b, user_holding_lp] = pre_states
let (post_states, chained_calls): (Vec<AccountPostState>, Vec<ChainedCall>) = match instruction
{
Instruction::NewDefinition {
token_a_amount,
token_b_amount,
amm_program_id,
} => {
let [
pool,
vault_a,
vault_b,
pool_definition_lp,
user_holding_a,
user_holding_b,
user_holding_lp,
] = pre_states
.try_into()
.expect("RemoveLiquidity instruction requires exactly seven accounts");
amm_program::remove::remove_liquidity(pool, vault_a, vault_b, pool_definition_lp, user_holding_a, user_holding_b, user_holding_lp, remove_liquidity_amount, min_amount_to_remove_token_a, min_amount_to_remove_token_b)
.expect("Transfer instruction requires exactly seven accounts");
amm_program::new_definition::new_definition(
pool,
vault_a,
vault_b,
pool_definition_lp,
user_holding_a,
user_holding_b,
user_holding_lp,
token_a_amount,
token_b_amount,
amm_program_id,
)
}
Instruction::AddLiquidity { min_amount_liquidity, max_amount_a, max_amount_b } => {
let [pool, vault_a, vault_b, pool_definition_lp, user_holding_a, user_holding_b, user_holding_lp] = pre_states
Instruction::AddLiquidity {
min_amount_liquidity,
max_amount_to_add_token_a,
max_amount_to_add_token_b,
} => {
let [
pool,
vault_a,
vault_b,
pool_definition_lp,
user_holding_a,
user_holding_b,
user_holding_lp,
] = pre_states
.try_into()
.expect("AddLiquidity instruction requires exactly seven accounts");
amm_program::add::add_liquidity(pool, vault_a, vault_b, pool_definition_lp, user_holding_a, user_holding_b, user_holding_lp, min_amount_liquidity, max_amount_a, max_amount_b)
.expect("Transfer instruction requires exactly seven accounts");
amm_program::add::add_liquidity(
pool,
vault_a,
vault_b,
pool_definition_lp,
user_holding_a,
user_holding_b,
user_holding_lp,
min_amount_liquidity,
max_amount_to_add_token_a,
max_amount_to_add_token_b,
)
}
Instruction::RemoveLiquidity {
remove_liquidity_amount,
min_amount_to_remove_token_a,
min_amount_to_remove_token_b,
} => {
let [
pool,
vault_a,
vault_b,
pool_definition_lp,
user_holding_a,
user_holding_b,
user_holding_lp,
] = pre_states
.try_into()
.expect("Transfer instruction requires exactly seven accounts");
amm_program::remove::remove_liquidity(
pool,
vault_a,
vault_b,
pool_definition_lp,
user_holding_a,
user_holding_b,
user_holding_lp,
remove_liquidity_amount,
min_amount_to_remove_token_a,
min_amount_to_remove_token_b,
)
}
Instruction::Swap {
swap_amount_in,
min_amount_out,
token_definition_id_in,
} => {
let [pool, vault_a, vault_b, user_holding_a, user_holding_b] = pre_states
.try_into()
.expect("Transfer instruction requires exactly five accounts");
amm_program::swap::swap(
pool,
vault_a,
vault_b,
user_holding_a,
user_holding_b,
swap_amount_in,
min_amount_out,
token_definition_id_in,
)
}
};
}
write_nssa_outputs_with_chained_call(instruction_words, pre_states_clone, post_states, chained_calls);
}
write_nssa_outputs_with_chained_call(
instruction_words,
pre_states_clone,
post_states,
chained_calls,
);
}

9
programs/amm/Cargo.toml
View File

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

10
programs/amm/core/Cargo.toml Normal file
View File

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

486
programs/amm/core/src/lib.rs
View File

@ -1,330 +1,248 @@
//! This crate contains core data structures and utilities for the Token Program.
//! This crate contains core data structures and utilities for the AMM Program.
use borsh::{BorshDeserialize, BorshSerialize};
use nssa_core::account::{AccountId, Data};
use nssa_core::{
account::{AccountId, Data},
program::{PdaSeed, ProgramId},
};
use serde::{Deserialize, Serialize};
pub const CURRENT_VERSION: u8 = 1;
// The AMM program has five functions (four directly accessible via instructions):
// 1. New AMM definition. Arguments to this function are:
// * Seven accounts: [amm_pool, vault_holding_a, vault_holding_b, pool_lp, user_holding_a,
// user_holding_b, user_holding_lp]. For new AMM Pool: amm_pool, vault_holding_a,
// vault_holding_b, pool_lp and user_holding_lp are default accounts. 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
// * PDA remark: Accounts amm_pool, vault_holding_a, vault_holding_b and pool_lp are PDA. The
// AccountId for these accounts must be computed using: amm_pool AccountId <-
// compute_pool_pda vault_holding_a, vault_holding_b <- compute_vault_pda pool_lp
// <-compute_liquidity_token_pda
// * Requires authorization: user_holding_a, user_holding_b
// * An instruction data of 65-bytes, indicating the initial amm reserves' balances and
// token_program_id with the following layout: [0x00 || array of balances (little-endian 16
// bytes) || AMM_PROGRAM_ID)]
// * Internally, calls compute_liquidity_token_pda_seed, compute_vault_pda_seed to authorize
// transfers.
// * Internally, calls compute_pool_da, compute_vault_pda and compute_vault_pda to check
// various AccountIds are correct.
// 3. Add liquidity Arguments to this function are:
// * Seven accounts: [amm_pool, vault_holding_a, vault_holding_b, pool_lp, user_holding_a,
// user_holding_a, user_holding_lp].
// * Requires authorization: user_holding_a, user_holding_b
// * An instruction data byte string of length 49, amounts for minimum amount of liquidity from
// add (min_amount_lp),
// * max amount added for each token (max_amount_a and max_amount_b); indicate [0x02 || array
// of of balances (little-endian 16 bytes)].
// * Internally, calls compute_liquidity_token_pda_seed to compute liquidity pool PDA seed.
// 4. Remove liquidity
// * Seven accounts: [amm_pool, vault_holding_a, vault_holding_b, pool_lp, user_holding_a,
// user_holding_a, user_holding_lp].
// * Requires authorization: user_holding_lp
// * An instruction data byte string of length 49, amounts for minimum amount of liquidity to
// redeem (balance_lp),
// * minimum balance of each token to remove (min_amount_a and min_amount_b); indicate [0x03 ||
// array of balances (little-endian 16 bytes)].
// * Internally, calls compute_vault_pda_seed to compute vault_a and vault_b's PDA seed.
/// AMM Program Instruction.
#[derive(Serialize, Deserialize)]
pub enum Instruction {
/// Create a new fungible token definition without metadata.
/// Initializes a new Pool (or re-initializes an inactive Pool).
///
/// Required accounts:
/// - Token Definition account (uninitialized),
/// - Token Holding account (uninitialized).
NewDefinition { name: String, total_supply: u128 },
/// - AMM Pool
/// - Vault Holding Account for Token A
/// - Vault Holding Account for Token B
/// - Pool Liquidity Token Definition
/// - User Holding Account for Token A (authorized)
/// - User Holding Account for Token B (authorized)
/// - User Holding Account for Pool Liquidity
NewDefinition {
token_a_amount: u128,
token_b_amount: u128,
amm_program_id: ProgramId,
},
/// Create a new fungible or non-fungible token definition with metadata.
/// Adds liquidity to the Pool
///
/// Required accounts:
/// - Token Definition account (uninitialized),
/// - Token Holding account (uninitialized),
/// - Token Metadata account (uninitialized).
NewDefinitionWithMetadata {
new_definition: NewTokenDefinition,
/// Boxed to avoid large enum variant size
metadata: Box<NewTokenMetadata>,
/// - AMM Pool (initialized)
/// - Vault Holding Account for Token A (initialized)
/// - Vault Holding Account for Token B (initialized)
/// - Pool Liquidity Token Definition (initialized)
/// - User Holding Account for Token A (authorized)
/// - User Holding Account for Token B (authorized)
/// - User Holding Account for Pool Liquidity
AddLiquidity {
min_amount_liquidity: u128,
max_amount_to_add_token_a: u128,
max_amount_to_add_token_b: u128,
},
/// Initialize a token holding account for a given token definition.
/// Removes liquidity from the Pool
///
/// Required accounts:
/// - Token Definition account (initialized),
/// - Token Holding account (uninitialized),
InitializeAccount,
/// - AMM Pool (initialized)
/// - Vault Holding Account for Token A (initialized)
/// - Vault Holding Account for Token B (initialized)
/// - Pool Liquidity Token Definition (initialized)
/// - User Holding Account for Token A (initialized)
/// - User Holding Account for Token B (initialized)
/// - User Holding Account for Pool Liquidity (authorized)
RemoveLiquidity {
remove_liquidity_amount: u128,
min_amount_to_remove_token_a: u128,
min_amount_to_remove_token_b: u128,
},
// 2. Swap assets Arguments to this function are:
// * Five accounts: [amm_pool, vault_holding_a, vault_holding_b, user_holding_a,
// user_holding_b].
// * Requires authorization: user holding account associated to TOKEN_DEFINITION_ID (either
// user_holding_a or user_holding_b)
// * An instruction data byte string of length 65, indicating which token type to swap,
// quantity of tokens put into the swap (of type TOKEN_DEFINITION_ID) and min_amount_out.
// [0x01 || amount (little-endian 16 bytes) || TOKEN_DEFINITION_ID].
// * Internally, calls swap logic.
// * Four accounts: [user_deposit, vault_deposit, vault_withdraw, user_withdraw].
// user_deposit and vault_deposit define deposit transaction. vault_withdraw and
// user_withdraw define withdraw transaction.
// * deposit_amount is the amount for user_deposit -> vault_deposit transfer.
// * reserve_amounts is the pool's reserves; used to compute the withdraw amount.
// * Outputs the token transfers as a Vec<ChainedCall> and the withdraw amount.
///TODO update description
/// Burn tokens from the holder's account.
/// Swap some quantity of Tokens (either Token A or Token B)
/// while maintaining the Pool constant product.
///
/// Required accounts:
/// - Token Definition account (initialized),
/// - Token Holding account (authorized).
Burn { amount_to_burn: u128 },
/*
fn add_liquidity(
pre_states: &[AccountWithMetadata],
balances: &[u128],
) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
if pre_states.len() != 7 {
panic!("Invalid number of input accounts");
}
let pool = &pre_states[0];
let vault_a = &pre_states[1];
let vault_b = &pre_states[2];
let pool_definition_lp = &pre_states[3];
let user_holding_a = &pre_states[4];
let user_holding_b = &pre_states[5];
let user_holding_lp = &pre_states[6];
let min_amount_lp = balances[0];
let max_amount_a = balances[1];
let max_amount_b = balances[2];
*/
///TODO: update for add
/// Mint new tokens to the holder's account.
///
/// Required accounts:
/// - Token Definition account (authorized),
/// - Token Holding account.
AddLiquidity { min_amount_liquidity: u128, max_amount_to_add_token_a: u128, max_amount_to_add_token_b: u128 },
/*
fn remove_liquidity(
pre_states: &[AccountWithMetadata],
amounts: &[u128],
) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
if pre_states.len() != 7 {
panic!("Invalid number of input accounts");
}
let pool = &pre_states[0];
let vault_a = &pre_states[1];
let vault_b = &pre_states[2];
let pool_definition_lp = &pre_states[3];
let user_holding_a = &pre_states[4];
let user_holding_b = &pre_states[5];
let user_holding_lp = &pre_states[6];
if amounts.len() != 3 {
panic!("Invalid number of balances");
}
let amount_lp = amounts[0];
let amount_min_a = amounts[1];
let amount_min_b = amounts[2];
*/
//TODO types
RemoveLiquidity { remove_liquidity_amount: u128, min_amount_to_remove_token_a: u128, min_amount_to_remove_token_b: u128 }
}
/*
#[derive(Serialize, Deserialize)]
pub enum NewTokenDefinition {
Fungible { name: String, total_supply: u128 },
NonFungible { name: String, print_balance: u128 },
}
#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub enum TokenDefinition {
Fungible {
name: String,
total_supply: u128,
metadata_id: Option<AccountId>,
},
NonFungible {
name: String,
metadata_id: AccountId,
/// - AMM Pool (initialized)
/// - Vault Holding Account for Token A (initialized)
/// - Vault Holding Account for Token B (initialized)
/// - User Holding Account for Token A
/// - User Holding Account for Token B
/// Either User Holding Account for Token A or Token B is authorized.
Swap {
swap_amount_in: u128,
min_amount_out: u128,
token_definition_id_in: AccountId,
},
}
impl TryFrom<&Data> for TokenDefinition {
type Error = std::io::Error;
const POOL_DEFINITION_DATA_SIZE: usize = 225;
fn try_from(data: &Data) -> Result<Self, Self::Error> {
TokenDefinition::try_from_slice(data.as_ref())
}
#[derive(Clone, Default)]
pub struct PoolDefinition {
pub definition_token_a_id: AccountId,
pub definition_token_b_id: AccountId,
pub vault_a_id: AccountId,
pub vault_b_id: AccountId,
pub liquidity_pool_id: AccountId,
pub liquidity_pool_supply: u128,
pub reserve_a: u128,
pub reserve_b: u128,
/// Fees are currently not used
pub fees: u128,
/// A pool becomes inactive (active = false)
/// once all of its liquidity has been removed (e.g., reserves are emptied and
/// liquidity_pool_supply = 0)
pub active: bool,
}
impl From<&TokenDefinition> for Data {
fn from(definition: &TokenDefinition) -> Self {
// Using size_of_val as size hint for Vec allocation
let mut data = Vec::with_capacity(std::mem::size_of_val(definition));
impl PoolDefinition {
pub fn into_data(self) -> Data {
let mut bytes = [0; POOL_DEFINITION_DATA_SIZE];
bytes[0..32].copy_from_slice(&self.definition_token_a_id.to_bytes());
bytes[32..64].copy_from_slice(&self.definition_token_b_id.to_bytes());
bytes[64..96].copy_from_slice(&self.vault_a_id.to_bytes());
bytes[96..128].copy_from_slice(&self.vault_b_id.to_bytes());
bytes[128..160].copy_from_slice(&self.liquidity_pool_id.to_bytes());
bytes[160..176].copy_from_slice(&self.liquidity_pool_supply.to_le_bytes());
bytes[176..192].copy_from_slice(&self.reserve_a.to_le_bytes());
bytes[192..208].copy_from_slice(&self.reserve_b.to_le_bytes());
bytes[208..224].copy_from_slice(&self.fees.to_le_bytes());
bytes[224] = self.active as u8;
BorshSerialize::serialize(definition, &mut data)
.expect("Serialization to Vec should not fail");
Data::try_from(data).expect("Token definition encoded data should fit into Data")
}
}
#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub enum TokenHolding {
Fungible {
definition_id: AccountId,
balance: u128,
},
NftMaster {
definition_id: AccountId,
/// The amount of printed copies left - 1 (1 reserved for master copy itself).
print_balance: u128,
},
NftPrintedCopy {
definition_id: AccountId,
/// Whether nft is owned by the holder.
owned: bool,
},
}
impl TokenHolding {
pub fn zeroized_clone_from(other: &Self) -> Self {
match other {
TokenHolding::Fungible { definition_id, .. } => TokenHolding::Fungible {
definition_id: *definition_id,
balance: 0,
},
TokenHolding::NftMaster { definition_id, .. } => TokenHolding::NftMaster {
definition_id: *definition_id,
print_balance: 0,
},
TokenHolding::NftPrintedCopy { definition_id, .. } => TokenHolding::NftPrintedCopy {
definition_id: *definition_id,
owned: false,
},
}
bytes
.to_vec()
.try_into()
.expect("225 bytes should fit into Data")
}
pub fn zeroized_from_definition(
definition_id: AccountId,
definition: &TokenDefinition,
) -> Self {
match definition {
TokenDefinition::Fungible { .. } => TokenHolding::Fungible {
definition_id,
balance: 0,
},
TokenDefinition::NonFungible { .. } => TokenHolding::NftPrintedCopy {
definition_id,
owned: false,
},
}
}
pub fn parse(data: &[u8]) -> Option<Self> {
if data.len() != POOL_DEFINITION_DATA_SIZE {
None
} else {
let definition_token_a_id = AccountId::new(data[0..32].try_into().expect("Parse data: The AMM program must be provided a valid AccountId for Token A definition"));
let definition_token_b_id = AccountId::new(data[32..64].try_into().expect("Parse data: The AMM program must be provided a valid AccountId for Vault B definition"));
let vault_a_id = AccountId::new(data[64..96].try_into().expect(
"Parse data: The AMM program must be provided a valid AccountId for Vault A",
));
let vault_b_id = AccountId::new(data[96..128].try_into().expect(
"Parse data: The AMM program must be provided a valid AccountId for Vault B",
));
let liquidity_pool_id = AccountId::new(data[128..160].try_into().expect("Parse data: The AMM program must be provided a valid AccountId for Token liquidity pool definition"));
let liquidity_pool_supply = u128::from_le_bytes(data[160..176].try_into().expect(
"Parse data: The AMM program must be provided a valid u128 for liquidity cap",
));
let reserve_a = u128::from_le_bytes(data[176..192].try_into().expect(
"Parse data: The AMM program must be provided a valid u128 for reserve A balance",
));
let reserve_b = u128::from_le_bytes(data[192..208].try_into().expect(
"Parse data: The AMM program must be provided a valid u128 for reserve B balance",
));
let fees = u128::from_le_bytes(
data[208..224]
.try_into()
.expect("Parse data: The AMM program must be provided a valid u128 for fees"),
);
pub fn definition_id(&self) -> AccountId {
match self {
TokenHolding::Fungible { definition_id, .. } => *definition_id,
TokenHolding::NftMaster { definition_id, .. } => *definition_id,
TokenHolding::NftPrintedCopy { definition_id, .. } => *definition_id,
let active = match data[224] {
0 => false,
1 => true,
_ => panic!("Parse data: The AMM program must be provided a valid bool for active"),
};
Some(Self {
definition_token_a_id,
definition_token_b_id,
vault_a_id,
vault_b_id,
liquidity_pool_id,
liquidity_pool_supply,
reserve_a,
reserve_b,
fees,
active,
})
}
}
}
impl TryFrom<&Data> for TokenHolding {
type Error = std::io::Error;
fn try_from(data: &Data) -> Result<Self, Self::Error> {
TokenHolding::try_from_slice(data.as_ref())
}
pub fn compute_pool_pda(
amm_program_id: ProgramId,
definition_token_a_id: AccountId,
definition_token_b_id: AccountId,
) -> AccountId {
AccountId::from((
&amm_program_id,
&compute_pool_pda_seed(definition_token_a_id, definition_token_b_id),
))
}
impl From<&TokenHolding> for Data {
fn from(holding: &TokenHolding) -> Self {
// Using size_of_val as size hint for Vec allocation
let mut data = Vec::with_capacity(std::mem::size_of_val(holding));
pub fn compute_pool_pda_seed(
definition_token_a_id: AccountId,
definition_token_b_id: AccountId,
) -> PdaSeed {
use risc0_zkvm::sha::{Impl, Sha256};
BorshSerialize::serialize(holding, &mut data)
.expect("Serialization to Vec should not fail");
let (token_1, token_2) = match definition_token_a_id
.value()
.cmp(definition_token_b_id.value())
{
std::cmp::Ordering::Less => (definition_token_b_id, definition_token_a_id),
std::cmp::Ordering::Greater => (definition_token_a_id, definition_token_b_id),
std::cmp::Ordering::Equal => panic!("Definitions match"),
};
Data::try_from(data).expect("Token holding encoded data should fit into Data")
}
let mut bytes = [0; 64];
bytes[0..32].copy_from_slice(&token_1.to_bytes());
bytes[32..].copy_from_slice(&token_2.to_bytes());
PdaSeed::new(
Impl::hash_bytes(&bytes)
.as_bytes()
.try_into()
.expect("Hash output must be exactly 32 bytes long"),
)
}
#[derive(Serialize, Deserialize)]
pub struct NewTokenMetadata {
pub uri: String,
pub creators: String,
pub fn compute_vault_pda(
amm_program_id: ProgramId,
pool_id: AccountId,
definition_token_id: AccountId,
) -> AccountId {
AccountId::from((
&amm_program_id,
&compute_vault_pda_seed(pool_id, definition_token_id),
))
}
#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct TokenMetadata {
pub version: u8,
pub definition_id: AccountId,
pub uri: String,
pub creators: String,
/// Block id
pub primary_sale_date: u64,
pub fn compute_vault_pda_seed(pool_id: AccountId, definition_token_id: AccountId) -> PdaSeed {
use risc0_zkvm::sha::{Impl, Sha256};
let mut bytes = [0; 64];
bytes[0..32].copy_from_slice(&pool_id.to_bytes());
bytes[32..].copy_from_slice(&definition_token_id.to_bytes());
PdaSeed::new(
Impl::hash_bytes(&bytes)
.as_bytes()
.try_into()
.expect("Hash output must be exactly 32 bytes long"),
)
}
impl TryFrom<&Data> for TokenMetadata {
type Error = std::io::Error;
fn try_from(data: &Data) -> Result<Self, Self::Error> {
TokenMetadata::try_from_slice(data.as_ref())
}
pub fn compute_liquidity_token_pda(amm_program_id: ProgramId, pool_id: AccountId) -> AccountId {
AccountId::from((&amm_program_id, &compute_liquidity_token_pda_seed(pool_id)))
}
impl From<&TokenMetadata> for Data {
fn from(metadata: &TokenMetadata) -> Self {
// Using size_of_val as size hint for Vec allocation
let mut data = Vec::with_capacity(std::mem::size_of_val(metadata));
pub fn compute_liquidity_token_pda_seed(pool_id: AccountId) -> PdaSeed {
use risc0_zkvm::sha::{Impl, Sha256};
BorshSerialize::serialize(metadata, &mut data)
.expect("Serialization to Vec should not fail");
let mut bytes = [0; 64];
bytes[0..32].copy_from_slice(&pool_id.to_bytes());
bytes[32..].copy_from_slice(&[0; 32]);
Data::try_from(data).expect("Token metadata encoded data should fit into Data")
}
}*/
PdaSeed::new(
Impl::hash_bytes(&bytes)
.as_bytes()
.try_into()
.expect("Hash output must be exactly 32 bytes long"),
)
}

169
programs/amm/src/add.rs
View File

@ -0,0 +1,169 @@
use amm_core::{PoolDefinition, compute_liquidity_token_pda_seed};
use nssa_core::{
account::AccountWithMetadata,
program::{AccountPostState, ChainedCall},
};
#[expect(clippy::too_many_arguments, reason = "TODO: Fix later")]
pub fn add_liquidity(
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
pool_definition_lp: AccountWithMetadata,
user_holding_a: AccountWithMetadata,
user_holding_b: AccountWithMetadata,
user_holding_lp: AccountWithMetadata,
min_amount_liquidity: u128,
max_amount_to_add_token_a: u128,
max_amount_to_add_token_b: u128,
) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
// 1. Fetch Pool state
let pool_def_data = PoolDefinition::parse(&pool.account.data)
.expect("Add liquidity: AMM Program expects valid Pool Definition Account");
if vault_a.account_id != pool_def_data.vault_a_id {
panic!("Vault A was not provided");
}
if pool_def_data.liquidity_pool_id != pool_definition_lp.account_id {
panic!("LP definition mismatch");
}
if vault_b.account_id != pool_def_data.vault_b_id {
panic!("Vault B was not provided");
}
if max_amount_to_add_token_a == 0 || max_amount_to_add_token_b == 0 {
panic!("Both max-balances must be nonzero");
}
if min_amount_liquidity == 0 {
panic!("Min-lp must be nonzero");
}
// 2. Determine deposit amount
let vault_b_token_holding = token_core::TokenHolding::try_from(&vault_b.account.data)
.expect("Add liquidity: AMM Program expects valid Token Holding Account for Vault B");
let token_core::TokenHolding::Fungible {
definition_id: _,
balance: vault_b_balance,
} = vault_b_token_holding
else {
panic!(
"Add liquidity: AMM Program expects valid Fungible Token Holding Account for Vault B"
);
};
let vault_a_token_holding = token_core::TokenHolding::try_from(&vault_a.account.data)
.expect("Add liquidity: AMM Program expects valid Token Holding Account for Vault A");
let token_core::TokenHolding::Fungible {
definition_id: _,
balance: vault_a_balance,
} = vault_a_token_holding
else {
panic!(
"Add liquidity: AMM Program expects valid Fungible Token Holding Account for Vault A"
);
};
if pool_def_data.reserve_a == 0 || pool_def_data.reserve_b == 0 {
panic!("Reserves must be nonzero");
}
if vault_a_balance < pool_def_data.reserve_a || vault_b_balance < pool_def_data.reserve_b {
panic!("Vaults' balances must be at least the reserve amounts");
}
// Calculate actual_amounts
let ideal_a: u128 =
(pool_def_data.reserve_a * max_amount_to_add_token_b) / pool_def_data.reserve_b;
let ideal_b: u128 =
(pool_def_data.reserve_b * max_amount_to_add_token_a) / pool_def_data.reserve_a;
let actual_amount_a = if ideal_a > max_amount_to_add_token_a {
max_amount_to_add_token_a
} else {
ideal_a
};
let actual_amount_b = if ideal_b > max_amount_to_add_token_b {
max_amount_to_add_token_b
} else {
ideal_b
};
// 3. Validate amounts
if max_amount_to_add_token_a < actual_amount_a || max_amount_to_add_token_b < actual_amount_b {
panic!("Actual trade amounts cannot exceed max_amounts");
}
if actual_amount_a == 0 || actual_amount_b == 0 {
panic!("A trade amount is 0");
}
// 4. Calculate LP to mint
let delta_lp = std::cmp::min(
pool_def_data.liquidity_pool_supply * actual_amount_a / pool_def_data.reserve_a,
pool_def_data.liquidity_pool_supply * actual_amount_b / pool_def_data.reserve_b,
);
if delta_lp == 0 {
panic!("Payable LP must be nonzero");
}
if delta_lp < min_amount_liquidity {
panic!("Payable LP is less than provided minimum LP amount");
}
// 5. Update pool account
let mut pool_post = pool.account.clone();
let pool_post_definition = PoolDefinition {
liquidity_pool_supply: pool_def_data.liquidity_pool_supply + delta_lp,
reserve_a: pool_def_data.reserve_a + actual_amount_a,
reserve_b: pool_def_data.reserve_b + actual_amount_b,
..pool_def_data
};
pool_post.data = pool_post_definition.into_data();
let token_program_id = user_holding_a.account.program_owner;
// Chain call for Token A (UserHoldingA -> Vault_A)
let call_token_a = ChainedCall::new(
token_program_id,
vec![user_holding_a.clone(), vault_a.clone()],
&token_core::Instruction::Transfer {
amount_to_transfer: actual_amount_a,
},
);
// Chain call for Token B (UserHoldingB -> Vault_B)
let call_token_b = ChainedCall::new(
token_program_id,
vec![user_holding_b.clone(), vault_b.clone()],
&token_core::Instruction::Transfer {
amount_to_transfer: actual_amount_b,
},
);
// Chain call for LP (mint new tokens for user_holding_lp)
let mut pool_definition_lp_auth = pool_definition_lp.clone();
pool_definition_lp_auth.is_authorized = true;
let call_token_lp = ChainedCall::new(
token_program_id,
vec![pool_definition_lp_auth.clone(), user_holding_lp.clone()],
&token_core::Instruction::Mint {
amount_to_mint: delta_lp,
},
)
.with_pda_seeds(vec![compute_liquidity_token_pda_seed(pool.account_id)]);
let chained_calls = vec![call_token_lp, call_token_b, call_token_a];
let post_states = vec![
AccountPostState::new(pool_post),
AccountPostState::new(vault_a.account.clone()),
AccountPostState::new(vault_b.account.clone()),
AccountPostState::new(pool_definition_lp.account.clone()),
AccountPostState::new(user_holding_a.account.clone()),
AccountPostState::new(user_holding_b.account.clone()),
AccountPostState::new(user_holding_lp.account.clone()),
];
(post_states, chained_calls)
}

10
programs/amm/src/lib.rs Normal file
View File

@ -0,0 +1,10 @@
//! The AMM Program implementation.
pub use token_core as core;
pub mod add;
pub mod new_definition;
pub mod remove;
pub mod swap;
mod tests;

156
programs/amm/src/new_definition.rs
View File

@ -0,0 +1,156 @@
use amm_core::{
PoolDefinition, compute_liquidity_token_pda, compute_liquidity_token_pda_seed,
compute_pool_pda, compute_vault_pda,
};
use nssa_core::{
account::{Account, AccountWithMetadata},
program::{AccountPostState, ChainedCall, ProgramId},
};
#[expect(clippy::too_many_arguments, reason = "TODO: Fix later")]
pub fn new_definition(
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
pool_definition_lp: AccountWithMetadata,
user_holding_a: AccountWithMetadata,
user_holding_b: AccountWithMetadata,
user_holding_lp: AccountWithMetadata,
token_a_amount: u128,
token_b_amount: u128,
amm_program_id: ProgramId,
) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
// Prevents pool constant coefficient (k) from being 0.
if token_a_amount == 0 || token_b_amount == 0 {
panic!("Balances must be nonzero")
}
// Verify token_a and token_b are different
let definition_token_a_id = token_core::TokenHolding::try_from(&user_holding_a.account.data)
.expect("New definition: AMM Program expects valid Token Holding account for Token A")
.definition_id();
let definition_token_b_id = token_core::TokenHolding::try_from(&user_holding_b.account.data)
.expect("New definition: AMM Program expects valid Token Holding account for Token B")
.definition_id();
// both instances of the same token program
let token_program = user_holding_a.account.program_owner;
if user_holding_b.account.program_owner != token_program {
panic!("User Token holdings must use the same Token Program");
}
if definition_token_a_id == definition_token_b_id {
panic!("Cannot set up a swap for a token with itself")
}
if pool.account_id
!= compute_pool_pda(amm_program_id, definition_token_a_id, definition_token_b_id)
{
panic!("Pool Definition Account ID does not match PDA");
}
if vault_a.account_id
!= compute_vault_pda(amm_program_id, pool.account_id, definition_token_a_id)
|| vault_b.account_id
!= compute_vault_pda(amm_program_id, pool.account_id, definition_token_b_id)
{
panic!("Vault ID does not match PDA");
}
if pool_definition_lp.account_id != compute_liquidity_token_pda(amm_program_id, pool.account_id)
{
panic!("Liquidity pool Token Definition Account ID does not match PDA");
}
// Verify that Pool Account is not active
let pool_account_data = if pool.account == Account::default() {
PoolDefinition::default()
} else {
PoolDefinition::parse(&pool.account.data).expect("AMM program expects a valid Pool account")
};
if pool_account_data.active {
panic!("Cannot initialize an active Pool Definition")
}
// LP Token minting calculation
// We assume LP is based on the initial deposit amount for Token_A.
// Update pool account
let mut pool_post = pool.account.clone();
let pool_post_definition = PoolDefinition {
definition_token_a_id,
definition_token_b_id,
vault_a_id: vault_a.account_id,
vault_b_id: vault_b.account_id,
liquidity_pool_id: pool_definition_lp.account_id,
liquidity_pool_supply: token_a_amount,
reserve_a: token_a_amount,
reserve_b: token_b_amount,
fees: 0u128, // TODO: we assume all fees are 0 for now.
active: true,
};
pool_post.data = pool_post_definition.into_data();
let pool_post: AccountPostState = if pool.account == Account::default() {
AccountPostState::new_claimed(pool_post.clone())
} else {
AccountPostState::new(pool_post.clone())
};
let token_program_id = user_holding_a.account.program_owner;
// Chain call for Token A (user_holding_a -> Vault_A)
let call_token_a = ChainedCall::new(
token_program_id,
vec![user_holding_a.clone(), vault_a.clone()],
&token_core::Instruction::Transfer {
amount_to_transfer: token_a_amount,
},
);
// Chain call for Token B (user_holding_b -> Vault_B)
let call_token_b = ChainedCall::new(
token_program_id,
vec![user_holding_b.clone(), vault_b.clone()],
&token_core::Instruction::Transfer {
amount_to_transfer: token_b_amount,
},
);
// Chain call for liquidity token (TokenLP definition -> User LP Holding)
let instruction = if pool.account == Account::default() {
token_core::Instruction::NewFungibleDefinition {
name: String::from("LP Token"),
total_supply: token_a_amount,
}
} else {
token_core::Instruction::Mint {
amount_to_mint: token_a_amount,
}
};
let mut pool_lp_auth = pool_definition_lp.clone();
pool_lp_auth.is_authorized = true;
let call_token_lp = ChainedCall::new(
token_program_id,
vec![pool_lp_auth.clone(), user_holding_lp.clone()],
&instruction,
)
.with_pda_seeds(vec![compute_liquidity_token_pda_seed(pool.account_id)]);
let chained_calls = vec![call_token_lp, call_token_b, call_token_a];
let post_states = vec![
pool_post.clone(),
AccountPostState::new(vault_a.account.clone()),
AccountPostState::new(vault_b.account.clone()),
AccountPostState::new(pool_definition_lp.account.clone()),
AccountPostState::new(user_holding_a.account.clone()),
AccountPostState::new(user_holding_b.account.clone()),
AccountPostState::new(user_holding_lp.account.clone()),
];
(post_states.clone(), chained_calls)
}

71
programs/amm/src/remove.rs
View File

@ -1,27 +1,22 @@
fn remove_liquidity(
pre_states: &[AccountWithMetadata],
amounts: &[u128],
use amm_core::{PoolDefinition, compute_liquidity_token_pda_seed, compute_vault_pda_seed};
use nssa_core::{
account::AccountWithMetadata,
program::{AccountPostState, ChainedCall},
};
#[expect(clippy::too_many_arguments, reason = "TODO: Fix later")]
pub fn remove_liquidity(
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
pool_definition_lp: AccountWithMetadata,
user_holding_a: AccountWithMetadata,
user_holding_b: AccountWithMetadata,
user_holding_lp: AccountWithMetadata,
remove_liquidity_amount: u128,
min_amount_to_remove_token_a: u128,
min_amount_to_remove_token_b: u128,
) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
if pre_states.len() != 7 {
panic!("Invalid number of input accounts");
}
let pool = &pre_states[0];
let vault_a = &pre_states[1];
let vault_b = &pre_states[2];
let pool_definition_lp = &pre_states[3];
let user_holding_a = &pre_states[4];
let user_holding_b = &pre_states[5];
let user_holding_lp = &pre_states[6];
if amounts.len() != 3 {
panic!("Invalid number of balances");
}
let amount_lp = amounts[0];
let amount_min_a = amounts[1];
let amount_min_b = amounts[2];
// 1. Fetch Pool state
let pool_def_data = PoolDefinition::parse(&pool.account.data)
.expect("Remove liquidity: AMM Program expects a valid Pool Definition Account");
@ -50,11 +45,11 @@ fn remove_liquidity(
running_vault_a.is_authorized = true;
running_vault_b.is_authorized = true;
if amount_min_a == 0 || amount_min_b == 0 {
if min_amount_to_remove_token_a == 0 || min_amount_to_remove_token_b == 0 {
panic!("Minimum withdraw amount must be nonzero");
}
if amount_lp == 0 {
if remove_liquidity_amount == 0 {
panic!("Liquidity amount must be nonzero");
}
@ -78,21 +73,21 @@ fn remove_liquidity(
}
let withdraw_amount_a =
(pool_def_data.reserve_a * amount_lp) / pool_def_data.liquidity_pool_supply;
(pool_def_data.reserve_a * remove_liquidity_amount) / pool_def_data.liquidity_pool_supply;
let withdraw_amount_b =
(pool_def_data.reserve_b * amount_lp) / pool_def_data.liquidity_pool_supply;
(pool_def_data.reserve_b * remove_liquidity_amount) / pool_def_data.liquidity_pool_supply;
// 3. Validate and slippage check
if withdraw_amount_a < amount_min_a {
if withdraw_amount_a < min_amount_to_remove_token_a {
panic!("Insufficient minimal withdraw amount (Token A) provided for liquidity amount");
}
if withdraw_amount_b < amount_min_b {
if withdraw_amount_b < min_amount_to_remove_token_b {
panic!("Insufficient minimal withdraw amount (Token B) provided for liquidity amount");
}
// 4. Calculate LP to reduce cap by
let delta_lp: u128 =
(pool_def_data.liquidity_pool_supply * amount_lp) / pool_def_data.liquidity_pool_supply;
let delta_lp: u128 = (pool_def_data.liquidity_pool_supply * remove_liquidity_amount)
/ pool_def_data.liquidity_pool_supply;
let active: bool = pool_def_data.liquidity_pool_supply - delta_lp != 0;
@ -150,13 +145,13 @@ fn remove_liquidity(
let post_states = vec![
AccountPostState::new(pool_post.clone()),
AccountPostState::new(pre_states[1].account.clone()),
AccountPostState::new(pre_states[2].account.clone()),
AccountPostState::new(pre_states[3].account.clone()),
AccountPostState::new(pre_states[4].account.clone()),
AccountPostState::new(pre_states[5].account.clone()),
AccountPostState::new(pre_states[6].account.clone()),
AccountPostState::new(vault_a.account.clone()),
AccountPostState::new(vault_b.account.clone()),
AccountPostState::new(pool_definition_lp.account.clone()),
AccountPostState::new(user_holding_a.account.clone()),
AccountPostState::new(user_holding_b.account.clone()),
AccountPostState::new(user_holding_lp.account.clone()),
];
(post_states, chained_calls)
}
}

177
programs/amm/src/swap.rs
View File

@ -0,0 +1,177 @@
pub use amm_core::{PoolDefinition, compute_liquidity_token_pda_seed, compute_vault_pda_seed};
use nssa_core::{
account::{AccountId, AccountWithMetadata},
program::{AccountPostState, ChainedCall},
};
#[expect(clippy::too_many_arguments, reason = "TODO: Fix later")]
pub fn swap(
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
user_holding_a: AccountWithMetadata,
user_holding_b: AccountWithMetadata,
swap_amount_in: u128,
min_amount_out: u128,
token_in_id: AccountId,
) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
// Verify vaults are in fact vaults
let pool_def_data = PoolDefinition::parse(&pool.account.data)
.expect("Swap: AMM Program expects a valid Pool Definition Account");
if !pool_def_data.active {
panic!("Pool is inactive");
}
if vault_a.account_id != pool_def_data.vault_a_id {
panic!("Vault A was not provided");
}
if vault_b.account_id != pool_def_data.vault_b_id {
panic!("Vault B was not provided");
}
// fetch pool reserves
// validates reserves is at least the vaults' balances
let vault_a_token_holding = token_core::TokenHolding::try_from(&vault_a.account.data)
.expect("Swap: AMM Program expects a valid Token Holding Account for Vault A");
let token_core::TokenHolding::Fungible {
definition_id: _,
balance: vault_a_balance,
} = vault_a_token_holding
else {
panic!("Swap: AMM Program expects a valid Fungible Token Holding Account for Vault A");
};
if vault_a_balance < pool_def_data.reserve_a {
panic!("Reserve for Token A exceeds vault balance");
}
let vault_b_token_holding = token_core::TokenHolding::try_from(&vault_b.account.data)
.expect("Swap: AMM Program expects a valid Token Holding Account for Vault B");
let token_core::TokenHolding::Fungible {
definition_id: _,
balance: vault_b_balance,
} = vault_b_token_holding
else {
panic!("Swap: AMM Program expects a valid Fungible Token Holding Account for Vault B");
};
if vault_b_balance < pool_def_data.reserve_b {
panic!("Reserve for Token B exceeds vault balance");
}
let (chained_calls, [deposit_a, withdraw_a], [deposit_b, withdraw_b]) =
if token_in_id == pool_def_data.definition_token_a_id {
let (chained_calls, deposit_a, withdraw_b) = swap_logic(
user_holding_a.clone(),
vault_a.clone(),
vault_b.clone(),
user_holding_b.clone(),
swap_amount_in,
min_amount_out,
&[pool_def_data.reserve_a, pool_def_data.reserve_b],
pool.account_id,
);
(chained_calls, [deposit_a, 0], [0, withdraw_b])
} else if token_in_id == pool_def_data.definition_token_b_id {
let (chained_calls, deposit_b, withdraw_a) = swap_logic(
user_holding_b.clone(),
vault_b.clone(),
vault_a.clone(),
user_holding_a.clone(),
swap_amount_in,
min_amount_out,
&[pool_def_data.reserve_b, pool_def_data.reserve_a],
pool.account_id,
);
(chained_calls, [0, withdraw_a], [deposit_b, 0])
} else {
panic!("AccountId is not a token type for the pool");
};
// Update pool account
let mut pool_post = pool.account.clone();
let pool_post_definition = PoolDefinition {
reserve_a: pool_def_data.reserve_a + deposit_a - withdraw_a,
reserve_b: pool_def_data.reserve_b + deposit_b - withdraw_b,
..pool_def_data
};
pool_post.data = pool_post_definition.into_data();
let post_states = vec![
AccountPostState::new(pool_post.clone()),
AccountPostState::new(vault_a.account.clone()),
AccountPostState::new(vault_b.account.clone()),
AccountPostState::new(user_holding_a.account.clone()),
AccountPostState::new(user_holding_b.account.clone()),
];
(post_states, chained_calls)
}
#[expect(clippy::too_many_arguments, reason = "TODO: Fix later")]
fn swap_logic(
user_deposit: AccountWithMetadata,
vault_deposit: AccountWithMetadata,
vault_withdraw: AccountWithMetadata,
user_withdraw: AccountWithMetadata,
swap_amount_in: u128,
min_amount_out: u128,
reserve_amounts: &[u128],
pool_id: AccountId,
) -> (Vec<ChainedCall>, u128, u128) {
let reserve_deposit_vault_amount = reserve_amounts[0];
let reserve_withdraw_vault_amount = reserve_amounts[1];
// Compute withdraw amount
// Maintains pool constant product
// k = pool_def_data.reserve_a * pool_def_data.reserve_b;
let withdraw_amount = (reserve_withdraw_vault_amount * swap_amount_in)
/ (reserve_deposit_vault_amount + swap_amount_in);
// Slippage check
if min_amount_out > withdraw_amount {
panic!("Withdraw amount is less than minimal amount out");
}
if withdraw_amount == 0 {
panic!("Withdraw amount should be nonzero");
}
let token_program_id = user_deposit.account.program_owner;
let mut chained_calls = Vec::new();
chained_calls.push(ChainedCall::new(
token_program_id,
vec![user_deposit, vault_deposit],
&token_core::Instruction::Transfer {
amount_to_transfer: swap_amount_in,
},
));
let mut vault_withdraw = vault_withdraw.clone();
vault_withdraw.is_authorized = true;
let pda_seed = compute_vault_pda_seed(
pool_id,
token_core::TokenHolding::try_from(&vault_withdraw.account.data)
.expect("Swap Logic: AMM Program expects valid token data")
.definition_id(),
);
chained_calls.push(
ChainedCall::new(
token_program_id,
vec![vault_withdraw, user_withdraw],
&token_core::Instruction::Transfer {
amount_to_transfer: withdraw_amount,
},
)
.with_pda_seeds(vec![pda_seed]),
);
(chained_calls, swap_amount_in, withdraw_amount)
}

1751
programs/amm/src/tests.rs Normal file
View File

File diff suppressed because it is too large Load Diff

1
wallet/Cargo.toml
View File

@ -9,6 +9,7 @@ nssa.workspace = true
common.workspace = true
key_protocol.workspace = true
token_core.workspace = true
amm_core.workspace = true
anyhow.workspace = true
serde_json.workspace = true

247
wallet/src/program_facades/amm.rs
View File

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