remove_defunct_pool: move accounts to ATAs initialized if needed

invariant-labs · Jan 21, 2025 · c94761c · c94761c
1 parent 6e91338
commit c94761c
Show file tree

Hide file tree

Showing 6 changed files with 68 additions and 39 deletions.
diff --git a/programs/invariant/Cargo.toml b/programs/invariant/Cargo.toml
@@ -21,7 +21,7 @@ all = []
 
 [dependencies]
 decimal = { path = "decimal" }
-anchor-lang = "0.21.0"
+anchor-lang = { version = "0.21.0", features = ['init-if-needed'] }
 anchor-spl = "0.21.0"
 integer-sqrt = "0.1.5"
 uint = "0.9.1"

diff --git a/programs/invariant/src/instructions/remove_defunct_pool.rs b/programs/invariant/src/instructions/remove_defunct_pool.rs
@@ -7,6 +7,7 @@ use crate::structs::State;
 use crate::ErrorCode::*;
 use crate::SEED;
 use anchor_lang::prelude::*;
+use anchor_spl::associated_token::AssociatedToken;
 use anchor_spl::token;
 use anchor_spl::token::CloseAccount;
 use anchor_spl::token::Token;
@@ -30,14 +31,16 @@ pub struct RemoveDefunctPool<'info> {
     pub tickmap: AccountLoader<'info, Tickmap>,
     pub token_x: Account<'info, Mint>,
     pub token_y: Account<'info, Mint>,
-    #[account(mut,
-        constraint = account_x.mint == token_x.key() @ InvalidMint,
-        constraint = &account_x.owner == admin.key @ InvalidOwner,
+    #[account(init_if_needed,
+        payer = admin,
+        associated_token::mint = token_x,
+        associated_token::authority = admin
     )]
     pub account_x: Box<Account<'info, TokenAccount>>,
-    #[account(mut,
-        constraint = account_y.mint == token_y.key() @ InvalidMint,
-        constraint = &account_y.owner == admin.key @ InvalidOwner
+    #[account(init_if_needed,
+        payer = admin,
+        associated_token::mint = token_y,
+        associated_token::authority = admin
     )]
     pub account_y: Box<Account<'info, TokenAccount>>,
     #[account(mut,
@@ -57,6 +60,9 @@ pub struct RemoveDefunctPool<'info> {
     #[account(constraint = &state.load()?.authority == program_authority.key @ InvalidAuthority)]
     pub program_authority: AccountInfo<'info>,
     pub token_program: Program<'info, Token>,
+    pub associated_token_program: Program<'info, AssociatedToken>,
+    pub rent: Sysvar<'info, Rent>,
+    pub system_program: Program<'info, System>,
 }
 
 impl<'info> interfaces::send_tokens::SendTokens<'info> for RemoveDefunctPool<'info> {

diff --git a/sdk/src/idl/invariant.ts b/sdk/src/idl/invariant.ts
@@ -998,6 +998,21 @@ export type Invariant = {
           "name": "tokenProgram",
           "isMut": false,
           "isSigner": false
+        },
+        {
+          "name": "associatedTokenProgram",
+          "isMut": false,
+          "isSigner": false
+        },
+        {
+          "name": "rent",
+          "isMut": false,
+          "isSigner": false
+        },
+        {
+          "name": "systemProgram",
+          "isMut": false,
+          "isSigner": false
         }
       ],
       "args": []
@@ -2592,6 +2607,21 @@ export const IDL: Invariant = {
           "name": "tokenProgram",
           "isMut": false,
           "isSigner": false
+        },
+        {
+          "name": "associatedTokenProgram",
+          "isMut": false,
+          "isSigner": false
+        },
+        {
+          "name": "rent",
+          "isMut": false,
+          "isSigner": false
+        },
+        {
+          "name": "systemProgram",
+          "isMut": false,
+          "isSigner": false
         }
       ],
       "args": []

diff --git a/sdk/src/market.ts b/sdk/src/market.ts
@@ -1,5 +1,5 @@
 import { BN, Program, utils, Provider } from '@project-serum/anchor'
