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

Merge pull request #306 from logos-blockchain/marvin/refactor-amm-program

Browse Source 
refactor amm program
This commit is contained in:
jonesmarvin8 2026-02-11 19:42:41 -05:00 committed by GitHub
commit 818fc5e601
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
23 changed files with 2872 additions and 3843 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
Cargo.lock generated
View File

@ -302,6 +302,25 @@ version = "0.2.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923"
[[package]]
name = "amm_core"
version = "0.1.0"
dependencies = [
"borsh",
"nssa_core",
"risc0-zkvm",
"serde",
]
[[package]]
name = "amm_program"
version = "0.1.0"
dependencies = [
"amm_core",
"nssa_core",
"token_core",
]
[[package]]
name = "android_system_properties"
version = "0.1.5"
@ -4808,8 +4827,8 @@ checksum = "38bf9645c8b145698bb0b18a4637dcacbc421ea49bef2317e4fd8065a387cf21"
name = "nssa"
version = "0.1.0"
dependencies = [
"amm_core",
"borsh",
"bytemuck",
"env_logger",
"hex",
"hex-literal 1.1.0",
@ -5339,6 +5358,8 @@ dependencies = [
name = "programs"
version = "0.1.0"
dependencies = [
"amm_core",
"amm_program",
"nssa_core",
"risc0-zkvm",
"serde",
@ -7754,6 +7775,7 @@ dependencies = [
name = "wallet"
version = "0.1.0"
dependencies = [
"amm_core",
"anyhow",
"async-stream",
"base58",
@ -7773,7 +7795,6 @@ dependencies = [
"nssa_core",
"optfield",
"rand 0.8.5",
"risc0-zkvm",
"serde",
"serde_json",
"sha2",

6
Cargo.toml
View File

@ -13,6 +13,10 @@ members = [
"common",
"nssa",
"nssa/core",
"programs/amm/core",
"programs/amm",
"programs/token/core",
"programs/token",
"sequencer_core",
"sequencer_rpc",
"sequencer_runner",
@ -50,6 +54,8 @@ wallet = { path = "wallet" }
wallet-ffi = { path = "wallet-ffi" }
token_core = { path = "programs/token/core" }
token_program = { path = "programs/token" }
amm_core = { path = "programs/amm/core" }
amm_program = { path = "programs/amm" }
test_program_methods = { path = "test_program_methods" }
bedrock_client = { path = "bedrock_client" }
indexer_core = { path = "indexer_core" }

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.

2
nssa/Cargo.toml
View File

@ -16,7 +16,6 @@ borsh.workspace = true
hex.workspace = true
secp256k1 = "0.31.1"
risc0-binfmt = "3.0.2"
bytemuck = "1.24.0"
log.workspace = true
[build-dependencies]
@ -25,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

236
nssa/src/state.rs
View File

@ -311,6 +311,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},
@ -2329,137 +2330,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 {
@ -2640,7 +2510,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(),
@ -2648,7 +2518,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 {
@ -2678,7 +2551,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(),
@ -2686,7 +2559,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(),
@ -2725,7 +2598,7 @@ pub mod tests {
Account {
program_owner: Program::amm().id(),
balance: 0u128,
data: PoolDefinition::into_data(PoolDefinition {
data: Data::from(&PoolDefinition {
definition_token_a_id: IdForTests::token_a_definition_id(),
definition_token_b_id: IdForTests::token_b_definition_id(),
vault_a_id: IdForTests::vault_a_id(),
@ -2844,7 +2717,7 @@ pub mod tests {
Account {
program_owner: Program::amm().id(),
balance: 0u128,
data: PoolDefinition::into_data(PoolDefinition {
data: Data::from(&PoolDefinition {
definition_token_a_id: IdForTests::token_a_definition_id(),
definition_token_b_id: IdForTests::token_b_definition_id(),
vault_a_id: IdForTests::vault_a_id(),
@ -2912,7 +2785,7 @@ pub mod tests {
Account {
program_owner: Program::amm().id(),
balance: 0u128,
data: PoolDefinition::into_data(PoolDefinition {
data: Data::from(&PoolDefinition {
definition_token_a_id: IdForTests::token_a_definition_id(),
definition_token_b_id: IdForTests::token_b_definition_id(),
vault_a_id: IdForTests::vault_a_id(),
@ -2980,7 +2853,7 @@ pub mod tests {
Account {
program_owner: Program::amm().id(),
balance: 0u128,
data: PoolDefinition::into_data(PoolDefinition {
data: Data::from(&PoolDefinition {
definition_token_a_id: IdForTests::token_a_definition_id(),
definition_token_b_id: IdForTests::token_b_definition_id(),
vault_a_id: IdForTests::vault_a_id(),
@ -3073,7 +2946,7 @@ pub mod tests {
Account {
program_owner: Program::amm().id(),
balance: 0u128,
data: PoolDefinition::into_data(PoolDefinition {
data: Data::from(&PoolDefinition {
definition_token_a_id: IdForTests::token_a_definition_id(),
definition_token_b_id: IdForTests::token_b_definition_id(),
vault_a_id: IdForTests::vault_a_id(),
@ -3179,7 +3052,7 @@ pub mod tests {
Account {
program_owner: Program::amm().id(),
balance: 0u128,
data: PoolDefinition::into_data(PoolDefinition {
data: Data::from(&PoolDefinition {
definition_token_a_id: IdForTests::token_a_definition_id(),
definition_token_b_id: IdForTests::token_b_definition_id(),
vault_a_id: IdForTests::vault_a_id(),
@ -3248,7 +3121,7 @@ pub mod tests {
Account {
program_owner: Program::amm().id(),
balance: 0u128,
data: PoolDefinition::into_data(PoolDefinition {
data: Data::from(&PoolDefinition {
definition_token_a_id: IdForTests::token_a_definition_id(),
definition_token_b_id: IdForTests::token_b_definition_id(),
vault_a_id: IdForTests::vault_a_id(),
@ -3277,11 +3150,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 =
@ -3347,11 +3215,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(),
@ -3424,12 +3292,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(),
@ -3509,12 +3376,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(),
@ -3582,12 +3448,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(),
@ -3646,11 +3511,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(),
@ -3708,11 +3573,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(),
@ -3759,12 +3624,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

@ -8,5 +8,7 @@ license = { workspace = true }
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 }

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

File diff suppressed because it is too large Load Diff

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

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

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

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

197
programs/amm/core/src/lib.rs Normal file
View File

@ -0,0 +1,197 @@
//! This crate contains core data structures and utilities for the AMM Program.
use borsh::{BorshDeserialize, BorshSerialize};
use nssa_core::{
account::{AccountId, Data},
program::{PdaSeed, ProgramId},
};
use serde::{Deserialize, Serialize};
/// AMM Program Instruction.
#[derive(Serialize, Deserialize)]
pub enum Instruction {
/// Initializes a new Pool (or re-initializes an inactive Pool).
///
/// Required accounts:
/// - 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,
},
/// Adds liquidity to the Pool
///
/// Required accounts:
/// - 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,
},
/// Removes liquidity from the Pool
///
/// Required accounts:
/// - 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,
},
/// Swap some quantity of Tokens (either Token A or Token B)
/// while maintaining the Pool constant product.
///
/// Required accounts:
/// - 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,
},
}
#[derive(Clone, Default, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
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 TryFrom<&Data> for PoolDefinition {
type Error = std::io::Error;
fn try_from(data: &Data) -> Result<Self, Self::Error> {
PoolDefinition::try_from_slice(data.as_ref())
}
}
impl From<&PoolDefinition> for Data {
fn from(definition: &PoolDefinition) -> Self {
// Using size_of_val as size hint for Vec allocation
let mut data = Vec::with_capacity(std::mem::size_of_val(definition));
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")
}
}
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),
))
}
pub fn compute_pool_pda_seed(
definition_token_a_id: AccountId,
definition_token_b_id: AccountId,
) -> PdaSeed {
use risc0_zkvm::sha::{Impl, Sha256};
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"),
};
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"),
)
}
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),
))
}
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"),
)
}
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)))
}
pub 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"),
)
}

178
programs/amm/src/add.rs Normal file
View File

@ -0,0 +1,178 @@
use std::num::NonZeroU128;
use amm_core::{PoolDefinition, compute_liquidity_token_pda_seed};
use nssa_core::{
account::{AccountWithMetadata, Data},
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: NonZeroU128,
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::try_from(&pool.account.data)
.expect("Add liquidity: AMM Program expects valid Pool Definition Account");
assert_eq!(
vault_a.account_id, pool_def_data.vault_a_id,
"Vault A was not provided"
);
assert_eq!(
pool_def_data.liquidity_pool_id, pool_definition_lp.account_id,
"LP definition mismatch"
);
assert_eq!(
vault_b.account_id, pool_def_data.vault_b_id,
"Vault B was not provided"
);
assert!(
max_amount_to_add_token_a != 0 && max_amount_to_add_token_b != 0,
"Both max-balances 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"
);
};
assert!(pool_def_data.reserve_a != 0, "Reserves must be nonzero");
assert!(pool_def_data.reserve_b != 0, "Reserves must be nonzero");
assert!(
vault_a_balance >= pool_def_data.reserve_a,
"Vaults' balances must be at least the reserve amounts"
);
assert!(
vault_b_balance >= pool_def_data.reserve_b,
"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
assert!(
max_amount_to_add_token_a >= actual_amount_a,
"Actual trade amounts cannot exceed max_amounts"
);
assert!(
max_amount_to_add_token_b >= actual_amount_b,
"Actual trade amounts cannot exceed max_amounts"
);
assert!(actual_amount_a != 0, "A trade amount is 0");
assert!(actual_amount_b != 0, "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,
);
assert!(delta_lp != 0, "Payable LP must be nonzero");
assert!(
delta_lp >= min_amount_liquidity.into(),
"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 = Data::from(&pool_post_definition);
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 amm_core as core;
pub mod add;
pub mod new_definition;
pub mod remove;
pub mod swap;
mod tests;

158
programs/amm/src/new_definition.rs Normal file
View File

@ -0,0 +1,158 @@
use std::num::NonZeroU128;
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, Data},
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: NonZeroU128,
token_b_amount: NonZeroU128,
amm_program_id: ProgramId,
) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
// 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;
assert_eq!(
user_holding_b.account.program_owner, token_program,
"User Token holdings must use the same Token Program"
);
assert!(
definition_token_a_id != definition_token_b_id,
"Cannot set up a swap for a token with itself"
);
assert_eq!(
pool.account_id,
compute_pool_pda(amm_program_id, definition_token_a_id, definition_token_b_id),
"Pool Definition Account ID does not match PDA"
);
assert_eq!(
vault_a.account_id,
compute_vault_pda(amm_program_id, pool.account_id, definition_token_a_id),
"Vault ID does not match PDA"
);
assert_eq!(
vault_b.account_id,
compute_vault_pda(amm_program_id, pool.account_id, definition_token_b_id),
"Vault ID does not match PDA"
);
assert_eq!(
pool_definition_lp.account_id,
compute_liquidity_token_pda(amm_program_id, pool.account_id),
"Liquidity pool Token Definition Account ID does not match PDA"
);
// TODO: return here
// Verify that Pool Account is not active
let pool_account_data = if pool.account == Account::default() {
PoolDefinition::default()
} else {
PoolDefinition::try_from(&pool.account.data)
.expect("AMM program expects a valid Pool account")
};
assert!(
!pool_account_data.active,
"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.into(),
reserve_a: token_a_amount.into(),
reserve_b: token_b_amount.into(),
fees: 0u128, // TODO: we assume all fees are 0 for now.
active: true,
};
pool_post.data = Data::from(&pool_post_definition);
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.into(),
},
);
// 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.into(),
},
);
// 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.into(),
}
} else {
token_core::Instruction::Mint {
amount_to_mint: token_a_amount.into(),
}
};
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)
}

166
programs/amm/src/remove.rs Normal file
View File

@ -0,0 +1,166 @@
use std::num::NonZeroU128;
use amm_core::{PoolDefinition, compute_liquidity_token_pda_seed, compute_vault_pda_seed};
use nssa_core::{
account::{AccountWithMetadata, Data},
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: NonZeroU128,
min_amount_to_remove_token_a: u128,
min_amount_to_remove_token_b: u128,
) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
let remove_liquidity_amount: u128 = remove_liquidity_amount.into();
// 1. Fetch Pool state
let pool_def_data = PoolDefinition::try_from(&pool.account.data)
.expect("Remove liquidity: AMM Program expects a valid Pool Definition Account");
assert!(pool_def_data.active, "Pool is inactive");
assert_eq!(
pool_def_data.liquidity_pool_id, pool_definition_lp.account_id,
"LP definition mismatch"
);
assert_eq!(
vault_a.account_id, pool_def_data.vault_a_id,
"Vault A was not provided"
);
assert_eq!(
vault_b.account_id, pool_def_data.vault_b_id,
"Vault B was not provided"
);
// Vault addresses do not need to be checked with PDA
// calculation for setting authorization since stored
// in the Pool Definition.
let mut running_vault_a = vault_a.clone();
let mut running_vault_b = vault_b.clone();
running_vault_a.is_authorized = true;
running_vault_b.is_authorized = true;
assert!(
min_amount_to_remove_token_a != 0,
"Minimum withdraw amount must be nonzero"
);
assert!(
min_amount_to_remove_token_b != 0,
"Minimum withdraw amount must be nonzero"
);
// 2. Compute withdrawal amounts
let user_holding_lp_data = token_core::TokenHolding::try_from(&user_holding_lp.account.data)
.expect("Remove liquidity: AMM Program expects a valid Token Account for liquidity token");
let token_core::TokenHolding::Fungible {
definition_id: _,
balance: user_lp_balance,
} = user_holding_lp_data
else {
panic!(
"Remove liquidity: AMM Program expects a valid Fungible Token Holding Account for liquidity token"
);
};
assert!(
user_lp_balance <= pool_def_data.liquidity_pool_supply,
"Invalid liquidity account provided"
);
assert_eq!(
user_holding_lp_data.definition_id(),
pool_def_data.liquidity_pool_id,
"Invalid liquidity account provided"
);
let withdraw_amount_a =
(pool_def_data.reserve_a * remove_liquidity_amount) / pool_def_data.liquidity_pool_supply;
let withdraw_amount_b =
(pool_def_data.reserve_b * remove_liquidity_amount) / pool_def_data.liquidity_pool_supply;
// 3. Validate and slippage check
assert!(
withdraw_amount_a >= min_amount_to_remove_token_a,
"Insufficient minimal withdraw amount (Token A) provided for liquidity amount"
);
assert!(
withdraw_amount_b >= min_amount_to_remove_token_b,
"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 * remove_liquidity_amount)
/ pool_def_data.liquidity_pool_supply;
let active: bool = pool_def_data.liquidity_pool_supply - delta_lp != 0;
// 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 - withdraw_amount_a,
reserve_b: pool_def_data.reserve_b - withdraw_amount_b,
active,
..pool_def_data.clone()
};
pool_post.data = Data::from(&pool_post_definition);
let token_program_id = user_holding_a.account.program_owner;
// Chaincall for Token A withdraw
let call_token_a = ChainedCall::new(
token_program_id,
vec![running_vault_a, user_holding_a.clone()],
&token_core::Instruction::Transfer {
amount_to_transfer: withdraw_amount_a,
},
)
.with_pda_seeds(vec![compute_vault_pda_seed(
pool.account_id,
pool_def_data.definition_token_a_id,
)]);
// Chaincall for Token B withdraw
let call_token_b = ChainedCall::new(
token_program_id,
vec![running_vault_b, user_holding_b.clone()],
&token_core::Instruction::Transfer {
amount_to_transfer: withdraw_amount_b,
},
)
.with_pda_seeds(vec![compute_vault_pda_seed(
pool.account_id,
pool_def_data.definition_token_b_id,
)]);
// Chaincall for LP adjustment
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, user_holding_lp.clone()],
&token_core::Instruction::Burn {
amount_to_burn: 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.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)
}

176
programs/amm/src/swap.rs Normal file
View File

@ -0,0 +1,176 @@
pub use amm_core::{PoolDefinition, compute_liquidity_token_pda_seed, compute_vault_pda_seed};
use nssa_core::{
account::{AccountId, AccountWithMetadata, Data},
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::try_from(&pool.account.data)
.expect("Swap: AMM Program expects a valid Pool Definition Account");
assert!(pool_def_data.active, "Pool is inactive");
assert_eq!(
vault_a.account_id, pool_def_data.vault_a_id,
"Vault A was not provided"
);
assert_eq!(
vault_b.account_id, pool_def_data.vault_b_id,
"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");
};
assert!(
vault_a_balance >= pool_def_data.reserve_a,
"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");
};
assert!(
vault_b_balance >= pool_def_data.reserve_b,
"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 = Data::from(&pool_post_definition);
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_deposit_vault_amount: u128,
reserve_withdraw_vault_amount: u128,
pool_id: AccountId,
) -> (Vec<ChainedCall>, u128, u128) {
// 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
assert!(
min_amount_out <= withdraw_amount,
"Withdraw amount is less than minimal amount out"
);
assert!(withdraw_amount != 0, "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)
}

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

File diff suppressed because it is too large Load Diff

2
wallet/Cargo.toml
View File

@ -10,6 +10,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
@ -27,7 +28,6 @@ rand.workspace = true
itertools.workspace = true
sha2.workspace = true
futures.workspace = true
risc0-zkvm.workspace = true
async-stream = "0.3.6"
indicatif = { version = "0.18.3", features = ["improved_unicode"] }
optfield = "0.4.0"

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)
}