begin refactoring

programs/amm/core/src/lib.rs

@@ -0,0 +1,330 @@
+//! This crate contains core data structures and utilities for the Token Program.
+
+use borsh::{BorshDeserialize, BorshSerialize};
+use nssa_core::account::{AccountId, Data};
+use serde::{Deserialize, Serialize};
+
+pub const CURRENT_VERSION: u8 = 1;
+
+
+// The AMM program has five functions (four directly accessible via instructions):
+// 1. New AMM definition. Arguments to this function are:
+//      * Seven accounts: [amm_pool, vault_holding_a, vault_holding_b, pool_lp, user_holding_a,
+//        user_holding_b, user_holding_lp]. For new AMM Pool: amm_pool, vault_holding_a,
+//        vault_holding_b, pool_lp and user_holding_lp are default accounts. amm_pool is a default
+//        account that will initiate the amm definition account values vault_holding_a is a token
+//        holding account for token a vault_holding_b is a token holding account for token b pool_lp
+//        is a token holding account for the pool's lp token user_holding_a is a token holding
+//        account for token a user_holding_b is a token holding account for token b user_holding_lp
+//        is a token holding account for lp token
+//      * PDA remark: Accounts amm_pool, vault_holding_a, vault_holding_b and pool_lp are PDA. The
+//        AccountId for these accounts must be computed using: amm_pool AccountId <-
+//        compute_pool_pda vault_holding_a, vault_holding_b <- compute_vault_pda pool_lp
+//        <-compute_liquidity_token_pda
+//      * Requires authorization: user_holding_a, user_holding_b
+//      * An instruction data of 65-bytes, indicating the initial amm reserves' balances and
+//        token_program_id with the following layout: [0x00 || array of balances (little-endian 16
+//        bytes) || AMM_PROGRAM_ID)]
+//      * Internally, calls compute_liquidity_token_pda_seed, compute_vault_pda_seed to authorize
+//        transfers.
+//      * Internally, calls compute_pool_da, compute_vault_pda and compute_vault_pda to check
+//        various AccountIds are correct.
+// 3. Add liquidity Arguments to this function are:
+//      * Seven accounts: [amm_pool, vault_holding_a, vault_holding_b, pool_lp, user_holding_a,
+//        user_holding_a, user_holding_lp].
+//      * Requires authorization: user_holding_a, user_holding_b
+//      * An instruction data byte string of length 49, amounts for minimum amount of liquidity from
+//        add (min_amount_lp),
+//      * max amount added for each token (max_amount_a and max_amount_b); indicate [0x02 || array
+//        of of balances (little-endian 16 bytes)].
+//      * Internally, calls compute_liquidity_token_pda_seed to compute liquidity pool PDA seed.
+// 4. Remove liquidity
+//      * Seven accounts: [amm_pool, vault_holding_a, vault_holding_b, pool_lp, user_holding_a,
+//        user_holding_a, user_holding_lp].
+//      * Requires authorization: user_holding_lp
+//      * An instruction data byte string of length 49, amounts for minimum amount of liquidity to
+//        redeem (balance_lp),
+//      * minimum balance of each token to remove (min_amount_a and min_amount_b); indicate [0x03 ||
+//        array of balances (little-endian 16 bytes)].
+//      * Internally, calls compute_vault_pda_seed to compute vault_a and vault_b's PDA seed.
+
+
+
+/// AMM Program Instruction.
+#[derive(Serialize, Deserialize)]
+pub enum Instruction {
+
+    /// Create a new fungible token definition without metadata.
+    ///
+    /// Required accounts:
+    /// - Token Definition account (uninitialized),
+    /// - Token Holding account (uninitialized).
+    NewDefinition { name: String, total_supply: u128 },
+
+    /// Create a new fungible or non-fungible token definition with metadata.
+    ///
+    /// Required accounts:
+    /// - Token Definition account (uninitialized),
+    /// - Token Holding account (uninitialized),
+    /// - Token Metadata account (uninitialized).
+    NewDefinitionWithMetadata {
+        new_definition: NewTokenDefinition,
+        /// Boxed to avoid large enum variant size
+        metadata: Box<NewTokenMetadata>,
+    },
+
+
+    
+    /// Initialize a token holding account for a given token definition.
+    ///
+    /// Required accounts:
+    /// - Token Definition account (initialized),
+    /// - Token Holding account (uninitialized),
+    InitializeAccount,
+
+
+// 2. Swap assets Arguments to this function are:
+//      * Five accounts: [amm_pool, vault_holding_a, vault_holding_b, user_holding_a,
+//        user_holding_b].
+//      * Requires authorization: user holding account associated to TOKEN_DEFINITION_ID (either
+//        user_holding_a or user_holding_b)
+//      * An instruction data byte string of length 65, indicating which token type to swap,
+//        quantity of tokens put into the swap (of type TOKEN_DEFINITION_ID) and min_amount_out.
+//        [0x01 || amount (little-endian 16 bytes) || TOKEN_DEFINITION_ID].
+//      * Internally, calls swap logic.
+//              * Four accounts: [user_deposit, vault_deposit, vault_withdraw, user_withdraw].
+//                user_deposit and vault_deposit define deposit transaction. vault_withdraw and
+//                user_withdraw define withdraw transaction.
+//              * deposit_amount is the amount for user_deposit -> vault_deposit transfer.
+//              * reserve_amounts is the pool's reserves; used to compute the withdraw amount.
+//              * Outputs the token transfers as a Vec<ChainedCall> and the withdraw amount.
+
+
+    ///TODO update description
+    /// Burn tokens from the holder's account.
+    ///
+    /// Required accounts:
+    /// - Token Definition account (initialized),
+    /// - Token Holding account (authorized).
+    Burn { amount_to_burn: u128 },
+
+
+/*
+fn add_liquidity(
+    pre_states: &[AccountWithMetadata],
+    balances: &[u128],
+) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
+    if pre_states.len() != 7 {
+        panic!("Invalid number of input accounts");
+    }
+
+    let pool = &pre_states[0];
+    let vault_a = &pre_states[1];
+    let vault_b = &pre_states[2];
+    let pool_definition_lp = &pre_states[3];
+    let user_holding_a = &pre_states[4];
+    let user_holding_b = &pre_states[5];
+    let user_holding_lp = &pre_states[6];
+
+        let min_amount_lp = balances[0];
+    let max_amount_a = balances[1];
+    let max_amount_b = balances[2];
+
+*/
+
+    ///TODO: update for add
+    /// Mint new tokens to the holder's account.
+    ///
+    /// Required accounts:
+    /// - Token Definition account (authorized),
+    /// - Token Holding account.
+    AddLiquidity { min_amount_liquidity: u128, max_amount_to_add_token_a: u128, max_amount_to_add_token_b: u128 },
+
+
+    /*
+    fn remove_liquidity(
+    pre_states: &[AccountWithMetadata],
+    amounts: &[u128],
+) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
+    if pre_states.len() != 7 {
+        panic!("Invalid number of input accounts");
+    }
+
+    let pool = &pre_states[0];
+    let vault_a = &pre_states[1];
+    let vault_b = &pre_states[2];
+    let pool_definition_lp = &pre_states[3];
+    let user_holding_a = &pre_states[4];
+    let user_holding_b = &pre_states[5];
+    let user_holding_lp = &pre_states[6];
+
+    if amounts.len() != 3 {
+        panic!("Invalid number of balances");
+    }
+
+    let amount_lp = amounts[0];
+    let amount_min_a = amounts[1];
+    let amount_min_b = amounts[2];
+    
+     */
+    //TODO types
+    RemoveLiquidity { remove_liquidity_amount: u128, min_amount_to_remove_token_a: u128, min_amount_to_remove_token_b: u128 }
+}
+
+/*
+
+#[derive(Serialize, Deserialize)]
+pub enum NewTokenDefinition {
+    Fungible { name: String, total_supply: u128 },
+    NonFungible { name: String, print_balance: u128 },
+}
+
+#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
+pub enum TokenDefinition {
+    Fungible {
+        name: String,
+        total_supply: u128,
+        metadata_id: Option<AccountId>,
+    },
+    NonFungible {
+        name: String,
+        metadata_id: AccountId,
+    },
+}
+
+impl TryFrom<&Data> for TokenDefinition {
+    type Error = std::io::Error;
+
+    fn try_from(data: &Data) -> Result<Self, Self::Error> {
+        TokenDefinition::try_from_slice(data.as_ref())
+    }
+}
+
+impl From<&TokenDefinition> for Data {
+    fn from(definition: &TokenDefinition) -> Self {
+        // Using size_of_val as size hint for Vec allocation
+        let mut data = Vec::with_capacity(std::mem::size_of_val(definition));
+
+        BorshSerialize::serialize(definition, &mut data)
+            .expect("Serialization to Vec should not fail");
+
+        Data::try_from(data).expect("Token definition encoded data should fit into Data")
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
+pub enum TokenHolding {
+    Fungible {
+        definition_id: AccountId,
+        balance: u128,
+    },
+    NftMaster {
+        definition_id: AccountId,
+        /// The amount of printed copies left - 1 (1 reserved for master copy itself).
+        print_balance: u128,
+    },
+    NftPrintedCopy {
+        definition_id: AccountId,
+        /// Whether nft is owned by the holder.
+        owned: bool,
+    },
+}
+
+impl TokenHolding {
+    pub fn zeroized_clone_from(other: &Self) -> Self {
+        match other {
+            TokenHolding::Fungible { definition_id, .. } => TokenHolding::Fungible {
+                definition_id: *definition_id,
+                balance: 0,
+            },
+            TokenHolding::NftMaster { definition_id, .. } => TokenHolding::NftMaster {
+                definition_id: *definition_id,
+                print_balance: 0,
+            },
+            TokenHolding::NftPrintedCopy { definition_id, .. } => TokenHolding::NftPrintedCopy {
+                definition_id: *definition_id,
+                owned: false,
+            },
+        }
+    }
+
+    pub fn zeroized_from_definition(
+        definition_id: AccountId,
+        definition: &TokenDefinition,
+    ) -> Self {
+        match definition {
+            TokenDefinition::Fungible { .. } => TokenHolding::Fungible {
+                definition_id,
+                balance: 0,
+            },
+            TokenDefinition::NonFungible { .. } => TokenHolding::NftPrintedCopy {
+                definition_id,
+                owned: false,
+            },
+        }
+    }
+
+    pub fn definition_id(&self) -> AccountId {
+        match self {
+            TokenHolding::Fungible { definition_id, .. } => *definition_id,
+            TokenHolding::NftMaster { definition_id, .. } => *definition_id,
+            TokenHolding::NftPrintedCopy { definition_id, .. } => *definition_id,
+        }
+    }
+}
+
+impl TryFrom<&Data> for TokenHolding {
+    type Error = std::io::Error;
+
+    fn try_from(data: &Data) -> Result<Self, Self::Error> {
+        TokenHolding::try_from_slice(data.as_ref())
+    }
+}
+
+impl From<&TokenHolding> for Data {
+    fn from(holding: &TokenHolding) -> Self {
+        // Using size_of_val as size hint for Vec allocation
+        let mut data = Vec::with_capacity(std::mem::size_of_val(holding));
+
+        BorshSerialize::serialize(holding, &mut data)
+            .expect("Serialization to Vec should not fail");
+
+        Data::try_from(data).expect("Token holding encoded data should fit into Data")
+    }
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct NewTokenMetadata {
+    pub uri: String,
+    pub creators: String,
+}
+
+#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
+pub struct TokenMetadata {
+    pub version: u8,
+    pub definition_id: AccountId,
+    pub uri: String,
+    pub creators: String,
+    /// Block id
+    pub primary_sale_date: u64,
+}
+
+impl TryFrom<&Data> for TokenMetadata {
+    type Error = std::io::Error;
+
+    fn try_from(data: &Data) -> Result<Self, Self::Error> {
+        TokenMetadata::try_from_slice(data.as_ref())
+    }
+}
+
+impl From<&TokenMetadata> for Data {
+    fn from(metadata: &TokenMetadata) -> Self {
+        // Using size_of_val as size hint for Vec allocation
+        let mut data = Vec::with_capacity(std::mem::size_of_val(metadata));
+
+        BorshSerialize::serialize(metadata, &mut data)
+            .expect("Serialization to Vec should not fail");
+
+        Data::try_from(data).expect("Token metadata encoded data should fit into Data")
+    }
+}*/