-import { Token, TOKEN_PROGRAM_ID } from '@solana/spl-token'
+import { ASSOCIATED_TOKEN_PROGRAM_ID, Token, TOKEN_PROGRAM_ID } from '@solana/spl-token'
 import {
   ComputeBudgetProgram,
   Connection,
@@ -39,6 +39,7 @@ import { Invariant, IDL } from './idl/invariant'
 import { DENOMINATOR, IWallet, Pair, signAndSend } from '.'
 import { getMarketAddress, Network } from './network'
 import { bs58 } from '@project-serum/anchor/dist/cjs/utils/bytes'
+import { getAssociatedTokenAddress } from './token'
 
 const POSITION_SEED = 'positionv1'
 const TICK_SEED = 'tickv1'
@@ -1351,12 +1352,15 @@ export class Market {
   }
 
   async removeDefunctPoolInstruction(removeDefunctPool: RemoveDefunctPool) {
-    const { pair, accountX, accountY } = removeDefunctPool
+    const { pair } = removeDefunctPool
     const adminPubkey = removeDefunctPool.admin ?? this.wallet.publicKey
     const { address: stateAddress } = await this.getStateAddress()
     const poolAddress = await pair.getAddress(this.program.programId)
     const pool = await this.getPool(pair)
 
+    const accountX = getAssociatedTokenAddress(adminPubkey, pair.tokenX)
+    const accountY = getAssociatedTokenAddress(adminPubkey, pair.tokenY)
+
     return this.program.instruction.removeDefunctPool({
       accounts: {
         state: stateAddress,
@@ -1370,7 +1374,10 @@ export class Market {
         reserveY: pool.tokenYReserve,
         admin: adminPubkey,
         programAuthority: this.programAuthority,
-        tokenProgram: TOKEN_PROGRAM_ID
+        tokenProgram: TOKEN_PROGRAM_ID,
+        associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
+        rent: SYSVAR_RENT_PUBKEY,
+        systemProgram: SystemProgram.programId
       }
     })
   }
@@ -1711,8 +1718,6 @@ export interface ChangeFeeReceiver {
 
 export interface RemoveDefunctPool {
   pair: Pair
-  accountX: PublicKey
-  accountY: PublicKey
   admin?: PublicKey
 }
 

diff --git a/sdk/src/token.ts b/sdk/src/token.ts
@@ -0,0 +1,11 @@
+import { TOKEN_PROGRAM_ID, ASSOCIATED_TOKEN_PROGRAM_ID } from '@solana/spl-token'
+import { PublicKey } from '@solana/web3.js'
+
+// function only available in higher versions of spl-token than current 0.1.8
+export function getAssociatedTokenAddress(owner: PublicKey, token: PublicKey) {
+  const [address] = PublicKey.findProgramAddressSync(
+    [owner.toBuffer(), TOKEN_PROGRAM_ID.toBuffer(), token.toBuffer()],
+    ASSOCIATED_TOKEN_PROGRAM_ID
+  )
+  return address
+}
diff --git a/tests/remove-pool.spec.ts b/tests/remove-pool.spec.ts
@@ -61,15 +61,9 @@ describe('remove pool', () => {
   })
 
   it('#removeDefunctPool()', async () => {
-    const adminTokenXAccount = await tokenX.createAccount(admin.publicKey)
-    const adminTokenYAccount = await tokenY.createAccount(admin.publicKey)
-
     const poolState = await market.getPool(pair)
 
-    await market.removeDefunctPool(
-      { pair, admin: admin.publicKey, accountX: adminTokenXAccount, accountY: adminTokenYAccount },
-      admin
-    )
+    await market.removeDefunctPool({ pair, admin: admin.publicKey }, admin)
     await assertThrowsAsync(market.getPool(pair), 'Error: Account does not exist')
     await assertThrowsAsync(
       market.program.account.tickmap.fetch(poolState.tickmap),
@@ -93,16 +87,11 @@ describe('remove pool', () => {
       initTick
     })
 
-    const positionOwnerTokenXAccount = await tokenX.createAccount(positionOwner.publicKey)
-    const positionOwnerTokenYAccount = await tokenY.createAccount(positionOwner.publicKey)
-
     await assertThrowsAsync(
       market.removeDefunctPool(
         {
           pair,
-          admin: positionOwner.publicKey,
-          accountX: positionOwnerTokenXAccount,
-          accountY: positionOwnerTokenYAccount
+          admin: positionOwner.publicKey
         },
         positionOwner
       ),
@@ -165,16 +154,11 @@ describe('remove pool', () => {
     )
     await market.getPosition(positionOwner.publicKey, positionIndex)
 
-    const adminTokenXAccount = await tokenX.createAccount(admin.publicKey)
-    const adminTokenYAccount = await tokenY.createAccount(admin.publicKey)
-
     await assertThrowsAsync(
       market.removeDefunctPool(
         {
           pair,
-          admin: admin.publicKey,
-          accountX: adminTokenXAccount,
-          accountY: adminTokenYAccount
+          admin: admin.publicKey
         },
         admin
       ),
@@ -196,16 +180,9 @@ describe('remove pool', () => {
       },
       positionOwner
     )
-
-    const adminTokenXAccount = await tokenX.createAccount(admin.publicKey)
-    const adminTokenYAccount = await tokenY.createAccount(admin.publicKey)
-
     const poolState = await market.getPool(pair)
 
-    await market.removeDefunctPool(
-      { pair, admin: admin.publicKey, accountX: adminTokenXAccount, accountY: adminTokenYAccount },
-      admin
-    )
+    await market.removeDefunctPool({ pair, admin: admin.publicKey }, admin)
     await assertThrowsAsync(market.getPool(pair), 'Error: Account does not exist')
     await assertThrowsAsync(
       market.program.account.tickmap.fetch(poolState.tickmap),