spl_token_swap/
processor.rs

1//! Program state processor
2
3use crate::constraints::{SwapConstraints, SWAP_CONSTRAINTS};
4use crate::{
5    curve::{
6        base::SwapCurve,
7        calculator::{RoundDirection, TradeDirection},
8        fees::Fees,
9    },
10    error::SwapError,
11    instruction::{
12        DepositAllTokenTypes, DepositSingleTokenTypeExactAmountIn, Initialize, Swap,
13        SwapInstruction, WithdrawAllTokenTypes, WithdrawSingleTokenTypeExactAmountOut,
14    },
15    state::{SwapState, SwapV1, SwapVersion},
16};
17use num_traits::FromPrimitive;
18use solana_program::{
19    account_info::{next_account_info, AccountInfo},
20    clock::Clock,
21    decode_error::DecodeError,
22    entrypoint::ProgramResult,
23    instruction::Instruction,
24    msg,
25    program::invoke_signed,
26    program_error::{PrintProgramError, ProgramError},
27    program_option::COption,
28    pubkey::Pubkey,
29    sysvar::Sysvar,
30};
31use spl_token_2022::{
32    check_spl_token_program_account,
33    error::TokenError,
34    extension::{
35        mint_close_authority::MintCloseAuthority, transfer_fee::TransferFeeConfig,
36        BaseStateWithExtensions, StateWithExtensions,
37    },
38    state::{Account, Mint},
39};
40use std::{convert::TryInto, error::Error};
41
42/// Program state handler.
43pub struct Processor {}
44impl Processor {
45    /// Unpacks a spl_token `Account`.
46    pub fn unpack_token_account(
47        account_info: &AccountInfo,
48        token_program_id: &Pubkey,
49    ) -> Result<Account, SwapError> {
50        if account_info.owner != token_program_id
51            && check_spl_token_program_account(account_info.owner).is_err()
52        {
53            Err(SwapError::IncorrectTokenProgramId)
54        } else {
55            StateWithExtensions::<Account>::unpack(&account_info.data.borrow())
56                .map(|a| a.base)
57                .map_err(|_| SwapError::ExpectedAccount)
58        }
59    }
60
61    /// Unpacks a spl_token `Mint`.
62    pub fn unpack_mint(
63        account_info: &AccountInfo,
64        token_program_id: &Pubkey,
65    ) -> Result<Mint, SwapError> {
66        if account_info.owner != token_program_id
67            && check_spl_token_program_account(account_info.owner).is_err()
68        {
69            Err(SwapError::IncorrectTokenProgramId)
70        } else {
71            StateWithExtensions::<Mint>::unpack(&account_info.data.borrow())
72                .map(|m| m.base)
73                .map_err(|_| SwapError::ExpectedMint)
74        }
75    }
76
77    /// Unpacks a spl_token `Mint` with extension data
78    pub fn unpack_mint_with_extensions<'a>(
79        account_data: &'a [u8],
80        owner: &Pubkey,
81        token_program_id: &Pubkey,
82    ) -> Result<StateWithExtensions<'a, Mint>, SwapError> {
83        if owner != token_program_id && check_spl_token_program_account(owner).is_err() {
84            Err(SwapError::IncorrectTokenProgramId)
85        } else {
86            StateWithExtensions::<Mint>::unpack(account_data).map_err(|_| SwapError::ExpectedMint)
87        }
88    }
89
90    /// Calculates the authority id by generating a program address.
91    pub fn authority_id(
92        program_id: &Pubkey,
93        my_info: &Pubkey,
94        bump_seed: u8,
95    ) -> Result<Pubkey, SwapError> {
96        Pubkey::create_program_address(&[&my_info.to_bytes()[..32], &[bump_seed]], program_id)
97            .or(Err(SwapError::InvalidProgramAddress))
98    }
99
100    /// Issue a spl_token `Burn` instruction.
101    pub fn token_burn<'a>(
102        swap: &Pubkey,
103        token_program: AccountInfo<'a>,
104        burn_account: AccountInfo<'a>,
105        mint: AccountInfo<'a>,
106        authority: AccountInfo<'a>,
107        bump_seed: u8,
108        amount: u64,
109    ) -> Result<(), ProgramError> {
110        let swap_bytes = swap.to_bytes();
111        let authority_signature_seeds = [&swap_bytes[..32], &[bump_seed]];
112        let signers = &[&authority_signature_seeds[..]];
113
114        let ix = spl_token_2022::instruction::burn(
115            token_program.key,
116            burn_account.key,
117            mint.key,
118            authority.key,
119            &[],
120            amount,
121        )?;
122
123        invoke_signed_wrapper::<TokenError>(
124            &ix,
125            &[burn_account, mint, authority, token_program],
126            signers,
127        )
128    }
129
130    /// Issue a spl_token `MintTo` instruction.
131    pub fn token_mint_to<'a>(
132        swap: &Pubkey,
133        token_program: AccountInfo<'a>,
134        mint: AccountInfo<'a>,
135        destination: AccountInfo<'a>,
136        authority: AccountInfo<'a>,
137        bump_seed: u8,
138        amount: u64,
139    ) -> Result<(), ProgramError> {
140        let swap_bytes = swap.to_bytes();
141        let authority_signature_seeds = [&swap_bytes[..32], &[bump_seed]];
142        let signers = &[&authority_signature_seeds[..]];
143        let ix = spl_token_2022::instruction::mint_to(
144            token_program.key,
145            mint.key,
146            destination.key,
147            authority.key,
148            &[],
149            amount,
150        )?;
151
152        invoke_signed_wrapper::<TokenError>(
153            &ix,
154            &[mint, destination, authority, token_program],
155            signers,
156        )
157    }
158
159    /// Issue a spl_token `Transfer` instruction.
160    #[allow(clippy::too_many_arguments)]
161    pub fn token_transfer<'a>(
162        swap: &Pubkey,
163        token_program: AccountInfo<'a>,
164        source: AccountInfo<'a>,
165        mint: AccountInfo<'a>,
166        destination: AccountInfo<'a>,
167        authority: AccountInfo<'a>,
168        bump_seed: u8,
169        amount: u64,
170        decimals: u8,
171    ) -> Result<(), ProgramError> {
172        let swap_bytes = swap.to_bytes();
173        let authority_signature_seeds = [&swap_bytes[..32], &[bump_seed]];
174        let signers = &[&authority_signature_seeds[..]];
175        let ix = spl_token_2022::instruction::transfer_checked(
176            token_program.key,
177            source.key,
178            mint.key,
179            destination.key,
180            authority.key,
181            &[],
182            amount,
183            decimals,
184        )?;
185        invoke_signed_wrapper::<TokenError>(
186            &ix,
187            &[source, mint, destination, authority, token_program],
188            signers,
189        )
190    }
191
192    #[allow(clippy::too_many_arguments)]
193    fn check_accounts(
194        token_swap: &dyn SwapState,
195        program_id: &Pubkey,
196        swap_account_info: &AccountInfo,
197        authority_info: &AccountInfo,
198        token_a_info: &AccountInfo,
199        token_b_info: &AccountInfo,
200        pool_mint_info: &AccountInfo,
201        pool_token_program_info: &AccountInfo,
202        user_token_a_info: Option<&AccountInfo>,
203        user_token_b_info: Option<&AccountInfo>,
204        pool_fee_account_info: Option<&AccountInfo>,
205    ) -> ProgramResult {
206        if swap_account_info.owner != program_id {
207            return Err(ProgramError::IncorrectProgramId);
208        }
209        if *authority_info.key
210            != Self::authority_id(program_id, swap_account_info.key, token_swap.bump_seed())?
211        {
212            return Err(SwapError::InvalidProgramAddress.into());
213        }
214        if *token_a_info.key != *token_swap.token_a_account() {
215            return Err(SwapError::IncorrectSwapAccount.into());
216        }
217        if *token_b_info.key != *token_swap.token_b_account() {
218            return Err(SwapError::IncorrectSwapAccount.into());
219        }
220        if *pool_mint_info.key != *token_swap.pool_mint() {
221            return Err(SwapError::IncorrectPoolMint.into());
222        }
223        if *pool_token_program_info.key != *token_swap.token_program_id() {
224            return Err(SwapError::IncorrectTokenProgramId.into());
225        }
226        if let Some(user_token_a_info) = user_token_a_info {
227            if token_a_info.key == user_token_a_info.key {
228                return Err(SwapError::InvalidInput.into());
229            }
230        }
231        if let Some(user_token_b_info) = user_token_b_info {
232            if token_b_info.key == user_token_b_info.key {
233                return Err(SwapError::InvalidInput.into());
234            }
235        }
236        if let Some(pool_fee_account_info) = pool_fee_account_info {
237            if *pool_fee_account_info.key != *token_swap.pool_fee_account() {
238                return Err(SwapError::IncorrectFeeAccount.into());
239            }
240        }
241        Ok(())
242    }
243
244    /// Processes an [Initialize](enum.Instruction.html).
245    pub fn process_initialize(
246        program_id: &Pubkey,
247        fees: Fees,
248        swap_curve: SwapCurve,
249        accounts: &[AccountInfo],
250        swap_constraints: &Option<SwapConstraints>,
251    ) -> ProgramResult {
252        let account_info_iter = &mut accounts.iter();
253        let swap_info = next_account_info(account_info_iter)?;
254        let authority_info = next_account_info(account_info_iter)?;
255        let token_a_info = next_account_info(account_info_iter)?;
256        let token_b_info = next_account_info(account_info_iter)?;
257        let pool_mint_info = next_account_info(account_info_iter)?;
258        let fee_account_info = next_account_info(account_info_iter)?;
259        let destination_info = next_account_info(account_info_iter)?;
260        let pool_token_program_info = next_account_info(account_info_iter)?;
261
262        let token_program_id = *pool_token_program_info.key;
263        if SwapVersion::is_initialized(&swap_info.data.borrow()) {
264            return Err(SwapError::AlreadyInUse.into());
265        }
266
267        let (swap_authority, bump_seed) =
268            Pubkey::find_program_address(&[&swap_info.key.to_bytes()], program_id);
269        if *authority_info.key != swap_authority {
270            return Err(SwapError::InvalidProgramAddress.into());
271        }
272        let token_a = Self::unpack_token_account(token_a_info, &token_program_id)?;
273        let token_b = Self::unpack_token_account(token_b_info, &token_program_id)?;
274        let fee_account = Self::unpack_token_account(fee_account_info, &token_program_id)?;
275        let destination = Self::unpack_token_account(destination_info, &token_program_id)?;
276        let pool_mint = {
277            let pool_mint_data = pool_mint_info.data.borrow();
278            let pool_mint = Self::unpack_mint_with_extensions(
279                &pool_mint_data,
280                pool_mint_info.owner,
281                &token_program_id,
282            )?;
283            if let Ok(extension) = pool_mint.get_extension::<MintCloseAuthority>() {
284                let close_authority: Option<Pubkey> = extension.close_authority.into();
285                if close_authority.is_some() {
286                    return Err(SwapError::InvalidCloseAuthority.into());
287                }
288            }
289            pool_mint.base
290        };
291        if *authority_info.key != token_a.owner {
292            return Err(SwapError::InvalidOwner.into());
293        }
294        if *authority_info.key != token_b.owner {
295            return Err(SwapError::InvalidOwner.into());
296        }
297        if *authority_info.key == destination.owner {
298            return Err(SwapError::InvalidOutputOwner.into());
299        }
300        if *authority_info.key == fee_account.owner {
301            return Err(SwapError::InvalidOutputOwner.into());
302        }
303        if COption::Some(*authority_info.key) != pool_mint.mint_authority {
304            return Err(SwapError::InvalidOwner.into());
305        }
306
307        if token_a.mint == token_b.mint {
308            return Err(SwapError::RepeatedMint.into());
309        }
310        swap_curve
311            .calculator
312            .validate_supply(token_a.amount, token_b.amount)?;
313        if token_a.delegate.is_some() {
314            return Err(SwapError::InvalidDelegate.into());
315        }
316        if token_b.delegate.is_some() {
317            return Err(SwapError::InvalidDelegate.into());
318        }
319        if token_a.close_authority.is_some() {
320            return Err(SwapError::InvalidCloseAuthority.into());
321        }
322        if token_b.close_authority.is_some() {
323            return Err(SwapError::InvalidCloseAuthority.into());
324        }
325
326        if pool_mint.supply != 0 {
327            return Err(SwapError::InvalidSupply.into());
328        }
329        if pool_mint.freeze_authority.is_some() {
330            return Err(SwapError::InvalidFreezeAuthority.into());
331        }
332        if *pool_mint_info.key != fee_account.mint {
333            return Err(SwapError::IncorrectPoolMint.into());
334        }
335
336        if let Some(swap_constraints) = swap_constraints {
337            let owner_key = swap_constraints
338                .owner_key
339                .parse::<Pubkey>()
340                .map_err(|_| SwapError::InvalidOwner)?;
341            if fee_account.owner != owner_key {
342                return Err(SwapError::InvalidOwner.into());
343            }
344            swap_constraints.validate_curve(&swap_curve)?;
345            swap_constraints.validate_fees(&fees)?;
346        }
347        fees.validate()?;
348        swap_curve.calculator.validate()?;
349
350        let initial_amount = swap_curve.calculator.new_pool_supply();
351
352        Self::token_mint_to(
353            swap_info.key,
354            pool_token_program_info.clone(),
355            pool_mint_info.clone(),
356            destination_info.clone(),
357            authority_info.clone(),
358            bump_seed,
359            to_u64(initial_amount)?,
360        )?;
361
362        let obj = SwapVersion::SwapV1(SwapV1 {
363            is_initialized: true,
364            bump_seed,
365            token_program_id,
366            token_a: *token_a_info.key,
367            token_b: *token_b_info.key,
368            pool_mint: *pool_mint_info.key,
369            token_a_mint: token_a.mint,
370            token_b_mint: token_b.mint,
371            pool_fee_account: *fee_account_info.key,
372            fees,
373            swap_curve,
374        });
375        SwapVersion::pack(obj, &mut swap_info.data.borrow_mut())?;
376        Ok(())
377    }
378
379    /// Processes an [Swap](enum.Instruction.html).
380    pub fn process_swap(
381        program_id: &Pubkey,
382        amount_in: u64,
383        minimum_amount_out: u64,
384        accounts: &[AccountInfo],
385    ) -> ProgramResult {
386        let account_info_iter = &mut accounts.iter();
387        let swap_info = next_account_info(account_info_iter)?;
388        let authority_info = next_account_info(account_info_iter)?;
389        let user_transfer_authority_info = next_account_info(account_info_iter)?;
390        let source_info = next_account_info(account_info_iter)?;
391        let swap_source_info = next_account_info(account_info_iter)?;
392        let swap_destination_info = next_account_info(account_info_iter)?;
393        let destination_info = next_account_info(account_info_iter)?;
394        let pool_mint_info = next_account_info(account_info_iter)?;
395        let pool_fee_account_info = next_account_info(account_info_iter)?;
396        let source_token_mint_info = next_account_info(account_info_iter)?;
397        let destination_token_mint_info = next_account_info(account_info_iter)?;
398        let source_token_program_info = next_account_info(account_info_iter)?;
399        let destination_token_program_info = next_account_info(account_info_iter)?;
400        let pool_token_program_info = next_account_info(account_info_iter)?;
401
402        if swap_info.owner != program_id {
403            return Err(ProgramError::IncorrectProgramId);
404        }
405        let token_swap = SwapVersion::unpack(&swap_info.data.borrow())?;
406
407        if *authority_info.key
408            != Self::authority_id(program_id, swap_info.key, token_swap.bump_seed())?
409        {
410            return Err(SwapError::InvalidProgramAddress.into());
411        }
412        if !(*swap_source_info.key == *token_swap.token_a_account()
413            || *swap_source_info.key == *token_swap.token_b_account())
414        {
415            return Err(SwapError::IncorrectSwapAccount.into());
416        }
417        if !(*swap_destination_info.key == *token_swap.token_a_account()
418            || *swap_destination_info.key == *token_swap.token_b_account())
419        {
420            return Err(SwapError::IncorrectSwapAccount.into());
421        }
422        if *swap_source_info.key == *swap_destination_info.key {
423            return Err(SwapError::InvalidInput.into());
424        }
425        if swap_source_info.key == source_info.key {
426            return Err(SwapError::InvalidInput.into());
427        }
428        if swap_destination_info.key == destination_info.key {
429            return Err(SwapError::InvalidInput.into());
430        }
431        if *pool_mint_info.key != *token_swap.pool_mint() {
432            return Err(SwapError::IncorrectPoolMint.into());
433        }
434        if *pool_fee_account_info.key != *token_swap.pool_fee_account() {
435            return Err(SwapError::IncorrectFeeAccount.into());
436        }
437        if *pool_token_program_info.key != *token_swap.token_program_id() {
438            return Err(SwapError::IncorrectTokenProgramId.into());
439        }
440
441        let source_account =
442            Self::unpack_token_account(swap_source_info, token_swap.token_program_id())?;
443        let dest_account =
444            Self::unpack_token_account(swap_destination_info, token_swap.token_program_id())?;
445        let pool_mint = Self::unpack_mint(pool_mint_info, token_swap.token_program_id())?;
446
447        // Take transfer fees into account for actual amount transferred in
448        let actual_amount_in = {
449            let source_mint_data = source_token_mint_info.data.borrow();
450            let source_mint = Self::unpack_mint_with_extensions(
451                &source_mint_data,
452                source_token_mint_info.owner,
453                token_swap.token_program_id(),
454            )?;
455
456            if let Ok(transfer_fee_config) = source_mint.get_extension::<TransferFeeConfig>() {
457                amount_in.saturating_sub(
458                    transfer_fee_config
459                        .calculate_epoch_fee(Clock::get()?.epoch, amount_in)
460                        .ok_or(SwapError::FeeCalculationFailure)?,
461                )
462            } else {
463                amount_in
464            }
465        };
466
467        // Calculate the trade amounts
468        let trade_direction = if *swap_source_info.key == *token_swap.token_a_account() {
469            TradeDirection::AtoB
470        } else {
471            TradeDirection::BtoA
472        };
473        let result = token_swap
474            .swap_curve()
475            .swap(
476                to_u128(actual_amount_in)?,
477                to_u128(source_account.amount)?,
478                to_u128(dest_account.amount)?,
479                trade_direction,
480                token_swap.fees(),
481            )
482            .ok_or(SwapError::ZeroTradingTokens)?;
483
484        // Re-calculate the source amount swapped based on what the curve says
485        let (source_transfer_amount, source_mint_decimals) = {
486            let source_amount_swapped = to_u64(result.source_amount_swapped)?;
487
488            let source_mint_data = source_token_mint_info.data.borrow();
489            let source_mint = Self::unpack_mint_with_extensions(
490                &source_mint_data,
491                source_token_mint_info.owner,
492                token_swap.token_program_id(),
493            )?;
494            let amount =
495                if let Ok(transfer_fee_config) = source_mint.get_extension::<TransferFeeConfig>() {
496                    source_amount_swapped.saturating_add(
497                        transfer_fee_config
498                            .calculate_inverse_epoch_fee(Clock::get()?.epoch, source_amount_swapped)
499                            .ok_or(SwapError::FeeCalculationFailure)?,
500                    )
501                } else {
502                    source_amount_swapped
503                };
504            (amount, source_mint.base.decimals)
505        };
506
507        let (destination_transfer_amount, destination_mint_decimals) = {
508            let destination_mint_data = destination_token_mint_info.data.borrow();
509            let destination_mint = Self::unpack_mint_with_extensions(
510                &destination_mint_data,
511                source_token_mint_info.owner,
512                token_swap.token_program_id(),
513            )?;
514            let amount_out = to_u64(result.destination_amount_swapped)?;
515            let amount_received = if let Ok(transfer_fee_config) =
516                destination_mint.get_extension::<TransferFeeConfig>()
517            {
518                amount_out.saturating_sub(
519                    transfer_fee_config
520                        .calculate_epoch_fee(Clock::get()?.epoch, amount_out)
521                        .ok_or(SwapError::FeeCalculationFailure)?,
522                )
523            } else {
524                amount_out
525            };
526            if amount_received < minimum_amount_out {
527                return Err(SwapError::ExceededSlippage.into());
528            }
529            (amount_out, destination_mint.base.decimals)
530        };
531
532        let (swap_token_a_amount, swap_token_b_amount) = match trade_direction {
533            TradeDirection::AtoB => (
534                result.new_swap_source_amount,
535                result.new_swap_destination_amount,
536            ),
537            TradeDirection::BtoA => (
538                result.new_swap_destination_amount,
539                result.new_swap_source_amount,
540            ),
541        };
542
543        Self::token_transfer(
544            swap_info.key,
545            source_token_program_info.clone(),
546            source_info.clone(),
547            source_token_mint_info.clone(),
548            swap_source_info.clone(),
549            user_transfer_authority_info.clone(),
550            token_swap.bump_seed(),
551            source_transfer_amount,
552            source_mint_decimals,
553        )?;
554
555        if result.owner_fee > 0 {
556            let mut pool_token_amount = token_swap
557                .swap_curve()
558                .calculator
559                .withdraw_single_token_type_exact_out(
560                    result.owner_fee,
561                    swap_token_a_amount,
562                    swap_token_b_amount,
563                    to_u128(pool_mint.supply)?,
564                    trade_direction,
565                    RoundDirection::Floor,
566                )
567                .ok_or(SwapError::FeeCalculationFailure)?;
568            // Allow error to fall through
569            if let Ok(host_fee_account_info) = next_account_info(account_info_iter) {
570                let host_fee_account = Self::unpack_token_account(
571                    host_fee_account_info,
572                    token_swap.token_program_id(),
573                )?;
574                if *pool_mint_info.key != host_fee_account.mint {
575                    return Err(SwapError::IncorrectPoolMint.into());
576                }
577                let host_fee = token_swap
578                    .fees()
579                    .host_fee(pool_token_amount)
580                    .ok_or(SwapError::FeeCalculationFailure)?;
581                if host_fee > 0 {
582                    pool_token_amount = pool_token_amount
583                        .checked_sub(host_fee)
584                        .ok_or(SwapError::FeeCalculationFailure)?;
585                    Self::token_mint_to(
586                        swap_info.key,
587                        pool_token_program_info.clone(),
588                        pool_mint_info.clone(),
589                        host_fee_account_info.clone(),
590                        authority_info.clone(),
591                        token_swap.bump_seed(),
592                        to_u64(host_fee)?,
593                    )?;
594                }
595            }
596            if token_swap
597                .check_pool_fee_info(pool_fee_account_info)
598                .is_ok()
599            {
600                Self::token_mint_to(
601                    swap_info.key,
602                    pool_token_program_info.clone(),
603                    pool_mint_info.clone(),
604                    pool_fee_account_info.clone(),
605                    authority_info.clone(),
606                    token_swap.bump_seed(),
607                    to_u64(pool_token_amount)?,
608                )?;
609            };
610        }
611
612        Self::token_transfer(
613            swap_info.key,
614            destination_token_program_info.clone(),
615            swap_destination_info.clone(),
616            destination_token_mint_info.clone(),
617            destination_info.clone(),
618            authority_info.clone(),
619            token_swap.bump_seed(),
620            destination_transfer_amount,
621            destination_mint_decimals,
622        )?;
623
624        Ok(())
625    }
626
627    /// Processes an [DepositAllTokenTypes](enum.Instruction.html).
628    pub fn process_deposit_all_token_types(
629        program_id: &Pubkey,
630        pool_token_amount: u64,
631        maximum_token_a_amount: u64,
632        maximum_token_b_amount: u64,
633        accounts: &[AccountInfo],
634    ) -> ProgramResult {
635        let account_info_iter = &mut accounts.iter();
636        let swap_info = next_account_info(account_info_iter)?;
637        let authority_info = next_account_info(account_info_iter)?;
638        let user_transfer_authority_info = next_account_info(account_info_iter)?;
639        let source_a_info = next_account_info(account_info_iter)?;
640        let source_b_info = next_account_info(account_info_iter)?;
641        let token_a_info = next_account_info(account_info_iter)?;
642        let token_b_info = next_account_info(account_info_iter)?;
643        let pool_mint_info = next_account_info(account_info_iter)?;
644        let dest_info = next_account_info(account_info_iter)?;
645        let token_a_mint_info = next_account_info(account_info_iter)?;
646        let token_b_mint_info = next_account_info(account_info_iter)?;
647        let token_a_program_info = next_account_info(account_info_iter)?;
648        let token_b_program_info = next_account_info(account_info_iter)?;
649        let pool_token_program_info = next_account_info(account_info_iter)?;
650
651        let token_swap = SwapVersion::unpack(&swap_info.data.borrow())?;
652        let calculator = &token_swap.swap_curve().calculator;
653        if !calculator.allows_deposits() {
654            return Err(SwapError::UnsupportedCurveOperation.into());
655        }
656        Self::check_accounts(
657            token_swap.as_ref(),
658            program_id,
659            swap_info,
660            authority_info,
661            token_a_info,
662            token_b_info,
663            pool_mint_info,
664            pool_token_program_info,
665            Some(source_a_info),
666            Some(source_b_info),
667            None,
668        )?;
669
670        let token_a = Self::unpack_token_account(token_a_info, token_swap.token_program_id())?;
671        let token_b = Self::unpack_token_account(token_b_info, token_swap.token_program_id())?;
672        let pool_mint = Self::unpack_mint(pool_mint_info, token_swap.token_program_id())?;
673        let current_pool_mint_supply = to_u128(pool_mint.supply)?;
674        let (pool_token_amount, pool_mint_supply) = if current_pool_mint_supply > 0 {
675            (to_u128(pool_token_amount)?, current_pool_mint_supply)
676        } else {
677            (calculator.new_pool_supply(), calculator.new_pool_supply())
678        };
679
680        let results = calculator
681            .pool_tokens_to_trading_tokens(
682                pool_token_amount,
683                pool_mint_supply,
684                to_u128(token_a.amount)?,
685                to_u128(token_b.amount)?,
686                RoundDirection::Ceiling,
687            )
688            .ok_or(SwapError::ZeroTradingTokens)?;
689        let token_a_amount = to_u64(results.token_a_amount)?;
690        if token_a_amount > maximum_token_a_amount {
691            return Err(SwapError::ExceededSlippage.into());
692        }
693        if token_a_amount == 0 {
694            return Err(SwapError::ZeroTradingTokens.into());
695        }
696        let token_b_amount = to_u64(results.token_b_amount)?;
697        if token_b_amount > maximum_token_b_amount {
698            return Err(SwapError::ExceededSlippage.into());
699        }
700        if token_b_amount == 0 {
701            return Err(SwapError::ZeroTradingTokens.into());
702        }
703
704        let pool_token_amount = to_u64(pool_token_amount)?;
705
706        Self::token_transfer(
707            swap_info.key,
708            token_a_program_info.clone(),
709            source_a_info.clone(),
710            token_a_mint_info.clone(),
711            token_a_info.clone(),
712            user_transfer_authority_info.clone(),
713            token_swap.bump_seed(),
714            token_a_amount,
715            Self::unpack_mint(token_a_mint_info, token_swap.token_program_id())?.decimals,
716        )?;
717        Self::token_transfer(
718            swap_info.key,
719            token_b_program_info.clone(),
720            source_b_info.clone(),
721            token_b_mint_info.clone(),
722            token_b_info.clone(),
723            user_transfer_authority_info.clone(),
724            token_swap.bump_seed(),
725            token_b_amount,
726            Self::unpack_mint(token_b_mint_info, token_swap.token_program_id())?.decimals,
727        )?;
728        Self::token_mint_to(
729            swap_info.key,
730            pool_token_program_info.clone(),
731            pool_mint_info.clone(),
732            dest_info.clone(),
733            authority_info.clone(),
734            token_swap.bump_seed(),
735            pool_token_amount,
736        )?;
737
738        Ok(())
739    }
740
741    /// Processes an [WithdrawAllTokenTypes](enum.Instruction.html).
742    pub fn process_withdraw_all_token_types(
743        program_id: &Pubkey,
744        pool_token_amount: u64,
745        minimum_token_a_amount: u64,
746        minimum_token_b_amount: u64,
747        accounts: &[AccountInfo],
748    ) -> ProgramResult {
749        let account_info_iter = &mut accounts.iter();
750        let swap_info = next_account_info(account_info_iter)?;
751        let authority_info = next_account_info(account_info_iter)?;
752        let user_transfer_authority_info = next_account_info(account_info_iter)?;
753        let pool_mint_info = next_account_info(account_info_iter)?;
754        let source_info = next_account_info(account_info_iter)?;
755        let token_a_info = next_account_info(account_info_iter)?;
756        let token_b_info = next_account_info(account_info_iter)?;
757        let dest_token_a_info = next_account_info(account_info_iter)?;
758        let dest_token_b_info = next_account_info(account_info_iter)?;
759        let pool_fee_account_info = next_account_info(account_info_iter)?;
760        let token_a_mint_info = next_account_info(account_info_iter)?;
761        let token_b_mint_info = next_account_info(account_info_iter)?;
762        let pool_token_program_info = next_account_info(account_info_iter)?;
763        let token_a_program_info = next_account_info(account_info_iter)?;
764        let token_b_program_info = next_account_info(account_info_iter)?;
765
766        let token_swap = SwapVersion::unpack(&swap_info.data.borrow())?;
767        Self::check_accounts(
768            token_swap.as_ref(),
769            program_id,
770            swap_info,
771            authority_info,
772            token_a_info,
773            token_b_info,
774            pool_mint_info,
775            pool_token_program_info,
776            Some(dest_token_a_info),
777            Some(dest_token_b_info),
778            Some(pool_fee_account_info),
779        )?;
780
781        let token_a = Self::unpack_token_account(token_a_info, token_swap.token_program_id())?;
782        let token_b = Self::unpack_token_account(token_b_info, token_swap.token_program_id())?;
783        let pool_mint = Self::unpack_mint(pool_mint_info, token_swap.token_program_id())?;
784
785        let calculator = &token_swap.swap_curve().calculator;
786
787        let withdraw_fee = match token_swap.check_pool_fee_info(pool_fee_account_info) {
788            Ok(_) => {
789                if *pool_fee_account_info.key == *source_info.key {
790                    // withdrawing from the fee account, don't assess withdraw fee
791                    0
792                } else {
793                    token_swap
794                        .fees()
795                        .owner_withdraw_fee(to_u128(pool_token_amount)?)
796                        .ok_or(SwapError::FeeCalculationFailure)?
797                }
798            }
799            Err(_) => 0,
800        };
801        let pool_token_amount = to_u128(pool_token_amount)?
802            .checked_sub(withdraw_fee)
803            .ok_or(SwapError::CalculationFailure)?;
804
805        let results = calculator
806            .pool_tokens_to_trading_tokens(
807                pool_token_amount,
808                to_u128(pool_mint.supply)?,
809                to_u128(token_a.amount)?,
810                to_u128(token_b.amount)?,
811                RoundDirection::Floor,
812            )
813            .ok_or(SwapError::ZeroTradingTokens)?;
814        let token_a_amount = to_u64(results.token_a_amount)?;
815        let token_a_amount = std::cmp::min(token_a.amount, token_a_amount);
816        if token_a_amount < minimum_token_a_amount {
817            return Err(SwapError::ExceededSlippage.into());
818        }
819        if token_a_amount == 0 && token_a.amount != 0 {
820            return Err(SwapError::ZeroTradingTokens.into());
821        }
822        let token_b_amount = to_u64(results.token_b_amount)?;
823        let token_b_amount = std::cmp::min(token_b.amount, token_b_amount);
824        if token_b_amount < minimum_token_b_amount {
825            return Err(SwapError::ExceededSlippage.into());
826        }
827        if token_b_amount == 0 && token_b.amount != 0 {
828            return Err(SwapError::ZeroTradingTokens.into());
829        }
830
831        if withdraw_fee > 0 {
832            Self::token_transfer(
833                swap_info.key,
834                pool_token_program_info.clone(),
835                source_info.clone(),
836                pool_mint_info.clone(),
837                pool_fee_account_info.clone(),
838                user_transfer_authority_info.clone(),
839                token_swap.bump_seed(),
840                to_u64(withdraw_fee)?,
841                pool_mint.decimals,
842            )?;
843        }
844        Self::token_burn(
845            swap_info.key,
846            pool_token_program_info.clone(),
847            source_info.clone(),
848            pool_mint_info.clone(),
849            user_transfer_authority_info.clone(),
850            token_swap.bump_seed(),
851            to_u64(pool_token_amount)?,
852        )?;
853
854        if token_a_amount > 0 {
855            Self::token_transfer(
856                swap_info.key,
857                token_a_program_info.clone(),
858                token_a_info.clone(),
859                token_a_mint_info.clone(),
860                dest_token_a_info.clone(),
861                authority_info.clone(),
862                token_swap.bump_seed(),
863                token_a_amount,
864                Self::unpack_mint(token_a_mint_info, token_swap.token_program_id())?.decimals,
865            )?;
866        }
867        if token_b_amount > 0 {
868            Self::token_transfer(
869                swap_info.key,
870                token_b_program_info.clone(),
871                token_b_info.clone(),
872                token_b_mint_info.clone(),
873                dest_token_b_info.clone(),
874                authority_info.clone(),
875                token_swap.bump_seed(),
876                token_b_amount,
877                Self::unpack_mint(token_b_mint_info, token_swap.token_program_id())?.decimals,
878            )?;
879        }
880        Ok(())
881    }
882
883    /// Processes DepositSingleTokenTypeExactAmountIn
884    pub fn process_deposit_single_token_type_exact_amount_in(
885        program_id: &Pubkey,
886        source_token_amount: u64,
887        minimum_pool_token_amount: u64,
888        accounts: &[AccountInfo],
889    ) -> ProgramResult {
890        let account_info_iter = &mut accounts.iter();
891        let swap_info = next_account_info(account_info_iter)?;
892        let authority_info = next_account_info(account_info_iter)?;
893        let user_transfer_authority_info = next_account_info(account_info_iter)?;
894        let source_info = next_account_info(account_info_iter)?;
895        let swap_token_a_info = next_account_info(account_info_iter)?;
896        let swap_token_b_info = next_account_info(account_info_iter)?;
897        let pool_mint_info = next_account_info(account_info_iter)?;
898        let destination_info = next_account_info(account_info_iter)?;
899        let source_token_mint_info = next_account_info(account_info_iter)?;
900        let source_token_program_info = next_account_info(account_info_iter)?;
901        let pool_token_program_info = next_account_info(account_info_iter)?;
902
903        let token_swap = SwapVersion::unpack(&swap_info.data.borrow())?;
904        let calculator = &token_swap.swap_curve().calculator;
905        if !calculator.allows_deposits() {
906            return Err(SwapError::UnsupportedCurveOperation.into());
907        }
908        let source_account =
909            Self::unpack_token_account(source_info, token_swap.token_program_id())?;
910        let swap_token_a =
911            Self::unpack_token_account(swap_token_a_info, token_swap.token_program_id())?;
912        let swap_token_b =
913            Self::unpack_token_account(swap_token_b_info, token_swap.token_program_id())?;
914
915        let trade_direction = if source_account.mint == swap_token_a.mint {
916            TradeDirection::AtoB
917        } else if source_account.mint == swap_token_b.mint {
918            TradeDirection::BtoA
919        } else {
920            return Err(SwapError::IncorrectSwapAccount.into());
921        };
922
923        let (source_a_info, source_b_info) = match trade_direction {
924            TradeDirection::AtoB => (Some(source_info), None),
925            TradeDirection::BtoA => (None, Some(source_info)),
926        };
927
928        Self::check_accounts(
929            token_swap.as_ref(),
930            program_id,
931            swap_info,
932            authority_info,
933            swap_token_a_info,
934            swap_token_b_info,
935            pool_mint_info,
936            pool_token_program_info,
937            source_a_info,
938            source_b_info,
939            None,
940        )?;
941
942        let pool_mint = Self::unpack_mint(pool_mint_info, token_swap.token_program_id())?;
943        let pool_mint_supply = to_u128(pool_mint.supply)?;
944        let pool_token_amount = if pool_mint_supply > 0 {
945            token_swap
946                .swap_curve()
947                .deposit_single_token_type(
948                    to_u128(source_token_amount)?,
949                    to_u128(swap_token_a.amount)?,
950                    to_u128(swap_token_b.amount)?,
951                    pool_mint_supply,
952                    trade_direction,
953                    token_swap.fees(),
954                )
955                .ok_or(SwapError::ZeroTradingTokens)?
956        } else {
957            calculator.new_pool_supply()
958        };
959
960        let pool_token_amount = to_u64(pool_token_amount)?;
961        if pool_token_amount < minimum_pool_token_amount {
962            return Err(SwapError::ExceededSlippage.into());
963        }
964        if pool_token_amount == 0 {
965            return Err(SwapError::ZeroTradingTokens.into());
966        }
967
968        match trade_direction {
969            TradeDirection::AtoB => {
970                Self::token_transfer(
971                    swap_info.key,
972                    source_token_program_info.clone(),
973                    source_info.clone(),
974                    source_token_mint_info.clone(),
975                    swap_token_a_info.clone(),
976                    user_transfer_authority_info.clone(),
977                    token_swap.bump_seed(),
978                    source_token_amount,
979                    Self::unpack_mint(source_token_mint_info, token_swap.token_program_id())?
980                        .decimals,
981                )?;
982            }
983            TradeDirection::BtoA => {
984                Self::token_transfer(
985                    swap_info.key,
986                    source_token_program_info.clone(),
987                    source_info.clone(),
988                    source_token_mint_info.clone(),
989                    swap_token_b_info.clone(),
990                    user_transfer_authority_info.clone(),
991                    token_swap.bump_seed(),
992                    source_token_amount,
993                    Self::unpack_mint(source_token_mint_info, token_swap.token_program_id())?
994                        .decimals,
995                )?;
996            }
997        }
998        Self::token_mint_to(
999            swap_info.key,
1000            pool_token_program_info.clone(),
1001            pool_mint_info.clone(),
1002            destination_info.clone(),
1003            authority_info.clone(),
1004            token_swap.bump_seed(),
1005            pool_token_amount,
1006        )?;
1007
1008        Ok(())
1009    }
1010
1011    /// Processes a [WithdrawSingleTokenTypeExactAmountOut](enum.Instruction.html).
1012    pub fn process_withdraw_single_token_type_exact_amount_out(
1013        program_id: &Pubkey,
1014        destination_token_amount: u64,
1015        maximum_pool_token_amount: u64,
1016        accounts: &[AccountInfo],
1017    ) -> ProgramResult {
1018        let account_info_iter = &mut accounts.iter();
1019        let swap_info = next_account_info(account_info_iter)?;
1020        let authority_info = next_account_info(account_info_iter)?;
1021        let user_transfer_authority_info = next_account_info(account_info_iter)?;
1022        let pool_mint_info = next_account_info(account_info_iter)?;
1023        let source_info = next_account_info(account_info_iter)?;
1024        let swap_token_a_info = next_account_info(account_info_iter)?;
1025        let swap_token_b_info = next_account_info(account_info_iter)?;
1026        let destination_info = next_account_info(account_info_iter)?;
1027        let pool_fee_account_info = next_account_info(account_info_iter)?;
1028        let destination_token_mint_info = next_account_info(account_info_iter)?;
1029        let pool_token_program_info = next_account_info(account_info_iter)?;
1030        let destination_token_program_info = next_account_info(account_info_iter)?;
1031
1032        let token_swap = SwapVersion::unpack(&swap_info.data.borrow())?;
1033        let destination_account =
1034            Self::unpack_token_account(destination_info, token_swap.token_program_id())?;
1035        let swap_token_a =
1036            Self::unpack_token_account(swap_token_a_info, token_swap.token_program_id())?;
1037        let swap_token_b =
1038            Self::unpack_token_account(swap_token_b_info, token_swap.token_program_id())?;
1039
1040        let trade_direction = if destination_account.mint == swap_token_a.mint {
1041            TradeDirection::AtoB
1042        } else if destination_account.mint == swap_token_b.mint {
1043            TradeDirection::BtoA
1044        } else {
1045            return Err(SwapError::IncorrectSwapAccount.into());
1046        };
1047
1048        let (destination_a_info, destination_b_info) = match trade_direction {
1049            TradeDirection::AtoB => (Some(destination_info), None),
1050            TradeDirection::BtoA => (None, Some(destination_info)),
1051        };
1052        Self::check_accounts(
1053            token_swap.as_ref(),
1054            program_id,
1055            swap_info,
1056            authority_info,
1057            swap_token_a_info,
1058            swap_token_b_info,
1059            pool_mint_info,
1060            pool_token_program_info,
1061            destination_a_info,
1062            destination_b_info,
1063            Some(pool_fee_account_info),
1064        )?;
1065
1066        let pool_mint = Self::unpack_mint(pool_mint_info, token_swap.token_program_id())?;
1067        let pool_mint_supply = to_u128(pool_mint.supply)?;
1068        let swap_token_a_amount = to_u128(swap_token_a.amount)?;
1069        let swap_token_b_amount = to_u128(swap_token_b.amount)?;
1070
1071        let burn_pool_token_amount = token_swap
1072            .swap_curve()
1073            .withdraw_single_token_type_exact_out(
1074                to_u128(destination_token_amount)?,
1075                swap_token_a_amount,
1076                swap_token_b_amount,
1077                pool_mint_supply,
1078                trade_direction,
1079                token_swap.fees(),
1080            )
1081            .ok_or(SwapError::ZeroTradingTokens)?;
1082
1083        let withdraw_fee = match token_swap.check_pool_fee_info(pool_fee_account_info) {
1084            Ok(_) => {
1085                if *pool_fee_account_info.key == *source_info.key {
1086                    // withdrawing from the fee account, don't assess withdraw fee
1087                    0
1088                } else {
1089                    token_swap
1090                        .fees()
1091                        .owner_withdraw_fee(burn_pool_token_amount)
1092                        .ok_or(SwapError::FeeCalculationFailure)?
1093                }
1094            }
1095            Err(_) => 0,
1096        };
1097        let pool_token_amount = burn_pool_token_amount
1098            .checked_add(withdraw_fee)
1099            .ok_or(SwapError::CalculationFailure)?;
1100
1101        if to_u64(pool_token_amount)? > maximum_pool_token_amount {
1102            return Err(SwapError::ExceededSlippage.into());
1103        }
1104        if pool_token_amount == 0 {
1105            return Err(SwapError::ZeroTradingTokens.into());
1106        }
1107
1108        if withdraw_fee > 0 {
1109            Self::token_transfer(
1110                swap_info.key,
1111                pool_token_program_info.clone(),
1112                source_info.clone(),
1113                pool_mint_info.clone(),
1114                pool_fee_account_info.clone(),
1115                user_transfer_authority_info.clone(),
1116                token_swap.bump_seed(),
1117                to_u64(withdraw_fee)?,
1118                pool_mint.decimals,
1119            )?;
1120        }
1121        Self::token_burn(
1122            swap_info.key,
1123            pool_token_program_info.clone(),
1124            source_info.clone(),
1125            pool_mint_info.clone(),
1126            user_transfer_authority_info.clone(),
1127            token_swap.bump_seed(),
1128            to_u64(burn_pool_token_amount)?,
1129        )?;
1130
1131        match trade_direction {
1132            TradeDirection::AtoB => {
1133                Self::token_transfer(
1134                    swap_info.key,
1135                    destination_token_program_info.clone(),
1136                    swap_token_a_info.clone(),
1137                    destination_token_mint_info.clone(),
1138                    destination_info.clone(),
1139                    authority_info.clone(),
1140                    token_swap.bump_seed(),
1141                    destination_token_amount,
1142                    Self::unpack_mint(destination_token_mint_info, token_swap.token_program_id())?
1143                        .decimals,
1144                )?;
1145            }
1146            TradeDirection::BtoA => {
1147                Self::token_transfer(
1148                    swap_info.key,
1149                    destination_token_program_info.clone(),
1150                    swap_token_b_info.clone(),
1151                    destination_token_mint_info.clone(),
1152                    destination_info.clone(),
1153                    authority_info.clone(),
1154                    token_swap.bump_seed(),
1155                    destination_token_amount,
1156                    Self::unpack_mint(destination_token_mint_info, token_swap.token_program_id())?
1157                        .decimals,
1158                )?;
1159            }
1160        }
1161
1162        Ok(())
1163    }
1164
1165    /// Processes an [Instruction](enum.Instruction.html).
1166    pub fn process(program_id: &Pubkey, accounts: &[AccountInfo], input: &[u8]) -> ProgramResult {
1167        Self::process_with_constraints(program_id, accounts, input, &SWAP_CONSTRAINTS)
1168    }
1169
1170    /// Processes an instruction given extra constraint
1171    pub fn process_with_constraints(
1172        program_id: &Pubkey,
1173        accounts: &[AccountInfo],
1174        input: &[u8],
1175        swap_constraints: &Option<SwapConstraints>,
1176    ) -> ProgramResult {
1177        let instruction = SwapInstruction::unpack(input)?;
1178        match instruction {
1179            SwapInstruction::Initialize(Initialize { fees, swap_curve }) => {
1180                msg!("Instruction: Init");
1181                Self::process_initialize(program_id, fees, swap_curve, accounts, swap_constraints)
1182            }
1183            SwapInstruction::Swap(Swap {
1184                amount_in,
1185                minimum_amount_out,
1186            }) => {
1187                msg!("Instruction: Swap");
1188                Self::process_swap(program_id, amount_in, minimum_amount_out, accounts)
1189            }
1190            SwapInstruction::DepositAllTokenTypes(DepositAllTokenTypes {
1191                pool_token_amount,
1192                maximum_token_a_amount,
1193                maximum_token_b_amount,
1194            }) => {
1195                msg!("Instruction: DepositAllTokenTypes");
1196                Self::process_deposit_all_token_types(
1197                    program_id,
1198                    pool_token_amount,
1199                    maximum_token_a_amount,
1200                    maximum_token_b_amount,
1201                    accounts,
1202                )
1203            }
1204            SwapInstruction::WithdrawAllTokenTypes(WithdrawAllTokenTypes {
1205                pool_token_amount,
1206                minimum_token_a_amount,
1207                minimum_token_b_amount,
1208            }) => {
1209                msg!("Instruction: WithdrawAllTokenTypes");
1210                Self::process_withdraw_all_token_types(
1211                    program_id,
1212                    pool_token_amount,
1213                    minimum_token_a_amount,
1214                    minimum_token_b_amount,
1215                    accounts,
1216                )
1217            }
1218            SwapInstruction::DepositSingleTokenTypeExactAmountIn(
1219                DepositSingleTokenTypeExactAmountIn {
1220                    source_token_amount,
1221                    minimum_pool_token_amount,
1222                },
1223            ) => {
1224                msg!("Instruction: DepositSingleTokenTypeExactAmountIn");
1225                Self::process_deposit_single_token_type_exact_amount_in(
1226                    program_id,
1227                    source_token_amount,
1228                    minimum_pool_token_amount,
1229                    accounts,
1230                )
1231            }
1232            SwapInstruction::WithdrawSingleTokenTypeExactAmountOut(
1233                WithdrawSingleTokenTypeExactAmountOut {
1234                    destination_token_amount,
1235                    maximum_pool_token_amount,
1236                },
1237            ) => {
1238                msg!("Instruction: WithdrawSingleTokenTypeExactAmountOut");
1239                Self::process_withdraw_single_token_type_exact_amount_out(
1240                    program_id,
1241                    destination_token_amount,
1242                    maximum_pool_token_amount,
1243                    accounts,
1244                )
1245            }
1246        }
1247    }
1248}
1249
1250fn to_u128(val: u64) -> Result<u128, SwapError> {
1251    val.try_into().map_err(|_| SwapError::ConversionFailure)
1252}
1253
1254fn to_u64(val: u128) -> Result<u64, SwapError> {
1255    val.try_into().map_err(|_| SwapError::ConversionFailure)
1256}
1257
1258fn invoke_signed_wrapper<T>(
1259    instruction: &Instruction,
1260    account_infos: &[AccountInfo],
1261    signers_seeds: &[&[&[u8]]],
1262) -> Result<(), ProgramError>
1263where
1264    T: 'static + PrintProgramError + DecodeError<T> + FromPrimitive + Error,
1265{
1266    invoke_signed(instruction, account_infos, signers_seeds).map_err(|err| {
1267        err.print::<T>();
1268        err
1269    })
1270}
1271
1272#[cfg(test)]
1273mod tests {
1274    use super::*;
1275    use crate::{
1276        curve::calculator::{CurveCalculator, INITIAL_SWAP_POOL_AMOUNT},
1277        curve::{
1278            base::CurveType, constant_price::ConstantPriceCurve,
1279            constant_product::ConstantProductCurve, offset::OffsetCurve,
1280        },
1281        instruction::{
1282            deposit_all_token_types, deposit_single_token_type_exact_amount_in, initialize, swap,
1283            withdraw_all_token_types, withdraw_single_token_type_exact_amount_out,
1284        },
1285    };
1286    use solana_program::{
1287        clock::Clock, entrypoint::SUCCESS, instruction::Instruction, program_pack::Pack,
1288        program_stubs, rent::Rent,
1289    };
1290    use solana_sdk::account::{
1291        create_account_for_test, create_is_signer_account_infos, Account as SolanaAccount,
1292    };
1293    use spl_token_2022::{
1294        error::TokenError,
1295        extension::{
1296            transfer_fee::{instruction::initialize_transfer_fee_config, TransferFee},
1297            ExtensionType,
1298        },
1299        instruction::{
1300            approve, close_account, freeze_account, initialize_account, initialize_immutable_owner,
1301            initialize_mint, initialize_mint_close_authority, mint_to, revoke, set_authority,
1302            AuthorityType,
1303        },
1304    };
1305    use std::sync::Arc;
1306    use test_case::test_case;
1307
1308    // Test program id for the swap program.
1309    const SWAP_PROGRAM_ID: Pubkey = Pubkey::new_from_array([2u8; 32]);
1310
1311    struct TestSyscallStubs {}
1312    impl program_stubs::SyscallStubs for TestSyscallStubs {
1313        fn sol_invoke_signed(
1314            &self,
1315            instruction: &Instruction,
1316            account_infos: &[AccountInfo],
1317            signers_seeds: &[&[&[u8]]],
1318        ) -> ProgramResult {
1319            msg!("TestSyscallStubs::sol_invoke_signed()");
1320
1321            let mut new_account_infos = vec![];
1322
1323            // mimic check for token program in accounts
1324            if !account_infos
1325                .iter()
1326                .any(|x| *x.key == spl_token::id() || *x.key == spl_token_2022::id())
1327            {
1328                return Err(ProgramError::InvalidAccountData);
1329            }
1330
1331            for meta in instruction.accounts.iter() {
1332                for account_info in account_infos.iter() {
1333                    if meta.pubkey == *account_info.key {
1334                        let mut new_account_info = account_info.clone();
1335                        for seeds in signers_seeds.iter() {
1336                            let signer =
1337                                Pubkey::create_program_address(seeds, &SWAP_PROGRAM_ID).unwrap();
1338                            if *account_info.key == signer {
1339                                new_account_info.is_signer = true;
1340                            }
1341                        }
1342                        new_account_infos.push(new_account_info);
1343                    }
1344                }
1345            }
1346
1347            if instruction.program_id == spl_token::id() {
1348                spl_token::processor::Processor::process(
1349                    &instruction.program_id,
1350                    &new_account_infos,
1351                    &instruction.data,
1352                )
1353            } else if instruction.program_id == spl_token_2022::id() {
1354                spl_token_2022::processor::Processor::process(
1355                    &instruction.program_id,
1356                    &new_account_infos,
1357                    &instruction.data,
1358                )
1359            } else {
1360                Err(ProgramError::IncorrectProgramId)
1361            }
1362        }
1363
1364        fn sol_get_clock_sysvar(&self, var_addr: *mut u8) -> u64 {
1365            unsafe {
1366                *(var_addr as *mut _ as *mut Clock) = Clock::default();
1367            }
1368            SUCCESS
1369        }
1370    }
1371
1372    fn test_syscall_stubs() {
1373        use std::sync::Once;
1374        static ONCE: Once = Once::new();
1375
1376        ONCE.call_once(|| {
1377            program_stubs::set_syscall_stubs(Box::new(TestSyscallStubs {}));
1378        });
1379    }
1380
1381    #[derive(Default)]
1382    struct SwapTransferFees {
1383        pool_token: TransferFee,
1384        token_a: TransferFee,
1385        token_b: TransferFee,
1386    }
1387
1388    struct SwapAccountInfo {
1389        bump_seed: u8,
1390        authority_key: Pubkey,
1391        fees: Fees,
1392        transfer_fees: SwapTransferFees,
1393        swap_curve: SwapCurve,
1394        swap_key: Pubkey,
1395        swap_account: SolanaAccount,
1396        pool_mint_key: Pubkey,
1397        pool_mint_account: SolanaAccount,
1398        pool_fee_key: Pubkey,
1399        pool_fee_account: SolanaAccount,
1400        pool_token_key: Pubkey,
1401        pool_token_account: SolanaAccount,
1402        token_a_key: Pubkey,
1403        token_a_account: SolanaAccount,
1404        token_a_mint_key: Pubkey,
1405        token_a_mint_account: SolanaAccount,
1406        token_b_key: Pubkey,
1407        token_b_account: SolanaAccount,
1408        token_b_mint_key: Pubkey,
1409        token_b_mint_account: SolanaAccount,
1410        pool_token_program_id: Pubkey,
1411        token_a_program_id: Pubkey,
1412        token_b_program_id: Pubkey,
1413    }
1414
1415    impl SwapAccountInfo {
1416        #[allow(clippy::too_many_arguments)]
1417        pub fn new(
1418            user_key: &Pubkey,
1419            fees: Fees,
1420            transfer_fees: SwapTransferFees,
1421            swap_curve: SwapCurve,
1422            token_a_amount: u64,
1423            token_b_amount: u64,
1424            pool_token_program_id: &Pubkey,
1425            token_a_program_id: &Pubkey,
1426            token_b_program_id: &Pubkey,
1427        ) -> Self {
1428            let swap_key = Pubkey::new_unique();
1429            let swap_account = SolanaAccount::new(0, SwapVersion::LATEST_LEN, &SWAP_PROGRAM_ID);
1430            let (authority_key, bump_seed) =
1431                Pubkey::find_program_address(&[&swap_key.to_bytes()[..]], &SWAP_PROGRAM_ID);
1432
1433            let (pool_mint_key, mut pool_mint_account) = create_mint(
1434                pool_token_program_id,
1435                &authority_key,
1436                None,
1437                None,
1438                &transfer_fees.pool_token,
1439            );
1440            let (pool_token_key, pool_token_account) = mint_token(
1441                pool_token_program_id,
1442                &pool_mint_key,
1443                &mut pool_mint_account,
1444                &authority_key,
1445                user_key,
1446                0,
1447            );
1448            let (pool_fee_key, pool_fee_account) = mint_token(
1449                pool_token_program_id,
1450                &pool_mint_key,
1451                &mut pool_mint_account,
1452                &authority_key,
1453                user_key,
1454                0,
1455            );
1456            let (token_a_mint_key, mut token_a_mint_account) = create_mint(
1457                token_a_program_id,
1458                user_key,
1459                None,
1460                None,
1461                &transfer_fees.token_a,
1462            );
1463            let (token_a_key, token_a_account) = mint_token(
1464                token_a_program_id,
1465                &token_a_mint_key,
1466                &mut token_a_mint_account,
1467                user_key,
1468                &authority_key,
1469                token_a_amount,
1470            );
1471            let (token_b_mint_key, mut token_b_mint_account) = create_mint(
1472                token_b_program_id,
1473                user_key,
1474                None,
1475                None,
1476                &transfer_fees.token_b,
1477            );
1478            let (token_b_key, token_b_account) = mint_token(
1479                token_b_program_id,
1480                &token_b_mint_key,
1481                &mut token_b_mint_account,
1482                user_key,
1483                &authority_key,
1484                token_b_amount,
1485            );
1486
1487            SwapAccountInfo {
1488                bump_seed,
1489                authority_key,
1490                fees,
1491                transfer_fees,
1492                swap_curve,
1493                swap_key,
1494                swap_account,
1495                pool_mint_key,
1496                pool_mint_account,
1497                pool_fee_key,
1498                pool_fee_account,
1499                pool_token_key,
1500                pool_token_account,
1501                token_a_key,
1502                token_a_account,
1503                token_a_mint_key,
1504                token_a_mint_account,
1505                token_b_key,
1506                token_b_account,
1507                token_b_mint_key,
1508                token_b_mint_account,
1509                pool_token_program_id: *pool_token_program_id,
1510                token_a_program_id: *token_a_program_id,
1511                token_b_program_id: *token_b_program_id,
1512            }
1513        }
1514
1515        pub fn initialize_swap(&mut self) -> ProgramResult {
1516            do_process_instruction(
1517                initialize(
1518                    &SWAP_PROGRAM_ID,
1519                    &self.pool_token_program_id,
1520                    &self.swap_key,
1521                    &self.authority_key,
1522                    &self.token_a_key,
1523                    &self.token_b_key,
1524                    &self.pool_mint_key,
1525                    &self.pool_fee_key,
1526                    &self.pool_token_key,
1527                    self.fees.clone(),
1528                    self.swap_curve.clone(),
1529                )
1530                .unwrap(),
1531                vec![
1532                    &mut self.swap_account,
1533                    &mut SolanaAccount::default(),
1534                    &mut self.token_a_account,
1535                    &mut self.token_b_account,
1536                    &mut self.pool_mint_account,
1537                    &mut self.pool_fee_account,
1538                    &mut self.pool_token_account,
1539                    &mut SolanaAccount::default(),
1540                ],
1541            )
1542        }
1543
1544        pub fn setup_token_accounts(
1545            &mut self,
1546            mint_owner: &Pubkey,
1547            account_owner: &Pubkey,
1548            a_amount: u64,
1549            b_amount: u64,
1550            pool_amount: u64,
1551        ) -> (
1552            Pubkey,
1553            SolanaAccount,
1554            Pubkey,
1555            SolanaAccount,
1556            Pubkey,
1557            SolanaAccount,
1558        ) {
1559            let (token_a_key, token_a_account) = mint_token(
1560                &self.token_a_program_id,
1561                &self.token_a_mint_key,
1562                &mut self.token_a_mint_account,
1563                mint_owner,
1564                account_owner,
1565                a_amount,
1566            );
1567            let (token_b_key, token_b_account) = mint_token(
1568                &self.token_b_program_id,
1569                &self.token_b_mint_key,
1570                &mut self.token_b_mint_account,
1571                mint_owner,
1572                account_owner,
1573                b_amount,
1574            );
1575            let (pool_key, pool_account) = mint_token(
1576                &self.pool_token_program_id,
1577                &self.pool_mint_key,
1578                &mut self.pool_mint_account,
1579                &self.authority_key,
1580                account_owner,
1581                pool_amount,
1582            );
1583            (
1584                token_a_key,
1585                token_a_account,
1586                token_b_key,
1587                token_b_account,
1588                pool_key,
1589                pool_account,
1590            )
1591        }
1592
1593        fn get_swap_key(&self, mint_key: &Pubkey) -> &Pubkey {
1594            if *mint_key == self.token_a_mint_key {
1595                &self.token_a_key
1596            } else if *mint_key == self.token_b_mint_key {
1597                &self.token_b_key
1598            } else {
1599                panic!("Could not find matching swap token account");
1600            }
1601        }
1602
1603        fn get_token_program_id(&self, account_key: &Pubkey) -> &Pubkey {
1604            if *account_key == self.token_a_key {
1605                &self.token_a_program_id
1606            } else if *account_key == self.token_b_key {
1607                &self.token_b_program_id
1608            } else {
1609                panic!("Could not find matching swap token account");
1610            }
1611        }
1612
1613        fn get_token_mint(&self, account_key: &Pubkey) -> (Pubkey, SolanaAccount) {
1614            if *account_key == self.token_a_key {
1615                (self.token_a_mint_key, self.token_a_mint_account.clone())
1616            } else if *account_key == self.token_b_key {
1617                (self.token_b_mint_key, self.token_b_mint_account.clone())
1618            } else {
1619                panic!("Could not find matching swap token account");
1620            }
1621        }
1622
1623        fn get_token_account(&self, account_key: &Pubkey) -> &SolanaAccount {
1624            if *account_key == self.token_a_key {
1625                &self.token_a_account
1626            } else if *account_key == self.token_b_key {
1627                &self.token_b_account
1628            } else {
1629                panic!("Could not find matching swap token account");
1630            }
1631        }
1632
1633        fn set_token_account(&mut self, account_key: &Pubkey, account: SolanaAccount) {
1634            if *account_key == self.token_a_key {
1635                self.token_a_account = account;
1636                return;
1637            } else if *account_key == self.token_b_key {
1638                self.token_b_account = account;
1639                return;
1640            }
1641            panic!("Could not find matching swap token account");
1642        }
1643
1644        #[allow(clippy::too_many_arguments)]
1645        pub fn swap(
1646            &mut self,
1647            user_key: &Pubkey,
1648            user_source_key: &Pubkey,
1649            user_source_account: &mut SolanaAccount,
1650            swap_source_key: &Pubkey,
1651            swap_destination_key: &Pubkey,
1652            user_destination_key: &Pubkey,
1653            user_destination_account: &mut SolanaAccount,
1654            amount_in: u64,
1655            minimum_amount_out: u64,
1656        ) -> ProgramResult {
1657            let user_transfer_key = Pubkey::new_unique();
1658            let source_token_program_id = self.get_token_program_id(swap_source_key);
1659            let destination_token_program_id = self.get_token_program_id(swap_destination_key);
1660            // approve moving from user source account
1661            do_process_instruction(
1662                approve(
1663                    source_token_program_id,
1664                    user_source_key,
1665                    &user_transfer_key,
1666                    user_key,
1667                    &[],
1668                    amount_in,
1669                )
1670                .unwrap(),
1671                vec![
1672                    user_source_account,
1673                    &mut SolanaAccount::default(),
1674                    &mut SolanaAccount::default(),
1675                ],
1676            )
1677            .unwrap();
1678
1679            let (source_mint_key, mut source_mint_account) = self.get_token_mint(swap_source_key);
1680            let (destination_mint_key, mut destination_mint_account) =
1681                self.get_token_mint(swap_destination_key);
1682            let mut swap_source_account = self.get_token_account(swap_source_key).clone();
1683            let mut swap_destination_account = self.get_token_account(swap_destination_key).clone();
1684
1685            // perform the swap
1686            do_process_instruction(
1687                swap(
1688                    &SWAP_PROGRAM_ID,
1689                    source_token_program_id,
1690                    destination_token_program_id,
1691                    &self.pool_token_program_id,
1692                    &self.swap_key,
1693                    &self.authority_key,
1694                    &user_transfer_key,
1695                    user_source_key,
1696                    swap_source_key,
1697                    swap_destination_key,
1698                    user_destination_key,
1699                    &self.pool_mint_key,
1700                    &self.pool_fee_key,
1701                    &source_mint_key,
1702                    &destination_mint_key,
1703                    None,
1704                    Swap {
1705                        amount_in,
1706                        minimum_amount_out,
1707                    },
1708                )
1709                .unwrap(),
1710                vec![
1711                    &mut self.swap_account,
1712                    &mut SolanaAccount::default(),
1713                    &mut SolanaAccount::default(),
1714                    user_source_account,
1715                    &mut swap_source_account,
1716                    &mut swap_destination_account,
1717                    user_destination_account,
1718                    &mut self.pool_mint_account,
1719                    &mut self.pool_fee_account,
1720                    &mut source_mint_account,
1721                    &mut destination_mint_account,
1722                    &mut SolanaAccount::default(),
1723                    &mut SolanaAccount::default(),
1724                    &mut SolanaAccount::default(),
1725                ],
1726            )?;
1727
1728            self.set_token_account(swap_source_key, swap_source_account);
1729            self.set_token_account(swap_destination_key, swap_destination_account);
1730
1731            Ok(())
1732        }
1733
1734        #[allow(clippy::too_many_arguments)]
1735        pub fn deposit_all_token_types(
1736            &mut self,
1737            depositor_key: &Pubkey,
1738            depositor_token_a_key: &Pubkey,
1739            depositor_token_a_account: &mut SolanaAccount,
1740            depositor_token_b_key: &Pubkey,
1741            depositor_token_b_account: &mut SolanaAccount,
1742            depositor_pool_key: &Pubkey,
1743            depositor_pool_account: &mut SolanaAccount,
1744            pool_token_amount: u64,
1745            maximum_token_a_amount: u64,
1746            maximum_token_b_amount: u64,
1747        ) -> ProgramResult {
1748            let user_transfer_authority = Pubkey::new_unique();
1749            let token_a_program_id = depositor_token_a_account.owner;
1750            do_process_instruction(
1751                approve(
1752                    &token_a_program_id,
1753                    depositor_token_a_key,
1754                    &user_transfer_authority,
1755                    depositor_key,
1756                    &[],
1757                    maximum_token_a_amount,
1758                )
1759                .unwrap(),
1760                vec![
1761                    depositor_token_a_account,
1762                    &mut SolanaAccount::default(),
1763                    &mut SolanaAccount::default(),
1764                ],
1765            )
1766            .unwrap();
1767
1768            let token_b_program_id = depositor_token_b_account.owner;
1769            do_process_instruction(
1770                approve(
1771                    &token_b_program_id,
1772                    depositor_token_b_key,
1773                    &user_transfer_authority,
1774                    depositor_key,
1775                    &[],
1776                    maximum_token_b_amount,
1777                )
1778                .unwrap(),
1779                vec![
1780                    depositor_token_b_account,
1781                    &mut SolanaAccount::default(),
1782                    &mut SolanaAccount::default(),
1783                ],
1784            )
1785            .unwrap();
1786
1787            let pool_token_program_id = depositor_pool_account.owner;
1788            do_process_instruction(
1789                deposit_all_token_types(
1790                    &SWAP_PROGRAM_ID,
1791                    &token_a_program_id,
1792                    &token_b_program_id,
1793                    &pool_token_program_id,
1794                    &self.swap_key,
1795                    &self.authority_key,
1796                    &user_transfer_authority,
1797                    depositor_token_a_key,
1798                    depositor_token_b_key,
1799                    &self.token_a_key,
1800                    &self.token_b_key,
1801                    &self.pool_mint_key,
1802                    depositor_pool_key,
1803                    &self.token_a_mint_key,
1804                    &self.token_b_mint_key,
1805                    DepositAllTokenTypes {
1806                        pool_token_amount,
1807                        maximum_token_a_amount,
1808                        maximum_token_b_amount,
1809                    },
1810                )
1811                .unwrap(),
1812                vec![
1813                    &mut self.swap_account,
1814                    &mut SolanaAccount::default(),
1815                    &mut SolanaAccount::default(),
1816                    depositor_token_a_account,
1817                    depositor_token_b_account,
1818                    &mut self.token_a_account,
1819                    &mut self.token_b_account,
1820                    &mut self.pool_mint_account,
1821                    depositor_pool_account,
1822                    &mut self.token_a_mint_account,
1823                    &mut self.token_b_mint_account,
1824                    &mut SolanaAccount::default(),
1825                    &mut SolanaAccount::default(),
1826                    &mut SolanaAccount::default(),
1827                ],
1828            )
1829        }
1830
1831        #[allow(clippy::too_many_arguments)]
1832        pub fn withdraw_all_token_types(
1833            &mut self,
1834            user_key: &Pubkey,
1835            pool_key: &Pubkey,
1836            pool_account: &mut SolanaAccount,
1837            token_a_key: &Pubkey,
1838            token_a_account: &mut SolanaAccount,
1839            token_b_key: &Pubkey,
1840            token_b_account: &mut SolanaAccount,
1841            pool_token_amount: u64,
1842            minimum_token_a_amount: u64,
1843            minimum_token_b_amount: u64,
1844        ) -> ProgramResult {
1845            let user_transfer_authority_key = Pubkey::new_unique();
1846            let pool_token_program_id = pool_account.owner;
1847            // approve user transfer authority to take out pool tokens
1848            do_process_instruction(
1849                approve(
1850                    &pool_token_program_id,
1851                    pool_key,
1852                    &user_transfer_authority_key,
1853                    user_key,
1854                    &[],
1855                    pool_token_amount,
1856                )
1857                .unwrap(),
1858                vec![
1859                    pool_account,
1860                    &mut SolanaAccount::default(),
1861                    &mut SolanaAccount::default(),
1862                ],
1863            )
1864            .unwrap();
1865
1866            // withdraw token a and b correctly
1867            let token_a_program_id = token_a_account.owner;
1868            let token_b_program_id = token_b_account.owner;
1869            do_process_instruction(
1870                withdraw_all_token_types(
1871                    &SWAP_PROGRAM_ID,
1872                    &pool_token_program_id,
1873                    &token_a_program_id,
1874                    &token_b_program_id,
1875                    &self.swap_key,
1876                    &self.authority_key,
1877                    &user_transfer_authority_key,
1878                    &self.pool_mint_key,
1879                    &self.pool_fee_key,
1880                    pool_key,
1881                    &self.token_a_key,
1882                    &self.token_b_key,
1883                    token_a_key,
1884                    token_b_key,
1885                    &self.token_a_mint_key,
1886                    &self.token_b_mint_key,
1887                    WithdrawAllTokenTypes {
1888                        pool_token_amount,
1889                        minimum_token_a_amount,
1890                        minimum_token_b_amount,
1891                    },
1892                )
1893                .unwrap(),
1894                vec![
1895                    &mut self.swap_account,
1896                    &mut SolanaAccount::default(),
1897                    &mut SolanaAccount::default(),
1898                    &mut self.pool_mint_account,
1899                    pool_account,
1900                    &mut self.token_a_account,
1901                    &mut self.token_b_account,
1902                    token_a_account,
1903                    token_b_account,
1904                    &mut self.pool_fee_account,
1905                    &mut self.token_a_mint_account,
1906                    &mut self.token_b_mint_account,
1907                    &mut SolanaAccount::default(),
1908                    &mut SolanaAccount::default(),
1909                    &mut SolanaAccount::default(),
1910                ],
1911            )
1912        }
1913
1914        #[allow(clippy::too_many_arguments)]
1915        pub fn deposit_single_token_type_exact_amount_in(
1916            &mut self,
1917            depositor_key: &Pubkey,
1918            deposit_account_key: &Pubkey,
1919            deposit_token_account: &mut SolanaAccount,
1920            deposit_pool_key: &Pubkey,
1921            deposit_pool_account: &mut SolanaAccount,
1922            source_token_amount: u64,
1923            minimum_pool_token_amount: u64,
1924        ) -> ProgramResult {
1925            let user_transfer_authority_key = Pubkey::new_unique();
1926            let source_token_program_id = deposit_token_account.owner;
1927            do_process_instruction(
1928                approve(
1929                    &source_token_program_id,
1930                    deposit_account_key,
1931                    &user_transfer_authority_key,
1932                    depositor_key,
1933                    &[],
1934                    source_token_amount,
1935                )
1936                .unwrap(),
1937                vec![
1938                    deposit_token_account,
1939                    &mut SolanaAccount::default(),
1940                    &mut SolanaAccount::default(),
1941                ],
1942            )
1943            .unwrap();
1944
1945            let source_mint_key =
1946                StateWithExtensions::<Account>::unpack(&deposit_token_account.data)
1947                    .unwrap()
1948                    .base
1949                    .mint;
1950            let swap_source_key = self.get_swap_key(&source_mint_key);
1951            let (source_mint_key, mut source_mint_account) = self.get_token_mint(swap_source_key);
1952
1953            let pool_token_program_id = deposit_pool_account.owner;
1954            do_process_instruction(
1955                deposit_single_token_type_exact_amount_in(
1956                    &SWAP_PROGRAM_ID,
1957                    &source_token_program_id,
1958                    &pool_token_program_id,
1959                    &self.swap_key,
1960                    &self.authority_key,
1961                    &user_transfer_authority_key,
1962                    deposit_account_key,
1963                    &self.token_a_key,
1964                    &self.token_b_key,
1965                    &self.pool_mint_key,
1966                    deposit_pool_key,
1967                    &source_mint_key,
1968                    DepositSingleTokenTypeExactAmountIn {
1969                        source_token_amount,
1970                        minimum_pool_token_amount,
1971                    },
1972                )
1973                .unwrap(),
1974                vec![
1975                    &mut self.swap_account,
1976                    &mut SolanaAccount::default(),
1977                    &mut SolanaAccount::default(),
1978                    deposit_token_account,
1979                    &mut self.token_a_account,
1980                    &mut self.token_b_account,
1981                    &mut self.pool_mint_account,
1982                    deposit_pool_account,
1983                    &mut source_mint_account,
1984                    &mut SolanaAccount::default(),
1985                    &mut SolanaAccount::default(),
1986                ],
1987            )
1988        }
1989
1990        #[allow(clippy::too_many_arguments)]
1991        pub fn withdraw_single_token_type_exact_amount_out(
1992            &mut self,
1993            user_key: &Pubkey,
1994            pool_key: &Pubkey,
1995            pool_account: &mut SolanaAccount,
1996            destination_key: &Pubkey,
1997            destination_account: &mut SolanaAccount,
1998            destination_token_amount: u64,
1999            maximum_pool_token_amount: u64,
2000        ) -> ProgramResult {
2001            let user_transfer_authority_key = Pubkey::new_unique();
2002            let pool_token_program_id = pool_account.owner;
2003            // approve user transfer authority to take out pool tokens
2004            do_process_instruction(
2005                approve(
2006                    &pool_token_program_id,
2007                    pool_key,
2008                    &user_transfer_authority_key,
2009                    user_key,
2010                    &[],
2011                    maximum_pool_token_amount,
2012                )
2013                .unwrap(),
2014                vec![
2015                    pool_account,
2016                    &mut SolanaAccount::default(),
2017                    &mut SolanaAccount::default(),
2018                ],
2019            )
2020            .unwrap();
2021
2022            let destination_mint_key =
2023                StateWithExtensions::<Account>::unpack(&destination_account.data)
2024                    .unwrap()
2025                    .base
2026                    .mint;
2027            let swap_destination_key = self.get_swap_key(&destination_mint_key);
2028            let (destination_mint_key, mut destination_mint_account) =
2029                self.get_token_mint(swap_destination_key);
2030
2031            let destination_token_program_id = destination_account.owner;
2032            do_process_instruction(
2033                withdraw_single_token_type_exact_amount_out(
2034                    &SWAP_PROGRAM_ID,
2035                    &pool_token_program_id,
2036                    &destination_token_program_id,
2037                    &self.swap_key,
2038                    &self.authority_key,
2039                    &user_transfer_authority_key,
2040                    &self.pool_mint_key,
2041                    &self.pool_fee_key,
2042                    pool_key,
2043                    &self.token_a_key,
2044                    &self.token_b_key,
2045                    destination_key,
2046                    &destination_mint_key,
2047                    WithdrawSingleTokenTypeExactAmountOut {
2048                        destination_token_amount,
2049                        maximum_pool_token_amount,
2050                    },
2051                )
2052                .unwrap(),
2053                vec![
2054                    &mut self.swap_account,
2055                    &mut SolanaAccount::default(),
2056                    &mut SolanaAccount::default(),
2057                    &mut self.pool_mint_account,
2058                    pool_account,
2059                    &mut self.token_a_account,
2060                    &mut self.token_b_account,
2061                    destination_account,
2062                    &mut self.pool_fee_account,
2063                    &mut destination_mint_account,
2064                    &mut SolanaAccount::default(),
2065                    &mut SolanaAccount::default(),
2066                ],
2067            )
2068        }
2069    }
2070
2071    fn mint_minimum_balance() -> u64 {
2072        Rent::default().minimum_balance(spl_token::state::Mint::get_packed_len())
2073    }
2074
2075    fn account_minimum_balance() -> u64 {
2076        Rent::default().minimum_balance(spl_token::state::Account::get_packed_len())
2077    }
2078
2079    fn do_process_instruction_with_fee_constraints(
2080        instruction: Instruction,
2081        accounts: Vec<&mut SolanaAccount>,
2082        swap_constraints: &Option<SwapConstraints>,
2083    ) -> ProgramResult {
2084        test_syscall_stubs();
2085
2086        // approximate the logic in the actual runtime which runs the instruction
2087        // and only updates accounts if the instruction is successful
2088        let mut account_clones = accounts.iter().map(|x| (*x).clone()).collect::<Vec<_>>();
2089        let mut meta = instruction
2090            .accounts
2091            .iter()
2092            .zip(account_clones.iter_mut())
2093            .map(|(account_meta, account)| (&account_meta.pubkey, account_meta.is_signer, account))
2094            .collect::<Vec<_>>();
2095        let mut account_infos = create_is_signer_account_infos(&mut meta);
2096        let res = if instruction.program_id == SWAP_PROGRAM_ID {
2097            Processor::process_with_constraints(
2098                &instruction.program_id,
2099                &account_infos,
2100                &instruction.data,
2101                swap_constraints,
2102            )
2103        } else if instruction.program_id == spl_token::id() {
2104            spl_token::processor::Processor::process(
2105                &instruction.program_id,
2106                &account_infos,
2107                &instruction.data,
2108            )
2109        } else if instruction.program_id == spl_token_2022::id() {
2110            spl_token_2022::processor::Processor::process(
2111                &instruction.program_id,
2112                &account_infos,
2113                &instruction.data,
2114            )
2115        } else {
2116            Err(ProgramError::IncorrectProgramId)
2117        };
2118
2119        if res.is_ok() {
2120            let mut account_metas = instruction
2121                .accounts
2122                .iter()
2123                .zip(accounts)
2124                .map(|(account_meta, account)| (&account_meta.pubkey, account))
2125                .collect::<Vec<_>>();
2126            for account_info in account_infos.iter_mut() {
2127                for account_meta in account_metas.iter_mut() {
2128                    if account_info.key == account_meta.0 {
2129                        let account = &mut account_meta.1;
2130                        account.owner = *account_info.owner;
2131                        account.lamports = **account_info.lamports.borrow();
2132                        account.data = account_info.data.borrow().to_vec();
2133                    }
2134                }
2135            }
2136        }
2137        res
2138    }
2139
2140    fn do_process_instruction(
2141        instruction: Instruction,
2142        accounts: Vec<&mut SolanaAccount>,
2143    ) -> ProgramResult {
2144        do_process_instruction_with_fee_constraints(instruction, accounts, &SWAP_CONSTRAINTS)
2145    }
2146
2147    fn mint_token(
2148        program_id: &Pubkey,
2149        mint_key: &Pubkey,
2150        mint_account: &mut SolanaAccount,
2151        mint_authority_key: &Pubkey,
2152        account_owner_key: &Pubkey,
2153        amount: u64,
2154    ) -> (Pubkey, SolanaAccount) {
2155        let account_key = Pubkey::new_unique();
2156        let space = if *program_id == spl_token_2022::id() {
2157            ExtensionType::get_account_len::<Account>(&[
2158                ExtensionType::ImmutableOwner,
2159                ExtensionType::TransferFeeAmount,
2160            ])
2161        } else {
2162            Account::get_packed_len()
2163        };
2164        let minimum_balance = Rent::default().minimum_balance(space);
2165        let mut account_account = SolanaAccount::new(minimum_balance, space, program_id);
2166        let mut mint_authority_account = SolanaAccount::default();
2167        let mut rent_sysvar_account = create_account_for_test(&Rent::free());
2168
2169        // no-ops in normal token, so we're good to run it either way
2170        do_process_instruction(
2171            initialize_immutable_owner(program_id, &account_key).unwrap(),
2172            vec![&mut account_account],
2173        )
2174        .unwrap();
2175
2176        do_process_instruction(
2177            initialize_account(program_id, &account_key, mint_key, account_owner_key).unwrap(),
2178            vec![
2179                &mut account_account,
2180                mint_account,
2181                &mut mint_authority_account,
2182                &mut rent_sysvar_account,
2183            ],
2184        )
2185        .unwrap();
2186
2187        if amount > 0 {
2188            do_process_instruction(
2189                mint_to(
2190                    program_id,
2191                    mint_key,
2192                    &account_key,
2193                    mint_authority_key,
2194                    &[],
2195                    amount,
2196                )
2197                .unwrap(),
2198                vec![
2199                    mint_account,
2200                    &mut account_account,
2201                    &mut mint_authority_account,
2202                ],
2203            )
2204            .unwrap();
2205        }
2206
2207        (account_key, account_account)
2208    }
2209
2210    fn create_mint(
2211        program_id: &Pubkey,
2212        authority_key: &Pubkey,
2213        freeze_authority: Option<&Pubkey>,
2214        close_authority: Option<&Pubkey>,
2215        fees: &TransferFee,
2216    ) -> (Pubkey, SolanaAccount) {
2217        let mint_key = Pubkey::new_unique();
2218        let space = if *program_id == spl_token_2022::id() {
2219            if close_authority.is_some() {
2220                ExtensionType::get_account_len::<Mint>(&[
2221                    ExtensionType::MintCloseAuthority,
2222                    ExtensionType::TransferFeeConfig,
2223                ])
2224            } else {
2225                ExtensionType::get_account_len::<Mint>(&[ExtensionType::TransferFeeConfig])
2226            }
2227        } else {
2228            Mint::get_packed_len()
2229        };
2230        let minimum_balance = Rent::default().minimum_balance(space);
2231        let mut mint_account = SolanaAccount::new(minimum_balance, space, program_id);
2232        let mut rent_sysvar_account = create_account_for_test(&Rent::free());
2233
2234        if *program_id == spl_token_2022::id() {
2235            if close_authority.is_some() {
2236                do_process_instruction(
2237                    initialize_mint_close_authority(program_id, &mint_key, close_authority)
2238                        .unwrap(),
2239                    vec![&mut mint_account],
2240                )
2241                .unwrap();
2242            }
2243            do_process_instruction(
2244                initialize_transfer_fee_config(
2245                    program_id,
2246                    &mint_key,
2247                    freeze_authority,
2248                    freeze_authority,
2249                    fees.transfer_fee_basis_points.into(),
2250                    fees.maximum_fee.into(),
2251                )
2252                .unwrap(),
2253                vec![&mut mint_account],
2254            )
2255            .unwrap();
2256        }
2257        do_process_instruction(
2258            initialize_mint(program_id, &mint_key, authority_key, freeze_authority, 2).unwrap(),
2259            vec![&mut mint_account, &mut rent_sysvar_account],
2260        )
2261        .unwrap();
2262
2263        (mint_key, mint_account)
2264    }
2265
2266    #[test_case(spl_token::id(); "token")]
2267    #[test_case(spl_token_2022::id(); "token-2022")]
2268    fn test_token_program_id_error(token_program_id: Pubkey) {
2269        test_syscall_stubs();
2270        let swap_key = Pubkey::new_unique();
2271        let mut mint = (Pubkey::new_unique(), SolanaAccount::default());
2272        let mut destination = (Pubkey::new_unique(), SolanaAccount::default());
2273        let token_program = (token_program_id, SolanaAccount::default());
2274        let (authority_key, bump_seed) =
2275            Pubkey::find_program_address(&[&swap_key.to_bytes()[..]], &SWAP_PROGRAM_ID);
2276        let mut authority = (authority_key, SolanaAccount::default());
2277        let swap_bytes = swap_key.to_bytes();
2278        let authority_signature_seeds = [&swap_bytes[..32], &[bump_seed]];
2279        let signers = &[&authority_signature_seeds[..]];
2280        let ix = mint_to(
2281            &token_program.0,
2282            &mint.0,
2283            &destination.0,
2284            &authority.0,
2285            &[],
2286            10,
2287        )
2288        .unwrap();
2289        let mint = (&mut mint).into();
2290        let destination = (&mut destination).into();
2291        let authority = (&mut authority).into();
2292
2293        let err = invoke_signed(&ix, &[mint, destination, authority], signers).unwrap_err();
2294        assert_eq!(err, ProgramError::InvalidAccountData);
2295    }
2296
2297    #[test_case(spl_token::id(); "token")]
2298    #[test_case(spl_token_2022::id(); "token-2022")]
2299    fn test_token_error(token_program_id: Pubkey) {
2300        test_syscall_stubs();
2301        let swap_key = Pubkey::new_unique();
2302        let mut mint = (
2303            Pubkey::new_unique(),
2304            SolanaAccount::new(
2305                mint_minimum_balance(),
2306                spl_token::state::Mint::get_packed_len(),
2307                &token_program_id,
2308            ),
2309        );
2310        let mut destination = (
2311            Pubkey::new_unique(),
2312            SolanaAccount::new(
2313                account_minimum_balance(),
2314                spl_token::state::Account::get_packed_len(),
2315                &token_program_id,
2316            ),
2317        );
2318        let mut token_program = (token_program_id, SolanaAccount::default());
2319        let (authority_key, bump_seed) =
2320            Pubkey::find_program_address(&[&swap_key.to_bytes()[..]], &SWAP_PROGRAM_ID);
2321        let mut authority = (authority_key, SolanaAccount::default());
2322        let swap_bytes = swap_key.to_bytes();
2323        let authority_signature_seeds = [&swap_bytes[..32], &[bump_seed]];
2324        let signers = &[&authority_signature_seeds[..]];
2325        let mut rent_sysvar = (
2326            Pubkey::new_unique(),
2327            create_account_for_test(&Rent::default()),
2328        );
2329        do_process_instruction(
2330            initialize_mint(
2331                &token_program.0,
2332                &mint.0,
2333                &authority.0,
2334                Some(&authority.0),
2335                2,
2336            )
2337            .unwrap(),
2338            vec![&mut mint.1, &mut rent_sysvar.1],
2339        )
2340        .unwrap();
2341        do_process_instruction(
2342            initialize_account(&token_program.0, &destination.0, &mint.0, &authority.0).unwrap(),
2343            vec![
2344                &mut destination.1,
2345                &mut mint.1,
2346                &mut authority.1,
2347                &mut rent_sysvar.1,
2348                &mut token_program.1,
2349            ],
2350        )
2351        .unwrap();
2352        do_process_instruction(
2353            freeze_account(&token_program.0, &destination.0, &mint.0, &authority.0, &[]).unwrap(),
2354            vec![
2355                &mut destination.1,
2356                &mut mint.1,
2357                &mut authority.1,
2358                &mut token_program.1,
2359            ],
2360        )
2361        .unwrap();
2362        let ix = mint_to(
2363            &token_program.0,
2364            &mint.0,
2365            &destination.0,
2366            &authority.0,
2367            &[],
2368            10,
2369        )
2370        .unwrap();
2371        let mint_info = (&mut mint).into();
2372        let destination_info = (&mut destination).into();
2373        let authority_info = (&mut authority).into();
2374        let token_program_info = (&mut token_program).into();
2375
2376        let err = invoke_signed_wrapper::<TokenError>(
2377            &ix,
2378            &[
2379                mint_info,
2380                destination_info,
2381                authority_info,
2382                token_program_info,
2383            ],
2384            signers,
2385        )
2386        .unwrap_err();
2387        assert_eq!(err, ProgramError::Custom(TokenError::AccountFrozen as u32));
2388    }
2389
2390    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
2391    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
2392    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
2393    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
2394    fn test_initialize(
2395        pool_token_program_id: Pubkey,
2396        token_a_program_id: Pubkey,
2397        token_b_program_id: Pubkey,
2398    ) {
2399        let user_key = Pubkey::new_unique();
2400        let trade_fee_numerator = 1;
2401        let trade_fee_denominator = 2;
2402        let owner_trade_fee_numerator = 1;
2403        let owner_trade_fee_denominator = 10;
2404        let owner_withdraw_fee_numerator = 1;
2405        let owner_withdraw_fee_denominator = 5;
2406        let host_fee_numerator = 20;
2407        let host_fee_denominator = 100;
2408        let fees = Fees {
2409            trade_fee_numerator,
2410            trade_fee_denominator,
2411            owner_trade_fee_numerator,
2412            owner_trade_fee_denominator,
2413            owner_withdraw_fee_numerator,
2414            owner_withdraw_fee_denominator,
2415            host_fee_numerator,
2416            host_fee_denominator,
2417        };
2418
2419        let token_a_amount = 1000;
2420        let token_b_amount = 2000;
2421        let pool_token_amount = 10;
2422        let curve_type = CurveType::ConstantProduct;
2423        let swap_curve = SwapCurve {
2424            curve_type,
2425            calculator: Arc::new(ConstantProductCurve {}),
2426        };
2427
2428        let mut accounts = SwapAccountInfo::new(
2429            &user_key,
2430            fees,
2431            SwapTransferFees::default(),
2432            swap_curve,
2433            token_a_amount,
2434            token_b_amount,
2435            &pool_token_program_id,
2436            &token_a_program_id,
2437            &token_b_program_id,
2438        );
2439
2440        // uninitialized token a account
2441        {
2442            let old_account = accounts.token_a_account;
2443            accounts.token_a_account = SolanaAccount::new(0, 0, &token_a_program_id);
2444            assert_eq!(
2445                Err(SwapError::ExpectedAccount.into()),
2446                accounts.initialize_swap()
2447            );
2448            accounts.token_a_account = old_account;
2449        }
2450
2451        // uninitialized token b account
2452        {
2453            let old_account = accounts.token_b_account;
2454            accounts.token_b_account = SolanaAccount::new(0, 0, &token_b_program_id);
2455            assert_eq!(
2456                Err(SwapError::ExpectedAccount.into()),
2457                accounts.initialize_swap()
2458            );
2459            accounts.token_b_account = old_account;
2460        }
2461
2462        // uninitialized pool mint
2463        {
2464            let old_account = accounts.pool_mint_account;
2465            accounts.pool_mint_account = SolanaAccount::new(0, 0, &pool_token_program_id);
2466            assert_eq!(
2467                Err(SwapError::ExpectedMint.into()),
2468                accounts.initialize_swap()
2469            );
2470            accounts.pool_mint_account = old_account;
2471        }
2472
2473        // token A account owner is not swap authority
2474        {
2475            let (_token_a_key, token_a_account) = mint_token(
2476                &token_a_program_id,
2477                &accounts.token_a_mint_key,
2478                &mut accounts.token_a_mint_account,
2479                &user_key,
2480                &user_key,
2481                0,
2482            );
2483            let old_account = accounts.token_a_account;
2484            accounts.token_a_account = token_a_account;
2485            assert_eq!(
2486                Err(SwapError::InvalidOwner.into()),
2487                accounts.initialize_swap()
2488            );
2489            accounts.token_a_account = old_account;
2490        }
2491
2492        // token B account owner is not swap authority
2493        {
2494            let (_token_b_key, token_b_account) = mint_token(
2495                &token_b_program_id,
2496                &accounts.token_b_mint_key,
2497                &mut accounts.token_b_mint_account,
2498                &user_key,
2499                &user_key,
2500                0,
2501            );
2502            let old_account = accounts.token_b_account;
2503            accounts.token_b_account = token_b_account;
2504            assert_eq!(
2505                Err(SwapError::InvalidOwner.into()),
2506                accounts.initialize_swap()
2507            );
2508            accounts.token_b_account = old_account;
2509        }
2510
2511        // pool token account owner is swap authority
2512        {
2513            let (_pool_token_key, pool_token_account) = mint_token(
2514                &pool_token_program_id,
2515                &accounts.pool_mint_key,
2516                &mut accounts.pool_mint_account,
2517                &accounts.authority_key,
2518                &accounts.authority_key,
2519                0,
2520            );
2521            let old_account = accounts.pool_token_account;
2522            accounts.pool_token_account = pool_token_account;
2523            assert_eq!(
2524                Err(SwapError::InvalidOutputOwner.into()),
2525                accounts.initialize_swap()
2526            );
2527            accounts.pool_token_account = old_account;
2528        }
2529
2530        // pool fee account owner is swap authority
2531        {
2532            let (_pool_fee_key, pool_fee_account) = mint_token(
2533                &pool_token_program_id,
2534                &accounts.pool_mint_key,
2535                &mut accounts.pool_mint_account,
2536                &accounts.authority_key,
2537                &accounts.authority_key,
2538                0,
2539            );
2540            let old_account = accounts.pool_fee_account;
2541            accounts.pool_fee_account = pool_fee_account;
2542            assert_eq!(
2543                Err(SwapError::InvalidOutputOwner.into()),
2544                accounts.initialize_swap()
2545            );
2546            accounts.pool_fee_account = old_account;
2547        }
2548
2549        // pool mint authority is not swap authority
2550        {
2551            let (_pool_mint_key, pool_mint_account) = create_mint(
2552                &pool_token_program_id,
2553                &user_key,
2554                None,
2555                None,
2556                &TransferFee::default(),
2557            );
2558            let old_mint = accounts.pool_mint_account;
2559            accounts.pool_mint_account = pool_mint_account;
2560            assert_eq!(
2561                Err(SwapError::InvalidOwner.into()),
2562                accounts.initialize_swap()
2563            );
2564            accounts.pool_mint_account = old_mint;
2565        }
2566
2567        // pool mint token has freeze authority
2568        {
2569            let (_pool_mint_key, pool_mint_account) = create_mint(
2570                &pool_token_program_id,
2571                &accounts.authority_key,
2572                Some(&user_key),
2573                None,
2574                &TransferFee::default(),
2575            );
2576            let old_mint = accounts.pool_mint_account;
2577            accounts.pool_mint_account = pool_mint_account;
2578            assert_eq!(
2579                Err(SwapError::InvalidFreezeAuthority.into()),
2580                accounts.initialize_swap()
2581            );
2582            accounts.pool_mint_account = old_mint;
2583        }
2584
2585        // pool mint token has close authority, only available in token-2022
2586        if pool_token_program_id == spl_token_2022::id() {
2587            let (_pool_mint_key, pool_mint_account) = create_mint(
2588                &pool_token_program_id,
2589                &accounts.authority_key,
2590                None,
2591                Some(&user_key),
2592                &TransferFee::default(),
2593            );
2594            let old_mint = accounts.pool_mint_account;
2595            accounts.pool_mint_account = pool_mint_account;
2596            assert_eq!(
2597                Err(SwapError::InvalidCloseAuthority.into()),
2598                accounts.initialize_swap()
2599            );
2600            accounts.pool_mint_account = old_mint;
2601        }
2602
2603        // token A account owned by wrong program
2604        {
2605            let (_token_a_key, mut token_a_account) = mint_token(
2606                &token_a_program_id,
2607                &accounts.token_a_mint_key,
2608                &mut accounts.token_a_mint_account,
2609                &user_key,
2610                &accounts.authority_key,
2611                token_a_amount,
2612            );
2613            token_a_account.owner = SWAP_PROGRAM_ID;
2614            let old_account = accounts.token_a_account;
2615            accounts.token_a_account = token_a_account;
2616            assert_eq!(
2617                Err(SwapError::IncorrectTokenProgramId.into()),
2618                accounts.initialize_swap()
2619            );
2620            accounts.token_a_account = old_account;
2621        }
2622
2623        // token B account owned by wrong program
2624        {
2625            let (_token_b_key, mut token_b_account) = mint_token(
2626                &token_b_program_id,
2627                &accounts.token_b_mint_key,
2628                &mut accounts.token_b_mint_account,
2629                &user_key,
2630                &accounts.authority_key,
2631                token_b_amount,
2632            );
2633            token_b_account.owner = SWAP_PROGRAM_ID;
2634            let old_account = accounts.token_b_account;
2635            accounts.token_b_account = token_b_account;
2636            assert_eq!(
2637                Err(SwapError::IncorrectTokenProgramId.into()),
2638                accounts.initialize_swap()
2639            );
2640            accounts.token_b_account = old_account;
2641        }
2642
2643        // empty token A account
2644        {
2645            let (_token_a_key, token_a_account) = mint_token(
2646                &token_a_program_id,
2647                &accounts.token_a_mint_key,
2648                &mut accounts.token_a_mint_account,
2649                &user_key,
2650                &accounts.authority_key,
2651                0,
2652            );
2653            let old_account = accounts.token_a_account;
2654            accounts.token_a_account = token_a_account;
2655            assert_eq!(
2656                Err(SwapError::EmptySupply.into()),
2657                accounts.initialize_swap()
2658            );
2659            accounts.token_a_account = old_account;
2660        }
2661
2662        // empty token B account
2663        {
2664            let (_token_b_key, token_b_account) = mint_token(
2665                &token_b_program_id,
2666                &accounts.token_b_mint_key,
2667                &mut accounts.token_b_mint_account,
2668                &user_key,
2669                &accounts.authority_key,
2670                0,
2671            );
2672            let old_account = accounts.token_b_account;
2673            accounts.token_b_account = token_b_account;
2674            assert_eq!(
2675                Err(SwapError::EmptySupply.into()),
2676                accounts.initialize_swap()
2677            );
2678            accounts.token_b_account = old_account;
2679        }
2680
2681        // invalid pool tokens
2682        {
2683            let old_mint = accounts.pool_mint_account;
2684            let old_pool_account = accounts.pool_token_account;
2685
2686            let (_pool_mint_key, pool_mint_account) = create_mint(
2687                &pool_token_program_id,
2688                &accounts.authority_key,
2689                None,
2690                None,
2691                &TransferFee::default(),
2692            );
2693            accounts.pool_mint_account = pool_mint_account;
2694
2695            let (_empty_pool_token_key, empty_pool_token_account) = mint_token(
2696                &pool_token_program_id,
2697                &accounts.pool_mint_key,
2698                &mut accounts.pool_mint_account,
2699                &accounts.authority_key,
2700                &user_key,
2701                0,
2702            );
2703
2704            let (_pool_token_key, pool_token_account) = mint_token(
2705                &pool_token_program_id,
2706                &accounts.pool_mint_key,
2707                &mut accounts.pool_mint_account,
2708                &accounts.authority_key,
2709                &user_key,
2710                pool_token_amount,
2711            );
2712
2713            // non-empty pool token account
2714            accounts.pool_token_account = pool_token_account;
2715            assert_eq!(
2716                Err(SwapError::InvalidSupply.into()),
2717                accounts.initialize_swap()
2718            );
2719
2720            // pool tokens already in circulation
2721            accounts.pool_token_account = empty_pool_token_account;
2722            assert_eq!(
2723                Err(SwapError::InvalidSupply.into()),
2724                accounts.initialize_swap()
2725            );
2726
2727            accounts.pool_mint_account = old_mint;
2728            accounts.pool_token_account = old_pool_account;
2729        }
2730
2731        // pool fee account has wrong mint
2732        {
2733            let (_pool_fee_key, pool_fee_account) = mint_token(
2734                &token_a_program_id,
2735                &accounts.token_a_mint_key,
2736                &mut accounts.token_a_mint_account,
2737                &user_key,
2738                &user_key,
2739                0,
2740            );
2741            let old_account = accounts.pool_fee_account;
2742            accounts.pool_fee_account = pool_fee_account;
2743            assert_eq!(
2744                Err(SwapError::IncorrectPoolMint.into()),
2745                accounts.initialize_swap()
2746            );
2747            accounts.pool_fee_account = old_account;
2748        }
2749
2750        // token A account is delegated
2751        {
2752            do_process_instruction(
2753                approve(
2754                    &token_a_program_id,
2755                    &accounts.token_a_key,
2756                    &user_key,
2757                    &accounts.authority_key,
2758                    &[],
2759                    1,
2760                )
2761                .unwrap(),
2762                vec![
2763                    &mut accounts.token_a_account,
2764                    &mut SolanaAccount::default(),
2765                    &mut SolanaAccount::default(),
2766                ],
2767            )
2768            .unwrap();
2769            assert_eq!(
2770                Err(SwapError::InvalidDelegate.into()),
2771                accounts.initialize_swap()
2772            );
2773
2774            do_process_instruction(
2775                revoke(
2776                    &token_a_program_id,
2777                    &accounts.token_a_key,
2778                    &accounts.authority_key,
2779                    &[],
2780                )
2781                .unwrap(),
2782                vec![&mut accounts.token_a_account, &mut SolanaAccount::default()],
2783            )
2784            .unwrap();
2785        }
2786
2787        // token B account is delegated
2788        {
2789            do_process_instruction(
2790                approve(
2791                    &token_b_program_id,
2792                    &accounts.token_b_key,
2793                    &user_key,
2794                    &accounts.authority_key,
2795                    &[],
2796                    1,
2797                )
2798                .unwrap(),
2799                vec![
2800                    &mut accounts.token_b_account,
2801                    &mut SolanaAccount::default(),
2802                    &mut SolanaAccount::default(),
2803                ],
2804            )
2805            .unwrap();
2806            assert_eq!(
2807                Err(SwapError::InvalidDelegate.into()),
2808                accounts.initialize_swap()
2809            );
2810
2811            do_process_instruction(
2812                revoke(
2813                    &token_b_program_id,
2814                    &accounts.token_b_key,
2815                    &accounts.authority_key,
2816                    &[],
2817                )
2818                .unwrap(),
2819                vec![&mut accounts.token_b_account, &mut SolanaAccount::default()],
2820            )
2821            .unwrap();
2822        }
2823
2824        // token A account has close authority
2825        {
2826            do_process_instruction(
2827                set_authority(
2828                    &token_a_program_id,
2829                    &accounts.token_a_key,
2830                    Some(&user_key),
2831                    AuthorityType::CloseAccount,
2832                    &accounts.authority_key,
2833                    &[],
2834                )
2835                .unwrap(),
2836                vec![&mut accounts.token_a_account, &mut SolanaAccount::default()],
2837            )
2838            .unwrap();
2839            assert_eq!(
2840                Err(SwapError::InvalidCloseAuthority.into()),
2841                accounts.initialize_swap()
2842            );
2843
2844            do_process_instruction(
2845                set_authority(
2846                    &token_a_program_id,
2847                    &accounts.token_a_key,
2848                    None,
2849                    AuthorityType::CloseAccount,
2850                    &user_key,
2851                    &[],
2852                )
2853                .unwrap(),
2854                vec![&mut accounts.token_a_account, &mut SolanaAccount::default()],
2855            )
2856            .unwrap();
2857        }
2858
2859        // token B account has close authority
2860        {
2861            do_process_instruction(
2862                set_authority(
2863                    &token_b_program_id,
2864                    &accounts.token_b_key,
2865                    Some(&user_key),
2866                    AuthorityType::CloseAccount,
2867                    &accounts.authority_key,
2868                    &[],
2869                )
2870                .unwrap(),
2871                vec![&mut accounts.token_b_account, &mut SolanaAccount::default()],
2872            )
2873            .unwrap();
2874            assert_eq!(
2875                Err(SwapError::InvalidCloseAuthority.into()),
2876                accounts.initialize_swap()
2877            );
2878
2879            do_process_instruction(
2880                set_authority(
2881                    &token_b_program_id,
2882                    &accounts.token_b_key,
2883                    None,
2884                    AuthorityType::CloseAccount,
2885                    &user_key,
2886                    &[],
2887                )
2888                .unwrap(),
2889                vec![&mut accounts.token_b_account, &mut SolanaAccount::default()],
2890            )
2891            .unwrap();
2892        }
2893
2894        // wrong token program id
2895        {
2896            let wrong_program_id = Pubkey::new_unique();
2897            assert_eq!(
2898                Err(ProgramError::IncorrectProgramId),
2899                do_process_instruction(
2900                    initialize(
2901                        &SWAP_PROGRAM_ID,
2902                        &wrong_program_id,
2903                        &accounts.swap_key,
2904                        &accounts.authority_key,
2905                        &accounts.token_a_key,
2906                        &accounts.token_b_key,
2907                        &accounts.pool_mint_key,
2908                        &accounts.pool_fee_key,
2909                        &accounts.pool_token_key,
2910                        accounts.fees.clone(),
2911                        accounts.swap_curve.clone(),
2912                    )
2913                    .unwrap(),
2914                    vec![
2915                        &mut accounts.swap_account,
2916                        &mut SolanaAccount::default(),
2917                        &mut accounts.token_a_account,
2918                        &mut accounts.token_b_account,
2919                        &mut accounts.pool_mint_account,
2920                        &mut accounts.pool_fee_account,
2921                        &mut accounts.pool_token_account,
2922                        &mut SolanaAccount::default(),
2923                    ],
2924                )
2925            );
2926        }
2927
2928        // create swap with same token A and B
2929        {
2930            let (_token_a_repeat_key, token_a_repeat_account) = mint_token(
2931                &token_a_program_id,
2932                &accounts.token_a_mint_key,
2933                &mut accounts.token_a_mint_account,
2934                &user_key,
2935                &accounts.authority_key,
2936                10,
2937            );
2938            let old_account = accounts.token_b_account;
2939            accounts.token_b_account = token_a_repeat_account;
2940            assert_eq!(
2941                Err(SwapError::RepeatedMint.into()),
2942                accounts.initialize_swap()
2943            );
2944            accounts.token_b_account = old_account;
2945        }
2946
2947        // create valid swap
2948        accounts.initialize_swap().unwrap();
2949
2950        // create invalid flat swap
2951        {
2952            let token_b_price = 0;
2953            let fees = Fees {
2954                trade_fee_numerator,
2955                trade_fee_denominator,
2956                owner_trade_fee_numerator,
2957                owner_trade_fee_denominator,
2958                owner_withdraw_fee_numerator,
2959                owner_withdraw_fee_denominator,
2960                host_fee_numerator,
2961                host_fee_denominator,
2962            };
2963            let swap_curve = SwapCurve {
2964                curve_type: CurveType::ConstantPrice,
2965                calculator: Arc::new(ConstantPriceCurve { token_b_price }),
2966            };
2967            let mut accounts = SwapAccountInfo::new(
2968                &user_key,
2969                fees,
2970                SwapTransferFees::default(),
2971                swap_curve,
2972                token_a_amount,
2973                token_b_amount,
2974                &pool_token_program_id,
2975                &token_a_program_id,
2976                &token_b_program_id,
2977            );
2978            assert_eq!(
2979                Err(SwapError::InvalidCurve.into()),
2980                accounts.initialize_swap()
2981            );
2982        }
2983
2984        // create valid flat swap
2985        {
2986            let fees = Fees {
2987                trade_fee_numerator,
2988                trade_fee_denominator,
2989                owner_trade_fee_numerator,
2990                owner_trade_fee_denominator,
2991                owner_withdraw_fee_numerator,
2992                owner_withdraw_fee_denominator,
2993                host_fee_numerator,
2994                host_fee_denominator,
2995            };
2996            let token_b_price = 10_000;
2997            let swap_curve = SwapCurve {
2998                curve_type: CurveType::ConstantPrice,
2999                calculator: Arc::new(ConstantPriceCurve { token_b_price }),
3000            };
3001            let mut accounts = SwapAccountInfo::new(
3002                &user_key,
3003                fees,
3004                SwapTransferFees::default(),
3005                swap_curve,
3006                token_a_amount,
3007                token_b_amount,
3008                &pool_token_program_id,
3009                &token_a_program_id,
3010                &token_b_program_id,
3011            );
3012            accounts.initialize_swap().unwrap();
3013        }
3014
3015        // create invalid offset swap
3016        {
3017            let token_b_offset = 0;
3018            let fees = Fees {
3019                trade_fee_numerator,
3020                trade_fee_denominator,
3021                owner_trade_fee_numerator,
3022                owner_trade_fee_denominator,
3023                owner_withdraw_fee_numerator,
3024                owner_withdraw_fee_denominator,
3025                host_fee_numerator,
3026                host_fee_denominator,
3027            };
3028            let swap_curve = SwapCurve {
3029                curve_type: CurveType::Offset,
3030                calculator: Arc::new(OffsetCurve { token_b_offset }),
3031            };
3032            let mut accounts = SwapAccountInfo::new(
3033                &user_key,
3034                fees,
3035                SwapTransferFees::default(),
3036                swap_curve,
3037                token_a_amount,
3038                token_b_amount,
3039                &pool_token_program_id,
3040                &token_a_program_id,
3041                &token_b_program_id,
3042            );
3043            assert_eq!(
3044                Err(SwapError::InvalidCurve.into()),
3045                accounts.initialize_swap()
3046            );
3047        }
3048
3049        // create valid offset swap
3050        {
3051            let token_b_offset = 10;
3052            let fees = Fees {
3053                trade_fee_numerator,
3054                trade_fee_denominator,
3055                owner_trade_fee_numerator,
3056                owner_trade_fee_denominator,
3057                owner_withdraw_fee_numerator,
3058                owner_withdraw_fee_denominator,
3059                host_fee_numerator,
3060                host_fee_denominator,
3061            };
3062            let swap_curve = SwapCurve {
3063                curve_type: CurveType::Offset,
3064                calculator: Arc::new(OffsetCurve { token_b_offset }),
3065            };
3066            let mut accounts = SwapAccountInfo::new(
3067                &user_key,
3068                fees,
3069                SwapTransferFees::default(),
3070                swap_curve,
3071                token_a_amount,
3072                token_b_amount,
3073                &pool_token_program_id,
3074                &token_a_program_id,
3075                &token_b_program_id,
3076            );
3077            accounts.initialize_swap().unwrap();
3078        }
3079
3080        // wrong owner key in constraint
3081        {
3082            let new_key = Pubkey::new_unique();
3083            let trade_fee_numerator = 25;
3084            let trade_fee_denominator = 10000;
3085            let owner_trade_fee_numerator = 5;
3086            let owner_trade_fee_denominator = 10000;
3087            let host_fee_numerator = 20;
3088            let host_fee_denominator = 100;
3089            let fees = Fees {
3090                trade_fee_numerator,
3091                trade_fee_denominator,
3092                owner_trade_fee_numerator,
3093                owner_trade_fee_denominator,
3094                owner_withdraw_fee_numerator,
3095                owner_withdraw_fee_denominator,
3096                host_fee_numerator,
3097                host_fee_denominator,
3098            };
3099            let curve = ConstantProductCurve {};
3100            let swap_curve = SwapCurve {
3101                curve_type: CurveType::ConstantProduct,
3102                calculator: Arc::new(curve),
3103            };
3104            let owner_key = &new_key.to_string();
3105            let valid_curve_types = &[CurveType::ConstantProduct];
3106            let constraints = Some(SwapConstraints {
3107                owner_key,
3108                valid_curve_types,
3109                fees: &fees,
3110            });
3111            let mut accounts = SwapAccountInfo::new(
3112                &user_key,
3113                fees.clone(),
3114                SwapTransferFees::default(),
3115                swap_curve,
3116                token_a_amount,
3117                token_b_amount,
3118                &pool_token_program_id,
3119                &token_a_program_id,
3120                &token_b_program_id,
3121            );
3122            assert_eq!(
3123                Err(SwapError::InvalidOwner.into()),
3124                do_process_instruction_with_fee_constraints(
3125                    initialize(
3126                        &SWAP_PROGRAM_ID,
3127                        &pool_token_program_id,
3128                        &accounts.swap_key,
3129                        &accounts.authority_key,
3130                        &accounts.token_a_key,
3131                        &accounts.token_b_key,
3132                        &accounts.pool_mint_key,
3133                        &accounts.pool_fee_key,
3134                        &accounts.pool_token_key,
3135                        accounts.fees.clone(),
3136                        accounts.swap_curve.clone(),
3137                    )
3138                    .unwrap(),
3139                    vec![
3140                        &mut accounts.swap_account,
3141                        &mut SolanaAccount::default(),
3142                        &mut accounts.token_a_account,
3143                        &mut accounts.token_b_account,
3144                        &mut accounts.pool_mint_account,
3145                        &mut accounts.pool_fee_account,
3146                        &mut accounts.pool_token_account,
3147                        &mut SolanaAccount::default(),
3148                    ],
3149                    &constraints,
3150                )
3151            );
3152        }
3153
3154        // wrong fee in constraint
3155        {
3156            let trade_fee_numerator = 25;
3157            let trade_fee_denominator = 10000;
3158            let owner_trade_fee_numerator = 5;
3159            let owner_trade_fee_denominator = 10000;
3160            let host_fee_numerator = 20;
3161            let host_fee_denominator = 100;
3162            let fees = Fees {
3163                trade_fee_numerator,
3164                trade_fee_denominator,
3165                owner_trade_fee_numerator,
3166                owner_trade_fee_denominator,
3167                owner_withdraw_fee_numerator,
3168                owner_withdraw_fee_denominator,
3169                host_fee_numerator,
3170                host_fee_denominator,
3171            };
3172            let curve = ConstantProductCurve {};
3173            let swap_curve = SwapCurve {
3174                curve_type: CurveType::ConstantProduct,
3175                calculator: Arc::new(curve),
3176            };
3177            let owner_key = &user_key.to_string();
3178            let valid_curve_types = &[CurveType::ConstantProduct];
3179            let constraints = Some(SwapConstraints {
3180                owner_key,
3181                valid_curve_types,
3182                fees: &fees,
3183            });
3184            let mut bad_fees = fees.clone();
3185            bad_fees.trade_fee_numerator = trade_fee_numerator - 1;
3186            let mut accounts = SwapAccountInfo::new(
3187                &user_key,
3188                bad_fees,
3189                SwapTransferFees::default(),
3190                swap_curve,
3191                token_a_amount,
3192                token_b_amount,
3193                &pool_token_program_id,
3194                &token_a_program_id,
3195                &token_b_program_id,
3196            );
3197            assert_eq!(
3198                Err(SwapError::InvalidFee.into()),
3199                do_process_instruction_with_fee_constraints(
3200                    initialize(
3201                        &SWAP_PROGRAM_ID,
3202                        &pool_token_program_id,
3203                        &accounts.swap_key,
3204                        &accounts.authority_key,
3205                        &accounts.token_a_key,
3206                        &accounts.token_b_key,
3207                        &accounts.pool_mint_key,
3208                        &accounts.pool_fee_key,
3209                        &accounts.pool_token_key,
3210                        accounts.fees.clone(),
3211                        accounts.swap_curve.clone(),
3212                    )
3213                    .unwrap(),
3214                    vec![
3215                        &mut accounts.swap_account,
3216                        &mut SolanaAccount::default(),
3217                        &mut accounts.token_a_account,
3218                        &mut accounts.token_b_account,
3219                        &mut accounts.pool_mint_account,
3220                        &mut accounts.pool_fee_account,
3221                        &mut accounts.pool_token_account,
3222                        &mut SolanaAccount::default(),
3223                    ],
3224                    &constraints,
3225                )
3226            );
3227        }
3228
3229        // create valid swap with constraints
3230        {
3231            let trade_fee_numerator = 25;
3232            let trade_fee_denominator = 10000;
3233            let owner_trade_fee_numerator = 5;
3234            let owner_trade_fee_denominator = 10000;
3235            let host_fee_numerator = 20;
3236            let host_fee_denominator = 100;
3237            let fees = Fees {
3238                trade_fee_numerator,
3239                trade_fee_denominator,
3240                owner_trade_fee_numerator,
3241                owner_trade_fee_denominator,
3242                owner_withdraw_fee_numerator,
3243                owner_withdraw_fee_denominator,
3244                host_fee_numerator,
3245                host_fee_denominator,
3246            };
3247            let curve = ConstantProductCurve {};
3248            let swap_curve = SwapCurve {
3249                curve_type: CurveType::ConstantProduct,
3250                calculator: Arc::new(curve),
3251            };
3252            let owner_key = &user_key.to_string();
3253            let valid_curve_types = &[CurveType::ConstantProduct];
3254            let constraints = Some(SwapConstraints {
3255                owner_key,
3256                valid_curve_types,
3257                fees: &fees,
3258            });
3259            let mut accounts = SwapAccountInfo::new(
3260                &user_key,
3261                fees.clone(),
3262                SwapTransferFees::default(),
3263                swap_curve,
3264                token_a_amount,
3265                token_b_amount,
3266                &pool_token_program_id,
3267                &token_a_program_id,
3268                &token_b_program_id,
3269            );
3270            do_process_instruction_with_fee_constraints(
3271                initialize(
3272                    &SWAP_PROGRAM_ID,
3273                    &pool_token_program_id,
3274                    &accounts.swap_key,
3275                    &accounts.authority_key,
3276                    &accounts.token_a_key,
3277                    &accounts.token_b_key,
3278                    &accounts.pool_mint_key,
3279                    &accounts.pool_fee_key,
3280                    &accounts.pool_token_key,
3281                    accounts.fees,
3282                    accounts.swap_curve.clone(),
3283                )
3284                .unwrap(),
3285                vec![
3286                    &mut accounts.swap_account,
3287                    &mut SolanaAccount::default(),
3288                    &mut accounts.token_a_account,
3289                    &mut accounts.token_b_account,
3290                    &mut accounts.pool_mint_account,
3291                    &mut accounts.pool_fee_account,
3292                    &mut accounts.pool_token_account,
3293                    &mut SolanaAccount::default(),
3294                ],
3295                &constraints,
3296            )
3297            .unwrap();
3298        }
3299
3300        // create again
3301        {
3302            assert_eq!(
3303                Err(SwapError::AlreadyInUse.into()),
3304                accounts.initialize_swap()
3305            );
3306        }
3307        let swap_state = SwapVersion::unpack(&accounts.swap_account.data).unwrap();
3308        assert!(swap_state.is_initialized());
3309        assert_eq!(swap_state.bump_seed(), accounts.bump_seed);
3310        assert_eq!(
3311            swap_state.swap_curve().curve_type,
3312            accounts.swap_curve.curve_type
3313        );
3314        assert_eq!(*swap_state.token_a_account(), accounts.token_a_key);
3315        assert_eq!(*swap_state.token_b_account(), accounts.token_b_key);
3316        assert_eq!(*swap_state.pool_mint(), accounts.pool_mint_key);
3317        assert_eq!(*swap_state.token_a_mint(), accounts.token_a_mint_key);
3318        assert_eq!(*swap_state.token_b_mint(), accounts.token_b_mint_key);
3319        assert_eq!(*swap_state.pool_fee_account(), accounts.pool_fee_key);
3320        let token_a =
3321            StateWithExtensions::<Account>::unpack(&accounts.token_a_account.data).unwrap();
3322        assert_eq!(token_a.base.amount, token_a_amount);
3323        let token_b =
3324            StateWithExtensions::<Account>::unpack(&accounts.token_b_account.data).unwrap();
3325        assert_eq!(token_b.base.amount, token_b_amount);
3326        let pool_account =
3327            StateWithExtensions::<Account>::unpack(&accounts.pool_token_account.data).unwrap();
3328        let pool_mint =
3329            StateWithExtensions::<Mint>::unpack(&accounts.pool_mint_account.data).unwrap();
3330        assert_eq!(pool_mint.base.supply, pool_account.base.amount);
3331    }
3332
3333    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
3334    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
3335    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
3336    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
3337    fn test_deposit(
3338        pool_token_program_id: Pubkey,
3339        token_a_program_id: Pubkey,
3340        token_b_program_id: Pubkey,
3341    ) {
3342        let user_key = Pubkey::new_unique();
3343        let depositor_key = Pubkey::new_unique();
3344        let trade_fee_numerator = 1;
3345        let trade_fee_denominator = 2;
3346        let owner_trade_fee_numerator = 1;
3347        let owner_trade_fee_denominator = 10;
3348        let owner_withdraw_fee_numerator = 1;
3349        let owner_withdraw_fee_denominator = 5;
3350        let host_fee_numerator = 20;
3351        let host_fee_denominator = 100;
3352
3353        let fees = Fees {
3354            trade_fee_numerator,
3355            trade_fee_denominator,
3356            owner_trade_fee_numerator,
3357            owner_trade_fee_denominator,
3358            owner_withdraw_fee_numerator,
3359            owner_withdraw_fee_denominator,
3360            host_fee_numerator,
3361            host_fee_denominator,
3362        };
3363
3364        let token_a_amount = 1000;
3365        let token_b_amount = 9000;
3366        let curve_type = CurveType::ConstantProduct;
3367        let swap_curve = SwapCurve {
3368            curve_type,
3369            calculator: Arc::new(ConstantProductCurve {}),
3370        };
3371
3372        let mut accounts = SwapAccountInfo::new(
3373            &user_key,
3374            fees,
3375            SwapTransferFees::default(),
3376            swap_curve,
3377            token_a_amount,
3378            token_b_amount,
3379            &pool_token_program_id,
3380            &token_a_program_id,
3381            &token_b_program_id,
3382        );
3383
3384        // depositing 10% of the current pool amount in token A and B means
3385        // that our pool tokens will be worth 1 / 10 of the current pool amount
3386        let pool_amount = INITIAL_SWAP_POOL_AMOUNT / 10;
3387        let deposit_a = token_a_amount / 10;
3388        let deposit_b = token_b_amount / 10;
3389
3390        // swap not initialized
3391        {
3392            let (
3393                token_a_key,
3394                mut token_a_account,
3395                token_b_key,
3396                mut token_b_account,
3397                pool_key,
3398                mut pool_account,
3399            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3400            assert_eq!(
3401                Err(ProgramError::UninitializedAccount),
3402                accounts.deposit_all_token_types(
3403                    &depositor_key,
3404                    &token_a_key,
3405                    &mut token_a_account,
3406                    &token_b_key,
3407                    &mut token_b_account,
3408                    &pool_key,
3409                    &mut pool_account,
3410                    pool_amount.try_into().unwrap(),
3411                    deposit_a,
3412                    deposit_b,
3413                )
3414            );
3415        }
3416
3417        accounts.initialize_swap().unwrap();
3418
3419        // wrong owner for swap account
3420        {
3421            let (
3422                token_a_key,
3423                mut token_a_account,
3424                token_b_key,
3425                mut token_b_account,
3426                pool_key,
3427                mut pool_account,
3428            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3429            let old_swap_account = accounts.swap_account;
3430            let mut wrong_swap_account = old_swap_account.clone();
3431            wrong_swap_account.owner = pool_token_program_id;
3432            accounts.swap_account = wrong_swap_account;
3433            assert_eq!(
3434                Err(ProgramError::IncorrectProgramId),
3435                accounts.deposit_all_token_types(
3436                    &depositor_key,
3437                    &token_a_key,
3438                    &mut token_a_account,
3439                    &token_b_key,
3440                    &mut token_b_account,
3441                    &pool_key,
3442                    &mut pool_account,
3443                    pool_amount.try_into().unwrap(),
3444                    deposit_a,
3445                    deposit_b,
3446                )
3447            );
3448            accounts.swap_account = old_swap_account;
3449        }
3450
3451        // wrong bump seed for authority_key
3452        {
3453            let (
3454                token_a_key,
3455                mut token_a_account,
3456                token_b_key,
3457                mut token_b_account,
3458                pool_key,
3459                mut pool_account,
3460            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3461            let old_authority = accounts.authority_key;
3462            let (bad_authority_key, _bump_seed) = Pubkey::find_program_address(
3463                &[&accounts.swap_key.to_bytes()[..]],
3464                &pool_token_program_id,
3465            );
3466            accounts.authority_key = bad_authority_key;
3467            assert_eq!(
3468                Err(SwapError::InvalidProgramAddress.into()),
3469                accounts.deposit_all_token_types(
3470                    &depositor_key,
3471                    &token_a_key,
3472                    &mut token_a_account,
3473                    &token_b_key,
3474                    &mut token_b_account,
3475                    &pool_key,
3476                    &mut pool_account,
3477                    pool_amount.try_into().unwrap(),
3478                    deposit_a,
3479                    deposit_b,
3480                )
3481            );
3482            accounts.authority_key = old_authority;
3483        }
3484
3485        // not enough token A
3486        {
3487            let (
3488                token_a_key,
3489                mut token_a_account,
3490                token_b_key,
3491                mut token_b_account,
3492                pool_key,
3493                mut pool_account,
3494            ) = accounts.setup_token_accounts(
3495                &user_key,
3496                &depositor_key,
3497                deposit_a / 2,
3498                deposit_b,
3499                0,
3500            );
3501            assert_eq!(
3502                Err(TokenError::InsufficientFunds.into()),
3503                accounts.deposit_all_token_types(
3504                    &depositor_key,
3505                    &token_a_key,
3506                    &mut token_a_account,
3507                    &token_b_key,
3508                    &mut token_b_account,
3509                    &pool_key,
3510                    &mut pool_account,
3511                    pool_amount.try_into().unwrap(),
3512                    deposit_a,
3513                    deposit_b,
3514                )
3515            );
3516        }
3517
3518        // not enough token B
3519        {
3520            let (
3521                token_a_key,
3522                mut token_a_account,
3523                token_b_key,
3524                mut token_b_account,
3525                pool_key,
3526                mut pool_account,
3527            ) = accounts.setup_token_accounts(
3528                &user_key,
3529                &depositor_key,
3530                deposit_a,
3531                deposit_b / 2,
3532                0,
3533            );
3534            assert_eq!(
3535                Err(TokenError::InsufficientFunds.into()),
3536                accounts.deposit_all_token_types(
3537                    &depositor_key,
3538                    &token_a_key,
3539                    &mut token_a_account,
3540                    &token_b_key,
3541                    &mut token_b_account,
3542                    &pool_key,
3543                    &mut pool_account,
3544                    pool_amount.try_into().unwrap(),
3545                    deposit_a,
3546                    deposit_b,
3547                )
3548            );
3549        }
3550
3551        // wrong swap token accounts
3552        {
3553            let (
3554                token_a_key,
3555                mut token_a_account,
3556                token_b_key,
3557                mut token_b_account,
3558                pool_key,
3559                mut pool_account,
3560            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3561            let expected_error: ProgramError = if token_a_account.owner == token_b_account.owner {
3562                TokenError::MintMismatch.into()
3563            } else {
3564                ProgramError::InvalidAccountData
3565            };
3566            assert_eq!(
3567                Err(expected_error),
3568                accounts.deposit_all_token_types(
3569                    &depositor_key,
3570                    &token_b_key,
3571                    &mut token_b_account,
3572                    &token_a_key,
3573                    &mut token_a_account,
3574                    &pool_key,
3575                    &mut pool_account,
3576                    pool_amount.try_into().unwrap(),
3577                    deposit_a,
3578                    deposit_b,
3579                )
3580            );
3581        }
3582
3583        // wrong pool token account
3584        {
3585            let (
3586                token_a_key,
3587                mut token_a_account,
3588                token_b_key,
3589                mut token_b_account,
3590                _pool_key,
3591                mut _pool_account,
3592            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3593            let (
3594                wrong_token_key,
3595                mut wrong_token_account,
3596                _token_b_key,
3597                mut _token_b_account,
3598                _pool_key,
3599                pool_account,
3600            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3601            let expected_error: ProgramError = if token_a_account.owner == pool_account.owner {
3602                TokenError::MintMismatch.into()
3603            } else {
3604                SwapError::IncorrectTokenProgramId.into()
3605            };
3606            assert_eq!(
3607                Err(expected_error),
3608                accounts.deposit_all_token_types(
3609                    &depositor_key,
3610                    &token_a_key,
3611                    &mut token_a_account,
3612                    &token_b_key,
3613                    &mut token_b_account,
3614                    &wrong_token_key,
3615                    &mut wrong_token_account,
3616                    pool_amount.try_into().unwrap(),
3617                    deposit_a,
3618                    deposit_b,
3619                )
3620            );
3621        }
3622
3623        // no approval
3624        {
3625            let (
3626                token_a_key,
3627                mut token_a_account,
3628                token_b_key,
3629                mut token_b_account,
3630                pool_key,
3631                mut pool_account,
3632            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3633            let user_transfer_authority_key = Pubkey::new_unique();
3634            assert_eq!(
3635                Err(TokenError::OwnerMismatch.into()),
3636                do_process_instruction(
3637                    deposit_all_token_types(
3638                        &SWAP_PROGRAM_ID,
3639                        &token_a_program_id,
3640                        &token_b_program_id,
3641                        &pool_token_program_id,
3642                        &accounts.swap_key,
3643                        &accounts.authority_key,
3644                        &user_transfer_authority_key,
3645                        &token_a_key,
3646                        &token_b_key,
3647                        &accounts.token_a_key,
3648                        &accounts.token_b_key,
3649                        &accounts.pool_mint_key,
3650                        &pool_key,
3651                        &accounts.token_a_mint_key,
3652                        &accounts.token_b_mint_key,
3653                        DepositAllTokenTypes {
3654                            pool_token_amount: pool_amount.try_into().unwrap(),
3655                            maximum_token_a_amount: deposit_a,
3656                            maximum_token_b_amount: deposit_b,
3657                        },
3658                    )
3659                    .unwrap(),
3660                    vec![
3661                        &mut accounts.swap_account,
3662                        &mut SolanaAccount::default(),
3663                        &mut SolanaAccount::default(),
3664                        &mut token_a_account,
3665                        &mut token_b_account,
3666                        &mut accounts.token_a_account,
3667                        &mut accounts.token_b_account,
3668                        &mut accounts.pool_mint_account,
3669                        &mut pool_account,
3670                        &mut accounts.token_a_mint_account,
3671                        &mut accounts.token_b_mint_account,
3672                        &mut SolanaAccount::default(),
3673                        &mut SolanaAccount::default(),
3674                        &mut SolanaAccount::default(),
3675                    ],
3676                )
3677            );
3678        }
3679
3680        // wrong token program id
3681        {
3682            let (
3683                token_a_key,
3684                mut token_a_account,
3685                token_b_key,
3686                mut token_b_account,
3687                pool_key,
3688                mut pool_account,
3689            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3690            let wrong_key = Pubkey::new_unique();
3691            assert_eq!(
3692                Err(SwapError::IncorrectTokenProgramId.into()),
3693                do_process_instruction(
3694                    deposit_all_token_types(
3695                        &SWAP_PROGRAM_ID,
3696                        &wrong_key,
3697                        &wrong_key,
3698                        &wrong_key,
3699                        &accounts.swap_key,
3700                        &accounts.authority_key,
3701                        &accounts.authority_key,
3702                        &token_a_key,
3703                        &token_b_key,
3704                        &accounts.token_a_key,
3705                        &accounts.token_b_key,
3706                        &accounts.pool_mint_key,
3707                        &pool_key,
3708                        &accounts.token_a_mint_key,
3709                        &accounts.token_b_mint_key,
3710                        DepositAllTokenTypes {
3711                            pool_token_amount: pool_amount.try_into().unwrap(),
3712                            maximum_token_a_amount: deposit_a,
3713                            maximum_token_b_amount: deposit_b,
3714                        },
3715                    )
3716                    .unwrap(),
3717                    vec![
3718                        &mut accounts.swap_account,
3719                        &mut SolanaAccount::default(),
3720                        &mut SolanaAccount::default(),
3721                        &mut token_a_account,
3722                        &mut token_b_account,
3723                        &mut accounts.token_a_account,
3724                        &mut accounts.token_b_account,
3725                        &mut accounts.pool_mint_account,
3726                        &mut pool_account,
3727                        &mut accounts.token_a_mint_account,
3728                        &mut accounts.token_b_mint_account,
3729                        &mut SolanaAccount::default(),
3730                        &mut SolanaAccount::default(),
3731                        &mut SolanaAccount::default(),
3732                    ],
3733                )
3734            );
3735        }
3736
3737        // wrong swap token accounts
3738        {
3739            let (
3740                token_a_key,
3741                mut token_a_account,
3742                token_b_key,
3743                mut token_b_account,
3744                pool_key,
3745                mut pool_account,
3746            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3747
3748            let old_a_key = accounts.token_a_key;
3749            let old_a_account = accounts.token_a_account;
3750
3751            accounts.token_a_key = token_a_key;
3752            accounts.token_a_account = token_a_account.clone();
3753
3754            // wrong swap token a account
3755            assert_eq!(
3756                Err(SwapError::IncorrectSwapAccount.into()),
3757                accounts.deposit_all_token_types(
3758                    &depositor_key,
3759                    &token_a_key,
3760                    &mut token_a_account,
3761                    &token_b_key,
3762                    &mut token_b_account,
3763                    &pool_key,
3764                    &mut pool_account,
3765                    pool_amount.try_into().unwrap(),
3766                    deposit_a,
3767                    deposit_b,
3768                )
3769            );
3770
3771            accounts.token_a_key = old_a_key;
3772            accounts.token_a_account = old_a_account;
3773
3774            let old_b_key = accounts.token_b_key;
3775            let old_b_account = accounts.token_b_account;
3776
3777            accounts.token_b_key = token_b_key;
3778            accounts.token_b_account = token_b_account.clone();
3779
3780            // wrong swap token b account
3781            assert_eq!(
3782                Err(SwapError::IncorrectSwapAccount.into()),
3783                accounts.deposit_all_token_types(
3784                    &depositor_key,
3785                    &token_a_key,
3786                    &mut token_a_account,
3787                    &token_b_key,
3788                    &mut token_b_account,
3789                    &pool_key,
3790                    &mut pool_account,
3791                    pool_amount.try_into().unwrap(),
3792                    deposit_a,
3793                    deposit_b,
3794                )
3795            );
3796
3797            accounts.token_b_key = old_b_key;
3798            accounts.token_b_account = old_b_account;
3799        }
3800
3801        // wrong mint
3802        {
3803            let (
3804                token_a_key,
3805                mut token_a_account,
3806                token_b_key,
3807                mut token_b_account,
3808                pool_key,
3809                mut pool_account,
3810            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3811            let (pool_mint_key, pool_mint_account) = create_mint(
3812                &pool_token_program_id,
3813                &accounts.authority_key,
3814                None,
3815                None,
3816                &TransferFee::default(),
3817            );
3818            let old_pool_key = accounts.pool_mint_key;
3819            let old_pool_account = accounts.pool_mint_account;
3820            accounts.pool_mint_key = pool_mint_key;
3821            accounts.pool_mint_account = pool_mint_account;
3822
3823            assert_eq!(
3824                Err(SwapError::IncorrectPoolMint.into()),
3825                accounts.deposit_all_token_types(
3826                    &depositor_key,
3827                    &token_a_key,
3828                    &mut token_a_account,
3829                    &token_b_key,
3830                    &mut token_b_account,
3831                    &pool_key,
3832                    &mut pool_account,
3833                    pool_amount.try_into().unwrap(),
3834                    deposit_a,
3835                    deposit_b,
3836                )
3837            );
3838
3839            accounts.pool_mint_key = old_pool_key;
3840            accounts.pool_mint_account = old_pool_account;
3841        }
3842
3843        // deposit 1 pool token fails beacuse it equates to 0 swap tokens
3844        {
3845            let (
3846                token_a_key,
3847                mut token_a_account,
3848                token_b_key,
3849                mut token_b_account,
3850                pool_key,
3851                mut pool_account,
3852            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3853            assert_eq!(
3854                Err(SwapError::ZeroTradingTokens.into()),
3855                accounts.deposit_all_token_types(
3856                    &depositor_key,
3857                    &token_a_key,
3858                    &mut token_a_account,
3859                    &token_b_key,
3860                    &mut token_b_account,
3861                    &pool_key,
3862                    &mut pool_account,
3863                    1,
3864                    deposit_a,
3865                    deposit_b,
3866                )
3867            );
3868        }
3869
3870        // slippage exceeded
3871        {
3872            let (
3873                token_a_key,
3874                mut token_a_account,
3875                token_b_key,
3876                mut token_b_account,
3877                pool_key,
3878                mut pool_account,
3879            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3880            // maximum A amount in too low
3881            assert_eq!(
3882                Err(SwapError::ExceededSlippage.into()),
3883                accounts.deposit_all_token_types(
3884                    &depositor_key,
3885                    &token_a_key,
3886                    &mut token_a_account,
3887                    &token_b_key,
3888                    &mut token_b_account,
3889                    &pool_key,
3890                    &mut pool_account,
3891                    pool_amount.try_into().unwrap(),
3892                    deposit_a / 10,
3893                    deposit_b,
3894                )
3895            );
3896            // maximum B amount in too low
3897            assert_eq!(
3898                Err(SwapError::ExceededSlippage.into()),
3899                accounts.deposit_all_token_types(
3900                    &depositor_key,
3901                    &token_a_key,
3902                    &mut token_a_account,
3903                    &token_b_key,
3904                    &mut token_b_account,
3905                    &pool_key,
3906                    &mut pool_account,
3907                    pool_amount.try_into().unwrap(),
3908                    deposit_a,
3909                    deposit_b / 10,
3910                )
3911            );
3912        }
3913
3914        // invalid input: can't use swap pool tokens as source
3915        {
3916            let (
3917                _token_a_key,
3918                _token_a_account,
3919                _token_b_key,
3920                _token_b_account,
3921                pool_key,
3922                mut pool_account,
3923            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3924            let swap_token_a_key = accounts.token_a_key;
3925            let mut swap_token_a_account = accounts.get_token_account(&swap_token_a_key).clone();
3926            let swap_token_b_key = accounts.token_b_key;
3927            let mut swap_token_b_account = accounts.get_token_account(&swap_token_b_key).clone();
3928            let authority_key = accounts.authority_key;
3929            assert_eq!(
3930                Err(SwapError::InvalidInput.into()),
3931                accounts.deposit_all_token_types(
3932                    &authority_key,
3933                    &swap_token_a_key,
3934                    &mut swap_token_a_account,
3935                    &swap_token_b_key,
3936                    &mut swap_token_b_account,
3937                    &pool_key,
3938                    &mut pool_account,
3939                    pool_amount.try_into().unwrap(),
3940                    deposit_a,
3941                    deposit_b,
3942                )
3943            );
3944        }
3945
3946        // correctly deposit
3947        {
3948            let (
3949                token_a_key,
3950                mut token_a_account,
3951                token_b_key,
3952                mut token_b_account,
3953                pool_key,
3954                mut pool_account,
3955            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3956            accounts
3957                .deposit_all_token_types(
3958                    &depositor_key,
3959                    &token_a_key,
3960                    &mut token_a_account,
3961                    &token_b_key,
3962                    &mut token_b_account,
3963                    &pool_key,
3964                    &mut pool_account,
3965                    pool_amount.try_into().unwrap(),
3966                    deposit_a,
3967                    deposit_b,
3968                )
3969                .unwrap();
3970
3971            let swap_token_a =
3972                StateWithExtensions::<Account>::unpack(&accounts.token_a_account.data).unwrap();
3973            assert_eq!(swap_token_a.base.amount, deposit_a + token_a_amount);
3974            let swap_token_b =
3975                StateWithExtensions::<Account>::unpack(&accounts.token_b_account.data).unwrap();
3976            assert_eq!(swap_token_b.base.amount, deposit_b + token_b_amount);
3977            let token_a = StateWithExtensions::<Account>::unpack(&token_a_account.data).unwrap();
3978            assert_eq!(token_a.base.amount, 0);
3979            let token_b = StateWithExtensions::<Account>::unpack(&token_b_account.data).unwrap();
3980            assert_eq!(token_b.base.amount, 0);
3981            let pool_account = StateWithExtensions::<Account>::unpack(&pool_account.data).unwrap();
3982            let swap_pool_account =
3983                StateWithExtensions::<Account>::unpack(&accounts.pool_token_account.data).unwrap();
3984            let pool_mint =
3985                StateWithExtensions::<Mint>::unpack(&accounts.pool_mint_account.data).unwrap();
3986            assert_eq!(
3987                pool_mint.base.supply,
3988                pool_account.base.amount + swap_pool_account.base.amount
3989            );
3990        }
3991    }
3992
3993    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
3994    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
3995    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
3996    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
3997    fn test_withdraw(
3998        pool_token_program_id: Pubkey,
3999        token_a_program_id: Pubkey,
4000        token_b_program_id: Pubkey,
4001    ) {
4002        let user_key = Pubkey::new_unique();
4003        let trade_fee_numerator = 1;
4004        let trade_fee_denominator = 2;
4005        let owner_trade_fee_numerator = 1;
4006        let owner_trade_fee_denominator = 10;
4007        let owner_withdraw_fee_numerator = 1;
4008        let owner_withdraw_fee_denominator = 5;
4009        let host_fee_numerator = 7;
4010        let host_fee_denominator = 100;
4011
4012        let fees = Fees {
4013            trade_fee_numerator,
4014            trade_fee_denominator,
4015            owner_trade_fee_numerator,
4016            owner_trade_fee_denominator,
4017            owner_withdraw_fee_numerator,
4018            owner_withdraw_fee_denominator,
4019            host_fee_numerator,
4020            host_fee_denominator,
4021        };
4022
4023        let token_a_amount = 1000;
4024        let token_b_amount = 2000;
4025        let curve_type = CurveType::ConstantProduct;
4026        let swap_curve = SwapCurve {
4027            curve_type,
4028            calculator: Arc::new(ConstantProductCurve {}),
4029        };
4030
4031        let withdrawer_key = Pubkey::new_unique();
4032        let initial_a = token_a_amount / 10;
4033        let initial_b = token_b_amount / 10;
4034        let initial_pool = swap_curve.calculator.new_pool_supply() / 10;
4035        let withdraw_amount = initial_pool / 4;
4036        let minimum_token_a_amount = initial_a / 40;
4037        let minimum_token_b_amount = initial_b / 40;
4038
4039        let mut accounts = SwapAccountInfo::new(
4040            &user_key,
4041            fees,
4042            SwapTransferFees::default(),
4043            swap_curve,
4044            token_a_amount,
4045            token_b_amount,
4046            &pool_token_program_id,
4047            &token_a_program_id,
4048            &token_b_program_id,
4049        );
4050
4051        // swap not initialized
4052        {
4053            let (
4054                token_a_key,
4055                mut token_a_account,
4056                token_b_key,
4057                mut token_b_account,
4058                pool_key,
4059                mut pool_account,
4060            ) = accounts.setup_token_accounts(&user_key, &withdrawer_key, initial_a, initial_b, 0);
4061            assert_eq!(
4062                Err(ProgramError::UninitializedAccount),
4063                accounts.withdraw_all_token_types(
4064                    &withdrawer_key,
4065                    &pool_key,
4066                    &mut pool_account,
4067                    &token_a_key,
4068                    &mut token_a_account,
4069                    &token_b_key,
4070                    &mut token_b_account,
4071                    withdraw_amount.try_into().unwrap(),
4072                    minimum_token_a_amount,
4073                    minimum_token_b_amount,
4074                )
4075            );
4076        }
4077
4078        accounts.initialize_swap().unwrap();
4079
4080        // wrong owner for swap account
4081        {
4082            let (
4083                token_a_key,
4084                mut token_a_account,
4085                token_b_key,
4086                mut token_b_account,
4087                pool_key,
4088                mut pool_account,
4089            ) = accounts.setup_token_accounts(&user_key, &withdrawer_key, initial_a, initial_b, 0);
4090            let old_swap_account = accounts.swap_account;
4091            let mut wrong_swap_account = old_swap_account.clone();
4092            wrong_swap_account.owner = pool_token_program_id;
4093            accounts.swap_account = wrong_swap_account;
4094            assert_eq!(
4095                Err(ProgramError::IncorrectProgramId),
4096                accounts.withdraw_all_token_types(
4097                    &withdrawer_key,
4098                    &pool_key,
4099                    &mut pool_account,
4100                    &token_a_key,
4101                    &mut token_a_account,
4102                    &token_b_key,
4103                    &mut token_b_account,
4104                    withdraw_amount.try_into().unwrap(),
4105                    minimum_token_a_amount,
4106                    minimum_token_b_amount,
4107                )
4108            );
4109            accounts.swap_account = old_swap_account;
4110        }
4111
4112        // wrong bump seed for authority_key
4113        {
4114            let (
4115                token_a_key,
4116                mut token_a_account,
4117                token_b_key,
4118                mut token_b_account,
4119                pool_key,
4120                mut pool_account,
4121            ) = accounts.setup_token_accounts(&user_key, &withdrawer_key, initial_a, initial_b, 0);
4122            let old_authority = accounts.authority_key;
4123            let (bad_authority_key, _bump_seed) = Pubkey::find_program_address(
4124                &[&accounts.swap_key.to_bytes()[..]],
4125                &pool_token_program_id,
4126            );
4127            accounts.authority_key = bad_authority_key;
4128            assert_eq!(
4129                Err(SwapError::InvalidProgramAddress.into()),
4130                accounts.withdraw_all_token_types(
4131                    &withdrawer_key,
4132                    &pool_key,
4133                    &mut pool_account,
4134                    &token_a_key,
4135                    &mut token_a_account,
4136                    &token_b_key,
4137                    &mut token_b_account,
4138                    withdraw_amount.try_into().unwrap(),
4139                    minimum_token_a_amount,
4140                    minimum_token_b_amount,
4141                )
4142            );
4143            accounts.authority_key = old_authority;
4144        }
4145
4146        // not enough pool tokens
4147        {
4148            let (
4149                token_a_key,
4150                mut token_a_account,
4151                token_b_key,
4152                mut token_b_account,
4153                pool_key,
4154                mut pool_account,
4155            ) = accounts.setup_token_accounts(
4156                &user_key,
4157                &withdrawer_key,
4158                initial_a,
4159                initial_b,
4160                to_u64(withdraw_amount).unwrap() / 2u64,
4161            );
4162            assert_eq!(
4163                Err(TokenError::InsufficientFunds.into()),
4164                accounts.withdraw_all_token_types(
4165                    &withdrawer_key,
4166                    &pool_key,
4167                    &mut pool_account,
4168                    &token_a_key,
4169                    &mut token_a_account,
4170                    &token_b_key,
4171                    &mut token_b_account,
4172                    withdraw_amount.try_into().unwrap(),
4173                    minimum_token_a_amount / 2,
4174                    minimum_token_b_amount / 2,
4175                )
4176            );
4177        }
4178
4179        // wrong token a / b accounts
4180        {
4181            let (
4182                token_a_key,
4183                mut token_a_account,
4184                token_b_key,
4185                mut token_b_account,
4186                pool_key,
4187                mut pool_account,
4188            ) = accounts.setup_token_accounts(
4189                &user_key,
4190                &withdrawer_key,
4191                initial_a,
4192                initial_b,
4193                withdraw_amount.try_into().unwrap(),
4194            );
4195            let expected_error: ProgramError = if token_a_account.owner == token_b_account.owner {
4196                TokenError::MintMismatch.into()
4197            } else {
4198                ProgramError::InvalidAccountData
4199            };
4200            assert_eq!(
4201                Err(expected_error),
4202                accounts.withdraw_all_token_types(
4203                    &withdrawer_key,
4204                    &pool_key,
4205                    &mut pool_account,
4206                    &token_b_key,
4207                    &mut token_b_account,
4208                    &token_a_key,
4209                    &mut token_a_account,
4210                    withdraw_amount.try_into().unwrap(),
4211                    minimum_token_a_amount,
4212                    minimum_token_b_amount,
4213                )
4214            );
4215        }
4216
4217        // wrong pool token account
4218        {
4219            let (
4220                token_a_key,
4221                mut token_a_account,
4222                token_b_key,
4223                mut token_b_account,
4224                _pool_key,
4225                _pool_account,
4226            ) = accounts.setup_token_accounts(
4227                &user_key,
4228                &withdrawer_key,
4229                initial_a,
4230                initial_b,
4231                withdraw_amount.try_into().unwrap(),
4232            );
4233            let (
4234                wrong_token_a_key,
4235                mut wrong_token_a_account,
4236                _token_b_key,
4237                _token_b_account,
4238                _pool_key,
4239                pool_account,
4240            ) = accounts.setup_token_accounts(
4241                &user_key,
4242                &withdrawer_key,
4243                withdraw_amount.try_into().unwrap(),
4244                initial_b,
4245                withdraw_amount.try_into().unwrap(),
4246            );
4247            let expected_error: ProgramError = if token_a_account.owner == pool_account.owner {
4248                TokenError::MintMismatch.into()
4249            } else {
4250                SwapError::IncorrectTokenProgramId.into()
4251            };
4252            assert_eq!(
4253                Err(expected_error),
4254                accounts.withdraw_all_token_types(
4255                    &withdrawer_key,
4256                    &wrong_token_a_key,
4257                    &mut wrong_token_a_account,
4258                    &token_a_key,
4259                    &mut token_a_account,
4260                    &token_b_key,
4261                    &mut token_b_account,
4262                    withdraw_amount.try_into().unwrap(),
4263                    minimum_token_a_amount,
4264                    minimum_token_b_amount,
4265                )
4266            );
4267        }
4268
4269        // wrong pool fee account
4270        {
4271            let (
4272                token_a_key,
4273                mut token_a_account,
4274                token_b_key,
4275                mut token_b_account,
4276                wrong_pool_key,
4277                wrong_pool_account,
4278            ) = accounts.setup_token_accounts(
4279                &user_key,
4280                &withdrawer_key,
4281                initial_a,
4282                initial_b,
4283                withdraw_amount.try_into().unwrap(),
4284            );
4285            let (
4286                _token_a_key,
4287                _token_a_account,
4288                _token_b_key,
4289                _token_b_account,
4290                pool_key,
4291                mut pool_account,
4292            ) = accounts.setup_token_accounts(
4293                &user_key,
4294                &withdrawer_key,
4295                initial_a,
4296                initial_b,
4297                withdraw_amount.try_into().unwrap(),
4298            );
4299            let old_pool_fee_account = accounts.pool_fee_account;
4300            let old_pool_fee_key = accounts.pool_fee_key;
4301            accounts.pool_fee_account = wrong_pool_account;
4302            accounts.pool_fee_key = wrong_pool_key;
4303            assert_eq!(
4304                Err(SwapError::IncorrectFeeAccount.into()),
4305                accounts.withdraw_all_token_types(
4306                    &withdrawer_key,
4307                    &pool_key,
4308                    &mut pool_account,
4309                    &token_a_key,
4310                    &mut token_a_account,
4311                    &token_b_key,
4312                    &mut token_b_account,
4313                    withdraw_amount.try_into().unwrap(),
4314                    minimum_token_a_amount,
4315                    minimum_token_b_amount,
4316                ),
4317            );
4318            accounts.pool_fee_account = old_pool_fee_account;
4319            accounts.pool_fee_key = old_pool_fee_key;
4320        }
4321
4322        // no approval
4323        {
4324            let (
4325                token_a_key,
4326                mut token_a_account,
4327                token_b_key,
4328                mut token_b_account,
4329                pool_key,
4330                mut pool_account,
4331            ) = accounts.setup_token_accounts(
4332                &user_key,
4333                &withdrawer_key,
4334                0,
4335                0,
4336                withdraw_amount.try_into().unwrap(),
4337            );
4338            let user_transfer_authority_key = Pubkey::new_unique();
4339            assert_eq!(
4340                Err(TokenError::OwnerMismatch.into()),
4341                do_process_instruction(
4342                    withdraw_all_token_types(
4343                        &SWAP_PROGRAM_ID,
4344                        &pool_token_program_id,
4345                        &token_a_program_id,
4346                        &token_b_program_id,
4347                        &accounts.swap_key,
4348                        &accounts.authority_key,
4349                        &user_transfer_authority_key,
4350                        &accounts.pool_mint_key,
4351                        &accounts.pool_fee_key,
4352                        &pool_key,
4353                        &accounts.token_a_key,
4354                        &accounts.token_b_key,
4355                        &token_a_key,
4356                        &token_b_key,
4357                        &accounts.token_a_mint_key,
4358                        &accounts.token_b_mint_key,
4359                        WithdrawAllTokenTypes {
4360                            pool_token_amount: withdraw_amount.try_into().unwrap(),
4361                            minimum_token_a_amount,
4362                            minimum_token_b_amount,
4363                        }
4364                    )
4365                    .unwrap(),
4366                    vec![
4367                        &mut accounts.swap_account,
4368                        &mut SolanaAccount::default(),
4369                        &mut SolanaAccount::default(),
4370                        &mut accounts.pool_mint_account,
4371                        &mut pool_account,
4372                        &mut accounts.token_a_account,
4373                        &mut accounts.token_b_account,
4374                        &mut token_a_account,
4375                        &mut token_b_account,
4376                        &mut accounts.pool_fee_account,
4377                        &mut accounts.token_a_mint_account,
4378                        &mut accounts.token_b_mint_account,
4379                        &mut SolanaAccount::default(),
4380                        &mut SolanaAccount::default(),
4381                        &mut SolanaAccount::default(),
4382                    ],
4383                )
4384            );
4385        }
4386
4387        // wrong token program id
4388        {
4389            let (
4390                token_a_key,
4391                mut token_a_account,
4392                token_b_key,
4393                mut token_b_account,
4394                pool_key,
4395                mut pool_account,
4396            ) = accounts.setup_token_accounts(
4397                &user_key,
4398                &withdrawer_key,
4399                initial_a,
4400                initial_b,
4401                withdraw_amount.try_into().unwrap(),
4402            );
4403            let wrong_key = Pubkey::new_unique();
4404            assert_eq!(
4405                Err(SwapError::IncorrectTokenProgramId.into()),
4406                do_process_instruction(
4407                    withdraw_all_token_types(
4408                        &SWAP_PROGRAM_ID,
4409                        &wrong_key,
4410                        &wrong_key,
4411                        &wrong_key,
4412                        &accounts.swap_key,
4413                        &accounts.authority_key,
4414                        &accounts.authority_key,
4415                        &accounts.pool_mint_key,
4416                        &accounts.pool_fee_key,
4417                        &pool_key,
4418                        &accounts.token_a_key,
4419                        &accounts.token_b_key,
4420                        &token_a_key,
4421                        &token_b_key,
4422                        &accounts.token_a_mint_key,
4423                        &accounts.token_b_mint_key,
4424                        WithdrawAllTokenTypes {
4425                            pool_token_amount: withdraw_amount.try_into().unwrap(),
4426                            minimum_token_a_amount,
4427                            minimum_token_b_amount,
4428                        },
4429                    )
4430                    .unwrap(),
4431                    vec![
4432                        &mut accounts.swap_account,
4433                        &mut SolanaAccount::default(),
4434                        &mut SolanaAccount::default(),
4435                        &mut accounts.pool_mint_account,
4436                        &mut pool_account,
4437                        &mut accounts.token_a_account,
4438                        &mut accounts.token_b_account,
4439                        &mut token_a_account,
4440                        &mut token_b_account,
4441                        &mut accounts.pool_fee_account,
4442                        &mut accounts.token_a_mint_account,
4443                        &mut accounts.token_b_mint_account,
4444                        &mut SolanaAccount::default(),
4445                        &mut SolanaAccount::default(),
4446                        &mut SolanaAccount::default(),
4447                    ],
4448                )
4449            );
4450        }
4451
4452        // wrong swap token accounts
4453        {
4454            let (
4455                token_a_key,
4456                mut token_a_account,
4457                token_b_key,
4458                mut token_b_account,
4459                pool_key,
4460                mut pool_account,
4461            ) = accounts.setup_token_accounts(
4462                &user_key,
4463                &withdrawer_key,
4464                initial_a,
4465                initial_b,
4466                initial_pool.try_into().unwrap(),
4467            );
4468
4469            let old_a_key = accounts.token_a_key;
4470            let old_a_account = accounts.token_a_account;
4471
4472            accounts.token_a_key = token_a_key;
4473            accounts.token_a_account = token_a_account.clone();
4474
4475            // wrong swap token a account
4476            assert_eq!(
4477                Err(SwapError::IncorrectSwapAccount.into()),
4478                accounts.withdraw_all_token_types(
4479                    &withdrawer_key,
4480                    &pool_key,
4481                    &mut pool_account,
4482                    &token_a_key,
4483                    &mut token_a_account,
4484                    &token_b_key,
4485                    &mut token_b_account,
4486                    withdraw_amount.try_into().unwrap(),
4487                    minimum_token_a_amount,
4488                    minimum_token_b_amount,
4489                )
4490            );
4491
4492            accounts.token_a_key = old_a_key;
4493            accounts.token_a_account = old_a_account;
4494
4495            let old_b_key = accounts.token_b_key;
4496            let old_b_account = accounts.token_b_account;
4497
4498            accounts.token_b_key = token_b_key;
4499            accounts.token_b_account = token_b_account.clone();
4500
4501            // wrong swap token b account
4502            assert_eq!(
4503                Err(SwapError::IncorrectSwapAccount.into()),
4504                accounts.withdraw_all_token_types(
4505                    &withdrawer_key,
4506                    &pool_key,
4507                    &mut pool_account,
4508                    &token_a_key,
4509                    &mut token_a_account,
4510                    &token_b_key,
4511                    &mut token_b_account,
4512                    withdraw_amount.try_into().unwrap(),
4513                    minimum_token_a_amount,
4514                    minimum_token_b_amount,
4515                )
4516            );
4517
4518            accounts.token_b_key = old_b_key;
4519            accounts.token_b_account = old_b_account;
4520        }
4521
4522        // wrong mint
4523        {
4524            let (
4525                token_a_key,
4526                mut token_a_account,
4527                token_b_key,
4528                mut token_b_account,
4529                pool_key,
4530                mut pool_account,
4531            ) = accounts.setup_token_accounts(
4532                &user_key,
4533                &withdrawer_key,
4534                initial_a,
4535                initial_b,
4536                initial_pool.try_into().unwrap(),
4537            );
4538            let (pool_mint_key, pool_mint_account) = create_mint(
4539                &pool_token_program_id,
4540                &accounts.authority_key,
4541                None,
4542                None,
4543                &TransferFee::default(),
4544            );
4545            let old_pool_key = accounts.pool_mint_key;
4546            let old_pool_account = accounts.pool_mint_account;
4547            accounts.pool_mint_key = pool_mint_key;
4548            accounts.pool_mint_account = pool_mint_account;
4549
4550            assert_eq!(
4551                Err(SwapError::IncorrectPoolMint.into()),
4552                accounts.withdraw_all_token_types(
4553                    &withdrawer_key,
4554                    &pool_key,
4555                    &mut pool_account,
4556                    &token_a_key,
4557                    &mut token_a_account,
4558                    &token_b_key,
4559                    &mut token_b_account,
4560                    withdraw_amount.try_into().unwrap(),
4561                    minimum_token_a_amount,
4562                    minimum_token_b_amount,
4563                )
4564            );
4565
4566            accounts.pool_mint_key = old_pool_key;
4567            accounts.pool_mint_account = old_pool_account;
4568        }
4569
4570        // withdrawing 1 pool token fails because it equates to 0 output tokens
4571        {
4572            let (
4573                token_a_key,
4574                mut token_a_account,
4575                token_b_key,
4576                mut token_b_account,
4577                pool_key,
4578                mut pool_account,
4579            ) = accounts.setup_token_accounts(
4580                &user_key,
4581                &withdrawer_key,
4582                initial_a,
4583                initial_b,
4584                initial_pool.try_into().unwrap(),
4585            );
4586            assert_eq!(
4587                Err(SwapError::ZeroTradingTokens.into()),
4588                accounts.withdraw_all_token_types(
4589                    &withdrawer_key,
4590                    &pool_key,
4591                    &mut pool_account,
4592                    &token_a_key,
4593                    &mut token_a_account,
4594                    &token_b_key,
4595                    &mut token_b_account,
4596                    1,
4597                    0,
4598                    0,
4599                )
4600            );
4601        }
4602
4603        // slippage exceeded
4604        {
4605            let (
4606                token_a_key,
4607                mut token_a_account,
4608                token_b_key,
4609                mut token_b_account,
4610                pool_key,
4611                mut pool_account,
4612            ) = accounts.setup_token_accounts(
4613                &user_key,
4614                &withdrawer_key,
4615                initial_a,
4616                initial_b,
4617                initial_pool.try_into().unwrap(),
4618            );
4619            // minimum A amount out too high
4620            assert_eq!(
4621                Err(SwapError::ExceededSlippage.into()),
4622                accounts.withdraw_all_token_types(
4623                    &withdrawer_key,
4624                    &pool_key,
4625                    &mut pool_account,
4626                    &token_a_key,
4627                    &mut token_a_account,
4628                    &token_b_key,
4629                    &mut token_b_account,
4630                    withdraw_amount.try_into().unwrap(),
4631                    minimum_token_a_amount * 10,
4632                    minimum_token_b_amount,
4633                )
4634            );
4635            // minimum B amount out too high
4636            assert_eq!(
4637                Err(SwapError::ExceededSlippage.into()),
4638                accounts.withdraw_all_token_types(
4639                    &withdrawer_key,
4640                    &pool_key,
4641                    &mut pool_account,
4642                    &token_a_key,
4643                    &mut token_a_account,
4644                    &token_b_key,
4645                    &mut token_b_account,
4646                    withdraw_amount.try_into().unwrap(),
4647                    minimum_token_a_amount,
4648                    minimum_token_b_amount * 10,
4649                )
4650            );
4651        }
4652
4653        // invalid input: can't use swap pool tokens as destination
4654        {
4655            let (
4656                token_a_key,
4657                mut token_a_account,
4658                token_b_key,
4659                mut token_b_account,
4660                pool_key,
4661                mut pool_account,
4662            ) = accounts.setup_token_accounts(
4663                &user_key,
4664                &withdrawer_key,
4665                initial_a,
4666                initial_b,
4667                initial_pool.try_into().unwrap(),
4668            );
4669            let swap_token_a_key = accounts.token_a_key;
4670            let mut swap_token_a_account = accounts.get_token_account(&swap_token_a_key).clone();
4671            assert_eq!(
4672                Err(SwapError::InvalidInput.into()),
4673                accounts.withdraw_all_token_types(
4674                    &withdrawer_key,
4675                    &pool_key,
4676                    &mut pool_account,
4677                    &swap_token_a_key,
4678                    &mut swap_token_a_account,
4679                    &token_b_key,
4680                    &mut token_b_account,
4681                    withdraw_amount.try_into().unwrap(),
4682                    minimum_token_a_amount,
4683                    minimum_token_b_amount,
4684                )
4685            );
4686            let swap_token_b_key = accounts.token_b_key;
4687            let mut swap_token_b_account = accounts.get_token_account(&swap_token_b_key).clone();
4688            assert_eq!(
4689                Err(SwapError::InvalidInput.into()),
4690                accounts.withdraw_all_token_types(
4691                    &withdrawer_key,
4692                    &pool_key,
4693                    &mut pool_account,
4694                    &token_a_key,
4695                    &mut token_a_account,
4696                    &swap_token_b_key,
4697                    &mut swap_token_b_account,
4698                    withdraw_amount.try_into().unwrap(),
4699                    minimum_token_a_amount,
4700                    minimum_token_b_amount,
4701                )
4702            );
4703        }
4704
4705        // correct withdrawal
4706        {
4707            let (
4708                token_a_key,
4709                mut token_a_account,
4710                token_b_key,
4711                mut token_b_account,
4712                pool_key,
4713                mut pool_account,
4714            ) = accounts.setup_token_accounts(
4715                &user_key,
4716                &withdrawer_key,
4717                initial_a,
4718                initial_b,
4719                initial_pool.try_into().unwrap(),
4720            );
4721
4722            accounts
4723                .withdraw_all_token_types(
4724                    &withdrawer_key,
4725                    &pool_key,
4726                    &mut pool_account,
4727                    &token_a_key,
4728                    &mut token_a_account,
4729                    &token_b_key,
4730                    &mut token_b_account,
4731                    withdraw_amount.try_into().unwrap(),
4732                    minimum_token_a_amount,
4733                    minimum_token_b_amount,
4734                )
4735                .unwrap();
4736
4737            let swap_token_a =
4738                StateWithExtensions::<Account>::unpack(&accounts.token_a_account.data).unwrap();
4739            let swap_token_b =
4740                StateWithExtensions::<Account>::unpack(&accounts.token_b_account.data).unwrap();
4741            let pool_mint =
4742                StateWithExtensions::<Mint>::unpack(&accounts.pool_mint_account.data).unwrap();
4743            let withdraw_fee = accounts.fees.owner_withdraw_fee(withdraw_amount).unwrap();
4744            let results = accounts
4745                .swap_curve
4746                .calculator
4747                .pool_tokens_to_trading_tokens(
4748                    withdraw_amount - withdraw_fee,
4749                    pool_mint.base.supply.try_into().unwrap(),
4750                    swap_token_a.base.amount.try_into().unwrap(),
4751                    swap_token_b.base.amount.try_into().unwrap(),
4752                    RoundDirection::Floor,
4753                )
4754                .unwrap();
4755            assert_eq!(
4756                swap_token_a.base.amount,
4757                token_a_amount - to_u64(results.token_a_amount).unwrap()
4758            );
4759            assert_eq!(
4760                swap_token_b.base.amount,
4761                token_b_amount - to_u64(results.token_b_amount).unwrap()
4762            );
4763            let token_a = StateWithExtensions::<Account>::unpack(&token_a_account.data).unwrap();
4764            assert_eq!(
4765                token_a.base.amount,
4766                initial_a + to_u64(results.token_a_amount).unwrap()
4767            );
4768            let token_b = StateWithExtensions::<Account>::unpack(&token_b_account.data).unwrap();
4769            assert_eq!(
4770                token_b.base.amount,
4771                initial_b + to_u64(results.token_b_amount).unwrap()
4772            );
4773            let pool_account = StateWithExtensions::<Account>::unpack(&pool_account.data).unwrap();
4774            assert_eq!(
4775                pool_account.base.amount,
4776                to_u64(initial_pool - withdraw_amount).unwrap()
4777            );
4778            let fee_account =
4779                StateWithExtensions::<Account>::unpack(&accounts.pool_fee_account.data).unwrap();
4780            assert_eq!(
4781                fee_account.base.amount,
4782                TryInto::<u64>::try_into(withdraw_fee).unwrap()
4783            );
4784        }
4785
4786        // correct withdrawal from fee account
4787        {
4788            let (
4789                token_a_key,
4790                mut token_a_account,
4791                token_b_key,
4792                mut token_b_account,
4793                _pool_key,
4794                mut _pool_account,
4795            ) = accounts.setup_token_accounts(&user_key, &withdrawer_key, 0, 0, 0);
4796
4797            let pool_fee_key = accounts.pool_fee_key;
4798            let mut pool_fee_account = accounts.pool_fee_account.clone();
4799            let fee_account =
4800                StateWithExtensions::<Account>::unpack(&pool_fee_account.data).unwrap();
4801            let pool_fee_amount = fee_account.base.amount;
4802
4803            accounts
4804                .withdraw_all_token_types(
4805                    &user_key,
4806                    &pool_fee_key,
4807                    &mut pool_fee_account,
4808                    &token_a_key,
4809                    &mut token_a_account,
4810                    &token_b_key,
4811                    &mut token_b_account,
4812                    pool_fee_amount,
4813                    0,
4814                    0,
4815                )
4816                .unwrap();
4817
4818            let swap_token_a =
4819                StateWithExtensions::<Account>::unpack(&accounts.token_a_account.data).unwrap();
4820            let swap_token_b =
4821                StateWithExtensions::<Account>::unpack(&accounts.token_b_account.data).unwrap();
4822            let pool_mint =
4823                StateWithExtensions::<Mint>::unpack(&accounts.pool_mint_account.data).unwrap();
4824            let results = accounts
4825                .swap_curve
4826                .calculator
4827                .pool_tokens_to_trading_tokens(
4828                    pool_fee_amount.try_into().unwrap(),
4829                    pool_mint.base.supply.try_into().unwrap(),
4830                    swap_token_a.base.amount.try_into().unwrap(),
4831                    swap_token_b.base.amount.try_into().unwrap(),
4832                    RoundDirection::Floor,
4833                )
4834                .unwrap();
4835            let token_a = StateWithExtensions::<Account>::unpack(&token_a_account.data).unwrap();
4836            assert_eq!(
4837                token_a.base.amount,
4838                TryInto::<u64>::try_into(results.token_a_amount).unwrap()
4839            );
4840            let token_b = StateWithExtensions::<Account>::unpack(&token_b_account.data).unwrap();
4841            assert_eq!(
4842                token_b.base.amount,
4843                TryInto::<u64>::try_into(results.token_b_amount).unwrap()
4844            );
4845        }
4846    }
4847
4848    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
4849    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
4850    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
4851    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
4852    fn test_deposit_one_exact_in(
4853        pool_token_program_id: Pubkey,
4854        token_a_program_id: Pubkey,
4855        token_b_program_id: Pubkey,
4856    ) {
4857        let user_key = Pubkey::new_unique();
4858        let depositor_key = Pubkey::new_unique();
4859        let trade_fee_numerator = 1;
4860        let trade_fee_denominator = 2;
4861        let owner_trade_fee_numerator = 1;
4862        let owner_trade_fee_denominator = 10;
4863        let owner_withdraw_fee_numerator = 1;
4864        let owner_withdraw_fee_denominator = 5;
4865        let host_fee_numerator = 20;
4866        let host_fee_denominator = 100;
4867
4868        let fees = Fees {
4869            trade_fee_numerator,
4870            trade_fee_denominator,
4871            owner_trade_fee_numerator,
4872            owner_trade_fee_denominator,
4873            owner_withdraw_fee_numerator,
4874            owner_withdraw_fee_denominator,
4875            host_fee_numerator,
4876            host_fee_denominator,
4877        };
4878
4879        let token_a_amount = 1000;
4880        let token_b_amount = 9000;
4881        let curve_type = CurveType::ConstantProduct;
4882        let swap_curve = SwapCurve {
4883            curve_type,
4884            calculator: Arc::new(ConstantProductCurve {}),
4885        };
4886
4887        let mut accounts = SwapAccountInfo::new(
4888            &user_key,
4889            fees,
4890            SwapTransferFees::default(),
4891            swap_curve,
4892            token_a_amount,
4893            token_b_amount,
4894            &pool_token_program_id,
4895            &token_a_program_id,
4896            &token_b_program_id,
4897        );
4898
4899        let deposit_a = token_a_amount / 10;
4900        let deposit_b = token_b_amount / 10;
4901        let pool_amount = to_u64(INITIAL_SWAP_POOL_AMOUNT / 100).unwrap();
4902
4903        // swap not initialized
4904        {
4905            let (
4906                token_a_key,
4907                mut token_a_account,
4908                _token_b_key,
4909                _token_b_account,
4910                pool_key,
4911                mut pool_account,
4912            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
4913            assert_eq!(
4914                Err(ProgramError::UninitializedAccount),
4915                accounts.deposit_single_token_type_exact_amount_in(
4916                    &depositor_key,
4917                    &token_a_key,
4918                    &mut token_a_account,
4919                    &pool_key,
4920                    &mut pool_account,
4921                    deposit_a,
4922                    pool_amount,
4923                )
4924            );
4925        }
4926
4927        accounts.initialize_swap().unwrap();
4928
4929        // wrong owner for swap account
4930        {
4931            let (
4932                token_a_key,
4933                mut token_a_account,
4934                _token_b_key,
4935                _token_b_account,
4936                pool_key,
4937                mut pool_account,
4938            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
4939            let old_swap_account = accounts.swap_account;
4940            let mut wrong_swap_account = old_swap_account.clone();
4941            wrong_swap_account.owner = pool_token_program_id;
4942            accounts.swap_account = wrong_swap_account;
4943            assert_eq!(
4944                Err(ProgramError::IncorrectProgramId),
4945                accounts.deposit_single_token_type_exact_amount_in(
4946                    &depositor_key,
4947                    &token_a_key,
4948                    &mut token_a_account,
4949                    &pool_key,
4950                    &mut pool_account,
4951                    deposit_a,
4952                    pool_amount,
4953                )
4954            );
4955            accounts.swap_account = old_swap_account;
4956        }
4957
4958        // wrong bump seed for authority_key
4959        {
4960            let (
4961                token_a_key,
4962                mut token_a_account,
4963                _token_b_key,
4964                _token_b_account,
4965                pool_key,
4966                mut pool_account,
4967            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
4968            let old_authority = accounts.authority_key;
4969            let (bad_authority_key, _bump_seed) = Pubkey::find_program_address(
4970                &[&accounts.swap_key.to_bytes()[..]],
4971                &pool_token_program_id,
4972            );
4973            accounts.authority_key = bad_authority_key;
4974            assert_eq!(
4975                Err(SwapError::InvalidProgramAddress.into()),
4976                accounts.deposit_single_token_type_exact_amount_in(
4977                    &depositor_key,
4978                    &token_a_key,
4979                    &mut token_a_account,
4980                    &pool_key,
4981                    &mut pool_account,
4982                    deposit_a,
4983                    pool_amount,
4984                )
4985            );
4986            accounts.authority_key = old_authority;
4987        }
4988
4989        // not enough token A / B
4990        {
4991            let (
4992                token_a_key,
4993                mut token_a_account,
4994                token_b_key,
4995                mut token_b_account,
4996                pool_key,
4997                mut pool_account,
4998            ) = accounts.setup_token_accounts(
4999                &user_key,
5000                &depositor_key,
5001                deposit_a / 2,
5002                deposit_b / 2,
5003                0,
5004            );
5005            assert_eq!(
5006                Err(TokenError::InsufficientFunds.into()),
5007                accounts.deposit_single_token_type_exact_amount_in(
5008                    &depositor_key,
5009                    &token_a_key,
5010                    &mut token_a_account,
5011                    &pool_key,
5012                    &mut pool_account,
5013                    deposit_a,
5014                    0,
5015                )
5016            );
5017            assert_eq!(
5018                Err(TokenError::InsufficientFunds.into()),
5019                accounts.deposit_single_token_type_exact_amount_in(
5020                    &depositor_key,
5021                    &token_b_key,
5022                    &mut token_b_account,
5023                    &pool_key,
5024                    &mut pool_account,
5025                    deposit_b,
5026                    0,
5027                )
5028            );
5029        }
5030
5031        // wrong pool token account
5032        {
5033            let (
5034                token_a_key,
5035                mut token_a_account,
5036                token_b_key,
5037                mut token_b_account,
5038                _pool_key,
5039                pool_account,
5040            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
5041            let expected_error: ProgramError = if token_b_account.owner == pool_account.owner {
5042                TokenError::MintMismatch.into()
5043            } else {
5044                SwapError::IncorrectTokenProgramId.into()
5045            };
5046            assert_eq!(
5047                Err(expected_error),
5048                accounts.deposit_single_token_type_exact_amount_in(
5049                    &depositor_key,
5050                    &token_a_key,
5051                    &mut token_a_account,
5052                    &token_b_key,
5053                    &mut token_b_account,
5054                    deposit_a,
5055                    pool_amount,
5056                )
5057            );
5058        }
5059
5060        // no approval
5061        {
5062            let (
5063                token_a_key,
5064                mut token_a_account,
5065                _token_b_key,
5066                _token_b_account,
5067                pool_key,
5068                mut pool_account,
5069            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
5070            let user_transfer_authority_key = Pubkey::new_unique();
5071            assert_eq!(
5072                Err(TokenError::OwnerMismatch.into()),
5073                do_process_instruction(
5074                    deposit_single_token_type_exact_amount_in(
5075                        &SWAP_PROGRAM_ID,
5076                        &token_a_program_id,
5077                        &pool_token_program_id,
5078                        &accounts.swap_key,
5079                        &accounts.authority_key,
5080                        &user_transfer_authority_key,
5081                        &token_a_key,
5082                        &accounts.token_a_key,
5083                        &accounts.token_b_key,
5084                        &accounts.pool_mint_key,
5085                        &pool_key,
5086                        &accounts.token_a_mint_key,
5087                        DepositSingleTokenTypeExactAmountIn {
5088                            source_token_amount: deposit_a,
5089                            minimum_pool_token_amount: pool_amount,
5090                        },
5091                    )
5092                    .unwrap(),
5093                    vec![
5094                        &mut accounts.swap_account,
5095                        &mut SolanaAccount::default(),
5096                        &mut SolanaAccount::default(),
5097                        &mut token_a_account,
5098                        &mut accounts.token_a_account,
5099                        &mut accounts.token_b_account,
5100                        &mut accounts.pool_mint_account,
5101                        &mut pool_account,
5102                        &mut accounts.token_a_mint_account,
5103                        &mut SolanaAccount::default(),
5104                        &mut SolanaAccount::default(),
5105                    ],
5106                )
5107            );
5108        }
5109
5110        // wrong token program id
5111        {
5112            let (
5113                token_a_key,
5114                mut token_a_account,
5115                _token_b_key,
5116                _token_b_account,
5117                pool_key,
5118                mut pool_account,
5119            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
5120            let wrong_key = Pubkey::new_unique();
5121            assert_eq!(
5122                Err(SwapError::IncorrectTokenProgramId.into()),
5123                do_process_instruction(
5124                    deposit_single_token_type_exact_amount_in(
5125                        &SWAP_PROGRAM_ID,
5126                        &wrong_key,
5127                        &wrong_key,
5128                        &accounts.swap_key,
5129                        &accounts.authority_key,
5130                        &accounts.authority_key,
5131                        &token_a_key,
5132                        &accounts.token_a_key,
5133                        &accounts.token_b_key,
5134                        &accounts.pool_mint_key,
5135                        &pool_key,
5136                        &accounts.token_a_mint_key,
5137                        DepositSingleTokenTypeExactAmountIn {
5138                            source_token_amount: deposit_a,
5139                            minimum_pool_token_amount: pool_amount,
5140                        },
5141                    )
5142                    .unwrap(),
5143                    vec![
5144                        &mut accounts.swap_account,
5145                        &mut SolanaAccount::default(),
5146                        &mut SolanaAccount::default(),
5147                        &mut token_a_account,
5148                        &mut accounts.token_a_account,
5149                        &mut accounts.token_b_account,
5150                        &mut accounts.pool_mint_account,
5151                        &mut pool_account,
5152                        &mut accounts.token_a_mint_account,
5153                        &mut SolanaAccount::default(),
5154                        &mut SolanaAccount::default(),
5155                    ],
5156                )
5157            );
5158        }
5159
5160        // wrong swap token accounts
5161        {
5162            let (
5163                token_a_key,
5164                mut token_a_account,
5165                token_b_key,
5166                token_b_account,
5167                pool_key,
5168                mut pool_account,
5169            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
5170
5171            let old_a_key = accounts.token_a_key;
5172            let old_a_account = accounts.token_a_account;
5173
5174            accounts.token_a_key = token_a_key;
5175            accounts.token_a_account = token_a_account.clone();
5176
5177            // wrong swap token a account
5178            assert_eq!(
5179                Err(SwapError::IncorrectSwapAccount.into()),
5180                accounts.deposit_single_token_type_exact_amount_in(
5181                    &depositor_key,
5182                    &token_a_key,
5183                    &mut token_a_account,
5184                    &pool_key,
5185                    &mut pool_account,
5186                    deposit_a,
5187                    pool_amount,
5188                )
5189            );
5190
5191            accounts.token_a_key = old_a_key;
5192            accounts.token_a_account = old_a_account;
5193
5194            let old_b_key = accounts.token_b_key;
5195            let old_b_account = accounts.token_b_account;
5196
5197            accounts.token_b_key = token_b_key;
5198            accounts.token_b_account = token_b_account;
5199
5200            // wrong swap token b account
5201            assert_eq!(
5202                Err(SwapError::IncorrectSwapAccount.into()),
5203                accounts.deposit_single_token_type_exact_amount_in(
5204                    &depositor_key,
5205                    &token_a_key,
5206                    &mut token_a_account,
5207                    &pool_key,
5208                    &mut pool_account,
5209                    deposit_a,
5210                    pool_amount,
5211                )
5212            );
5213
5214            accounts.token_b_key = old_b_key;
5215            accounts.token_b_account = old_b_account;
5216        }
5217
5218        // wrong mint
5219        {
5220            let (
5221                token_a_key,
5222                mut token_a_account,
5223                _token_b_key,
5224                _token_b_account,
5225                pool_key,
5226                mut pool_account,
5227            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
5228            let (pool_mint_key, pool_mint_account) = create_mint(
5229                &pool_token_program_id,
5230                &accounts.authority_key,
5231                None,
5232                None,
5233                &TransferFee::default(),
5234            );
5235            let old_pool_key = accounts.pool_mint_key;
5236            let old_pool_account = accounts.pool_mint_account;
5237            accounts.pool_mint_key = pool_mint_key;
5238            accounts.pool_mint_account = pool_mint_account;
5239
5240            assert_eq!(
5241                Err(SwapError::IncorrectPoolMint.into()),
5242                accounts.deposit_single_token_type_exact_amount_in(
5243                    &depositor_key,
5244                    &token_a_key,
5245                    &mut token_a_account,
5246                    &pool_key,
5247                    &mut pool_account,
5248                    deposit_a,
5249                    pool_amount,
5250                )
5251            );
5252
5253            accounts.pool_mint_key = old_pool_key;
5254            accounts.pool_mint_account = old_pool_account;
5255        }
5256
5257        // slippage exceeded
5258        {
5259            let (
5260                token_a_key,
5261                mut token_a_account,
5262                token_b_key,
5263                mut token_b_account,
5264                pool_key,
5265                mut pool_account,
5266            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
5267            // minimum pool amount too high
5268            assert_eq!(
5269                Err(SwapError::ExceededSlippage.into()),
5270                accounts.deposit_single_token_type_exact_amount_in(
5271                    &depositor_key,
5272                    &token_a_key,
5273                    &mut token_a_account,
5274                    &pool_key,
5275                    &mut pool_account,
5276                    deposit_a / 10,
5277                    pool_amount,
5278                )
5279            );
5280            // minimum pool amount too high
5281            assert_eq!(
5282                Err(SwapError::ExceededSlippage.into()),
5283                accounts.deposit_single_token_type_exact_amount_in(
5284                    &depositor_key,
5285                    &token_b_key,
5286                    &mut token_b_account,
5287                    &pool_key,
5288                    &mut pool_account,
5289                    deposit_b / 10,
5290                    pool_amount,
5291                )
5292            );
5293        }
5294
5295        // invalid input: can't use swap pool tokens as source
5296        {
5297            let (
5298                _token_a_key,
5299                _token_a_account,
5300                _token_b_key,
5301                _token_b_account,
5302                pool_key,
5303                mut pool_account,
5304            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
5305            let swap_token_a_key = accounts.token_a_key;
5306            let mut swap_token_a_account = accounts.get_token_account(&swap_token_a_key).clone();
5307            let swap_token_b_key = accounts.token_b_key;
5308            let mut swap_token_b_account = accounts.get_token_account(&swap_token_b_key).clone();
5309            let authority_key = accounts.authority_key;
5310            assert_eq!(
5311                Err(SwapError::InvalidInput.into()),
5312                accounts.deposit_single_token_type_exact_amount_in(
5313                    &authority_key,
5314                    &swap_token_a_key,
5315                    &mut swap_token_a_account,
5316                    &pool_key,
5317                    &mut pool_account,
5318                    deposit_a,
5319                    pool_amount,
5320                )
5321            );
5322            assert_eq!(
5323                Err(SwapError::InvalidInput.into()),
5324                accounts.deposit_single_token_type_exact_amount_in(
5325                    &authority_key,
5326                    &swap_token_b_key,
5327                    &mut swap_token_b_account,
5328                    &pool_key,
5329                    &mut pool_account,
5330                    deposit_b,
5331                    pool_amount,
5332                )
5333            );
5334        }
5335
5336        // correctly deposit
5337        {
5338            let (
5339                token_a_key,
5340                mut token_a_account,
5341                token_b_key,
5342                mut token_b_account,
5343                pool_key,
5344                mut pool_account,
5345            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
5346            accounts
5347                .deposit_single_token_type_exact_amount_in(
5348                    &depositor_key,
5349                    &token_a_key,
5350                    &mut token_a_account,
5351                    &pool_key,
5352                    &mut pool_account,
5353                    deposit_a,
5354                    pool_amount,
5355                )
5356                .unwrap();
5357
5358            let swap_token_a =
5359                StateWithExtensions::<Account>::unpack(&accounts.token_a_account.data).unwrap();
5360            assert_eq!(swap_token_a.base.amount, deposit_a + token_a_amount);
5361
5362            let token_a = StateWithExtensions::<Account>::unpack(&token_a_account.data).unwrap();
5363            assert_eq!(token_a.base.amount, 0);
5364
5365            accounts
5366                .deposit_single_token_type_exact_amount_in(
5367                    &depositor_key,
5368                    &token_b_key,
5369                    &mut token_b_account,
5370                    &pool_key,
5371                    &mut pool_account,
5372                    deposit_b,
5373                    pool_amount,
5374                )
5375                .unwrap();
5376            let swap_token_b =
5377                StateWithExtensions::<Account>::unpack(&accounts.token_b_account.data).unwrap();
5378            assert_eq!(swap_token_b.base.amount, deposit_b + token_b_amount);
5379
5380            let token_b = StateWithExtensions::<Account>::unpack(&token_b_account.data).unwrap();
5381            assert_eq!(token_b.base.amount, 0);
5382
5383            let pool_account = StateWithExtensions::<Account>::unpack(&pool_account.data).unwrap();
5384            let swap_pool_account =
5385                StateWithExtensions::<Account>::unpack(&accounts.pool_token_account.data).unwrap();
5386            let pool_mint =
5387                StateWithExtensions::<Mint>::unpack(&accounts.pool_mint_account.data).unwrap();
5388            assert_eq!(
5389                pool_mint.base.supply,
5390                pool_account.base.amount + swap_pool_account.base.amount
5391            );
5392        }
5393    }
5394
5395    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
5396    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
5397    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
5398    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
5399    fn test_withdraw_one_exact_out(
5400        pool_token_program_id: Pubkey,
5401        token_a_program_id: Pubkey,
5402        token_b_program_id: Pubkey,
5403    ) {
5404        let user_key = Pubkey::new_unique();
5405        let trade_fee_numerator = 1;
5406        let trade_fee_denominator = 2;
5407        let owner_trade_fee_numerator = 1;
5408        let owner_trade_fee_denominator = 10;
5409        let owner_withdraw_fee_numerator = 1;
5410        let owner_withdraw_fee_denominator = 5;
5411        let host_fee_numerator = 7;
5412        let host_fee_denominator = 100;
5413
5414        let fees = Fees {
5415            trade_fee_numerator,
5416            trade_fee_denominator,
5417            owner_trade_fee_numerator,
5418            owner_trade_fee_denominator,
5419            owner_withdraw_fee_numerator,
5420            owner_withdraw_fee_denominator,
5421            host_fee_numerator,
5422            host_fee_denominator,
5423        };
5424
5425        let token_a_amount = 100_000;
5426        let token_b_amount = 200_000;
5427        let curve_type = CurveType::ConstantProduct;
5428        let swap_curve = SwapCurve {
5429            curve_type,
5430            calculator: Arc::new(ConstantProductCurve {}),
5431        };
5432
5433        let withdrawer_key = Pubkey::new_unique();
5434        let initial_a = token_a_amount / 10;
5435        let initial_b = token_b_amount / 10;
5436        let initial_pool = swap_curve.calculator.new_pool_supply() / 10;
5437        let maximum_pool_token_amount = to_u64(initial_pool / 4).unwrap();
5438        let destination_a_amount = initial_a / 40;
5439        let destination_b_amount = initial_b / 40;
5440
5441        let mut accounts = SwapAccountInfo::new(
5442            &user_key,
5443            fees,
5444            SwapTransferFees::default(),
5445            swap_curve,
5446            token_a_amount,
5447            token_b_amount,
5448            &pool_token_program_id,
5449            &token_a_program_id,
5450            &token_b_program_id,
5451        );
5452
5453        // swap not initialized
5454        {
5455            let (
5456                token_a_key,
5457                mut token_a_account,
5458                _token_b_key,
5459                _token_b_account,
5460                pool_key,
5461                mut pool_account,
5462            ) = accounts.setup_token_accounts(&user_key, &withdrawer_key, initial_a, initial_b, 0);
5463            assert_eq!(
5464                Err(ProgramError::UninitializedAccount),
5465                accounts.withdraw_single_token_type_exact_amount_out(
5466                    &withdrawer_key,
5467                    &pool_key,
5468                    &mut pool_account,
5469                    &token_a_key,
5470                    &mut token_a_account,
5471                    destination_a_amount,
5472                    maximum_pool_token_amount,
5473                )
5474            );
5475        }
5476
5477        accounts.initialize_swap().unwrap();
5478
5479        // wrong owner for swap account
5480        {
5481            let (
5482                token_a_key,
5483                mut token_a_account,
5484                _token_b_key,
5485                _token_b_account,
5486                pool_key,
5487                mut pool_account,
5488            ) = accounts.setup_token_accounts(&user_key, &withdrawer_key, initial_a, initial_b, 0);
5489            let old_swap_account = accounts.swap_account;
5490            let mut wrong_swap_account = old_swap_account.clone();
5491            wrong_swap_account.owner = pool_token_program_id;
5492            accounts.swap_account = wrong_swap_account;
5493            assert_eq!(
5494                Err(ProgramError::IncorrectProgramId),
5495                accounts.withdraw_single_token_type_exact_amount_out(
5496                    &withdrawer_key,
5497                    &pool_key,
5498                    &mut pool_account,
5499                    &token_a_key,
5500                    &mut token_a_account,
5501                    destination_a_amount,
5502                    maximum_pool_token_amount,
5503                )
5504            );
5505            accounts.swap_account = old_swap_account;
5506        }
5507
5508        // wrong bump seed for authority_key
5509        {
5510            let (
5511                _token_a_key,
5512                _token_a_account,
5513                token_b_key,
5514                mut token_b_account,
5515                pool_key,
5516                mut pool_account,
5517            ) = accounts.setup_token_accounts(&user_key, &withdrawer_key, initial_a, initial_b, 0);
5518            let old_authority = accounts.authority_key;
5519            let (bad_authority_key, _bump_seed) = Pubkey::find_program_address(
5520                &[&accounts.swap_key.to_bytes()[..]],
5521                &pool_token_program_id,
5522            );
5523            accounts.authority_key = bad_authority_key;
5524            assert_eq!(
5525                Err(SwapError::InvalidProgramAddress.into()),
5526                accounts.withdraw_single_token_type_exact_amount_out(
5527                    &withdrawer_key,
5528                    &pool_key,
5529                    &mut pool_account,
5530                    &token_b_key,
5531                    &mut token_b_account,
5532                    destination_b_amount,
5533                    maximum_pool_token_amount,
5534                )
5535            );
5536            accounts.authority_key = old_authority;
5537        }
5538
5539        // not enough pool tokens
5540        {
5541            let (
5542                _token_a_key,
5543                _token_a_account,
5544                token_b_key,
5545                mut token_b_account,
5546                pool_key,
5547                mut pool_account,
5548            ) = accounts.setup_token_accounts(
5549                &user_key,
5550                &withdrawer_key,
5551                initial_a,
5552                initial_b,
5553                maximum_pool_token_amount / 1000,
5554            );
5555            assert_eq!(
5556                Err(TokenError::InsufficientFunds.into()),
5557                accounts.withdraw_single_token_type_exact_amount_out(
5558                    &withdrawer_key,
5559                    &pool_key,
5560                    &mut pool_account,
5561                    &token_b_key,
5562                    &mut token_b_account,
5563                    destination_b_amount,
5564                    maximum_pool_token_amount,
5565                )
5566            );
5567        }
5568
5569        // wrong pool token account
5570        {
5571            let (
5572                token_a_key,
5573                mut token_a_account,
5574                token_b_key,
5575                mut token_b_account,
5576                _pool_key,
5577                pool_account,
5578            ) = accounts.setup_token_accounts(
5579                &user_key,
5580                &withdrawer_key,
5581                maximum_pool_token_amount,
5582                initial_b,
5583                maximum_pool_token_amount,
5584            );
5585            let expected_error: ProgramError = if token_a_account.owner == pool_account.owner {
5586                TokenError::MintMismatch.into()
5587            } else {
5588                SwapError::IncorrectTokenProgramId.into()
5589            };
5590            assert_eq!(
5591                Err(expected_error),
5592                accounts.withdraw_single_token_type_exact_amount_out(
5593                    &withdrawer_key,
5594                    &token_a_key,
5595                    &mut token_a_account,
5596                    &token_b_key,
5597                    &mut token_b_account,
5598                    destination_b_amount,
5599                    maximum_pool_token_amount,
5600                )
5601            );
5602        }
5603
5604        // wrong pool fee account
5605        {
5606            let (
5607                token_a_key,
5608                mut token_a_account,
5609                _token_b_key,
5610                _token_b_account,
5611                wrong_pool_key,
5612                wrong_pool_account,
5613            ) = accounts.setup_token_accounts(
5614                &user_key,
5615                &withdrawer_key,
5616                initial_a,
5617                initial_b,
5618                maximum_pool_token_amount,
5619            );
5620            let (
5621                _token_a_key,
5622                _token_a_account,
5623                _token_b_key,
5624                _token_b_account,
5625                pool_key,
5626                mut pool_account,
5627            ) = accounts.setup_token_accounts(
5628                &user_key,
5629                &withdrawer_key,
5630                initial_a,
5631                initial_b,
5632                maximum_pool_token_amount,
5633            );
5634            let old_pool_fee_account = accounts.pool_fee_account;
5635            let old_pool_fee_key = accounts.pool_fee_key;
5636            accounts.pool_fee_account = wrong_pool_account;
5637            accounts.pool_fee_key = wrong_pool_key;
5638            assert_eq!(
5639                Err(SwapError::IncorrectFeeAccount.into()),
5640                accounts.withdraw_single_token_type_exact_amount_out(
5641                    &withdrawer_key,
5642                    &pool_key,
5643                    &mut pool_account,
5644                    &token_a_key,
5645                    &mut token_a_account,
5646                    destination_a_amount,
5647                    maximum_pool_token_amount,
5648                )
5649            );
5650            accounts.pool_fee_account = old_pool_fee_account;
5651            accounts.pool_fee_key = old_pool_fee_key;
5652        }
5653
5654        // no approval
5655        {
5656            let (
5657                token_a_key,
5658                mut token_a_account,
5659                _token_b_key,
5660                _token_b_account,
5661                pool_key,
5662                mut pool_account,
5663            ) = accounts.setup_token_accounts(
5664                &user_key,
5665                &withdrawer_key,
5666                0,
5667                0,
5668                maximum_pool_token_amount,
5669            );
5670            let user_transfer_authority_key = Pubkey::new_unique();
5671            assert_eq!(
5672                Err(TokenError::OwnerMismatch.into()),
5673                do_process_instruction(
5674                    withdraw_single_token_type_exact_amount_out(
5675                        &SWAP_PROGRAM_ID,
5676                        &pool_token_program_id,
5677                        &token_a_program_id,
5678                        &accounts.swap_key,
5679                        &accounts.authority_key,
5680                        &user_transfer_authority_key,
5681                        &accounts.pool_mint_key,
5682                        &accounts.pool_fee_key,
5683                        &pool_key,
5684                        &accounts.token_a_key,
5685                        &accounts.token_b_key,
5686                        &token_a_key,
5687                        &accounts.token_a_mint_key,
5688                        WithdrawSingleTokenTypeExactAmountOut {
5689                            destination_token_amount: destination_a_amount,
5690                            maximum_pool_token_amount,
5691                        }
5692                    )
5693                    .unwrap(),
5694                    vec![
5695                        &mut accounts.swap_account,
5696                        &mut SolanaAccount::default(),
5697                        &mut SolanaAccount::default(),
5698                        &mut accounts.pool_mint_account,
5699                        &mut pool_account,
5700                        &mut accounts.token_a_account,
5701                        &mut accounts.token_b_account,
5702                        &mut token_a_account,
5703                        &mut accounts.pool_fee_account,
5704                        &mut accounts.token_a_mint_account,
5705                        &mut SolanaAccount::default(),
5706                        &mut SolanaAccount::default(),
5707                    ],
5708                )
5709            );
5710        }
5711
5712        // wrong token program id
5713        {
5714            let (
5715                token_a_key,
5716                mut token_a_account,
5717                _token_b_key,
5718                _token_b_account,
5719                pool_key,
5720                mut pool_account,
5721            ) = accounts.setup_token_accounts(
5722                &user_key,
5723                &withdrawer_key,
5724                initial_a,
5725                initial_b,
5726                maximum_pool_token_amount,
5727            );
5728            let wrong_key = Pubkey::new_unique();
5729            assert_eq!(
5730                Err(SwapError::IncorrectTokenProgramId.into()),
5731                do_process_instruction(
5732                    withdraw_single_token_type_exact_amount_out(
5733                        &SWAP_PROGRAM_ID,
5734                        &wrong_key,
5735                        &wrong_key,
5736                        &accounts.swap_key,
5737                        &accounts.authority_key,
5738                        &accounts.authority_key,
5739                        &accounts.pool_mint_key,
5740                        &accounts.pool_fee_key,
5741                        &pool_key,
5742                        &accounts.token_a_key,
5743                        &accounts.token_b_key,
5744                        &token_a_key,
5745                        &accounts.token_a_mint_key,
5746                        WithdrawSingleTokenTypeExactAmountOut {
5747                            destination_token_amount: destination_a_amount,
5748                            maximum_pool_token_amount,
5749                        }
5750                    )
5751                    .unwrap(),
5752                    vec![
5753                        &mut accounts.swap_account,
5754                        &mut SolanaAccount::default(),
5755                        &mut SolanaAccount::default(),
5756                        &mut accounts.pool_mint_account,
5757                        &mut pool_account,
5758                        &mut accounts.token_a_account,
5759                        &mut accounts.token_b_account,
5760                        &mut token_a_account,
5761                        &mut accounts.pool_fee_account,
5762                        &mut accounts.token_a_mint_account,
5763                        &mut SolanaAccount::default(),
5764                        &mut SolanaAccount::default(),
5765                    ],
5766                )
5767            );
5768        }
5769
5770        // wrong swap token accounts
5771        {
5772            let (
5773                token_a_key,
5774                mut token_a_account,
5775                token_b_key,
5776                mut token_b_account,
5777                pool_key,
5778                mut pool_account,
5779            ) = accounts.setup_token_accounts(
5780                &user_key,
5781                &withdrawer_key,
5782                initial_a,
5783                initial_b,
5784                initial_pool.try_into().unwrap(),
5785            );
5786
5787            let old_a_key = accounts.token_a_key;
5788            let old_a_account = accounts.token_a_account;
5789
5790            accounts.token_a_key = token_a_key;
5791            accounts.token_a_account = token_a_account.clone();
5792
5793            // wrong swap token a account
5794            assert_eq!(
5795                Err(SwapError::IncorrectSwapAccount.into()),
5796                accounts.withdraw_single_token_type_exact_amount_out(
5797                    &withdrawer_key,
5798                    &pool_key,
5799                    &mut pool_account,
5800                    &token_a_key,
5801                    &mut token_a_account,
5802                    destination_a_amount,
5803                    maximum_pool_token_amount,
5804                )
5805            );
5806
5807            accounts.token_a_key = old_a_key;
5808            accounts.token_a_account = old_a_account;
5809
5810            let old_b_key = accounts.token_b_key;
5811            let old_b_account = accounts.token_b_account;
5812
5813            accounts.token_b_key = token_b_key;
5814            accounts.token_b_account = token_b_account.clone();
5815
5816            // wrong swap token b account
5817            assert_eq!(
5818                Err(SwapError::IncorrectSwapAccount.into()),
5819                accounts.withdraw_single_token_type_exact_amount_out(
5820                    &withdrawer_key,
5821                    &pool_key,
5822                    &mut pool_account,
5823                    &token_b_key,
5824                    &mut token_b_account,
5825                    destination_b_amount,
5826                    maximum_pool_token_amount,
5827                )
5828            );
5829
5830            accounts.token_b_key = old_b_key;
5831            accounts.token_b_account = old_b_account;
5832        }
5833
5834        // wrong mint
5835        {
5836            let (
5837                token_a_key,
5838                mut token_a_account,
5839                _token_b_key,
5840                _token_b_account,
5841                pool_key,
5842                mut pool_account,
5843            ) = accounts.setup_token_accounts(
5844                &user_key,
5845                &withdrawer_key,
5846                initial_a,
5847                initial_b,
5848                initial_pool.try_into().unwrap(),
5849            );
5850            let (pool_mint_key, pool_mint_account) = create_mint(
5851                &pool_token_program_id,
5852                &accounts.authority_key,
5853                None,
5854                None,
5855                &TransferFee::default(),
5856            );
5857            let old_pool_key = accounts.pool_mint_key;
5858            let old_pool_account = accounts.pool_mint_account;
5859            accounts.pool_mint_key = pool_mint_key;
5860            accounts.pool_mint_account = pool_mint_account;
5861
5862            assert_eq!(
5863                Err(SwapError::IncorrectPoolMint.into()),
5864                accounts.withdraw_single_token_type_exact_amount_out(
5865                    &withdrawer_key,
5866                    &pool_key,
5867                    &mut pool_account,
5868                    &token_a_key,
5869                    &mut token_a_account,
5870                    destination_a_amount,
5871                    maximum_pool_token_amount,
5872                )
5873            );
5874
5875            accounts.pool_mint_key = old_pool_key;
5876            accounts.pool_mint_account = old_pool_account;
5877        }
5878
5879        // slippage exceeded
5880        {
5881            let (
5882                token_a_key,
5883                mut token_a_account,
5884                token_b_key,
5885                mut token_b_account,
5886                pool_key,
5887                mut pool_account,
5888            ) = accounts.setup_token_accounts(
5889                &user_key,
5890                &withdrawer_key,
5891                initial_a,
5892                initial_b,
5893                maximum_pool_token_amount,
5894            );
5895
5896            // maximum pool token amount too low
5897            assert_eq!(
5898                Err(SwapError::ExceededSlippage.into()),
5899                accounts.withdraw_single_token_type_exact_amount_out(
5900                    &withdrawer_key,
5901                    &pool_key,
5902                    &mut pool_account,
5903                    &token_a_key,
5904                    &mut token_a_account,
5905                    destination_a_amount,
5906                    maximum_pool_token_amount / 1000,
5907                )
5908            );
5909            assert_eq!(
5910                Err(SwapError::ExceededSlippage.into()),
5911                accounts.withdraw_single_token_type_exact_amount_out(
5912                    &withdrawer_key,
5913                    &pool_key,
5914                    &mut pool_account,
5915                    &token_b_key,
5916                    &mut token_b_account,
5917                    destination_b_amount,
5918                    maximum_pool_token_amount / 1000,
5919                )
5920            );
5921        }
5922
5923        // invalid input: can't use swap pool tokens as destination
5924        {
5925            let (
5926                _token_a_key,
5927                _token_a_account,
5928                _token_b_key,
5929                _token_b_account,
5930                pool_key,
5931                mut pool_account,
5932            ) = accounts.setup_token_accounts(
5933                &user_key,
5934                &withdrawer_key,
5935                initial_a,
5936                initial_b,
5937                maximum_pool_token_amount,
5938            );
5939            let swap_token_a_key = accounts.token_a_key;
5940            let mut swap_token_a_account = accounts.get_token_account(&swap_token_a_key).clone();
5941            assert_eq!(
5942                Err(SwapError::InvalidInput.into()),
5943                accounts.withdraw_single_token_type_exact_amount_out(
5944                    &withdrawer_key,
5945                    &pool_key,
5946                    &mut pool_account,
5947                    &swap_token_a_key,
5948                    &mut swap_token_a_account,
5949                    destination_a_amount,
5950                    maximum_pool_token_amount,
5951                )
5952            );
5953            let swap_token_b_key = accounts.token_b_key;
5954            let mut swap_token_b_account = accounts.get_token_account(&swap_token_b_key).clone();
5955            assert_eq!(
5956                Err(SwapError::InvalidInput.into()),
5957                accounts.withdraw_single_token_type_exact_amount_out(
5958                    &withdrawer_key,
5959                    &pool_key,
5960                    &mut pool_account,
5961                    &swap_token_b_key,
5962                    &mut swap_token_b_account,
5963                    destination_b_amount,
5964                    maximum_pool_token_amount,
5965                )
5966            );
5967        }
5968
5969        // correct withdrawal
5970        {
5971            let (
5972                token_a_key,
5973                mut token_a_account,
5974                _token_b_key,
5975                _token_b_account,
5976                pool_key,
5977                mut pool_account,
5978            ) = accounts.setup_token_accounts(
5979                &user_key,
5980                &withdrawer_key,
5981                initial_a,
5982                initial_b,
5983                initial_pool.try_into().unwrap(),
5984            );
5985
5986            let swap_token_a =
5987                StateWithExtensions::<Account>::unpack(&accounts.token_a_account.data).unwrap();
5988            let swap_token_b =
5989                StateWithExtensions::<Account>::unpack(&accounts.token_b_account.data).unwrap();
5990            let pool_mint =
5991                StateWithExtensions::<Mint>::unpack(&accounts.pool_mint_account.data).unwrap();
5992
5993            let pool_token_amount = accounts
5994                .swap_curve
5995                .withdraw_single_token_type_exact_out(
5996                    destination_a_amount.try_into().unwrap(),
5997                    swap_token_a.base.amount.try_into().unwrap(),
5998                    swap_token_b.base.amount.try_into().unwrap(),
5999                    pool_mint.base.supply.try_into().unwrap(),
6000                    TradeDirection::AtoB,
6001                    &accounts.fees,
6002                )
6003                .unwrap();
6004            let withdraw_fee = accounts.fees.owner_withdraw_fee(pool_token_amount).unwrap();
6005
6006            accounts
6007                .withdraw_single_token_type_exact_amount_out(
6008                    &withdrawer_key,
6009                    &pool_key,
6010                    &mut pool_account,
6011                    &token_a_key,
6012                    &mut token_a_account,
6013                    destination_a_amount,
6014                    maximum_pool_token_amount,
6015                )
6016                .unwrap();
6017
6018            let swap_token_a =
6019                StateWithExtensions::<Account>::unpack(&accounts.token_a_account.data).unwrap();
6020
6021            assert_eq!(
6022                swap_token_a.base.amount,
6023                token_a_amount - destination_a_amount
6024            );
6025            let token_a = StateWithExtensions::<Account>::unpack(&token_a_account.data).unwrap();
6026            assert_eq!(token_a.base.amount, initial_a + destination_a_amount);
6027
6028            let pool_account = StateWithExtensions::<Account>::unpack(&pool_account.data).unwrap();
6029            assert_eq!(
6030                pool_account.base.amount,
6031                to_u64(initial_pool - pool_token_amount - withdraw_fee).unwrap()
6032            );
6033            let fee_account =
6034                StateWithExtensions::<Account>::unpack(&accounts.pool_fee_account.data).unwrap();
6035            assert_eq!(fee_account.base.amount, to_u64(withdraw_fee).unwrap());
6036        }
6037
6038        // correct withdrawal from fee account
6039        {
6040            let (
6041                token_a_key,
6042                mut token_a_account,
6043                _token_b_key,
6044                _token_b_account,
6045                _pool_key,
6046                _pool_account,
6047            ) = accounts.setup_token_accounts(&user_key, &withdrawer_key, initial_a, initial_b, 0);
6048
6049            let fee_a_amount = 2;
6050            let pool_fee_key = accounts.pool_fee_key;
6051            let mut pool_fee_account = accounts.pool_fee_account.clone();
6052            let fee_account =
6053                StateWithExtensions::<Account>::unpack(&pool_fee_account.data).unwrap();
6054            let pool_fee_amount = fee_account.base.amount;
6055
6056            let swap_token_a =
6057                StateWithExtensions::<Account>::unpack(&accounts.token_a_account.data).unwrap();
6058
6059            let token_a_amount = swap_token_a.base.amount;
6060            accounts
6061                .withdraw_single_token_type_exact_amount_out(
6062                    &user_key,
6063                    &pool_fee_key,
6064                    &mut pool_fee_account,
6065                    &token_a_key,
6066                    &mut token_a_account,
6067                    fee_a_amount,
6068                    pool_fee_amount,
6069                )
6070                .unwrap();
6071
6072            let swap_token_a =
6073                StateWithExtensions::<Account>::unpack(&accounts.token_a_account.data).unwrap();
6074
6075            assert_eq!(swap_token_a.base.amount, token_a_amount - fee_a_amount);
6076            let token_a = StateWithExtensions::<Account>::unpack(&token_a_account.data).unwrap();
6077            assert_eq!(token_a.base.amount, initial_a + fee_a_amount);
6078        }
6079    }
6080
6081    #[allow(clippy::too_many_arguments)]
6082    fn check_valid_swap_curve(
6083        fees: Fees,
6084        transfer_fees: SwapTransferFees,
6085        curve_type: CurveType,
6086        calculator: Arc<dyn CurveCalculator + Send + Sync>,
6087        token_a_amount: u64,
6088        token_b_amount: u64,
6089        pool_token_program_id: &Pubkey,
6090        token_a_program_id: &Pubkey,
6091        token_b_program_id: &Pubkey,
6092    ) {
6093        let user_key = Pubkey::new_unique();
6094        let swapper_key = Pubkey::new_unique();
6095
6096        let swap_curve = SwapCurve {
6097            curve_type,
6098            calculator,
6099        };
6100
6101        let mut accounts = SwapAccountInfo::new(
6102            &user_key,
6103            fees.clone(),
6104            transfer_fees,
6105            swap_curve.clone(),
6106            token_a_amount,
6107            token_b_amount,
6108            pool_token_program_id,
6109            token_a_program_id,
6110            token_b_program_id,
6111        );
6112        let initial_a = token_a_amount / 5;
6113        let initial_b = token_b_amount / 5;
6114        accounts.initialize_swap().unwrap();
6115
6116        let swap_token_a_key = accounts.token_a_key;
6117        let swap_token_b_key = accounts.token_b_key;
6118
6119        let (
6120            token_a_key,
6121            mut token_a_account,
6122            token_b_key,
6123            mut token_b_account,
6124            _pool_key,
6125            _pool_account,
6126        ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
6127        // swap one way
6128        let a_to_b_amount = initial_a / 10;
6129        let minimum_token_b_amount = 0;
6130        let pool_mint =
6131            StateWithExtensions::<Mint>::unpack(&accounts.pool_mint_account.data).unwrap();
6132        let initial_supply = pool_mint.base.supply;
6133        accounts
6134            .swap(
6135                &swapper_key,
6136                &token_a_key,
6137                &mut token_a_account,
6138                &swap_token_a_key,
6139                &swap_token_b_key,
6140                &token_b_key,
6141                &mut token_b_account,
6142                a_to_b_amount,
6143                minimum_token_b_amount,
6144            )
6145            .unwrap();
6146
6147        // tweak values based on transfer fees assessed
6148        let token_a_fee = accounts
6149            .transfer_fees
6150            .token_a
6151            .calculate_fee(a_to_b_amount)
6152            .unwrap();
6153        let actual_a_to_b_amount = a_to_b_amount - token_a_fee;
6154        let results = swap_curve
6155            .swap(
6156                actual_a_to_b_amount.try_into().unwrap(),
6157                token_a_amount.try_into().unwrap(),
6158                token_b_amount.try_into().unwrap(),
6159                TradeDirection::AtoB,
6160                &fees,
6161            )
6162            .unwrap();
6163
6164        let swap_token_a =
6165            StateWithExtensions::<Account>::unpack(&accounts.token_a_account.data).unwrap();
6166        let token_a_amount = swap_token_a.base.amount;
6167        assert_eq!(
6168            token_a_amount,
6169            TryInto::<u64>::try_into(results.new_swap_source_amount).unwrap()
6170        );
6171        let token_a = StateWithExtensions::<Account>::unpack(&token_a_account.data).unwrap();
6172        assert_eq!(token_a.base.amount, initial_a - a_to_b_amount);
6173
6174        let swap_token_b =
6175            StateWithExtensions::<Account>::unpack(&accounts.token_b_account.data).unwrap();
6176        let token_b_amount = swap_token_b.base.amount;
6177        assert_eq!(
6178            token_b_amount,
6179            TryInto::<u64>::try_into(results.new_swap_destination_amount).unwrap()
6180        );
6181        let token_b = StateWithExtensions::<Account>::unpack(&token_b_account.data).unwrap();
6182        assert_eq!(
6183            token_b.base.amount,
6184            initial_b + to_u64(results.destination_amount_swapped).unwrap()
6185        );
6186
6187        let first_fee = if results.owner_fee > 0 {
6188            swap_curve
6189                .calculator
6190                .withdraw_single_token_type_exact_out(
6191                    results.owner_fee,
6192                    token_a_amount.try_into().unwrap(),
6193                    token_b_amount.try_into().unwrap(),
6194                    initial_supply.try_into().unwrap(),
6195                    TradeDirection::AtoB,
6196                    RoundDirection::Floor,
6197                )
6198                .unwrap()
6199        } else {
6200            0
6201        };
6202        let fee_account =
6203            StateWithExtensions::<Account>::unpack(&accounts.pool_fee_account.data).unwrap();
6204        assert_eq!(
6205            fee_account.base.amount,
6206            TryInto::<u64>::try_into(first_fee).unwrap()
6207        );
6208
6209        let first_swap_amount = results.destination_amount_swapped;
6210
6211        // swap the other way
6212        let pool_mint =
6213            StateWithExtensions::<Mint>::unpack(&accounts.pool_mint_account.data).unwrap();
6214        let initial_supply = pool_mint.base.supply;
6215
6216        let b_to_a_amount = initial_b / 10;
6217        let minimum_a_amount = 0;
6218        accounts
6219            .swap(
6220                &swapper_key,
6221                &token_b_key,
6222                &mut token_b_account,
6223                &swap_token_b_key,
6224                &swap_token_a_key,
6225                &token_a_key,
6226                &mut token_a_account,
6227                b_to_a_amount,
6228                minimum_a_amount,
6229            )
6230            .unwrap();
6231
6232        let mut results = swap_curve
6233            .swap(
6234                b_to_a_amount.try_into().unwrap(),
6235                token_b_amount.try_into().unwrap(),
6236                token_a_amount.try_into().unwrap(),
6237                TradeDirection::BtoA,
6238                &fees,
6239            )
6240            .unwrap();
6241        // tweak values based on transfer fees assessed
6242        let token_a_fee = accounts
6243            .transfer_fees
6244            .token_a
6245            .calculate_fee(results.destination_amount_swapped.try_into().unwrap())
6246            .unwrap();
6247        results.destination_amount_swapped -= token_a_fee as u128;
6248
6249        let swap_token_a =
6250            StateWithExtensions::<Account>::unpack(&accounts.token_a_account.data).unwrap();
6251        let token_a_amount = swap_token_a.base.amount;
6252        assert_eq!(
6253            token_a_amount,
6254            TryInto::<u64>::try_into(results.new_swap_destination_amount).unwrap()
6255        );
6256        let token_a = StateWithExtensions::<Account>::unpack(&token_a_account.data).unwrap();
6257        assert_eq!(
6258            token_a.base.amount,
6259            initial_a - a_to_b_amount + to_u64(results.destination_amount_swapped).unwrap()
6260        );
6261
6262        let swap_token_b =
6263            StateWithExtensions::<Account>::unpack(&accounts.token_b_account.data).unwrap();
6264        let token_b_amount = swap_token_b.base.amount;
6265        assert_eq!(
6266            token_b_amount,
6267            TryInto::<u64>::try_into(results.new_swap_source_amount).unwrap()
6268        );
6269        let token_b = StateWithExtensions::<Account>::unpack(&token_b_account.data).unwrap();
6270        assert_eq!(
6271            token_b.base.amount,
6272            initial_b + to_u64(first_swap_amount).unwrap()
6273                - to_u64(results.source_amount_swapped).unwrap()
6274        );
6275
6276        let second_fee = if results.owner_fee > 0 {
6277            swap_curve
6278                .calculator
6279                .withdraw_single_token_type_exact_out(
6280                    results.owner_fee,
6281                    token_a_amount.try_into().unwrap(),
6282                    token_b_amount.try_into().unwrap(),
6283                    initial_supply.try_into().unwrap(),
6284                    TradeDirection::BtoA,
6285                    RoundDirection::Floor,
6286                )
6287                .unwrap()
6288        } else {
6289            0
6290        };
6291        let fee_account =
6292            StateWithExtensions::<Account>::unpack(&accounts.pool_fee_account.data).unwrap();
6293        assert_eq!(
6294            fee_account.base.amount,
6295            to_u64(first_fee + second_fee).unwrap()
6296        );
6297    }
6298
6299    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
6300    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
6301    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
6302    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
6303    fn test_valid_swap_curve_all_fees(
6304        pool_token_program_id: Pubkey,
6305        token_a_program_id: Pubkey,
6306        token_b_program_id: Pubkey,
6307    ) {
6308        // All fees
6309        let trade_fee_numerator = 1;
6310        let trade_fee_denominator = 10;
6311        let owner_trade_fee_numerator = 1;
6312        let owner_trade_fee_denominator = 30;
6313        let owner_withdraw_fee_numerator = 1;
6314        let owner_withdraw_fee_denominator = 30;
6315        let host_fee_numerator = 20;
6316        let host_fee_denominator = 100;
6317        let fees = Fees {
6318            trade_fee_numerator,
6319            trade_fee_denominator,
6320            owner_trade_fee_numerator,
6321            owner_trade_fee_denominator,
6322            owner_withdraw_fee_numerator,
6323            owner_withdraw_fee_denominator,
6324            host_fee_numerator,
6325            host_fee_denominator,
6326        };
6327
6328        let token_a_amount = 10_000_000_000;
6329        let token_b_amount = 50_000_000_000;
6330
6331        check_valid_swap_curve(
6332            fees.clone(),
6333            SwapTransferFees::default(),
6334            CurveType::ConstantProduct,
6335            Arc::new(ConstantProductCurve {}),
6336            token_a_amount,
6337            token_b_amount,
6338            &pool_token_program_id,
6339            &token_a_program_id,
6340            &token_b_program_id,
6341        );
6342        let token_b_price = 1;
6343        check_valid_swap_curve(
6344            fees.clone(),
6345            SwapTransferFees::default(),
6346            CurveType::ConstantPrice,
6347            Arc::new(ConstantPriceCurve { token_b_price }),
6348            token_a_amount,
6349            token_b_amount,
6350            &pool_token_program_id,
6351            &token_a_program_id,
6352            &token_b_program_id,
6353        );
6354        let token_b_offset = 10_000_000_000;
6355        check_valid_swap_curve(
6356            fees,
6357            SwapTransferFees::default(),
6358            CurveType::Offset,
6359            Arc::new(OffsetCurve { token_b_offset }),
6360            token_a_amount,
6361            token_b_amount,
6362            &pool_token_program_id,
6363            &token_a_program_id,
6364            &token_b_program_id,
6365        );
6366    }
6367
6368    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
6369    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
6370    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
6371    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
6372    fn test_valid_swap_curve_trade_fee_only(
6373        pool_token_program_id: Pubkey,
6374        token_a_program_id: Pubkey,
6375        token_b_program_id: Pubkey,
6376    ) {
6377        let trade_fee_numerator = 1;
6378        let trade_fee_denominator = 10;
6379        let owner_trade_fee_numerator = 0;
6380        let owner_trade_fee_denominator = 0;
6381        let owner_withdraw_fee_numerator = 0;
6382        let owner_withdraw_fee_denominator = 0;
6383        let host_fee_numerator = 0;
6384        let host_fee_denominator = 0;
6385        let fees = Fees {
6386            trade_fee_numerator,
6387            trade_fee_denominator,
6388            owner_trade_fee_numerator,
6389            owner_trade_fee_denominator,
6390            owner_withdraw_fee_numerator,
6391            owner_withdraw_fee_denominator,
6392            host_fee_numerator,
6393            host_fee_denominator,
6394        };
6395
6396        let token_a_amount = 10_000_000_000;
6397        let token_b_amount = 50_000_000_000;
6398
6399        check_valid_swap_curve(
6400            fees.clone(),
6401            SwapTransferFees::default(),
6402            CurveType::ConstantProduct,
6403            Arc::new(ConstantProductCurve {}),
6404            token_a_amount,
6405            token_b_amount,
6406            &pool_token_program_id,
6407            &token_a_program_id,
6408            &token_b_program_id,
6409        );
6410        let token_b_price = 10_000;
6411        check_valid_swap_curve(
6412            fees.clone(),
6413            SwapTransferFees::default(),
6414            CurveType::ConstantPrice,
6415            Arc::new(ConstantPriceCurve { token_b_price }),
6416            token_a_amount,
6417            token_b_amount / token_b_price,
6418            &pool_token_program_id,
6419            &token_a_program_id,
6420            &token_b_program_id,
6421        );
6422        let token_b_offset = 1;
6423        check_valid_swap_curve(
6424            fees,
6425            SwapTransferFees::default(),
6426            CurveType::Offset,
6427            Arc::new(OffsetCurve { token_b_offset }),
6428            token_a_amount,
6429            token_b_amount,
6430            &pool_token_program_id,
6431            &token_a_program_id,
6432            &token_b_program_id,
6433        );
6434    }
6435
6436    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
6437    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
6438    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
6439    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
6440    fn test_valid_swap_with_fee_constraints(
6441        pool_token_program_id: Pubkey,
6442        token_a_program_id: Pubkey,
6443        token_b_program_id: Pubkey,
6444    ) {
6445        let owner_key = Pubkey::new_unique();
6446
6447        let trade_fee_numerator = 1;
6448        let trade_fee_denominator = 10;
6449        let owner_trade_fee_numerator = 1;
6450        let owner_trade_fee_denominator = 30;
6451        let owner_withdraw_fee_numerator = 1;
6452        let owner_withdraw_fee_denominator = 30;
6453        let host_fee_numerator = 10;
6454        let host_fee_denominator = 100;
6455
6456        let token_a_amount = 1_000_000;
6457        let token_b_amount = 5_000_000;
6458
6459        let fees = Fees {
6460            trade_fee_numerator,
6461            trade_fee_denominator,
6462            owner_trade_fee_numerator,
6463            owner_trade_fee_denominator,
6464            owner_withdraw_fee_numerator,
6465            owner_withdraw_fee_denominator,
6466            host_fee_numerator,
6467            host_fee_denominator,
6468        };
6469
6470        let curve = ConstantProductCurve {};
6471        let swap_curve = SwapCurve {
6472            curve_type: CurveType::ConstantProduct,
6473            calculator: Arc::new(curve),
6474        };
6475
6476        let owner_key_str = &owner_key.to_string();
6477        let valid_curve_types = &[CurveType::ConstantProduct];
6478        let constraints = Some(SwapConstraints {
6479            owner_key: owner_key_str,
6480            valid_curve_types,
6481            fees: &fees,
6482        });
6483        let mut accounts = SwapAccountInfo::new(
6484            &owner_key,
6485            fees.clone(),
6486            SwapTransferFees::default(),
6487            swap_curve,
6488            token_a_amount,
6489            token_b_amount,
6490            &pool_token_program_id,
6491            &token_a_program_id,
6492            &token_b_program_id,
6493        );
6494
6495        // initialize swap
6496        do_process_instruction_with_fee_constraints(
6497            initialize(
6498                &SWAP_PROGRAM_ID,
6499                &pool_token_program_id,
6500                &accounts.swap_key,
6501                &accounts.authority_key,
6502                &accounts.token_a_key,
6503                &accounts.token_b_key,
6504                &accounts.pool_mint_key,
6505                &accounts.pool_fee_key,
6506                &accounts.pool_token_key,
6507                accounts.fees.clone(),
6508                accounts.swap_curve.clone(),
6509            )
6510            .unwrap(),
6511            vec![
6512                &mut accounts.swap_account,
6513                &mut SolanaAccount::default(),
6514                &mut accounts.token_a_account,
6515                &mut accounts.token_b_account,
6516                &mut accounts.pool_mint_account,
6517                &mut accounts.pool_fee_account,
6518                &mut accounts.pool_token_account,
6519                &mut SolanaAccount::default(),
6520            ],
6521            &constraints,
6522        )
6523        .unwrap();
6524
6525        let authority_key = accounts.authority_key;
6526
6527        let (
6528            token_a_key,
6529            mut token_a_account,
6530            token_b_key,
6531            mut token_b_account,
6532            pool_key,
6533            mut pool_account,
6534        ) = accounts.setup_token_accounts(
6535            &owner_key,
6536            &authority_key,
6537            token_a_amount,
6538            token_b_amount,
6539            0,
6540        );
6541
6542        let amount_in = token_a_amount / 2;
6543        let minimum_amount_out = 0;
6544
6545        // perform the swap
6546        do_process_instruction_with_fee_constraints(
6547            swap(
6548                &SWAP_PROGRAM_ID,
6549                &token_a_program_id,
6550                &token_b_program_id,
6551                &pool_token_program_id,
6552                &accounts.swap_key,
6553                &accounts.authority_key,
6554                &accounts.authority_key,
6555                &token_a_key,
6556                &accounts.token_a_key,
6557                &accounts.token_b_key,
6558                &token_b_key,
6559                &accounts.pool_mint_key,
6560                &accounts.pool_fee_key,
6561                &accounts.token_a_mint_key,
6562                &accounts.token_b_mint_key,
6563                Some(&pool_key),
6564                Swap {
6565                    amount_in,
6566                    minimum_amount_out,
6567                },
6568            )
6569            .unwrap(),
6570            vec![
6571                &mut accounts.swap_account,
6572                &mut SolanaAccount::default(),
6573                &mut SolanaAccount::default(),
6574                &mut token_a_account,
6575                &mut accounts.token_a_account,
6576                &mut accounts.token_b_account,
6577                &mut token_b_account,
6578                &mut accounts.pool_mint_account,
6579                &mut accounts.pool_fee_account,
6580                &mut accounts.token_a_mint_account,
6581                &mut accounts.token_b_mint_account,
6582                &mut SolanaAccount::default(),
6583                &mut SolanaAccount::default(),
6584                &mut SolanaAccount::default(),
6585                &mut pool_account,
6586            ],
6587            &constraints,
6588        )
6589        .unwrap();
6590
6591        // check that fees were taken in the host fee account
6592        let host_fee_account = StateWithExtensions::<Account>::unpack(&pool_account.data).unwrap();
6593        let owner_fee_account =
6594            StateWithExtensions::<Account>::unpack(&accounts.pool_fee_account.data).unwrap();
6595        let total_fee = owner_fee_account.base.amount * host_fee_denominator
6596            / (host_fee_denominator - host_fee_numerator);
6597        assert_eq!(
6598            total_fee,
6599            host_fee_account.base.amount + owner_fee_account.base.amount
6600        );
6601    }
6602
6603    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
6604    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
6605    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
6606    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
6607    fn test_invalid_swap(
6608        pool_token_program_id: Pubkey,
6609        token_a_program_id: Pubkey,
6610        token_b_program_id: Pubkey,
6611    ) {
6612        let user_key = Pubkey::new_unique();
6613        let swapper_key = Pubkey::new_unique();
6614        let trade_fee_numerator = 1;
6615        let trade_fee_denominator = 4;
6616        let owner_trade_fee_numerator = 1;
6617        let owner_trade_fee_denominator = 10;
6618        let owner_withdraw_fee_numerator = 1;
6619        let owner_withdraw_fee_denominator = 5;
6620        let host_fee_numerator = 9;
6621        let host_fee_denominator = 100;
6622        let fees = Fees {
6623            trade_fee_numerator,
6624            trade_fee_denominator,
6625            owner_trade_fee_numerator,
6626            owner_trade_fee_denominator,
6627            owner_withdraw_fee_numerator,
6628            owner_withdraw_fee_denominator,
6629            host_fee_numerator,
6630            host_fee_denominator,
6631        };
6632
6633        let token_a_amount = 1000;
6634        let token_b_amount = 5000;
6635        let curve_type = CurveType::ConstantProduct;
6636        let swap_curve = SwapCurve {
6637            curve_type,
6638            calculator: Arc::new(ConstantProductCurve {}),
6639        };
6640        let mut accounts = SwapAccountInfo::new(
6641            &user_key,
6642            fees,
6643            SwapTransferFees::default(),
6644            swap_curve,
6645            token_a_amount,
6646            token_b_amount,
6647            &pool_token_program_id,
6648            &token_a_program_id,
6649            &token_b_program_id,
6650        );
6651
6652        let initial_a = token_a_amount / 5;
6653        let initial_b = token_b_amount / 5;
6654        let minimum_token_b_amount = initial_b / 2;
6655
6656        let swap_token_a_key = accounts.token_a_key;
6657        let swap_token_b_key = accounts.token_b_key;
6658
6659        // swap not initialized
6660        {
6661            let (
6662                token_a_key,
6663                mut token_a_account,
6664                token_b_key,
6665                mut token_b_account,
6666                _pool_key,
6667                _pool_account,
6668            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
6669            assert_eq!(
6670                Err(ProgramError::UninitializedAccount),
6671                accounts.swap(
6672                    &swapper_key,
6673                    &token_a_key,
6674                    &mut token_a_account,
6675                    &swap_token_a_key,
6676                    &swap_token_b_key,
6677                    &token_b_key,
6678                    &mut token_b_account,
6679                    initial_a,
6680                    minimum_token_b_amount,
6681                )
6682            );
6683        }
6684
6685        accounts.initialize_swap().unwrap();
6686
6687        // wrong swap account program id
6688        {
6689            let (
6690                token_a_key,
6691                mut token_a_account,
6692                token_b_key,
6693                mut token_b_account,
6694                _pool_key,
6695                _pool_account,
6696            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
6697            let old_swap_account = accounts.swap_account;
6698            let mut wrong_swap_account = old_swap_account.clone();
6699            wrong_swap_account.owner = pool_token_program_id;
6700            accounts.swap_account = wrong_swap_account;
6701            assert_eq!(
6702                Err(ProgramError::IncorrectProgramId),
6703                accounts.swap(
6704                    &swapper_key,
6705                    &token_a_key,
6706                    &mut token_a_account,
6707                    &swap_token_a_key,
6708                    &swap_token_b_key,
6709                    &token_b_key,
6710                    &mut token_b_account,
6711                    initial_a,
6712                    minimum_token_b_amount,
6713                )
6714            );
6715            accounts.swap_account = old_swap_account;
6716        }
6717
6718        // wrong bump seed
6719        {
6720            let (
6721                token_a_key,
6722                mut token_a_account,
6723                token_b_key,
6724                mut token_b_account,
6725                _pool_key,
6726                _pool_account,
6727            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
6728            let old_authority = accounts.authority_key;
6729            let (bad_authority_key, _bump_seed) = Pubkey::find_program_address(
6730                &[&accounts.swap_key.to_bytes()[..]],
6731                &pool_token_program_id,
6732            );
6733            accounts.authority_key = bad_authority_key;
6734            assert_eq!(
6735                Err(SwapError::InvalidProgramAddress.into()),
6736                accounts.swap(
6737                    &swapper_key,
6738                    &token_a_key,
6739                    &mut token_a_account,
6740                    &swap_token_a_key,
6741                    &swap_token_b_key,
6742                    &token_b_key,
6743                    &mut token_b_account,
6744                    initial_a,
6745                    minimum_token_b_amount,
6746                )
6747            );
6748            accounts.authority_key = old_authority;
6749        }
6750
6751        // wrong token program id
6752        {
6753            let (
6754                token_a_key,
6755                mut token_a_account,
6756                token_b_key,
6757                mut token_b_account,
6758                _pool_key,
6759                _pool_account,
6760            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
6761            let wrong_program_id = Pubkey::new_unique();
6762            assert_eq!(
6763                Err(SwapError::IncorrectTokenProgramId.into()),
6764                do_process_instruction(
6765                    swap(
6766                        &SWAP_PROGRAM_ID,
6767                        &wrong_program_id,
6768                        &wrong_program_id,
6769                        &wrong_program_id,
6770                        &accounts.swap_key,
6771                        &accounts.authority_key,
6772                        &accounts.authority_key,
6773                        &token_a_key,
6774                        &accounts.token_a_key,
6775                        &accounts.token_b_key,
6776                        &token_b_key,
6777                        &accounts.pool_mint_key,
6778                        &accounts.pool_fee_key,
6779                        &accounts.token_a_mint_key,
6780                        &accounts.token_b_mint_key,
6781                        None,
6782                        Swap {
6783                            amount_in: initial_a,
6784                            minimum_amount_out: minimum_token_b_amount,
6785                        },
6786                    )
6787                    .unwrap(),
6788                    vec![
6789                        &mut accounts.swap_account,
6790                        &mut SolanaAccount::default(),
6791                        &mut SolanaAccount::default(),
6792                        &mut token_a_account,
6793                        &mut accounts.token_a_account,
6794                        &mut accounts.token_b_account,
6795                        &mut token_b_account,
6796                        &mut accounts.pool_mint_account,
6797                        &mut accounts.pool_fee_account,
6798                        &mut accounts.token_a_mint_account,
6799                        &mut accounts.token_b_mint_account,
6800                        &mut SolanaAccount::default(),
6801                        &mut SolanaAccount::default(),
6802                        &mut SolanaAccount::default(),
6803                    ],
6804                ),
6805            );
6806        }
6807
6808        // not enough token a to swap
6809        {
6810            let (
6811                token_a_key,
6812                mut token_a_account,
6813                token_b_key,
6814                mut token_b_account,
6815                _pool_key,
6816                _pool_account,
6817            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
6818            assert_eq!(
6819                Err(TokenError::InsufficientFunds.into()),
6820                accounts.swap(
6821                    &swapper_key,
6822                    &token_a_key,
6823                    &mut token_a_account,
6824                    &swap_token_a_key,
6825                    &swap_token_b_key,
6826                    &token_b_key,
6827                    &mut token_b_account,
6828                    initial_a * 2,
6829                    minimum_token_b_amount * 2,
6830                )
6831            );
6832        }
6833
6834        // wrong swap token A / B accounts
6835        {
6836            let (
6837                token_a_key,
6838                mut token_a_account,
6839                token_b_key,
6840                mut token_b_account,
6841                _pool_key,
6842                _pool_account,
6843            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
6844            let user_transfer_key = Pubkey::new_unique();
6845            assert_eq!(
6846                Err(SwapError::IncorrectSwapAccount.into()),
6847                do_process_instruction(
6848                    swap(
6849                        &SWAP_PROGRAM_ID,
6850                        &token_a_program_id,
6851                        &token_b_program_id,
6852                        &pool_token_program_id,
6853                        &accounts.swap_key,
6854                        &accounts.authority_key,
6855                        &user_transfer_key,
6856                        &token_a_key,
6857                        &token_a_key,
6858                        &token_b_key,
6859                        &token_b_key,
6860                        &accounts.pool_mint_key,
6861                        &accounts.pool_fee_key,
6862                        &accounts.token_a_mint_key,
6863                        &accounts.token_b_mint_key,
6864                        None,
6865                        Swap {
6866                            amount_in: initial_a,
6867                            minimum_amount_out: minimum_token_b_amount,
6868                        },
6869                    )
6870                    .unwrap(),
6871                    vec![
6872                        &mut accounts.swap_account,
6873                        &mut SolanaAccount::default(),
6874                        &mut SolanaAccount::default(),
6875                        &mut token_a_account.clone(),
6876                        &mut token_a_account,
6877                        &mut token_b_account.clone(),
6878                        &mut token_b_account,
6879                        &mut accounts.pool_mint_account,
6880                        &mut accounts.pool_fee_account,
6881                        &mut accounts.token_a_mint_account,
6882                        &mut accounts.token_b_mint_account,
6883                        &mut SolanaAccount::default(),
6884                        &mut SolanaAccount::default(),
6885                        &mut SolanaAccount::default(),
6886                    ],
6887                ),
6888            );
6889        }
6890
6891        // wrong user token A / B accounts
6892        {
6893            let (
6894                token_a_key,
6895                mut token_a_account,
6896                token_b_key,
6897                mut token_b_account,
6898                _pool_key,
6899                _pool_account,
6900            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
6901            assert_eq!(
6902                Err(TokenError::MintMismatch.into()),
6903                accounts.swap(
6904                    &swapper_key,
6905                    &token_b_key,
6906                    &mut token_b_account,
6907                    &swap_token_a_key,
6908                    &swap_token_b_key,
6909                    &token_a_key,
6910                    &mut token_a_account,
6911                    initial_a,
6912                    minimum_token_b_amount,
6913                )
6914            );
6915        }
6916
6917        // swap from a to a
6918        {
6919            let (
6920                token_a_key,
6921                mut token_a_account,
6922                _token_b_key,
6923                _token_b_account,
6924                _pool_key,
6925                _pool_account,
6926            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
6927            assert_eq!(
6928                Err(SwapError::InvalidInput.into()),
6929                accounts.swap(
6930                    &swapper_key,
6931                    &token_a_key,
6932                    &mut token_a_account.clone(),
6933                    &swap_token_a_key,
6934                    &swap_token_a_key,
6935                    &token_a_key,
6936                    &mut token_a_account,
6937                    initial_a,
6938                    minimum_token_b_amount,
6939                )
6940            );
6941        }
6942
6943        // incorrect mint provided
6944        {
6945            let (
6946                token_a_key,
6947                mut token_a_account,
6948                token_b_key,
6949                mut token_b_account,
6950                _pool_key,
6951                _pool_account,
6952            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
6953            let (pool_mint_key, pool_mint_account) = create_mint(
6954                &pool_token_program_id,
6955                &accounts.authority_key,
6956                None,
6957                None,
6958                &TransferFee::default(),
6959            );
6960            let old_pool_key = accounts.pool_mint_key;
6961            let old_pool_account = accounts.pool_mint_account;
6962            accounts.pool_mint_key = pool_mint_key;
6963            accounts.pool_mint_account = pool_mint_account;
6964
6965            assert_eq!(
6966                Err(SwapError::IncorrectPoolMint.into()),
6967                accounts.swap(
6968                    &swapper_key,
6969                    &token_a_key,
6970                    &mut token_a_account,
6971                    &swap_token_a_key,
6972                    &swap_token_b_key,
6973                    &token_b_key,
6974                    &mut token_b_account,
6975                    initial_a,
6976                    minimum_token_b_amount,
6977                )
6978            );
6979
6980            accounts.pool_mint_key = old_pool_key;
6981            accounts.pool_mint_account = old_pool_account;
6982        }
6983
6984        // incorrect fee account provided
6985        {
6986            let (
6987                token_a_key,
6988                mut token_a_account,
6989                token_b_key,
6990                mut token_b_account,
6991                wrong_pool_key,
6992                wrong_pool_account,
6993            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
6994            let old_pool_fee_account = accounts.pool_fee_account;
6995            let old_pool_fee_key = accounts.pool_fee_key;
6996            accounts.pool_fee_account = wrong_pool_account;
6997            accounts.pool_fee_key = wrong_pool_key;
6998            assert_eq!(
6999                Err(SwapError::IncorrectFeeAccount.into()),
7000                accounts.swap(
7001                    &swapper_key,
7002                    &token_a_key,
7003                    &mut token_a_account,
7004                    &swap_token_a_key,
7005                    &swap_token_b_key,
7006                    &token_b_key,
7007                    &mut token_b_account,
7008                    initial_a,
7009                    minimum_token_b_amount,
7010                )
7011            );
7012            accounts.pool_fee_account = old_pool_fee_account;
7013            accounts.pool_fee_key = old_pool_fee_key;
7014        }
7015
7016        // no approval
7017        {
7018            let (
7019                token_a_key,
7020                mut token_a_account,
7021                token_b_key,
7022                mut token_b_account,
7023                _pool_key,
7024                _pool_account,
7025            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
7026            let user_transfer_key = Pubkey::new_unique();
7027            assert_eq!(
7028                Err(TokenError::OwnerMismatch.into()),
7029                do_process_instruction(
7030                    swap(
7031                        &SWAP_PROGRAM_ID,
7032                        &token_a_program_id,
7033                        &token_b_program_id,
7034                        &pool_token_program_id,
7035                        &accounts.swap_key,
7036                        &accounts.authority_key,
7037                        &user_transfer_key,
7038                        &token_a_key,
7039                        &accounts.token_a_key,
7040                        &accounts.token_b_key,
7041                        &token_b_key,
7042                        &accounts.pool_mint_key,
7043                        &accounts.pool_fee_key,
7044                        &accounts.token_a_mint_key,
7045                        &accounts.token_b_mint_key,
7046                        None,
7047                        Swap {
7048                            amount_in: initial_a,
7049                            minimum_amount_out: minimum_token_b_amount,
7050                        },
7051                    )
7052                    .unwrap(),
7053                    vec![
7054                        &mut accounts.swap_account,
7055                        &mut SolanaAccount::default(),
7056                        &mut SolanaAccount::default(),
7057                        &mut token_a_account,
7058                        &mut accounts.token_a_account,
7059                        &mut accounts.token_b_account,
7060                        &mut token_b_account,
7061                        &mut accounts.pool_mint_account,
7062                        &mut accounts.pool_fee_account,
7063                        &mut accounts.token_a_mint_account,
7064                        &mut accounts.token_b_mint_account,
7065                        &mut SolanaAccount::default(),
7066                        &mut SolanaAccount::default(),
7067                        &mut SolanaAccount::default(),
7068                    ],
7069                ),
7070            );
7071        }
7072
7073        // output token value 0
7074        {
7075            let (
7076                token_a_key,
7077                mut token_a_account,
7078                token_b_key,
7079                mut token_b_account,
7080                _pool_key,
7081                _pool_account,
7082            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
7083            assert_eq!(
7084                Err(SwapError::ZeroTradingTokens.into()),
7085                accounts.swap(
7086                    &swapper_key,
7087                    &token_b_key,
7088                    &mut token_b_account,
7089                    &swap_token_b_key,
7090                    &swap_token_a_key,
7091                    &token_a_key,
7092                    &mut token_a_account,
7093                    1,
7094                    1,
7095                )
7096            );
7097        }
7098
7099        // slippage exceeded: minimum out amount too high
7100        {
7101            let (
7102                token_a_key,
7103                mut token_a_account,
7104                token_b_key,
7105                mut token_b_account,
7106                _pool_key,
7107                _pool_account,
7108            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
7109            assert_eq!(
7110                Err(SwapError::ExceededSlippage.into()),
7111                accounts.swap(
7112                    &swapper_key,
7113                    &token_a_key,
7114                    &mut token_a_account,
7115                    &swap_token_a_key,
7116                    &swap_token_b_key,
7117                    &token_b_key,
7118                    &mut token_b_account,
7119                    initial_a,
7120                    minimum_token_b_amount * 2,
7121                )
7122            );
7123        }
7124
7125        // invalid input: can't use swap pool as user source / dest
7126        {
7127            let (
7128                token_a_key,
7129                mut token_a_account,
7130                token_b_key,
7131                mut token_b_account,
7132                _pool_key,
7133                _pool_account,
7134            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
7135            let mut swap_token_a_account = accounts.get_token_account(&swap_token_a_key).clone();
7136            let authority_key = accounts.authority_key;
7137            assert_eq!(
7138                Err(SwapError::InvalidInput.into()),
7139                accounts.swap(
7140                    &authority_key,
7141                    &swap_token_a_key,
7142                    &mut swap_token_a_account,
7143                    &swap_token_a_key,
7144                    &swap_token_b_key,
7145                    &token_b_key,
7146                    &mut token_b_account,
7147                    initial_a,
7148                    minimum_token_b_amount,
7149                )
7150            );
7151            let mut swap_token_b_account = accounts.get_token_account(&swap_token_b_key).clone();
7152            assert_eq!(
7153                Err(SwapError::InvalidInput.into()),
7154                accounts.swap(
7155                    &swapper_key,
7156                    &token_a_key,
7157                    &mut token_a_account,
7158                    &swap_token_a_key,
7159                    &swap_token_b_key,
7160                    &swap_token_b_key,
7161                    &mut swap_token_b_account,
7162                    initial_a,
7163                    minimum_token_b_amount,
7164                )
7165            );
7166        }
7167
7168        // still correct: constraint specified, no host fee account
7169        {
7170            let authority_key = accounts.authority_key;
7171            let (
7172                token_a_key,
7173                mut token_a_account,
7174                token_b_key,
7175                mut token_b_account,
7176                _pool_key,
7177                _pool_account,
7178            ) = accounts.setup_token_accounts(&user_key, &authority_key, initial_a, initial_b, 0);
7179            let owner_key = &swapper_key.to_string();
7180            let fees = Fees {
7181                trade_fee_numerator,
7182                trade_fee_denominator,
7183                owner_trade_fee_numerator,
7184                owner_trade_fee_denominator,
7185                owner_withdraw_fee_numerator,
7186                owner_withdraw_fee_denominator,
7187                host_fee_numerator,
7188                host_fee_denominator,
7189            };
7190            let constraints = Some(SwapConstraints {
7191                owner_key,
7192                valid_curve_types: &[],
7193                fees: &fees,
7194            });
7195            do_process_instruction_with_fee_constraints(
7196                swap(
7197                    &SWAP_PROGRAM_ID,
7198                    &token_a_program_id,
7199                    &token_b_program_id,
7200                    &pool_token_program_id,
7201                    &accounts.swap_key,
7202                    &accounts.authority_key,
7203                    &accounts.authority_key,
7204                    &token_a_key,
7205                    &accounts.token_a_key,
7206                    &accounts.token_b_key,
7207                    &token_b_key,
7208                    &accounts.pool_mint_key,
7209                    &accounts.pool_fee_key,
7210                    &accounts.token_a_mint_key,
7211                    &accounts.token_b_mint_key,
7212                    None,
7213                    Swap {
7214                        amount_in: initial_a,
7215                        minimum_amount_out: minimum_token_b_amount,
7216                    },
7217                )
7218                .unwrap(),
7219                vec![
7220                    &mut accounts.swap_account,
7221                    &mut SolanaAccount::default(),
7222                    &mut SolanaAccount::default(),
7223                    &mut token_a_account,
7224                    &mut accounts.token_a_account,
7225                    &mut accounts.token_b_account,
7226                    &mut token_b_account,
7227                    &mut accounts.pool_mint_account,
7228                    &mut accounts.pool_fee_account,
7229                    &mut accounts.token_a_mint_account,
7230                    &mut accounts.token_b_mint_account,
7231                    &mut SolanaAccount::default(),
7232                    &mut SolanaAccount::default(),
7233                    &mut SolanaAccount::default(),
7234                ],
7235                &constraints,
7236            )
7237            .unwrap();
7238        }
7239
7240        // invalid mint for host fee account
7241        {
7242            let authority_key = accounts.authority_key;
7243            let (
7244                token_a_key,
7245                mut token_a_account,
7246                token_b_key,
7247                mut token_b_account,
7248                _pool_key,
7249                _pool_account,
7250            ) = accounts.setup_token_accounts(&user_key, &authority_key, initial_a, initial_b, 0);
7251            let (
7252                bad_token_a_key,
7253                mut bad_token_a_account,
7254                _token_b_key,
7255                mut _token_b_account,
7256                _pool_key,
7257                _pool_account,
7258            ) = accounts.setup_token_accounts(&user_key, &authority_key, initial_a, initial_b, 0);
7259            let owner_key = &swapper_key.to_string();
7260            let fees = Fees {
7261                trade_fee_numerator,
7262                trade_fee_denominator,
7263                owner_trade_fee_numerator,
7264                owner_trade_fee_denominator,
7265                owner_withdraw_fee_numerator,
7266                owner_withdraw_fee_denominator,
7267                host_fee_numerator,
7268                host_fee_denominator,
7269            };
7270            let constraints = Some(SwapConstraints {
7271                owner_key,
7272                valid_curve_types: &[],
7273                fees: &fees,
7274            });
7275            assert_eq!(
7276                Err(SwapError::IncorrectPoolMint.into()),
7277                do_process_instruction_with_fee_constraints(
7278                    swap(
7279                        &SWAP_PROGRAM_ID,
7280                        &token_a_program_id,
7281                        &token_b_program_id,
7282                        &pool_token_program_id,
7283                        &accounts.swap_key,
7284                        &accounts.authority_key,
7285                        &accounts.authority_key,
7286                        &token_a_key,
7287                        &accounts.token_a_key,
7288                        &accounts.token_b_key,
7289                        &token_b_key,
7290                        &accounts.pool_mint_key,
7291                        &accounts.pool_fee_key,
7292                        &accounts.token_a_mint_key,
7293                        &accounts.token_b_mint_key,
7294                        Some(&bad_token_a_key),
7295                        Swap {
7296                            amount_in: initial_a,
7297                            minimum_amount_out: 0,
7298                        },
7299                    )
7300                    .unwrap(),
7301                    vec![
7302                        &mut accounts.swap_account,
7303                        &mut SolanaAccount::default(),
7304                        &mut SolanaAccount::default(),
7305                        &mut token_a_account,
7306                        &mut accounts.token_a_account,
7307                        &mut accounts.token_b_account,
7308                        &mut token_b_account,
7309                        &mut accounts.pool_mint_account,
7310                        &mut accounts.pool_fee_account,
7311                        &mut accounts.token_a_mint_account,
7312                        &mut accounts.token_b_mint_account,
7313                        &mut SolanaAccount::default(),
7314                        &mut SolanaAccount::default(),
7315                        &mut SolanaAccount::default(),
7316                        &mut bad_token_a_account,
7317                    ],
7318                    &constraints,
7319                ),
7320            );
7321        }
7322    }
7323
7324    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
7325    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
7326    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
7327    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
7328    fn test_overdraw_offset_curve(
7329        pool_token_program_id: Pubkey,
7330        token_a_program_id: Pubkey,
7331        token_b_program_id: Pubkey,
7332    ) {
7333        let trade_fee_numerator = 1;
7334        let trade_fee_denominator = 10;
7335        let owner_trade_fee_numerator = 1;
7336        let owner_trade_fee_denominator = 30;
7337        let owner_withdraw_fee_numerator = 1;
7338        let owner_withdraw_fee_denominator = 30;
7339        let host_fee_numerator = 10;
7340        let host_fee_denominator = 100;
7341
7342        let token_a_amount = 1_000_000_000;
7343        let token_b_amount = 0;
7344        let fees = Fees {
7345            trade_fee_numerator,
7346            trade_fee_denominator,
7347            owner_trade_fee_numerator,
7348            owner_trade_fee_denominator,
7349            owner_withdraw_fee_numerator,
7350            owner_withdraw_fee_denominator,
7351            host_fee_numerator,
7352            host_fee_denominator,
7353        };
7354
7355        let token_b_offset = 2_000_000;
7356        let swap_curve = SwapCurve {
7357            curve_type: CurveType::Offset,
7358            calculator: Arc::new(OffsetCurve { token_b_offset }),
7359        };
7360        let user_key = Pubkey::new_unique();
7361        let swapper_key = Pubkey::new_unique();
7362
7363        let mut accounts = SwapAccountInfo::new(
7364            &user_key,
7365            fees,
7366            SwapTransferFees::default(),
7367            swap_curve,
7368            token_a_amount,
7369            token_b_amount,
7370            &pool_token_program_id,
7371            &token_a_program_id,
7372            &token_b_program_id,
7373        );
7374
7375        accounts.initialize_swap().unwrap();
7376
7377        let swap_token_a_key = accounts.token_a_key;
7378        let swap_token_b_key = accounts.token_b_key;
7379        let initial_a = 500_000;
7380        let initial_b = 1_000;
7381
7382        let (
7383            token_a_key,
7384            mut token_a_account,
7385            token_b_key,
7386            mut token_b_account,
7387            _pool_key,
7388            _pool_account,
7389        ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
7390
7391        // swap a to b way, fails, there's no liquidity
7392        let a_to_b_amount = initial_a;
7393        let minimum_token_b_amount = 0;
7394
7395        assert_eq!(
7396            Err(SwapError::ZeroTradingTokens.into()),
7397            accounts.swap(
7398                &swapper_key,
7399                &token_a_key,
7400                &mut token_a_account,
7401                &swap_token_a_key,
7402                &swap_token_b_key,
7403                &token_b_key,
7404                &mut token_b_account,
7405                a_to_b_amount,
7406                minimum_token_b_amount,
7407            )
7408        );
7409
7410        // swap b to a, succeeds at offset price
7411        let b_to_a_amount = initial_b;
7412        let minimum_token_a_amount = 0;
7413        accounts
7414            .swap(
7415                &swapper_key,
7416                &token_b_key,
7417                &mut token_b_account,
7418                &swap_token_b_key,
7419                &swap_token_a_key,
7420                &token_a_key,
7421                &mut token_a_account,
7422                b_to_a_amount,
7423                minimum_token_a_amount,
7424            )
7425            .unwrap();
7426
7427        // try a to b again, succeeds due to new liquidity
7428        accounts
7429            .swap(
7430                &swapper_key,
7431                &token_a_key,
7432                &mut token_a_account,
7433                &swap_token_a_key,
7434                &swap_token_b_key,
7435                &token_b_key,
7436                &mut token_b_account,
7437                a_to_b_amount,
7438                minimum_token_b_amount,
7439            )
7440            .unwrap();
7441
7442        // try a to b again, fails due to no more liquidity
7443        assert_eq!(
7444            Err(SwapError::ZeroTradingTokens.into()),
7445            accounts.swap(
7446                &swapper_key,
7447                &token_a_key,
7448                &mut token_a_account,
7449                &swap_token_a_key,
7450                &swap_token_b_key,
7451                &token_b_key,
7452                &mut token_b_account,
7453                a_to_b_amount,
7454                minimum_token_b_amount,
7455            )
7456        );
7457
7458        // Try to deposit, fails because deposits are not allowed for offset
7459        // curve swaps
7460        {
7461            let initial_a = 100;
7462            let initial_b = 100;
7463            let pool_amount = 100;
7464            let (
7465                token_a_key,
7466                mut token_a_account,
7467                token_b_key,
7468                mut token_b_account,
7469                pool_key,
7470                mut pool_account,
7471            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
7472            assert_eq!(
7473                Err(SwapError::UnsupportedCurveOperation.into()),
7474                accounts.deposit_all_token_types(
7475                    &swapper_key,
7476                    &token_a_key,
7477                    &mut token_a_account,
7478                    &token_b_key,
7479                    &mut token_b_account,
7480                    &pool_key,
7481                    &mut pool_account,
7482                    pool_amount,
7483                    initial_a,
7484                    initial_b,
7485                )
7486            );
7487        }
7488    }
7489
7490    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
7491    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
7492    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
7493    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
7494    fn test_withdraw_all_offset_curve(
7495        pool_token_program_id: Pubkey,
7496        token_a_program_id: Pubkey,
7497        token_b_program_id: Pubkey,
7498    ) {
7499        let trade_fee_numerator = 1;
7500        let trade_fee_denominator = 10;
7501        let owner_trade_fee_numerator = 1;
7502        let owner_trade_fee_denominator = 30;
7503        let owner_withdraw_fee_numerator = 0;
7504        let owner_withdraw_fee_denominator = 30;
7505        let host_fee_numerator = 10;
7506        let host_fee_denominator = 100;
7507
7508        let token_a_amount = 1_000_000_000;
7509        let token_b_amount = 10;
7510        let fees = Fees {
7511            trade_fee_numerator,
7512            trade_fee_denominator,
7513            owner_trade_fee_numerator,
7514            owner_trade_fee_denominator,
7515            owner_withdraw_fee_numerator,
7516            owner_withdraw_fee_denominator,
7517            host_fee_numerator,
7518            host_fee_denominator,
7519        };
7520
7521        let token_b_offset = 2_000_000;
7522        let swap_curve = SwapCurve {
7523            curve_type: CurveType::Offset,
7524            calculator: Arc::new(OffsetCurve { token_b_offset }),
7525        };
7526        let total_pool = swap_curve.calculator.new_pool_supply();
7527        let user_key = Pubkey::new_unique();
7528        let withdrawer_key = Pubkey::new_unique();
7529
7530        let mut accounts = SwapAccountInfo::new(
7531            &user_key,
7532            fees,
7533            SwapTransferFees::default(),
7534            swap_curve,
7535            token_a_amount,
7536            token_b_amount,
7537            &pool_token_program_id,
7538            &token_a_program_id,
7539            &token_b_program_id,
7540        );
7541
7542        accounts.initialize_swap().unwrap();
7543
7544        let (
7545            token_a_key,
7546            mut token_a_account,
7547            token_b_key,
7548            mut token_b_account,
7549            _pool_key,
7550            _pool_account,
7551        ) = accounts.setup_token_accounts(&user_key, &withdrawer_key, 0, 0, 0);
7552
7553        let pool_key = accounts.pool_token_key;
7554        let mut pool_account = accounts.pool_token_account.clone();
7555
7556        // WithdrawAllTokenTypes takes all tokens for A and B.
7557        // The curve's calculation for token B will say to transfer
7558        // `token_b_offset + token_b_amount`, but only `token_b_amount` will be
7559        // moved.
7560        accounts
7561            .withdraw_all_token_types(
7562                &user_key,
7563                &pool_key,
7564                &mut pool_account,
7565                &token_a_key,
7566                &mut token_a_account,
7567                &token_b_key,
7568                &mut token_b_account,
7569                total_pool.try_into().unwrap(),
7570                0,
7571                0,
7572            )
7573            .unwrap();
7574
7575        let token_a = StateWithExtensions::<Account>::unpack(&token_a_account.data).unwrap();
7576        assert_eq!(token_a.base.amount, token_a_amount);
7577        let token_b = StateWithExtensions::<Account>::unpack(&token_b_account.data).unwrap();
7578        assert_eq!(token_b.base.amount, token_b_amount);
7579        let swap_token_a =
7580            StateWithExtensions::<Account>::unpack(&accounts.token_a_account.data).unwrap();
7581        assert_eq!(swap_token_a.base.amount, 0);
7582        let swap_token_b =
7583            StateWithExtensions::<Account>::unpack(&accounts.token_b_account.data).unwrap();
7584        assert_eq!(swap_token_b.base.amount, 0);
7585    }
7586
7587    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
7588    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
7589    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
7590    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
7591    fn test_withdraw_all_constant_price_curve(
7592        pool_token_program_id: Pubkey,
7593        token_a_program_id: Pubkey,
7594        token_b_program_id: Pubkey,
7595    ) {
7596        let trade_fee_numerator = 1;
7597        let trade_fee_denominator = 10;
7598        let owner_trade_fee_numerator = 1;
7599        let owner_trade_fee_denominator = 30;
7600        let owner_withdraw_fee_numerator = 0;
7601        let owner_withdraw_fee_denominator = 30;
7602        let host_fee_numerator = 10;
7603        let host_fee_denominator = 100;
7604
7605        // initialize "unbalanced", so that withdrawing all will have some issues
7606        // A: 1_000_000_000
7607        // B: 2_000_000_000 (1_000 * 2_000_000)
7608        let swap_token_a_amount = 1_000_000_000;
7609        let swap_token_b_amount = 1_000;
7610        let token_b_price = 2_000_000;
7611        let fees = Fees {
7612            trade_fee_numerator,
7613            trade_fee_denominator,
7614            owner_trade_fee_numerator,
7615            owner_trade_fee_denominator,
7616            owner_withdraw_fee_numerator,
7617            owner_withdraw_fee_denominator,
7618            host_fee_numerator,
7619            host_fee_denominator,
7620        };
7621
7622        let swap_curve = SwapCurve {
7623            curve_type: CurveType::ConstantPrice,
7624            calculator: Arc::new(ConstantPriceCurve { token_b_price }),
7625        };
7626        let total_pool = swap_curve.calculator.new_pool_supply();
7627        let user_key = Pubkey::new_unique();
7628        let withdrawer_key = Pubkey::new_unique();
7629
7630        let mut accounts = SwapAccountInfo::new(
7631            &user_key,
7632            fees,
7633            SwapTransferFees::default(),
7634            swap_curve,
7635            swap_token_a_amount,
7636            swap_token_b_amount,
7637            &pool_token_program_id,
7638            &token_a_program_id,
7639            &token_b_program_id,
7640        );
7641
7642        accounts.initialize_swap().unwrap();
7643
7644        let (
7645            token_a_key,
7646            mut token_a_account,
7647            token_b_key,
7648            mut token_b_account,
7649            _pool_key,
7650            _pool_account,
7651        ) = accounts.setup_token_accounts(&user_key, &withdrawer_key, 0, 0, 0);
7652
7653        let pool_key = accounts.pool_token_key;
7654        let mut pool_account = accounts.pool_token_account.clone();
7655
7656        // WithdrawAllTokenTypes will not take all token A and B, since their
7657        // ratio is unbalanced.  It will try to take 1_500_000_000 worth of
7658        // each token, which means 1_500_000_000 token A, and 750 token B.
7659        // With no slippage, this will leave 250 token B in the pool.
7660        assert_eq!(
7661            Err(SwapError::ExceededSlippage.into()),
7662            accounts.withdraw_all_token_types(
7663                &user_key,
7664                &pool_key,
7665                &mut pool_account,
7666                &token_a_key,
7667                &mut token_a_account,
7668                &token_b_key,
7669                &mut token_b_account,
7670                total_pool.try_into().unwrap(),
7671                swap_token_a_amount,
7672                swap_token_b_amount,
7673            )
7674        );
7675
7676        accounts
7677            .withdraw_all_token_types(
7678                &user_key,
7679                &pool_key,
7680                &mut pool_account,
7681                &token_a_key,
7682                &mut token_a_account,
7683                &token_b_key,
7684                &mut token_b_account,
7685                total_pool.try_into().unwrap(),
7686                0,
7687                0,
7688            )
7689            .unwrap();
7690
7691        let token_a = StateWithExtensions::<Account>::unpack(&token_a_account.data).unwrap();
7692        assert_eq!(token_a.base.amount, swap_token_a_amount);
7693        let token_b = StateWithExtensions::<Account>::unpack(&token_b_account.data).unwrap();
7694        assert_eq!(token_b.base.amount, 750);
7695        let swap_token_a =
7696            StateWithExtensions::<Account>::unpack(&accounts.token_a_account.data).unwrap();
7697        assert_eq!(swap_token_a.base.amount, 0);
7698        let swap_token_b =
7699            StateWithExtensions::<Account>::unpack(&accounts.token_b_account.data).unwrap();
7700        assert_eq!(swap_token_b.base.amount, 250);
7701
7702        // deposit now, not enough to cover the tokens already in there
7703        let token_b_amount = 10;
7704        let token_a_amount = token_b_amount * token_b_price;
7705        let (
7706            token_a_key,
7707            mut token_a_account,
7708            token_b_key,
7709            mut token_b_account,
7710            pool_key,
7711            mut pool_account,
7712        ) = accounts.setup_token_accounts(
7713            &user_key,
7714            &withdrawer_key,
7715            token_a_amount,
7716            token_b_amount,
7717            0,
7718        );
7719
7720        assert_eq!(
7721            Err(SwapError::ExceededSlippage.into()),
7722            accounts.deposit_all_token_types(
7723                &withdrawer_key,
7724                &token_a_key,
7725                &mut token_a_account,
7726                &token_b_key,
7727                &mut token_b_account,
7728                &pool_key,
7729                &mut pool_account,
7730                1, // doesn't matter
7731                token_a_amount,
7732                token_b_amount,
7733            )
7734        );
7735
7736        // deposit enough tokens, success!
7737        let token_b_amount = 125;
7738        let token_a_amount = token_b_amount * token_b_price;
7739        let (
7740            token_a_key,
7741            mut token_a_account,
7742            token_b_key,
7743            mut token_b_account,
7744            pool_key,
7745            mut pool_account,
7746        ) = accounts.setup_token_accounts(
7747            &user_key,
7748            &withdrawer_key,
7749            token_a_amount,
7750            token_b_amount,
7751            0,
7752        );
7753
7754        accounts
7755            .deposit_all_token_types(
7756                &withdrawer_key,
7757                &token_a_key,
7758                &mut token_a_account,
7759                &token_b_key,
7760                &mut token_b_account,
7761                &pool_key,
7762                &mut pool_account,
7763                1, // doesn't matter
7764                token_a_amount,
7765                token_b_amount,
7766            )
7767            .unwrap();
7768    }
7769
7770    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
7771    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
7772    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
7773    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
7774    fn test_deposits_allowed_single_token(
7775        pool_token_program_id: Pubkey,
7776        token_a_program_id: Pubkey,
7777        token_b_program_id: Pubkey,
7778    ) {
7779        let trade_fee_numerator = 1;
7780        let trade_fee_denominator = 10;
7781        let owner_trade_fee_numerator = 1;
7782        let owner_trade_fee_denominator = 30;
7783        let owner_withdraw_fee_numerator = 0;
7784        let owner_withdraw_fee_denominator = 30;
7785        let host_fee_numerator = 10;
7786        let host_fee_denominator = 100;
7787
7788        let token_a_amount = 1_000_000;
7789        let token_b_amount = 0;
7790        let fees = Fees {
7791            trade_fee_numerator,
7792            trade_fee_denominator,
7793            owner_trade_fee_numerator,
7794            owner_trade_fee_denominator,
7795            owner_withdraw_fee_numerator,
7796            owner_withdraw_fee_denominator,
7797            host_fee_numerator,
7798            host_fee_denominator,
7799        };
7800
7801        let token_b_offset = 2_000_000;
7802        let swap_curve = SwapCurve {
7803            curve_type: CurveType::Offset,
7804            calculator: Arc::new(OffsetCurve { token_b_offset }),
7805        };
7806        let creator_key = Pubkey::new_unique();
7807        let depositor_key = Pubkey::new_unique();
7808
7809        let mut accounts = SwapAccountInfo::new(
7810            &creator_key,
7811            fees,
7812            SwapTransferFees::default(),
7813            swap_curve,
7814            token_a_amount,
7815            token_b_amount,
7816            &pool_token_program_id,
7817            &token_a_program_id,
7818            &token_b_program_id,
7819        );
7820
7821        accounts.initialize_swap().unwrap();
7822
7823        let initial_a = 1_000_000;
7824        let initial_b = 2_000_000;
7825        let (
7826            _depositor_token_a_key,
7827            _depositor_token_a_account,
7828            depositor_token_b_key,
7829            mut depositor_token_b_account,
7830            depositor_pool_key,
7831            mut depositor_pool_account,
7832        ) = accounts.setup_token_accounts(&creator_key, &depositor_key, initial_a, initial_b, 0);
7833
7834        assert_eq!(
7835            Err(SwapError::UnsupportedCurveOperation.into()),
7836            accounts.deposit_single_token_type_exact_amount_in(
7837                &depositor_key,
7838                &depositor_token_b_key,
7839                &mut depositor_token_b_account,
7840                &depositor_pool_key,
7841                &mut depositor_pool_account,
7842                initial_b,
7843                0,
7844            )
7845        );
7846    }
7847
7848    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
7849    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
7850    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
7851    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
7852    fn test_withdraw_with_invalid_fee_account(
7853        pool_token_program_id: Pubkey,
7854        token_a_program_id: Pubkey,
7855        token_b_program_id: Pubkey,
7856    ) {
7857        let user_key = Pubkey::new_unique();
7858
7859        let fees = Fees {
7860            trade_fee_numerator: 1,
7861            trade_fee_denominator: 2,
7862            owner_trade_fee_numerator: 1,
7863            owner_trade_fee_denominator: 10,
7864            owner_withdraw_fee_numerator: 1,
7865            owner_withdraw_fee_denominator: 5,
7866            host_fee_numerator: 7,
7867            host_fee_denominator: 100,
7868        };
7869
7870        let token_a_amount = 1000;
7871        let token_b_amount = 2000;
7872        let swap_curve = SwapCurve {
7873            curve_type: CurveType::ConstantProduct,
7874            calculator: Arc::new(ConstantProductCurve {}),
7875        };
7876
7877        let withdrawer_key = Pubkey::new_unique();
7878        let initial_a = token_a_amount / 10;
7879        let initial_b = token_b_amount / 10;
7880        let initial_pool = swap_curve.calculator.new_pool_supply() / 10;
7881        let withdraw_amount = initial_pool / 4;
7882        let minimum_token_a_amount = initial_a / 40;
7883        let minimum_token_b_amount = initial_b / 40;
7884
7885        let mut accounts = SwapAccountInfo::new(
7886            &user_key,
7887            fees,
7888            SwapTransferFees::default(),
7889            swap_curve,
7890            token_a_amount,
7891            token_b_amount,
7892            &pool_token_program_id,
7893            &token_a_program_id,
7894            &token_b_program_id,
7895        );
7896
7897        accounts.initialize_swap().unwrap();
7898
7899        let (
7900            token_a_key,
7901            mut token_a_account,
7902            token_b_key,
7903            mut token_b_account,
7904            pool_key,
7905            mut pool_account,
7906        ) = accounts.setup_token_accounts(
7907            &user_key,
7908            &withdrawer_key,
7909            initial_a,
7910            initial_b,
7911            initial_pool.try_into().unwrap(),
7912        );
7913
7914        let destination_key = Pubkey::new_unique();
7915        let mut destination = SolanaAccount::new(
7916            account_minimum_balance(),
7917            Account::get_packed_len(),
7918            &withdrawer_key,
7919        );
7920
7921        do_process_instruction(
7922            close_account(
7923                &pool_token_program_id,
7924                &accounts.pool_fee_key,
7925                &destination_key,
7926                &user_key,
7927                &[],
7928            )
7929            .unwrap(),
7930            vec![
7931                &mut accounts.pool_fee_account,
7932                &mut destination,
7933                &mut SolanaAccount::default(),
7934            ],
7935        )
7936        .unwrap();
7937
7938        let user_transfer_authority_key = Pubkey::new_unique();
7939        let pool_token_amount = withdraw_amount.try_into().unwrap();
7940
7941        do_process_instruction(
7942            approve(
7943                &pool_token_program_id,
7944                &pool_key,
7945                &user_transfer_authority_key,
7946                &withdrawer_key,
7947                &[],
7948                pool_token_amount,
7949            )
7950            .unwrap(),
7951            vec![
7952                &mut pool_account,
7953                &mut SolanaAccount::default(),
7954                &mut SolanaAccount::default(),
7955            ],
7956        )
7957        .unwrap();
7958
7959        do_process_instruction(
7960            withdraw_all_token_types(
7961                &SWAP_PROGRAM_ID,
7962                &pool_token_program_id,
7963                &token_a_program_id,
7964                &token_b_program_id,
7965                &accounts.swap_key,
7966                &accounts.authority_key,
7967                &user_transfer_authority_key,
7968                &accounts.pool_mint_key,
7969                &accounts.pool_fee_key,
7970                &pool_key,
7971                &accounts.token_a_key,
7972                &accounts.token_b_key,
7973                &token_a_key,
7974                &token_b_key,
7975                &accounts.token_a_mint_key,
7976                &accounts.token_b_mint_key,
7977                WithdrawAllTokenTypes {
7978                    pool_token_amount,
7979                    minimum_token_a_amount,
7980                    minimum_token_b_amount,
7981                },
7982            )
7983            .unwrap(),
7984            vec![
7985                &mut accounts.swap_account,
7986                &mut SolanaAccount::default(),
7987                &mut SolanaAccount::default(),
7988                &mut accounts.pool_mint_account,
7989                &mut pool_account,
7990                &mut accounts.token_a_account,
7991                &mut accounts.token_b_account,
7992                &mut token_a_account,
7993                &mut token_b_account,
7994                &mut accounts.pool_fee_account,
7995                &mut accounts.token_a_mint_account,
7996                &mut accounts.token_b_mint_account,
7997                &mut SolanaAccount::default(),
7998                &mut SolanaAccount::default(),
7999                &mut SolanaAccount::default(),
8000            ],
8001        )
8002        .unwrap();
8003    }
8004
8005    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
8006    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
8007    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
8008    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
8009    fn test_withdraw_one_exact_out_with_invalid_fee_account(
8010        pool_token_program_id: Pubkey,
8011        token_a_program_id: Pubkey,
8012        token_b_program_id: Pubkey,
8013    ) {
8014        let user_key = Pubkey::new_unique();
8015
8016        let fees = Fees {
8017            trade_fee_numerator: 1,
8018            trade_fee_denominator: 2,
8019            owner_trade_fee_numerator: 1,
8020            owner_trade_fee_denominator: 10,
8021            owner_withdraw_fee_numerator: 1,
8022            owner_withdraw_fee_denominator: 5,
8023            host_fee_numerator: 7,
8024            host_fee_denominator: 100,
8025        };
8026
8027        let token_a_amount = 1000;
8028        let token_b_amount = 2000;
8029        let swap_curve = SwapCurve {
8030            curve_type: CurveType::ConstantProduct,
8031            calculator: Arc::new(ConstantProductCurve {}),
8032        };
8033
8034        let withdrawer_key = Pubkey::new_unique();
8035        let initial_a = token_a_amount / 10;
8036        let initial_b = token_b_amount / 10;
8037        let initial_pool = swap_curve.calculator.new_pool_supply() / 10;
8038        let maximum_pool_token_amount = to_u64(initial_pool / 4).unwrap();
8039        let destination_a_amount = initial_a / 40;
8040
8041        let mut accounts = SwapAccountInfo::new(
8042            &user_key,
8043            fees,
8044            SwapTransferFees::default(),
8045            swap_curve,
8046            token_a_amount,
8047            token_b_amount,
8048            &pool_token_program_id,
8049            &token_a_program_id,
8050            &token_b_program_id,
8051        );
8052
8053        accounts.initialize_swap().unwrap();
8054
8055        let (
8056            token_a_key,
8057            mut token_a_account,
8058            _token_b_key,
8059            _token_b_account,
8060            pool_key,
8061            mut pool_account,
8062        ) = accounts.setup_token_accounts(
8063            &user_key,
8064            &withdrawer_key,
8065            initial_a,
8066            initial_b,
8067            initial_pool.try_into().unwrap(),
8068        );
8069
8070        let destination_key = Pubkey::new_unique();
8071        let mut destination = SolanaAccount::new(
8072            account_minimum_balance(),
8073            Account::get_packed_len(),
8074            &withdrawer_key,
8075        );
8076
8077        do_process_instruction(
8078            close_account(
8079                &pool_token_program_id,
8080                &accounts.pool_fee_key,
8081                &destination_key,
8082                &user_key,
8083                &[],
8084            )
8085            .unwrap(),
8086            vec![
8087                &mut accounts.pool_fee_account,
8088                &mut destination,
8089                &mut SolanaAccount::default(),
8090            ],
8091        )
8092        .unwrap();
8093
8094        let user_transfer_authority_key = Pubkey::new_unique();
8095
8096        do_process_instruction(
8097            approve(
8098                &pool_token_program_id,
8099                &pool_key,
8100                &user_transfer_authority_key,
8101                &withdrawer_key,
8102                &[],
8103                maximum_pool_token_amount,
8104            )
8105            .unwrap(),
8106            vec![
8107                &mut pool_account,
8108                &mut SolanaAccount::default(),
8109                &mut SolanaAccount::default(),
8110            ],
8111        )
8112        .unwrap();
8113
8114        do_process_instruction(
8115            withdraw_single_token_type_exact_amount_out(
8116                &SWAP_PROGRAM_ID,
8117                &pool_token_program_id,
8118                &token_a_program_id,
8119                &accounts.swap_key,
8120                &accounts.authority_key,
8121                &user_transfer_authority_key,
8122                &accounts.pool_mint_key,
8123                &accounts.pool_fee_key,
8124                &pool_key,
8125                &accounts.token_a_key,
8126                &accounts.token_b_key,
8127                &token_a_key,
8128                &accounts.token_a_mint_key,
8129                WithdrawSingleTokenTypeExactAmountOut {
8130                    destination_token_amount: destination_a_amount,
8131                    maximum_pool_token_amount,
8132                },
8133            )
8134            .unwrap(),
8135            vec![
8136                &mut accounts.swap_account,
8137                &mut SolanaAccount::default(),
8138                &mut SolanaAccount::default(),
8139                &mut accounts.pool_mint_account,
8140                &mut pool_account,
8141                &mut accounts.token_a_account,
8142                &mut accounts.token_b_account,
8143                &mut token_a_account,
8144                &mut accounts.pool_fee_account,
8145                &mut accounts.token_a_mint_account,
8146                &mut SolanaAccount::default(),
8147                &mut SolanaAccount::default(),
8148            ],
8149        )
8150        .unwrap();
8151    }
8152
8153    #[test_case(spl_token::id(), spl_token::id(), spl_token::id(); "all-token")]
8154    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
8155    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
8156    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
8157    fn test_valid_swap_with_invalid_fee_account(
8158        pool_token_program_id: Pubkey,
8159        token_a_program_id: Pubkey,
8160        token_b_program_id: Pubkey,
8161    ) {
8162        let owner_key = &Pubkey::new_unique();
8163
8164        let token_a_amount = 1_000_000;
8165        let token_b_amount = 5_000_000;
8166
8167        let fees = Fees {
8168            trade_fee_numerator: 1,
8169            trade_fee_denominator: 10,
8170            owner_trade_fee_numerator: 1,
8171            owner_trade_fee_denominator: 30,
8172            owner_withdraw_fee_numerator: 1,
8173            owner_withdraw_fee_denominator: 30,
8174            host_fee_numerator: 10,
8175            host_fee_denominator: 100,
8176        };
8177
8178        let swap_curve = SwapCurve {
8179            curve_type: CurveType::ConstantProduct,
8180            calculator: Arc::new(ConstantProductCurve {}),
8181        };
8182
8183        let owner_key_str = &owner_key.to_string();
8184        let constraints = Some(SwapConstraints {
8185            owner_key: owner_key_str,
8186            valid_curve_types: &[CurveType::ConstantProduct],
8187            fees: &fees,
8188        });
8189        let mut accounts = SwapAccountInfo::new(
8190            owner_key,
8191            fees.clone(),
8192            SwapTransferFees::default(),
8193            swap_curve,
8194            token_a_amount,
8195            token_b_amount,
8196            &pool_token_program_id,
8197            &token_a_program_id,
8198            &token_b_program_id,
8199        );
8200
8201        do_process_instruction_with_fee_constraints(
8202            initialize(
8203                &SWAP_PROGRAM_ID,
8204                &pool_token_program_id,
8205                &accounts.swap_key,
8206                &accounts.authority_key,
8207                &accounts.token_a_key,
8208                &accounts.token_b_key,
8209                &accounts.pool_mint_key,
8210                &accounts.pool_fee_key,
8211                &accounts.pool_token_key,
8212                accounts.fees.clone(),
8213                accounts.swap_curve.clone(),
8214            )
8215            .unwrap(),
8216            vec![
8217                &mut accounts.swap_account,
8218                &mut SolanaAccount::default(),
8219                &mut accounts.token_a_account,
8220                &mut accounts.token_b_account,
8221                &mut accounts.pool_mint_account,
8222                &mut accounts.pool_fee_account,
8223                &mut accounts.pool_token_account,
8224                &mut SolanaAccount::default(),
8225            ],
8226            &constraints,
8227        )
8228        .unwrap();
8229
8230        let authority_key = accounts.authority_key;
8231
8232        let (
8233            token_a_key,
8234            mut token_a_account,
8235            token_b_key,
8236            mut token_b_account,
8237            pool_key,
8238            mut pool_account,
8239        ) = accounts.setup_token_accounts(
8240            owner_key,
8241            &authority_key,
8242            token_a_amount,
8243            token_b_amount,
8244            0,
8245        );
8246
8247        let destination_key = Pubkey::new_unique();
8248        let mut destination = SolanaAccount::new(
8249            account_minimum_balance(),
8250            Account::get_packed_len(),
8251            owner_key,
8252        );
8253
8254        do_process_instruction(
8255            close_account(
8256                &pool_token_program_id,
8257                &accounts.pool_fee_key,
8258                &destination_key,
8259                owner_key,
8260                &[],
8261            )
8262            .unwrap(),
8263            vec![
8264                &mut accounts.pool_fee_account,
8265                &mut destination,
8266                &mut SolanaAccount::default(),
8267            ],
8268        )
8269        .unwrap();
8270
8271        do_process_instruction_with_fee_constraints(
8272            swap(
8273                &SWAP_PROGRAM_ID,
8274                &token_a_program_id,
8275                &token_b_program_id,
8276                &pool_token_program_id,
8277                &accounts.swap_key,
8278                &accounts.authority_key,
8279                &accounts.authority_key,
8280                &token_a_key,
8281                &accounts.token_a_key,
8282                &accounts.token_b_key,
8283                &token_b_key,
8284                &accounts.pool_mint_key,
8285                &accounts.pool_fee_key,
8286                &accounts.token_a_mint_key,
8287                &accounts.token_b_mint_key,
8288                Some(&pool_key),
8289                Swap {
8290                    amount_in: token_a_amount / 2,
8291                    minimum_amount_out: 0,
8292                },
8293            )
8294            .unwrap(),
8295            vec![
8296                &mut accounts.swap_account,
8297                &mut SolanaAccount::default(),
8298                &mut SolanaAccount::default(),
8299                &mut token_a_account,
8300                &mut accounts.token_a_account,
8301                &mut accounts.token_b_account,
8302                &mut token_b_account,
8303                &mut accounts.pool_mint_account,
8304                &mut accounts.pool_fee_account,
8305                &mut accounts.token_a_mint_account,
8306                &mut accounts.token_b_mint_account,
8307                &mut SolanaAccount::default(),
8308                &mut SolanaAccount::default(),
8309                &mut SolanaAccount::default(),
8310                &mut pool_account,
8311            ],
8312            &constraints,
8313        )
8314        .unwrap();
8315    }
8316
8317    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token_2022::id(); "all-token-2022")]
8318    #[test_case(spl_token::id(), spl_token_2022::id(), spl_token_2022::id(); "mixed-pool-token")]
8319    #[test_case(spl_token_2022::id(), spl_token_2022::id(), spl_token::id(); "mixed-pool-token-2022")]
8320    fn test_swap_curve_with_transfer_fees(
8321        pool_token_program_id: Pubkey,
8322        token_a_program_id: Pubkey,
8323        token_b_program_id: Pubkey,
8324    ) {
8325        // All fees
8326        let trade_fee_numerator = 1;
8327        let trade_fee_denominator = 10;
8328        let owner_trade_fee_numerator = 1;
8329        let owner_trade_fee_denominator = 30;
8330        let owner_withdraw_fee_numerator = 1;
8331        let owner_withdraw_fee_denominator = 30;
8332        let host_fee_numerator = 20;
8333        let host_fee_denominator = 100;
8334        let fees = Fees {
8335            trade_fee_numerator,
8336            trade_fee_denominator,
8337            owner_trade_fee_numerator,
8338            owner_trade_fee_denominator,
8339            owner_withdraw_fee_numerator,
8340            owner_withdraw_fee_denominator,
8341            host_fee_numerator,
8342            host_fee_denominator,
8343        };
8344
8345        let token_a_amount = 10_000_000_000;
8346        let token_b_amount = 50_000_000_000;
8347
8348        check_valid_swap_curve(
8349            fees,
8350            SwapTransferFees {
8351                pool_token: TransferFee::default(),
8352                token_a: TransferFee {
8353                    epoch: 0.into(),
8354                    transfer_fee_basis_points: 100.into(),
8355                    maximum_fee: 1_000_000_000.into(),
8356                },
8357                token_b: TransferFee::default(),
8358            },
8359            CurveType::ConstantProduct,
8360            Arc::new(ConstantProductCurve {}),
8361            token_a_amount,
8362            token_b_amount,
8363            &pool_token_program_id,
8364            &token_a_program_id,
8365            &token_b_program_id,
8366        );
8367    }
8368}