chore: add support for precalculated recipient token addresses in TokenTransferBuilder

hrishikeshjain · hrishikeshjain · commit 6d621503afd8 · 2025-04-25T19:22:02.000+05:30
Ticket: COIN-3141
diff --git a/modules/sdk-coin-sol/src/lib/ataInitializationBuilder.ts b/modules/sdk-coin-sol/src/lib/ataInitializationBuilder.ts
@@ -130,6 +130,7 @@ export class AtaInitializationBuilder extends TransactionBuilder {
     if (this._tokenAssociateRecipients.length === 0) {
       assert(this._mint && this._tokenName, 'Mint must be set before building the transaction');
       this._owner = this._owner || this._sender;
+
       this._tokenAssociateRecipients.push({
         ownerAddress: this._owner,
         tokenName: this._tokenName,
@@ -143,7 +144,13 @@ export class AtaInitializationBuilder extends TransactionBuilder {
         if (!token) {
           throw new BuildTransactionError('Invalid transaction: invalid token name, got: ' + recipient.tokenName);
         }
-        const ataPk = await getAssociatedTokenAccountAddress(token.tokenAddress, recipient.ownerAddress);
+
+        // Use the provided ataAddress if it exists, otherwise calculate it
+        let ataPk = recipient.ataAddress;
+        if (!ataPk) {
+          ataPk = await getAssociatedTokenAccountAddress(token.tokenAddress, recipient.ownerAddress);
+        }
+
         this._instructionsData.push({
           type: InstructionBuilderTypes.CreateAssociatedTokenAccount,
           params: {
diff --git a/modules/sdk-coin-sol/src/lib/iface.ts b/modules/sdk-coin-sol/src/lib/iface.ts
@@ -175,4 +175,5 @@ export interface TransactionExplanation extends BaseTransactionExplanation {
 export class TokenAssociateRecipient {
   ownerAddress: string;
   tokenName: string;
+  ataAddress?: string;
 }
diff --git a/modules/sdk-coin-sol/src/lib/tokenTransferBuilder.ts b/modules/sdk-coin-sol/src/lib/tokenTransferBuilder.ts
@@ -54,6 +54,7 @@ export class TokenTransferBuilder extends TransactionBuilder {
         this._createAtaParams.push({
           ownerAddress: ataInitInstruction.params.ownerAddress,
           tokenName: ataInitInstruction.params.tokenName,
+          ataAddress: ataInitInstruction.params.ataAddress,
         });
       }
     }
@@ -127,14 +128,18 @@ export class TokenTransferBuilder extends TransactionBuilder {
       uniqueCreateAtaParams.map(async (recipient: TokenAssociateRecipient): Promise<AtaInit> => {
         const coin = getSolTokenFromTokenName(recipient.tokenName);
         assert(coin instanceof SolCoin);
-        const recipientTokenAddress = await getAssociatedTokenAccountAddress(coin.tokenAddress, recipient.ownerAddress);
+        // Use the provided ataAddress if it exists, otherwise calculate it
+        let ataAddress = recipient.ataAddress;
+        if (!ataAddress) {
+          ataAddress = await getAssociatedTokenAccountAddress(coin.tokenAddress, recipient.ownerAddress);
+        }
         return {
           type: InstructionBuilderTypes.CreateAssociatedTokenAccount,
           params: {
             ownerAddress: recipient.ownerAddress,
             tokenName: coin.name,
             mintAddress: coin.tokenAddress,
-            ataAddress: recipientTokenAddress,
+            ataAddress,
             payerAddress: this._sender,
           },
         };
diff --git a/modules/sdk-coin-sol/src/lib/transactionBuilder.ts b/modules/sdk-coin-sol/src/lib/transactionBuilder.ts
@@ -25,7 +25,7 @@ import { KeyPair } from '.';
 import { InstructionBuilderTypes } from './constants';
 import { solInstructionFactory } from './solInstructionFactory';
 import assert from 'assert';
-import { DurableNonceParams, InstructionParams, Memo, Nonce, Transfer } from './iface';
+import { DurableNonceParams, InstructionParams, Memo, Nonce, SetPriorityFee, Transfer } from './iface';
 import { instructionParamsFactory } from './instructionParamsFactory';
 
 export abstract class TransactionBuilder extends BaseTransactionBuilder {
@@ -94,7 +94,7 @@ export abstract class TransactionBuilder extends BaseTransactionBuilder {
 
       // If prio fee instruction exists, set the priority fee variable
       if (instruction.type === InstructionBuilderTypes.SetPriorityFee) {
-        const priorityFeeInstructionsData = filteredPriorityFeeInstructionsData[0];
+        const priorityFeeInstructionsData = filteredPriorityFeeInstructionsData[0] as SetPriorityFee;
         this.setPriorityFee({ amount: Number(priorityFeeInstructionsData.params.fee) });
       }
     }
diff --git a/modules/sdk-coin-sol/src/lib/transferBuilderV2.ts b/modules/sdk-coin-sol/src/lib/transferBuilderV2.ts
@@ -59,6 +59,7 @@ export class TransferBuilderV2 extends TransactionBuilder {
         this._createAtaParams.push({
           ownerAddress: ataInitInstruction.params.ownerAddress,
           tokenName: ataInitInstruction.params.tokenName,
+          ataAddress: ataInitInstruction.params.ataAddress,
         });
       }
     }
diff --git a/modules/sdk-coin-sol/test/unit/transactionBuilder/ataInitBuilder.ts b/modules/sdk-coin-sol/test/unit/transactionBuilder/ataInitBuilder.ts
@@ -251,10 +251,12 @@ describe('Sol Associated Token Account Builder', () => {
       {
         ownerAddress: sender.pubkey,
         tokenName: mint,
+        ataAddress: sender.ataPubkey,
       },
       {
         ownerAddress: ownerPubkeys.pubkey,
         tokenName: 'sol:ray',
+        ataAddress: 'ACEuzYtR4gBFt6HLQTYisg2T7k8Vh4ss1SpnqmbVQSNy',
       },
     ];
     const multiAtaInitBuilder = (recipients) => {
@@ -325,6 +327,7 @@ describe('Sol Associated Token Account Builder', () => {
         const duplicateRecipient = {
           ownerAddress: sender.pubkey,
           tokenName: mint,
+          ataAddress: sender.ataPubkey,
         };
         const txBuilder = multiAtaInitBuilder(recipients);
         should(() => txBuilder.enableToken(duplicateRecipient)).throwError(
@@ -336,6 +339,7 @@ describe('Sol Associated Token Account Builder', () => {
         const errorMintRecipient = {
           ownerAddress: ownerPubkeys.pubkey,
           tokenName: 'invalidToken',
+          ataAddress: ownerPubkeys.ataPubkey,
         };
         const txBuilder = multiAtaInitBuilder(recipients);
         should(() => txBuilder.enableToken(errorMintRecipient)).throwError(
@@ -356,6 +360,7 @@ describe('Sol Associated Token Account Builder', () => {
         txBuilder.enableToken({
           ownerAddress: account.pub,
           tokenName: mint,
+          ataAddress: sender.ataPubkey,
         });
         txBuilder.sign({ key: account.prv });
         await txBuilder.build().should.rejectedWith('Invalid transaction: missing nonce blockhash');
@@ -367,6 +372,7 @@ describe('Sol Associated Token Account Builder', () => {
         txBuilder.enableToken({
           ownerAddress: account.pub,
           tokenName: mint,
+          ataAddress: sender.ataPubkey,
         });
         txBuilder.sign({ key: account.prv });
         await txBuilder.build().should.rejectedWith('Invalid transaction: missing sender');
@@ -394,6 +400,7 @@ describe('Sol Associated Token Account Builder', () => {
         const invalidOwner = {
           ownerAddress: 'invalid owner',
           tokenName: mint,
+          ataAddress: sender.ataPubkey,
         };
         const txBuilder = multiAtaInitBuilder(recipients);
         should(() => txBuilder.enableToken(invalidOwner)).throwError(
diff --git a/modules/sdk-coin-sol/test/unit/transactionBuilder/tokenTransferBuilder.ts b/modules/sdk-coin-sol/test/unit/transactionBuilder/tokenTransferBuilder.ts
@@ -355,7 +355,11 @@ describe('Sol Token Transfer Builder', () => {
       txBuilder.sender(walletPK);
       txBuilder.send({ address: otherAccount.pub, amount, tokenName: nameUSDC });
       txBuilder.memo(memo);
-      txBuilder.createAssociatedTokenAccount({ ownerAddress: otherAccount.pub, tokenName: nameUSDC });
+      txBuilder.createAssociatedTokenAccount({
+        ownerAddress: otherAccount.pub,
+        tokenName: nameUSDC,
+        ataAddress: ataAddress,
+      });
       const prioFeeMicroLamports = '0';
       const priorityFee: FeeOptions = {
         amount: prioFeeMicroLamports,
@@ -414,9 +418,21 @@ describe('Sol Token Transfer Builder', () => {
       txBuilder.send({ address: account1.pub, amount, tokenName: nameUSDC });
       txBuilder.send({ address: account2.pub, amount, tokenName: nameUSDC });
       txBuilder.memo(memo);
-      txBuilder.createAssociatedTokenAccount({ ownerAddress: otherAccount.pub, tokenName: nameUSDC });
-      txBuilder.createAssociatedTokenAccount({ ownerAddress: account1.pub, tokenName: nameUSDC });
-      txBuilder.createAssociatedTokenAccount({ ownerAddress: account2.pub, tokenName: nameUSDC });
+      txBuilder.createAssociatedTokenAccount({
+        ownerAddress: otherAccount.pub,
+        tokenName: nameUSDC,
+        ataAddress: ataAddress,
+      });
+      txBuilder.createAssociatedTokenAccount({
+        ownerAddress: account1.pub,
+        tokenName: nameUSDC,
+        ataAddress: ataAddress1,
+      });
+      txBuilder.createAssociatedTokenAccount({
+        ownerAddress: account2.pub,
+        tokenName: nameUSDC,
+        ataAddress: ataAddress2,
+      });
       txBuilder.setPriorityFee(priorityFee);
       const tx = await txBuilder.build();
       tx.inputs.length.should.equal(3);
@@ -519,9 +535,21 @@ describe('Sol Token Transfer Builder', () => {
       txBuilder.send({ address: otherAccount.pub, amount, tokenName: nameUSDC });
       txBuilder.send({ address: otherAccount.pub, amount, tokenName: nameUSDC });
       txBuilder.memo(memo);
-      txBuilder.createAssociatedTokenAccount({ ownerAddress: otherAccount.pub, tokenName: nameUSDC });
-      txBuilder.createAssociatedTokenAccount({ ownerAddress: otherAccount.pub, tokenName: nameUSDC });
-      txBuilder.createAssociatedTokenAccount({ ownerAddress: otherAccount.pub, tokenName: nameUSDC });
+      txBuilder.createAssociatedTokenAccount({
+        ownerAddress: otherAccount.pub,
+        tokenName: nameUSDC,
+        ataAddress: ataAddress,
+      });
+      txBuilder.createAssociatedTokenAccount({
+        ownerAddress: otherAccount.pub,
+        tokenName: nameUSDC,
+        ataAddress: ataAddress,
+      });
+      txBuilder.createAssociatedTokenAccount({
+        ownerAddress: otherAccount.pub,
+        tokenName: nameUSDC,
+        ataAddress: ataAddress,
+      });
       txBuilder.setPriorityFee(priorityFee);
       const tx = await txBuilder.build();
       tx.inputs.length.should.equal(3);
@@ -649,15 +677,23 @@ describe('Sol Token Transfer Builder', () => {
     it('for invalid ownerAddress', () => {
       const txBuilder = tokenTransferBuilder();
       should(() =>
-        txBuilder.createAssociatedTokenAccount({ ownerAddress: invalidPubKey, tokenName: nameUSDC })
+        txBuilder.createAssociatedTokenAccount({
+          ownerAddress: invalidPubKey,
+          tokenName: nameUSDC,
+          ataAddress: ataAddress,
+        })
       ).throwError('Invalid or missing ownerAddress, got: ' + invalidPubKey);
     });
 
     it('for invalid tokenName', () => {
       const invalidTokenName = 'tsol:random';
       const txBuilder = tokenTransferBuilder();
       should(() =>
-        txBuilder.createAssociatedTokenAccount({ ownerAddress: nonceAccount.pub, tokenName: invalidTokenName })
+        txBuilder.createAssociatedTokenAccount({
+          ownerAddress: nonceAccount.pub,
+          tokenName: invalidTokenName,
+          ataAddress: ataAddress,
+        })
       ).throwError('Invalid token name, got: ' + invalidTokenName);
     });
   });

Original file line number	Diff line number	Diff line change
`@@ -175,4 +175,5 @@ export interface TransactionExplanation extends BaseTransactionExplanation {`
`175`	`175`	`export class TokenAssociateRecipient {`
`176`	`176`	`ownerAddress: string;`
`177`	`177`	`tokenName: string;`
	`178`	`+ ataAddress?: string;`
`178`	`179`	`}`
Original file line number	Diff line number	Diff line change
`@@ -59,6 +59,7 @@ export class TransferBuilderV2 extends TransactionBuilder {`
`59`	`59`	`this._createAtaParams.push({`
`60`	`60`	`ownerAddress: ataInitInstruction.params.ownerAddress,`
`61`	`61`	`tokenName: ataInitInstruction.params.tokenName,`
	`62`	`+ ataAddress: ataInitInstruction.params.ataAddress,`
`62`	`63`	`});`
`63`	`64`	`}`
`64`	`65`	`}`