lez-programs/programs/amm/core/src/lib.rs

//! This crate contains core data structures and utilities for the AMM Program.

use borsh::{BorshDeserialize, BorshSerialize};
use nssa_core::{
    account::{AccountId, AccountWithMetadata, Data},
    program::{PdaSeed, ProgramId},
};
use serde::{Deserialize, Serialize};
use spel_framework_macros::account_type;

// These stable seed bytes are part of the PDA derivation scheme and must stay unchanged for
// compatibility.
const LIQUIDITY_TOKEN_PDA_SEED: [u8; 32] = [0; 32];
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
    /// in the LP-lock holding PDA; the caller receives `initial_lp - MINIMUM_LIQUIDITY`.
    ///
    /// Required accounts:
    /// - AMM Pool
    /// - Vault Holding Account for Token A
    /// - Vault Holding Account for Token B
    /// - Pool Liquidity Token Definition
    /// - LP Lock Holding Account, derived as `compute_lp_lock_holding_pda(self_program_id,
    ///   pool.account_id)`
    /// - User Holding Account for Token A (authorized)
    /// - User Holding Account for Token B (authorized)
    /// - User Holding Account for Pool Liquidity (authorized when uninitialized)
    NewDefinition {
        token_a_amount: u128,
        token_b_amount: u128,
        fees: u128,
        /// Unix timestamp (milliseconds) after which this transaction is invalid.
        deadline: u64,
    },

    /// 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,
        /// Unix timestamp (milliseconds) after which this transaction is invalid.
        deadline: u64,
    },

    /// 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,
        /// Unix timestamp (milliseconds) after which this transaction is invalid.
        deadline: u64,
    },

    /// 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 is authorized.
    SwapExactInput {
        swap_amount_in: u128,
        min_amount_out: u128,
        token_definition_id_in: AccountId,
        /// Unix timestamp (milliseconds) after which this transaction is invalid.
        deadline: u64,
    },

    /// Swap tokens specifying the exact desired output amount,
    /// 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 is authorized.
    SwapExactOutput {
        exact_amount_out: u128,
        max_amount_in: u128,
        token_definition_id_in: AccountId,
        /// Unix timestamp (milliseconds) after which this transaction is invalid.
        deadline: u64,
    },

    /// Sync pool reserves with current vault balances.
    ///
    /// Required accounts:
    /// - AMM Pool (initialized, with LP supply at or above minimum liquidity)
    /// - Vault Holding Account for Token A (initialized)
    /// - Vault Holding Account for Token B (initialized)
    SyncReserves,
}

pub const MINIMUM_LIQUIDITY: u128 = 1_000;

#[account_type]
#[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,
    /// Total LP supply tracked by the pool. After initialization it includes the permanently
    /// locked `MINIMUM_LIQUIDITY`; a zero supply means the pool is uninitialized
    pub liquidity_pool_supply: u128,
    pub reserve_a: u128,
    pub reserve_b: u128,
    /// Fee tier in basis points.
    pub fees: u128,
}

pub const FEE_BPS_DENOMINATOR: u128 = 10_000;
pub const FEE_TIER_BPS_1: u128 = 1;
pub const FEE_TIER_BPS_5: u128 = 5;
pub const FEE_TIER_BPS_30: u128 = 30;
pub const FEE_TIER_BPS_100: u128 = 100;

pub fn is_supported_fee_tier(fees: u128) -> bool {
    matches!(
        fees,
        FEE_TIER_BPS_1 | FEE_TIER_BPS_5 | FEE_TIER_BPS_30 | FEE_TIER_BPS_100
    )
}

pub fn assert_supported_fee_tier(fees: u128) {
    assert!(
        is_supported_fee_tier(fees),
        "Fee tier must be one of 1, 5, 30, or 100 basis points"
    );
}

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

/// 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,
    definition_token_b_id: AccountId,
) -> AccountId {
    AccountId::for_public_pda(
        &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];
    let (token_1_bytes, token_2_bytes) = bytes.split_at_mut(32);
    token_1_bytes.copy_from_slice(&token_1.to_bytes());
    token_2_bytes.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::for_public_pda(
        &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];
    let (pool_bytes, definition_bytes) = bytes.split_at_mut(32);
    pool_bytes.copy_from_slice(&pool_id.to_bytes());
    definition_bytes.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::for_public_pda(&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];
    let (pool_bytes, seed_bytes) = bytes.split_at_mut(32);
    pool_bytes.copy_from_slice(&pool_id.to_bytes());
    seed_bytes.copy_from_slice(&LIQUIDITY_TOKEN_PDA_SEED);

    PdaSeed::new(
        Impl::hash_bytes(&bytes)
            .as_bytes()
            .try_into()
            .expect("Hash output must be exactly 32 bytes long"),
    )
}

pub fn compute_lp_lock_holding_pda(amm_program_id: ProgramId, pool_id: AccountId) -> AccountId {
    AccountId::for_public_pda(&amm_program_id, &compute_lp_lock_holding_pda_seed(pool_id))
}

pub fn compute_lp_lock_holding_pda_seed(pool_id: AccountId) -> PdaSeed {
    use risc0_zkvm::sha::{Impl, Sha256};

    let mut bytes = [0; 64];
    let (pool_bytes, seed_bytes) = bytes.split_at_mut(32);
    pool_bytes.copy_from_slice(&pool_id.to_bytes());
    seed_bytes.copy_from_slice(&LP_LOCK_HOLDING_PDA_SEED);

    PdaSeed::new(
        Impl::hash_bytes(&bytes)
            .as_bytes()
            .try_into()
            .expect("Hash output must be exactly 32 bytes long"),
    )
}

fn read_fungible_holding(account: &AccountWithMetadata, context: &str) -> (AccountId, u128) {
    let token_holding = token_core::TokenHolding::try_from(&account.account.data)
        .unwrap_or_else(|_| panic!("{context}: AMM Program expects a valid Token Holding Account"));

    let token_core::TokenHolding::Fungible {
        definition_id,
        balance,
    } = token_holding
    else {
        panic!("{context}: AMM Program expects a valid Fungible Token Holding Account");
    };

    (definition_id, balance)
}

pub fn read_vault_fungible_balances(
    context: &str,
    vault_a: &AccountWithMetadata,
    vault_b: &AccountWithMetadata,
) -> (u128, u128) {
    let vault_a_context = format!("{context}: Vault A");
    let vault_b_context = format!("{context}: Vault B");
    let (_, vault_a_balance) = read_fungible_holding(vault_a, &vault_a_context);
    let (_, vault_b_balance) = read_fungible_holding(vault_b, &vault_b_context);

    (vault_a_balance, vault_b_balance)
}