fix(svm): H-01 optimize claiming relayer refunds

programs/svm-spoke/src/instructions/refund_claims.rs

@@ -10,16 +10,21 @@ use crate::{
 };
 
 #[derive(Accounts)]
-#[instruction(mint: Pubkey, refund_address: Pubkey)]
 pub struct InitializeClaimAccount<'info> {
     #[account(mut)]
     pub signer: Signer<'info>,
 
+    /// CHECK: This is only used for claim_account PDA derivation and it is up to the caller to ensure it is valid.
+    pub mint: UncheckedAccount<'info>,
+
+    /// CHECK: This is only used for claim_account PDA derivation and it is up to the caller to ensure it is valid.
+    pub refund_address: UncheckedAccount<'info>,
+
     #[account(
         init,
         payer = signer,
         space = DISCRIMINATOR_SIZE + ClaimAccount::INIT_SPACE,
-        seeds = [b"claim_account", mint.as_ref(), refund_address.as_ref()],
+        seeds = [b"claim_account", mint.key().as_ref(), refund_address.key().as_ref()],
         bump
     )]
     pub claim_account: Account<'info, ClaimAccount>,
@@ -120,11 +125,16 @@ pub fn claim_relayer_refund(ctx: Context<ClaimRelayerRefund>) -> Result<()> {
 // relayer refunds were executed with ATA after initializing the claim account. In such cases, the initializer should be
 // able to close the claim account manually.
 #[derive(Accounts)]
-#[instruction(mint: Pubkey, refund_address: Pubkey)]
 pub struct CloseClaimAccount<'info> {
     #[account(mut, address = claim_account.initializer @ SvmError::InvalidClaimInitializer)]
     pub signer: Signer<'info>,
 
+    /// CHECK: This is only used for claim_account PDA derivation and it is up to the caller to ensure it is valid.
+    pub mint: UncheckedAccount<'info>,
+
+    /// CHECK: This is only used for claim_account PDA derivation and it is up to the caller to ensure it is valid.
+    pub refund_address: UncheckedAccount<'info>,
+
     #[account(
         mut,
         close = signer,

programs/svm-spoke/src/lib.rs

@@ -634,18 +634,12 @@ pub mod svm_spoke {
     ///
     /// ### Required Accounts:
     /// - signer (Signer): The account that pays for the transaction and initializes the claim account.
+    /// - mint: The mint associated with the claim account.
+    /// - refund_address: The refund address associated with the claim account.
     /// - claim_account (Writable): The newly created claim account PDA to store claim data for this associated mint.
     ///   Seed: ["claim_account",mint,refund_address].
     /// - system_program: The system program required for account creation.
-    ///
-    /// ### Parameters:
-    /// - _mint: The public key of the mint associated with the claim account.
-    /// - _refund_address: The public key of the refund address associated with the claim account.
-    pub fn initialize_claim_account(
-        ctx: Context<InitializeClaimAccount>,
-        _mint: Pubkey,
-        _refund_address: Pubkey,
-    ) -> Result<()> {
+    pub fn initialize_claim_account(ctx: Context<InitializeClaimAccount>) -> Result<()> {
         instructions::initialize_claim_account(ctx)
     }
 
@@ -657,16 +651,10 @@ pub mod svm_spoke {
     ///
     /// ### Required Accounts:
     /// - signer (Signer): The account that authorizes the closure. Must be the initializer of the claim account.
+    /// - mint: The mint associated with the claim account.
+    /// - refund_address: The refund address associated with the claim account.
     /// - claim_account (Writable): The claim account PDA to be closed. Seed: ["claim_account",mint,refund_address].
-    ///
-    /// ### Parameters:
-    /// - _mint: The public key of the mint associated with the claim account.
-    /// - _refund_address: The public key of the refund address associated with the claim account.
-    pub fn close_claim_account(
-        ctx: Context<CloseClaimAccount>,
-        _mint: Pubkey,           // Only used in account constraints.
-        _refund_address: Pubkey, // Only used in account constraints.
-    ) -> Result<()> {
+    pub fn close_claim_account(ctx: Context<CloseClaimAccount>) -> Result<()> {
         instructions::close_claim_account(ctx)
     }
 

test/svm/SvmSpoke.Bundle.ts

@@ -7,7 +7,7 @@ import {
   mintTo,
   TOKEN_PROGRAM_ID,
 } from "@solana/spl-token";
-import { ComputeBudgetProgram, Keypair, PublicKey } from "@solana/web3.js";
+import { ComputeBudgetProgram, Keypair, PublicKey, TransactionInstruction } from "@solana/web3.js";
 import { assert } from "chai";
 import * as crypto from "crypto";
 import { ethers } from "ethers";
@@ -1025,7 +1025,7 @@ describe("svm_spoke.bundle", () => {
         } else if (!testConfig.deferredRefunds && testConfig.atomicAccountCreation) {
           refundAccounts.push(tokenAccount);
         } else {
-          await program.methods.initializeClaimAccount(mint, tokenOwner).rpc();
+          await program.methods.initializeClaimAccount().accounts({ mint, refundAddress: tokenOwner }).rpc();
           refundAccounts.push(claimAccount);
         }
 
@@ -1407,8 +1407,8 @@ describe("svm_spoke.bundle", () => {
           [Buffer.from("claim_account"), mint.toBuffer(), relayerB.publicKey.toBuffer()],
           program.programId
         );
-        await program.methods.initializeClaimAccount(mint, relayerA.publicKey).rpc();
-        await program.methods.initializeClaimAccount(mint, relayerB.publicKey).rpc();
+        await program.methods.initializeClaimAccount().accounts({ mint, refundAddress: relayerA.publicKey }).rpc();
+        await program.methods.initializeClaimAccount().accounts({ mint, refundAddress: relayerB.publicKey }).rpc();
       }
 
       // Prepare leaf using token accounts.
@@ -1742,4 +1742,220 @@ describe("svm_spoke.bundle", () => {
       assert.include(err.toString(), "Invalid Merkle proof", "Expected merkle verification to fail");
     }
   });
+
+  describe("Execute Max multiple refunds with claims", async () => {
+    const executeMaxRefundClaims = async (testConfig: {
+      solanaDistributions: number;
+      useAddressLookup: boolean;
+      separatePhases: boolean;
+    }) => {
+      // Add leaves for other EVM chains to have non-empty proofs array to ensure we don't run out of memory when processing.
+      const evmDistributions = 100; // This would fit in 7 proof array elements.
+
+      const refundAddresses: web3.PublicKey[] = []; // These are relayer authority addresses used in leaf building.
+      const claimAccounts: web3.PublicKey[] = [];
+      const tokenAccounts: web3.PublicKey[] = [];
+      const refundAmounts: BN[] = [];
+      const initializeInstructions: TransactionInstruction[] = [];
+      const claimInstructions: TransactionInstruction[] = [];
+
+      for (let i = 0; i < testConfig.solanaDistributions; i++) {
+        // Create the token account.
+        const tokenOwner = Keypair.generate().publicKey;
+        const tokenAccount = (await getOrCreateAssociatedTokenAccount(connection, payer, mint, tokenOwner)).address;
+        refundAddresses.push(tokenOwner);
+        tokenAccounts.push(tokenAccount);
+
+        const [claimAccount] = PublicKey.findProgramAddressSync(
+          [Buffer.from("claim_account"), mint.toBuffer(), tokenOwner.toBuffer()],
+          program.programId
+        );
+
+        // Create instruction to initialize claim account.
+        initializeInstructions.push(
+          await program.methods.initializeClaimAccount().accounts({ mint, refundAddress: tokenOwner }).instruction()
+        );
+        claimAccounts.push(claimAccount);
+
+        refundAmounts.push(new BN(randomBigInt(2).toString()));
+
+        // Create instruction to claim refund to the token account.
+        const claimRelayerRefundAccounts = {
+          signer: owner,
+          initializer: owner,
+          state,
+          vault,
+          mint,
+          tokenAccount,
+          refundAddress: tokenOwner,
+          claimAccount,
+          tokenProgram: TOKEN_PROGRAM_ID,
+          program: program.programId,
+        };
+        claimInstructions.push(
+          await program.methods.claimRelayerRefund().accounts(claimRelayerRefundAccounts).instruction()
+        );
+      }
+
+      const { relayerRefundLeaves, merkleTree } = buildRelayerRefundMerkleTree({
+        totalEvmDistributions: evmDistributions,
+        totalSolanaDistributions: testConfig.solanaDistributions,
+        mixLeaves: false,
+        chainId: chainId.toNumber(),
+        mint,
+        svmRelayers: refundAddresses,
+        svmRefundAmounts: refundAmounts,
+      });
+
+      const root = merkleTree.getRoot();
+      const proof = merkleTree.getProof(relayerRefundLeaves[0]);
+      const leaf = relayerRefundLeaves[0] as RelayerRefundLeafSolana;
+
+      const stateAccountData = await program.account.state.fetch(state);
+      const rootBundleId = stateAccountData.rootBundleId;
+
+      const rootBundleIdBuffer = Buffer.alloc(4);
+      rootBundleIdBuffer.writeUInt32LE(rootBundleId);
+      const seeds = [Buffer.from("root_bundle"), seed.toArrayLike(Buffer, "le", 8), rootBundleIdBuffer];
+      const [rootBundle] = PublicKey.findProgramAddressSync(seeds, program.programId);
+
+      // Relay root bundle
+      const relayRootBundleAccounts = { state, rootBundle, signer: owner, payer: owner, program: program.programId };
+      await program.methods.relayRootBundle(Array.from(root), Array.from(root)).accounts(relayRootBundleAccounts).rpc();
+
+      // Verify valid leaf
+      const proofAsNumbers = proof.map((p) => Array.from(p));
+
+      const [instructionParams] = PublicKey.findProgramAddressSync(
+        [Buffer.from("instruction_params"), owner.toBuffer()],
+        program.programId
+      );
+
+      const executeAccounts = {
+        instructionParams,
+        state,
+        rootBundle: rootBundle,
+        signer: owner,
+        vault,
+        tokenProgram: TOKEN_PROGRAM_ID,
+        mint,
+        transferLiability,
+        program: program.programId,
+      };
+
+      const executeRemainingAccounts = claimAccounts.map((account) => ({
+        pubkey: account,
+        isWritable: true,
+        isSigner: false,
+      }));
+
+      // Build the instruction to execute relayer refund leaf and write its instruction args to the data account.
+      await loadExecuteRelayerRefundLeafParams(program, owner, stateAccountData.rootBundleId, leaf, proofAsNumbers);
+
+      const executeInstruction = await program.methods
+        .executeRelayerRefundLeafDeferred()
+        .accounts(executeAccounts)
+        .remainingAccounts(executeRemainingAccounts)
+        .instruction();
+
+      // Initialize, execute and claim depending on the chosen method.
+      const instructions = [...initializeInstructions, executeInstruction, ...claimInstructions];
+      if (!testConfig.separatePhases) {
+        // Pack all instructions in one transaction.
+        if (testConfig.useAddressLookup)
+          await sendTransactionWithLookupTable(
+            connection,
+            instructions,
+            (anchor.AnchorProvider.env().wallet as anchor.Wallet).payer
+          );
+        else
+          await web3.sendAndConfirmTransaction(
+            connection,
+            new web3.Transaction().add(...instructions),
+            [(anchor.AnchorProvider.env().wallet as anchor.Wallet).payer],
+            {
+              commitment: "confirmed",
+            }
+          );
+      } else {
+        // Send claim account initialization, execution and claim in separate transactions.
+        if (testConfig.useAddressLookup) {
+          await sendTransactionWithLookupTable(
+            connection,
+            initializeInstructions,
+            (anchor.AnchorProvider.env().wallet as anchor.Wallet).payer
+          );
+          await sendTransactionWithLookupTable(
+            connection,
+            [executeInstruction],
+            (anchor.AnchorProvider.env().wallet as anchor.Wallet).payer
+          );
+          await sendTransactionWithLookupTable(
+            connection,
+            claimInstructions,
+            (anchor.AnchorProvider.env().wallet as anchor.Wallet).payer
+          );
+        } else {
+          await web3.sendAndConfirmTransaction(
+            connection,
+            new web3.Transaction().add(...initializeInstructions),
+            [(anchor.AnchorProvider.env().wallet as anchor.Wallet).payer],
+            {
+              commitment: "confirmed",
+            }
+          );
+          await web3.sendAndConfirmTransaction(
+            connection,
+            new web3.Transaction().add(executeInstruction),
+            [(anchor.AnchorProvider.env().wallet as anchor.Wallet).payer],
+            {
+              commitment: "confirmed",
+            }
+          );
+          await web3.sendAndConfirmTransaction(
+            connection,
+            new web3.Transaction().add(...claimInstructions),
+            [(anchor.AnchorProvider.env().wallet as anchor.Wallet).payer],
+            {
+              commitment: "confirmed",
+            }
+          );
+        }
+      }
+
+      // Verify all refund token account balances.
+      const refundBalances = await Promise.all(
+        tokenAccounts.map(async (account) => {
+          return (await connection.getTokenAccountBalance(account)).value.amount;
+        })
+      );
+      refundBalances.forEach((balance, i) => {
+        assertSE(balance, refundAmounts[i].toString(), `Refund account ${i} balance should match refund amount`);
+      });
+    };
+
+    it("Execute Max multiple refunds with claims in one legacy transaction", async () => {
+      // Larger amount would hit transaction message size limit.
+      const solanaDistributions = 5;
+      await executeMaxRefundClaims({ solanaDistributions, useAddressLookup: false, separatePhases: false });
+    });
+
+    it("Execute Max multiple refunds with claims in one versioned transaction", async () => {
+      // Larger amount would hit maximum instruction trace length limit.
+      const solanaDistributions = 12;
+      await executeMaxRefundClaims({ solanaDistributions, useAddressLookup: true, separatePhases: false });
+    });
+
+    it("Execute Max multiple refunds with claims in separate phase legacy transactions", async () => {
+      // Larger amount would hit transaction message size limit.
+      const solanaDistributions = 7;
+      await executeMaxRefundClaims({ solanaDistributions, useAddressLookup: false, separatePhases: true });
+    });
+
+    it("Execute Max multiple refunds with claims in separate phase versioned transactions", async () => {
+      // Larger amount would hit maximum instruction trace length limit.
+      const solanaDistributions = 21;
+      await executeMaxRefundClaims({ solanaDistributions, useAddressLookup: true, separatePhases: true });
+    });
+  });
 });

test/svm/SvmSpoke.RefundClaims.ts

@@ -50,8 +50,8 @@ describe("svm_spoke.refund_claims", () => {
 
   const initializeClaimAccount = async (initializer = claimInitializer) => {
     const initializeClaimAccountIx = await program.methods
-      .initializeClaimAccount(mint, relayer.publicKey)
-      .accounts({ signer: initializer.publicKey })
+      .initializeClaimAccount()
+      .accounts({ signer: initializer.publicKey, mint, refundAddress: relayer.publicKey })
       .instruction();
     await web3.sendAndConfirmTransaction(connection, new web3.Transaction().add(initializeClaimAccountIx), [
       initializer,
@@ -289,7 +289,10 @@ describe("svm_spoke.refund_claims", () => {
 
     // Should not be able to close the claim account from default wallet as the initializer was different.
     try {
-      await program.methods.closeClaimAccount(mint, relayer.publicKey).accounts({ signer: payer.publicKey }).rpc();
+      await program.methods
+        .closeClaimAccount()
+        .accounts({ signer: payer.publicKey, mint, refundAddress: relayer.publicKey })
+        .rpc();
       assert.fail("Closing claim account from different initializer should fail");
     } catch (error: any) {
       assert.instanceOf(error, AnchorError);
@@ -302,8 +305,8 @@ describe("svm_spoke.refund_claims", () => {
 
     // Close the claim account from initializer before executing relayer refunds.
     await program.methods
-      .closeClaimAccount(mint, relayer.publicKey)
-      .accounts({ signer: claimInitializer.publicKey })
+      .closeClaimAccount()
+      .accounts({ signer: claimInitializer.publicKey, mint, refundAddress: relayer.publicKey })
       .signers([claimInitializer])
       .rpc();
 
@@ -324,8 +327,8 @@ describe("svm_spoke.refund_claims", () => {
     // It should be not possible to close the claim account with non-zero refund liability.
     try {
       await program.methods
-        .closeClaimAccount(mint, relayer.publicKey)
-        .accounts({ signer: claimInitializer.publicKey })
+        .closeClaimAccount()
+        .accounts({ signer: claimInitializer.publicKey, mint, refundAddress: relayer.publicKey })
         .signers([claimInitializer])
         .rpc();
       assert.fail("Closing claim account with non-zero refund liability should fail");