feat(amm): add Initialize instruction with config-gated chained calls

Introduce a singleton AMM configuration account, a PDA derived from the
constant "CONFIG" seed, created once via a new `Initialize` instruction.
The config stores the Token Program ID the AMM issues every chained call
to, replacing the previous behavior of trusting the program owner of a
caller-supplied holding.

The config account's existence is the Program's initialization gate: the
chained-call instructions (new_definition, add_liquidity, remove_liquidity,
swap_exact_input, swap_exact_output) now take the config as their first
account, validate it against `compute_config_pda(self_program_id)`, and
read the Token Program ID from it on demand — rejecting calls until the
Program is initialized. Vaults and user holdings are asserted to match the
configured Token Program. sync_reserves is left ungated, as it cannot act
on a pool that could not have existed before initialization.

- amm_core: AmmConfig type, compute_config_pda/_seed, Initialize variant
- amm: initialize.rs + config threading through chained-call instructions
- guest: initialize instruction; config + self_program_id on gated calls
- tests: config fixtures, init-gate unit tests, end-to-end Initialize VM test
This commit is contained in:
r4bbit 2026-06-17 16:29:30 +02:00
parent e8fe634a2c
commit 3624ea1451
11 changed files with 724 additions and 61 deletions

View File

@ -2,9 +2,32 @@
"version": "0.1.0",
"name": "amm",
"instructions": [
{
"name": "initialize",
"accounts": [
{
"name": "config",
"writable": false,
"signer": false,
"init": false
}
],
"args": [
{
"name": "token_program_id",
"type": "program_id"
}
]
},
{
"name": "new_definition",
"accounts": [
{
"name": "config",
"writable": false,
"signer": false,
"init": false
},
{
"name": "pool",
"writable": false,
@ -76,6 +99,12 @@
{
"name": "add_liquidity",
"accounts": [
{
"name": "config",
"writable": false,
"signer": false,
"init": false
},
{
"name": "pool",
"writable": false,
@ -141,6 +170,12 @@
{
"name": "remove_liquidity",
"accounts": [
{
"name": "config",
"writable": false,
"signer": false,
"init": false
},
{
"name": "pool",
"writable": false,
@ -206,6 +241,12 @@
{
"name": "swap_exact_input",
"accounts": [
{
"name": "config",
"writable": false,
"signer": false,
"init": false
},
{
"name": "pool",
"writable": false,
@ -259,6 +300,12 @@
{
"name": "swap_exact_output",
"accounts": [
{
"name": "config",
"writable": false,
"signer": false,
"init": false
},
{
"name": "pool",
"writable": false,
@ -379,6 +426,18 @@
]
}
},
{
"name": "AmmConfig",
"type": {
"kind": "struct",
"fields": [
{
"name": "token_program_id",
"type": "program_id"
}
]
}
},
{
"name": "TokenDefinition",
"type": {

View File

@ -16,6 +16,21 @@ const LP_LOCK_HOLDING_PDA_SEED: [u8; 32] = [1; 32];
/// AMM Program Instruction.
#[derive(Serialize, Deserialize)]
pub enum Instruction {
/// Initializes the AMM Program by creating its singleton configuration account.
///
/// The configuration account is a PDA derived from the constant `"CONFIG"` seed
/// (`compute_config_pda(self_program_id)`). It stores the Token Program ID that the AMM
/// uses for every chained call. The Program must be initialized via this instruction before
/// any pool can be created or interacted with — the other instructions read the Token
/// Program ID from this account and reject calls when it does not yet exist.
///
/// Required accounts:
/// - AMM Config Account, uninitialized, derived as `compute_config_pda(self_program_id)`
Initialize {
/// Program ID of the Token Program the AMM will issue chained calls to.
token_program_id: ProgramId,
},
/// Initializes a new Pool (or re-initializes an existing zero-supply Pool).
///
/// On initialization, `MINIMUM_LIQUIDITY` LP tokens are permanently locked
@ -177,6 +192,61 @@ impl From<&PoolDefinition> for Data {
}
}
/// Singleton configuration account for the AMM Program.
///
/// Stored at the PDA derived from the constant `"CONFIG"` seed
/// (`compute_config_pda(amm_program_id)`). Created once via the `Initialize` instruction; its
/// existence is the Program's "initialized" flag. Every chained-call instruction reads
/// `token_program_id` from here instead of trusting the program owner of a caller-supplied
/// account.
#[account_type]
#[derive(Clone, Default, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct AmmConfig {
/// Program ID of the Token Program the AMM issues chained calls to.
pub token_program_id: ProgramId,
}
impl TryFrom<&Data> for AmmConfig {
type Error = std::io::Error;
fn try_from(data: &Data) -> Result<Self, Self::Error> {
AmmConfig::try_from_slice(data.as_ref())
}
}
impl From<&AmmConfig> for Data {
fn from(config: &AmmConfig) -> Self {
let mut data = Vec::with_capacity(std::mem::size_of_val(config));
BorshSerialize::serialize(config, &mut data).expect("Serialization to Vec should not fail");
Data::try_from(data).expect("AMM config encoded data should fit into Data")
}
}
// Stable seed marker for the singleton config PDA. The literal `"CONFIG"` bytes are hashed into
// the 32-byte seed; this must stay unchanged for address compatibility.
const CONFIG_PDA_SEED: &[u8] = b"CONFIG";
/// Derives the [`AccountId`] of the AMM Program's singleton config PDA.
#[must_use]
pub fn compute_config_pda(amm_program_id: ProgramId) -> AccountId {
AccountId::for_public_pda(&amm_program_id, &compute_config_pda_seed())
}
/// Derives the [`PdaSeed`] of the AMM Program's singleton config PDA from the `"CONFIG"` bytes.
#[must_use]
pub fn compute_config_pda_seed() -> PdaSeed {
use risc0_zkvm::sha::{Impl, Sha256};
PdaSeed::new(
Impl::hash_bytes(CONFIG_PDA_SEED)
.as_bytes()
.try_into()
.expect("Hash output must be exactly 32 bytes long"),
)
}
pub fn compute_pool_pda(
amm_program_id: ProgramId,
definition_token_a_id: AccountId,

View File

@ -6,6 +6,7 @@ use spel_framework::prelude::*;
use spel_framework::context::ProgramContext;
use nssa_core::{
account::{AccountId, AccountWithMetadata},
program::ProgramId,
};
#[cfg(not(test))]
@ -19,15 +20,31 @@ mod amm {
)]
use super::*;
/// Initializes the AMM Program by creating its singleton config account.
///
/// Expected accounts:
/// 1. `config` — uninitialized config PDA derived from `compute_config_pda(self_program_id)`.
#[instruction]
pub fn initialize(
ctx: ProgramContext,
config: AccountWithMetadata,
token_program_id: ProgramId,
) -> SpelResult {
let post_states =
amm_program::initialize::initialize(config, token_program_id, ctx.self_program_id);
Ok(spel_framework::SpelOutput::execute(post_states, vec![]))
}
/// Initializes a new Pool (or re-initializes an existing zero-supply Pool).
/// A fresh user LP holding must be explicitly authorized by the caller.
#[expect(
clippy::too_many_arguments,
reason = "instruction interface requires explicit pool, vault, mint, lock, and user accounts"
reason = "instruction interface requires explicit config, pool, vault, mint, lock, and user accounts"
)]
#[instruction]
pub fn new_definition(
ctx: ProgramContext,
config: AccountWithMetadata,
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
@ -42,6 +59,7 @@ mod amm {
deadline: u64,
) -> SpelResult {
let (post_states, chained_calls) = amm_program::new_definition::new_definition(
config,
pool,
vault_a,
vault_b,
@ -66,6 +84,8 @@ mod amm {
)]
#[instruction]
pub fn add_liquidity(
ctx: ProgramContext,
config: AccountWithMetadata,
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
@ -79,6 +99,7 @@ mod amm {
deadline: u64,
) -> SpelResult {
let (post_states, chained_calls) = amm_program::add::add_liquidity(
config,
pool,
vault_a,
vault_b,
@ -89,6 +110,7 @@ mod amm {
NonZeroU128::new(min_amount_liquidity).expect("min_amount_liquidity must be nonzero"),
max_amount_to_add_token_a,
max_amount_to_add_token_b,
ctx.self_program_id,
);
Ok(spel_framework::SpelOutput::execute(post_states, chained_calls)
.with_timestamp_validity_window(..deadline))
@ -101,6 +123,8 @@ mod amm {
)]
#[instruction]
pub fn remove_liquidity(
ctx: ProgramContext,
config: AccountWithMetadata,
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
@ -114,6 +138,7 @@ mod amm {
deadline: u64,
) -> SpelResult {
let (post_states, chained_calls) = amm_program::remove::remove_liquidity(
config,
pool,
vault_a,
vault_b,
@ -125,6 +150,7 @@ mod amm {
.expect("remove_liquidity_amount must be nonzero"),
min_amount_to_remove_token_a,
min_amount_to_remove_token_b,
ctx.self_program_id,
);
Ok(spel_framework::SpelOutput::execute(post_states, chained_calls)
.with_timestamp_validity_window(..deadline))
@ -137,6 +163,8 @@ mod amm {
)]
#[instruction]
pub fn swap_exact_input(
ctx: ProgramContext,
config: AccountWithMetadata,
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
@ -148,6 +176,7 @@ mod amm {
deadline: u64,
) -> SpelResult {
let (post_states, chained_calls) = amm_program::swap::swap_exact_input(
config,
pool,
vault_a,
vault_b,
@ -156,6 +185,7 @@ mod amm {
swap_amount_in,
min_amount_out,
token_definition_id_in,
ctx.self_program_id,
);
Ok(spel_framework::SpelOutput::execute(post_states, chained_calls)
.with_timestamp_validity_window(..deadline))
@ -168,6 +198,8 @@ mod amm {
)]
#[instruction]
pub fn swap_exact_output(
ctx: ProgramContext,
config: AccountWithMetadata,
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
@ -179,6 +211,7 @@ mod amm {
deadline: u64,
) -> SpelResult {
let (post_states, chained_calls) = amm_program::swap::swap_exact_output(
config,
pool,
vault_a,
vault_b,
@ -187,6 +220,7 @@ mod amm {
exact_amount_out,
max_amount_in,
token_definition_id_in,
ctx.self_program_id,
);
Ok(spel_framework::SpelOutput::execute(post_states, chained_calls)
.with_timestamp_validity_window(..deadline))

View File

@ -1,12 +1,12 @@
use std::num::NonZeroU128;
use amm_core::{
assert_supported_fee_tier, compute_liquidity_token_pda_seed, read_vault_fungible_balances,
PoolDefinition,
assert_supported_fee_tier, compute_config_pda, compute_liquidity_token_pda_seed,
read_vault_fungible_balances, AmmConfig, PoolDefinition,
};
use nssa_core::{
account::{AccountWithMetadata, Data},
program::{AccountPostState, ChainedCall},
program::{AccountPostState, ChainedCall, ProgramId},
};
#[expect(
@ -14,6 +14,7 @@ use nssa_core::{
reason = "instruction surface passes explicit pool, vault, and user accounts"
)]
pub fn add_liquidity(
config: AccountWithMetadata,
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
@ -24,7 +25,19 @@ pub fn add_liquidity(
min_amount_liquidity: NonZeroU128,
max_amount_to_add_token_a: u128,
max_amount_to_add_token_b: u128,
amm_program_id: ProgramId,
) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
// The Token Program is taken from the config account, not trusted from a caller-supplied
// holding. Validating the config PDA is also the Program's initialization gate.
assert_eq!(
config.account_id,
compute_config_pda(amm_program_id),
"Add liquidity: AMM config Account ID does not match PDA"
);
let token_program_id = AmmConfig::try_from(&config.account.data)
.expect("Add liquidity: AMM Program must be initialized before use")
.token_program_id;
// 1. Fetch Pool state
let pool_def_data = PoolDefinition::try_from(&pool.account.data)
.expect("Add liquidity: AMM Program expects valid Pool Definition Account");
@ -45,14 +58,21 @@ pub fn add_liquidity(
"Vault B was not provided"
);
let token_program_id = vault_a.account.program_owner;
assert_eq!(
vault_a.account.program_owner, token_program_id,
"Vault A must be owned by the configured Token Program"
);
assert_eq!(
vault_b.account.program_owner, token_program_id,
"Vault B must be owned by the configured Token Program"
);
assert_eq!(
user_holding_a.account.program_owner, token_program_id,
"User Token A holding must be owned by the vault's Token Program"
"User Token A holding must be owned by the configured Token Program"
);
assert_eq!(
user_holding_b.account.program_owner, token_program_id,
"User Token B holding must be owned by the vault's Token Program"
"User Token B holding must be owned by the configured Token Program"
);
assert!(
@ -187,6 +207,7 @@ pub fn add_liquidity(
let chained_calls = vec![call_token_lp, call_token_b, call_token_a];
let post_states = vec![
AccountPostState::new(config.account.clone()),
AccountPostState::new(pool_post),
AccountPostState::new(vault_a.account.clone()),
AccountPostState::new(vault_b.account.clone()),

View File

@ -0,0 +1,97 @@
use amm_core::{compute_config_pda, compute_config_pda_seed, AmmConfig};
use nssa_core::{
account::{Account, AccountWithMetadata, Data},
program::{AccountPostState, Claim, ProgramId},
};
/// Initializes the AMM Program by creating its singleton configuration account.
///
/// The config account is a PDA derived from the constant `"CONFIG"` seed
/// (`compute_config_pda(amm_program_id)`) and stores `token_program_id`, the Token Program the
/// AMM issues every chained call to. Its existence is the Program's "initialized" flag: the
/// chained-call instructions read the Token Program ID from it and reject calls until it exists.
///
/// # Panics
/// Panics if:
/// - `config.account_id` does not match `compute_config_pda(amm_program_id)`.
/// - `config.account` is not the default (the Program is already initialized).
pub fn initialize(
config: AccountWithMetadata,
token_program_id: ProgramId,
amm_program_id: ProgramId,
) -> Vec<AccountPostState> {
assert_eq!(
config.account_id,
compute_config_pda(amm_program_id),
"Initialize: AMM config Account ID does not match PDA"
);
assert_eq!(
config.account,
Account::default(),
"Initialize: AMM config account must be uninitialized"
);
let mut config_post = config.account.clone();
config_post.data = Data::from(&AmmConfig { token_program_id });
vec![AccountPostState::new_claimed(
config_post,
Claim::Pda(compute_config_pda_seed()),
)]
}
#[cfg(test)]
mod tests {
use amm_core::compute_config_pda;
use nssa_core::account::{AccountId, Nonce};
use super::*;
const AMM_PROGRAM_ID: ProgramId = [42; 8];
const TOKEN_PROGRAM_ID: ProgramId = [15; 8];
fn config_uninit() -> AccountWithMetadata {
AccountWithMetadata {
account: Account::default(),
is_authorized: false,
account_id: compute_config_pda(AMM_PROGRAM_ID),
}
}
#[test]
fn returns_single_pda_claimed_post_state() {
let post_states = initialize(config_uninit(), TOKEN_PROGRAM_ID, AMM_PROGRAM_ID);
assert_eq!(post_states.len(), 1);
assert_eq!(
post_states[0].required_claim(),
Some(Claim::Pda(compute_config_pda_seed()))
);
}
#[test]
fn stores_token_program_id() {
let post_states = initialize(config_uninit(), TOKEN_PROGRAM_ID, AMM_PROGRAM_ID);
let config = AmmConfig::try_from(&post_states[0].account().data)
.expect("post state must contain a valid AmmConfig");
assert_eq!(config.token_program_id, TOKEN_PROGRAM_ID);
}
#[test]
#[should_panic(expected = "AMM config Account ID does not match PDA")]
fn wrong_config_account_id_panics() {
let mut wrong = config_uninit();
wrong.account_id = AccountId::new([0; 32]);
initialize(wrong, TOKEN_PROGRAM_ID, AMM_PROGRAM_ID);
}
#[test]
#[should_panic(expected = "AMM config account must be uninitialized")]
fn already_initialized_config_panics() {
let mut initialized = config_uninit();
initialized.account.data = Data::from(&AmmConfig {
token_program_id: TOKEN_PROGRAM_ID,
});
initialized.account.nonce = Nonce(0);
initialize(initialized, TOKEN_PROGRAM_ID, AMM_PROGRAM_ID);
}
}

View File

@ -3,6 +3,7 @@
pub use amm_core as core;
pub mod add;
pub mod initialize;
pub mod new_definition;
pub mod remove;
pub mod swap;

View File

@ -1,10 +1,10 @@
use std::num::NonZeroU128;
use amm_core::{
assert_supported_fee_tier, compute_liquidity_token_pda, compute_liquidity_token_pda_seed,
compute_lp_lock_holding_pda, compute_lp_lock_holding_pda_seed, compute_pool_pda,
compute_pool_pda_seed, compute_vault_pda, compute_vault_pda_seed, PoolDefinition,
MINIMUM_LIQUIDITY,
assert_supported_fee_tier, compute_config_pda, compute_liquidity_token_pda,
compute_liquidity_token_pda_seed, compute_lp_lock_holding_pda,
compute_lp_lock_holding_pda_seed, compute_pool_pda, compute_pool_pda_seed, compute_vault_pda,
compute_vault_pda_seed, AmmConfig, PoolDefinition, MINIMUM_LIQUIDITY,
};
use nssa_core::{
account::{Account, AccountWithMetadata, Data},
@ -17,6 +17,7 @@ use token_core::TokenDefinition;
reason = "instruction surface passes explicit pool, vault, mint, lock, and user accounts"
)]
pub fn new_definition(
config: AccountWithMetadata,
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
@ -37,12 +38,24 @@ pub fn new_definition(
.expect("New definition: AMM Program expects valid Token Holding account for Token B")
.definition_id();
let token_program = user_holding_a.account.program_owner;
// both instances of the same token program
// The Token Program is taken from the config account, not trusted from a caller-supplied
// holding. Validating the config PDA is also the Program's initialization gate.
assert_eq!(
user_holding_b.account.program_owner, token_program,
"User Token holdings must use the same Token Program"
config.account_id,
compute_config_pda(amm_program_id),
"New definition: AMM config Account ID does not match PDA"
);
let token_program_id = AmmConfig::try_from(&config.account.data)
.expect("New definition: AMM Program must be initialized before use")
.token_program_id;
assert_eq!(
user_holding_a.account.program_owner, token_program_id,
"User Token A holding must be owned by the configured Token Program"
);
assert_eq!(
user_holding_b.account.program_owner, token_program_id,
"User Token B holding must be owned by the configured Token Program"
);
// Verify token_a and token_b are different
assert!(
@ -124,8 +137,6 @@ pub fn new_definition(
)),
);
let token_program_id = user_holding_a.account.program_owner;
// Chain call for Token A (user_holding_a -> Vault_A)
let mut vault_a_authorized = vault_a.clone();
vault_a_authorized.is_authorized = true;
@ -199,6 +210,7 @@ pub fn new_definition(
];
let post_states = vec![
AccountPostState::new(config.account.clone()),
pool_post.clone(),
AccountPostState::new(vault_a.account.clone()),
AccountPostState::new(vault_b.account.clone()),

View File

@ -1,12 +1,12 @@
use std::num::NonZeroU128;
use amm_core::{
assert_supported_fee_tier, compute_liquidity_token_pda_seed, compute_vault_pda_seed,
PoolDefinition, MINIMUM_LIQUIDITY,
assert_supported_fee_tier, compute_config_pda, compute_liquidity_token_pda_seed,
compute_vault_pda_seed, AmmConfig, PoolDefinition, MINIMUM_LIQUIDITY,
};
use nssa_core::{
account::{AccountWithMetadata, Data},
program::{AccountPostState, ChainedCall},
program::{AccountPostState, ChainedCall, ProgramId},
};
#[expect(
@ -14,6 +14,7 @@ use nssa_core::{
reason = "instruction surface passes explicit pool, vault, and user accounts"
)]
pub fn remove_liquidity(
config: AccountWithMetadata,
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
@ -24,9 +25,21 @@ pub fn remove_liquidity(
remove_liquidity_amount: NonZeroU128,
min_amount_to_remove_token_a: u128,
min_amount_to_remove_token_b: u128,
amm_program_id: ProgramId,
) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
let remove_liquidity_amount: u128 = remove_liquidity_amount.into();
// The Token Program is taken from the config account, not trusted from a caller-supplied
// holding. Validating the config PDA is also the Program's initialization gate.
assert_eq!(
config.account_id,
compute_config_pda(amm_program_id),
"Remove liquidity: AMM config Account ID does not match PDA"
);
let token_program_id = AmmConfig::try_from(&config.account.data)
.expect("Remove liquidity: AMM Program must be initialized before use")
.token_program_id;
// 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");
@ -49,14 +62,21 @@ pub fn remove_liquidity(
"Vault B was not provided"
);
let token_program_id = vault_a.account.program_owner;
assert_eq!(
vault_a.account.program_owner, token_program_id,
"Vault A must be owned by the configured Token Program"
);
assert_eq!(
vault_b.account.program_owner, token_program_id,
"Vault B must be owned by the configured Token Program"
);
assert_eq!(
user_holding_a.account.program_owner, token_program_id,
"User Token A holding must be owned by the vault's Token Program"
"User Token A holding must be owned by the configured Token Program"
);
assert_eq!(
user_holding_b.account.program_owner, token_program_id,
"User Token B holding must be owned by the vault's Token Program"
"User Token B holding must be owned by the configured Token Program"
);
// Vault addresses do not need to be checked with PDA
@ -204,6 +224,7 @@ pub fn remove_liquidity(
let chained_calls = vec![call_token_lp, call_token_b, call_token_a];
let post_states = vec![
AccountPostState::new(config.account.clone()),
AccountPostState::new(pool_post.clone()),
AccountPostState::new(vault_a.account.clone()),
AccountPostState::new(vault_b.account.clone()),

View File

@ -1,10 +1,11 @@
use amm_core::{
assert_supported_fee_tier, read_vault_fungible_balances, FEE_BPS_DENOMINATOR, MINIMUM_LIQUIDITY,
assert_supported_fee_tier, compute_config_pda, read_vault_fungible_balances, AmmConfig,
FEE_BPS_DENOMINATOR, MINIMUM_LIQUIDITY,
};
pub use amm_core::{compute_liquidity_token_pda_seed, compute_vault_pda_seed, PoolDefinition};
use nssa_core::{
account::{AccountId, AccountWithMetadata, Data},
program::{AccountPostState, ChainedCall},
program::{AccountPostState, ChainedCall, ProgramId},
};
/// Validates swap setup: checks pool liquidity is ready, vaults match, and reserves are sufficient.
@ -55,6 +56,7 @@ fn validate_swap_setup(
reason = "consistent with codebase style"
)]
fn create_swap_post_states(
config: AccountWithMetadata,
pool: AccountWithMetadata,
pool_def_data: PoolDefinition,
vault_a: AccountWithMetadata,
@ -86,6 +88,7 @@ fn create_swap_post_states(
pool_post.data = Data::from(&pool_post_definition);
vec![
AccountPostState::new(config.account),
AccountPostState::new(pool_post),
AccountPostState::new(vault_a.account),
AccountPostState::new(vault_b.account),
@ -100,6 +103,7 @@ fn create_swap_post_states(
)]
#[must_use]
pub fn swap_exact_input(
config: AccountWithMetadata,
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
@ -108,17 +112,35 @@ pub fn swap_exact_input(
swap_amount_in: u128,
min_amount_out: u128,
token_in_id: AccountId,
amm_program_id: ProgramId,
) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
let pool_def_data = validate_swap_setup(&pool, &vault_a, &vault_b);
let token_program_id = vault_a.account.program_owner;
// The Token Program is taken from the config account, not trusted from a caller-supplied
// account. Validating the config PDA is also the Program's initialization gate.
assert_eq!(
config.account_id,
compute_config_pda(amm_program_id),
"Swap exact input: AMM config Account ID does not match PDA"
);
let token_program_id = AmmConfig::try_from(&config.account.data)
.expect("Swap exact input: AMM Program must be initialized before use")
.token_program_id;
assert_eq!(
vault_a.account.program_owner, token_program_id,
"Vault A must be owned by the configured Token Program"
);
assert_eq!(
vault_b.account.program_owner, token_program_id,
"Vault B must be owned by the configured Token Program"
);
assert_eq!(
user_holding_a.account.program_owner, token_program_id,
"User Token A holding must be owned by the vault's Token Program"
"User Token A holding must be owned by the configured Token Program"
);
assert_eq!(
user_holding_b.account.program_owner, token_program_id,
"User Token B holding must be owned by the vault's Token Program"
"User Token B holding must be owned by the configured Token Program"
);
let (chained_calls, [deposit_a, withdraw_a], [deposit_b, withdraw_b]) =
@ -157,6 +179,7 @@ pub fn swap_exact_input(
};
let post_states = create_swap_post_states(
config,
pool,
pool_def_data,
vault_a,
@ -262,6 +285,7 @@ fn swap_logic(
)]
#[must_use]
pub fn swap_exact_output(
config: AccountWithMetadata,
pool: AccountWithMetadata,
vault_a: AccountWithMetadata,
vault_b: AccountWithMetadata,
@ -270,17 +294,35 @@ pub fn swap_exact_output(
exact_amount_out: u128,
max_amount_in: u128,
token_in_id: AccountId,
amm_program_id: ProgramId,
) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
let pool_def_data = validate_swap_setup(&pool, &vault_a, &vault_b);
let token_program_id = vault_a.account.program_owner;
// The Token Program is taken from the config account, not trusted from a caller-supplied
// account. Validating the config PDA is also the Program's initialization gate.
assert_eq!(
config.account_id,
compute_config_pda(amm_program_id),
"Swap exact output: AMM config Account ID does not match PDA"
);
let token_program_id = AmmConfig::try_from(&config.account.data)
.expect("Swap exact output: AMM Program must be initialized before use")
.token_program_id;
assert_eq!(
vault_a.account.program_owner, token_program_id,
"Vault A must be owned by the configured Token Program"
);
assert_eq!(
vault_b.account.program_owner, token_program_id,
"Vault B must be owned by the configured Token Program"
);
assert_eq!(
user_holding_a.account.program_owner, token_program_id,
"User Token A holding must be owned by the vault's Token Program"
"User Token A holding must be owned by the configured Token Program"
);
assert_eq!(
user_holding_b.account.program_owner, token_program_id,
"User Token B holding must be owned by the vault's Token Program"
"User Token B holding must be owned by the configured Token Program"
);
let (chained_calls, [deposit_a, withdraw_a], [deposit_b, withdraw_b]) =
@ -319,6 +361,7 @@ pub fn swap_exact_output(
};
let post_states = create_swap_post_states(
config,
pool,
pool_def_data,
vault_a,

File diff suppressed because it is too large Load Diff

View File

@ -43,6 +43,10 @@ impl Ids {
amm_methods::AMM_ID
}
fn config() -> AccountId {
amm_core::compute_config_pda(Self::amm_program())
}
fn token_a_definition() -> AccountId {
AccountId::new([3; 32])
}
@ -283,6 +287,17 @@ impl Balances {
}
impl Accounts {
fn config() -> Account {
Account {
program_owner: Ids::amm_program(),
balance: 0_u128,
data: Data::from(&amm_core::AmmConfig {
token_program_id: Ids::token_program(),
}),
nonce: Nonce(0),
}
}
fn user_a_holding() -> Account {
Account {
program_owner: Ids::token_program(),
@ -914,6 +929,7 @@ fn deploy_programs(state: &mut V03State) {
fn state_for_amm_tests() -> V03State {
let mut state = V03State::new_with_genesis_accounts(&[], vec![], 0);
deploy_programs(&mut state);
state.force_insert_account(Ids::config(), Accounts::config());
state.force_insert_account(Ids::pool_definition(), Accounts::pool_definition_init());
state.force_insert_account(
Ids::token_a_definition(),
@ -938,6 +954,7 @@ fn state_for_amm_tests() -> V03State {
fn state_for_amm_tests_with_new_def() -> V03State {
let mut state = V03State::new_with_genesis_accounts(&[], vec![], 0);
deploy_programs(&mut state);
state.force_insert_account(Ids::config(), Accounts::config());
state.force_insert_account(
Ids::token_a_definition(),
Accounts::token_a_definition_account(),
@ -977,6 +994,7 @@ fn try_execute_new_definition(
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1032,6 +1050,7 @@ fn execute_swap_a_to_b(state: &mut V03State, swap_amount_in: u128, min_amount_ou
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1061,6 +1080,7 @@ fn execute_swap_b_to_a(state: &mut V03State, swap_amount_in: u128, min_amount_ou
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1095,6 +1115,7 @@ fn execute_add_liquidity(
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1135,6 +1156,7 @@ fn execute_remove_liquidity(
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1154,6 +1176,26 @@ fn execute_remove_liquidity(
state.transition_from_public_transaction(&tx, 0, 0).unwrap();
}
#[cfg(test)]
fn execute_initialize(state: &mut V03State) {
let instruction = amm_core::Instruction::Initialize {
token_program_id: Ids::token_program(),
};
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![Ids::config()],
vec![],
instruction,
)
.unwrap();
let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
let tx = PublicTransaction::new(message, witness_set);
state.transition_from_public_transaction(&tx, 0, 0).unwrap();
}
fn fungible_balance(account: &Account) -> u128 {
let holding = TokenHolding::try_from(&account.data).expect("expected token holding");
let TokenHolding::Fungible {
@ -1185,6 +1227,26 @@ fn fungible_total_supply(account: &Account) -> u128 {
total_supply
}
#[test]
fn amm_initialize_creates_config_account() {
let mut state = V03State::new_with_genesis_accounts(&[], vec![], 0);
deploy_programs(&mut state);
// Before initialization the config PDA does not exist.
assert_eq!(state.get_account_by_id(Ids::config()), Account::default());
execute_initialize(&mut state);
// Initialization creates the config PDA, owned by the AMM program.
let config_account = state.get_account_by_id(Ids::config());
assert_eq!(config_account, Accounts::config());
// Explicitly assert the stored Token Program ID round-trips from the instruction argument.
let config = amm_core::AmmConfig::try_from(&config_account.data)
.expect("config account must hold a valid AmmConfig");
assert_eq!(config.token_program_id, Ids::token_program());
}
#[test]
fn amm_remove_liquidity() {
let mut state = state_for_amm_tests();
@ -1199,6 +1261,7 @@ fn amm_remove_liquidity() {
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1262,6 +1325,7 @@ fn amm_remove_liquidity_insufficient_user_lp_fails() {
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1487,6 +1551,7 @@ fn amm_add_liquidity() {
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1550,6 +1615,7 @@ fn amm_swap_b_to_a() {
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1602,6 +1668,7 @@ fn amm_swap_a_to_b() {
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1708,6 +1775,7 @@ fn amm_swap_rejects_expired_deadline() {
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1744,6 +1812,7 @@ fn amm_swap_exact_output_rejects_expired_deadline() {
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1780,6 +1849,7 @@ fn amm_add_liquidity_rejects_expired_deadline() {
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1822,6 +1892,7 @@ fn amm_remove_liquidity_rejects_expired_deadline() {
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
@ -1860,6 +1931,7 @@ fn amm_new_definition_rejects_expired_deadline() {
let message = public_transaction::Message::try_new(
Ids::amm_program(),
vec![
Ids::config(),
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),