programs/amm/src/remove.rs

@@ -0,0 +1,162 @@
+fn remove_liquidity(
+    pre_states: &[AccountWithMetadata],
+    amounts: &[u128],
+) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
+    if pre_states.len() != 7 {
+        panic!("Invalid number of input accounts");
+    }
+
+    let pool = &pre_states[0];
+    let vault_a = &pre_states[1];
+    let vault_b = &pre_states[2];
+    let pool_definition_lp = &pre_states[3];
+    let user_holding_a = &pre_states[4];
+    let user_holding_b = &pre_states[5];
+    let user_holding_lp = &pre_states[6];
+
+    if amounts.len() != 3 {
+        panic!("Invalid number of balances");
+    }
+
+    let amount_lp = amounts[0];
+    let amount_min_a = amounts[1];
+    let amount_min_b = amounts[2];
+
+    // 1. Fetch Pool state
+    let pool_def_data = PoolDefinition::parse(&pool.account.data)
+        .expect("Remove liquidity: AMM Program expects a valid Pool Definition Account");
+
+    if !pool_def_data.active {
+        panic!("Pool is inactive");
+    }
+
+    if pool_def_data.liquidity_pool_id != pool_definition_lp.account_id {
+        panic!("LP definition mismatch");
+    }
+
+    if vault_a.account_id != pool_def_data.vault_a_id {
+        panic!("Vault A was not provided");
+    }
+
+    if vault_b.account_id != pool_def_data.vault_b_id {
+        panic!("Vault B was not provided");
+    }
+
+    // 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;
+
+    if amount_min_a == 0 || amount_min_b == 0 {
+        panic!("Minimum withdraw amount must be nonzero");
+    }
+
+    if amount_lp == 0 {
+        panic!("Liquidity 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"
+        );
+    };
+
+    if user_lp_balance > pool_def_data.liquidity_pool_supply
+        || user_holding_lp_data.definition_id() != pool_def_data.liquidity_pool_id
+    {
+        panic!("Invalid liquidity account provided");
+    }
+
+    let withdraw_amount_a =
+        (pool_def_data.reserve_a * amount_lp) / pool_def_data.liquidity_pool_supply;
+    let withdraw_amount_b =
+        (pool_def_data.reserve_b * amount_lp) / pool_def_data.liquidity_pool_supply;
+
+    // 3. Validate and slippage check
+    if withdraw_amount_a < amount_min_a {
+        panic!("Insufficient minimal withdraw amount (Token A) provided for liquidity amount");
+    }
+    if withdraw_amount_b < amount_min_b {
+        panic!("Insufficient minimal withdraw amount (Token B) provided for liquidity amount");
+    }
+
+    // 4. Calculate LP to reduce cap by
+    let delta_lp: u128 =
+        (pool_def_data.liquidity_pool_supply * amount_lp) / pool_def_data.liquidity_pool_supply;
+
+    let 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 = pool_post_definition.into_data();
+
+    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(pre_states[1].account.clone()),
+        AccountPostState::new(pre_states[2].account.clone()),
+        AccountPostState::new(pre_states[3].account.clone()),
+        AccountPostState::new(pre_states[4].account.clone()),
+        AccountPostState::new(pre_states[5].account.clone()),
+        AccountPostState::new(pre_states[6].account.clone()),
+    ];
+
+    (post_states, chained_calls)
+}