Handle Token-2022 transfer fees in Rust quote SDK

rust-sdk/Cargo.toml

@@ -8,7 +8,8 @@ authors = ["minh <minh@raccoons.dev>"]
 [dependencies]
 anyhow = "1.0.71"
 cp-amm = { path = "../programs/cp-amm" }
-ruint = "1.3.0"
+spl-token-2022 = "6.0.0"
 
 [dev-dependencies]
-bytemuck = { workspace = true}
+bytemuck = { workspace = true}
+spl-pod = "0.5.1"

rust-sdk/src/quote.rs

@@ -0,0 +1,306 @@
+use anyhow::{anyhow, ensure, Context, Result};
+use cp_amm::{
+    params::swap::TradeDirection,
+    state::{fee::FeeMode, Pool, SwapResult},
+    ActivationType,
+};
+use spl_token_2022::extension::{
+    transfer_fee::{TransferFee, TransferFeeConfig},
+    BaseStateWithExtensions, StateWithExtensions,
+};
+use spl_token_2022::state::Mint;
+
+/// Off-chain quote for a potential swap.
+#[derive(Debug, PartialEq)]
+pub struct Quote {
+    /// Resulting pool math prior to Token-2022 transfer fees being applied.
+    pub swap_result: SwapResult,
+    /// Amount that effectively reaches the pool after deducting input transfer fees.
+    pub effective_amount_in: u64,
+    /// Portion of the provided amount collected as input transfer fee.
+    pub input_transfer_fee: u64,
+    /// Amount the taker receives after output transfer fees are withheld.
+    pub effective_amount_out: u64,
+    /// Portion of the pool's output withheld as transfer fee.
+    pub output_transfer_fee: u64,
+}
+
+/// Transfer-fee metadata extracted for each pool token mint.
+#[derive(Clone, Debug, PartialEq)]
+pub struct MintTransferFees {
+    token_a: Option<TransferFee>,
+    token_b: Option<TransferFee>,
+    token_a_known: bool,
+    token_b_known: bool,
+}
+
+impl Default for MintTransferFees {
+    fn default() -> Self {
+        Self {
+            token_a: None,
+            token_b: None,
+            token_a_known: false,
+            token_b_known: false,
+        }
+    }
+}
+
+impl MintTransferFees {
+    /// Create a transfer-fee bundle when both token mints have already been inspected.
+    pub fn new(token_a: Option<TransferFee>, token_b: Option<TransferFee>) -> Self {
+        Self {
+            token_a,
+            token_b,
+            token_a_known: true,
+            token_b_known: true,
+        }
+    }
+
+    pub fn input_fee(&self, a_to_b: bool) -> Option<&TransferFee> {
+        if a_to_b {
+            self.token_a.as_ref()
+        } else {
+            self.token_b.as_ref()
+        }
+    }
+
+    pub fn output_fee(&self, a_to_b: bool) -> Option<&TransferFee> {
+        if a_to_b {
+            self.token_b.as_ref()
+        } else {
+            self.token_a.as_ref()
+        }
+    }
+
+    fn input_known(&self, a_to_b: bool) -> bool {
+        if a_to_b {
+            self.token_a_known
+        } else {
+            self.token_b_known
+        }
+    }
+
+    fn output_known(&self, a_to_b: bool) -> bool {
+        if a_to_b {
+            self.token_b_known
+        } else {
+            self.token_a_known
+        }
+    }
+
+    /// Construct transfer-fee metadata by parsing on-chain mint accounts.
+    pub fn from_pool_mints(
+        pool: &Pool,
+        current_epoch: u64,
+        token_a_mint_account: Option<&[u8]>,
+        token_b_mint_account: Option<&[u8]>,
+    ) -> Result<Self> {
+        let token_a_flag = TokenProgramFlag::try_from(pool.token_a_flag)?;
+        let token_b_flag = TokenProgramFlag::try_from(pool.token_b_flag)?;
+
+        let (token_a, token_a_known) = match token_a_flag {
+            TokenProgramFlag::TokenProgram => (None, true),
+            TokenProgramFlag::TokenProgram2022 => {
+                let mint_data = token_a_mint_account
+                    .context("token A mint account data required for Token-2022 pools")?;
+                (load_transfer_fee(mint_data, current_epoch)?, true)
+            }
+        };
+
+        let (token_b, token_b_known) = match token_b_flag {
+            TokenProgramFlag::TokenProgram => (None, true),
+            TokenProgramFlag::TokenProgram2022 => {
+                let mint_data = token_b_mint_account
+                    .context("token B mint account data required for Token-2022 pools")?;
+                (load_transfer_fee(mint_data, current_epoch)?, true)
+            }
+        };
+
+        Ok(Self {
+            token_a,
+            token_b,
+            token_a_known,
+            token_b_known,
+        })
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum TokenProgramFlag {
+    TokenProgram,
+    TokenProgram2022,
+}
+
+impl TryFrom<u8> for TokenProgramFlag {
+    type Error = anyhow::Error;
+
+    fn try_from(value: u8) -> Result<Self> {
+        match value {
+            0 => Ok(Self::TokenProgram),
+            1 => Ok(Self::TokenProgram2022),
+            _ => Err(anyhow!("unknown token program flag: {value}")),
+        }
+    }
+}
+
+/// Compute an off-chain quote, mirroring on-chain transfer-fee behaviour.
+///
+/// The `transfer_fees` argument should be populated with [`MintTransferFees::from_pool_mints`]
+/// or another source that reflects the current epoch's Token-2022 transfer rules. Quotes for
+/// Token-2022 pools error out when the relevant mint data has not been provided, preventing
+/// callers from accidentally ignoring the additional fees enforced on-chain.
+pub fn get_quote(
+    pool: &Pool,
+    current_timestamp: u64,
+    current_slot: u64,
+    actual_amount_in: u64,
+    a_to_b: bool,
+    has_referral: bool,
+    transfer_fees: &MintTransferFees,
+) -> Result<Quote> {
+    ensure!(actual_amount_in > 0, "amount is zero");
+
+    let activation_type =
+        ActivationType::try_from(pool.activation_type).context("invalid activation type")?;
+
+    let current_point = match activation_type {
+        ActivationType::Slot => current_slot,
+        ActivationType::Timestamp => current_timestamp,
+    };
+
+    ensure!(
+        pool.pool_status == 0 && pool.activation_point <= current_point,
+        "Swap is disabled",
+    );
+
+    let trade_direction = if a_to_b {
+        TradeDirection::AtoB
+    } else {
+        TradeDirection::BtoA
+    };
+
+    let (input_fee, output_fee) = transfer_fees_for_direction(pool, transfer_fees, trade_direction)?;
+
+    let TransferFeeAmount {
+        amount: effective_amount_in,
+        transfer_fee: input_transfer_fee,
+    } = calculate_transfer_fee_excluded_amount(input_fee, actual_amount_in)?;
+
+    ensure!(effective_amount_in > 0, "amount is zero after transfer fees");
+
+    let fee_mode = &FeeMode::get_fee_mode(pool.collect_fee_mode, trade_direction, has_referral)?;
+
+    let swap_result = pool.get_swap_result(
+        effective_amount_in,
+        fee_mode,
+        trade_direction,
+        current_point,
+    )?;
+
+    let TransferFeeAmount {
+        amount: effective_amount_out,
+        transfer_fee: output_transfer_fee,
+    } = calculate_transfer_fee_excluded_amount(output_fee, swap_result.output_amount)?;
+
+    Ok(Quote {
+        swap_result,
+        effective_amount_in,
+        input_transfer_fee,
+        effective_amount_out,
+        output_transfer_fee,
+    })
+}
+
+fn transfer_fees_for_direction<'a>(
+    pool: &Pool,
+    transfer_fees: &'a MintTransferFees,
+    trade_direction: TradeDirection,
+) -> Result<(Option<&'a TransferFee>, Option<&'a TransferFee>)> {
+    match trade_direction {
+        TradeDirection::AtoB => Ok((
+            require_transfer_fee(
+                TokenProgramFlag::try_from(pool.token_a_flag)?,
+                transfer_fees.input_known(true),
+                transfer_fees.input_fee(true),
+                "token A",
+            )?,
+            require_transfer_fee(
+                TokenProgramFlag::try_from(pool.token_b_flag)?,
+                transfer_fees.output_known(true),
+                transfer_fees.output_fee(true),
+                "token B",
+            )?,
+        )),
+        TradeDirection::BtoA => Ok((
+            require_transfer_fee(
+                TokenProgramFlag::try_from(pool.token_b_flag)?,
+                transfer_fees.input_known(false),
+                transfer_fees.input_fee(false),
+                "token B",
+            )?,
+            require_transfer_fee(
+                TokenProgramFlag::try_from(pool.token_a_flag)?,
+                transfer_fees.output_known(false),
+                transfer_fees.output_fee(false),
+                "token A",
+            )?,
+        )),
+    }
+}
+
+fn require_transfer_fee<'a>(
+    flag: TokenProgramFlag,
+    is_known: bool,
+    transfer_fee: Option<&'a TransferFee>,
+    label: &str,
+) -> Result<Option<&'a TransferFee>> {
+    match flag {
+        TokenProgramFlag::TokenProgram => Ok(None),
+        TokenProgramFlag::TokenProgram2022 => {
+            ensure!(
+                is_known,
+                "missing transfer-fee data for {label} Token-2022 mint"
+            );
+            Ok(transfer_fee)
+        }
+    }
+}
+
+#[derive(Debug, PartialEq)]
+struct TransferFeeAmount {
+    amount: u64,
+    transfer_fee: u64,
+}
+
+fn calculate_transfer_fee_excluded_amount(
+    transfer_fee: Option<&TransferFee>,
+    transfer_fee_included_amount: u64,
+) -> Result<TransferFeeAmount> {
+    if let Some(transfer_fee) = transfer_fee {
+        let transfer_fee_amount = transfer_fee
+            .calculate_fee(transfer_fee_included_amount)
+            .ok_or_else(|| anyhow!("transfer fee calculation overflow"))?;
+        let transfer_fee_excluded_amount = transfer_fee_included_amount
+            .checked_sub(transfer_fee_amount)
+            .ok_or_else(|| anyhow!("transfer fee exceeds provided amount"))?;
+        Ok(TransferFeeAmount {
+            amount: transfer_fee_excluded_amount,
+            transfer_fee: transfer_fee_amount,
+        })
+    } else {
+        Ok(TransferFeeAmount {
+            amount: transfer_fee_included_amount,
+            transfer_fee: 0,
+        })
+    }
+}
+
+fn load_transfer_fee(mint_account_data: &[u8], epoch: u64) -> Result<Option<TransferFee>> {
+    let mint = StateWithExtensions::<Mint>::unpack(mint_account_data)?;
+    if let Ok(config) = mint.get_extension::<TransferFeeConfig>() {
+        Ok(Some(config.get_epoch_fee(epoch).clone()))
+    } else {
+        Ok(None)
+    }
+}
+