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