feat: lookuptable support + batch compress (#1087)

* add lut

* cleanup

* cleanup

---------

Co-authored-by: Swen <swen.schaeferjohann@code.berlin>

Author: SwenSchaeferjohann

cli/src/commands/balance/index.ts

@@ -1,8 +1,6 @@
 import { Command, Flags } from "@oclif/core";
-import { CustomLoader, getSolanaRpcUrl, rpc } from "../../utils/utils";
+import { CustomLoader, rpc } from "../../utils/utils";
 import { PublicKey } from "@solana/web3.js";
-import { getTestRpc } from "@lightprotocol/stateless.js";
-import { WasmFactory } from "@lightprotocol/hasher.rs";
 
 class BalanceCommand extends Command {
   static summary = "Get balance";

js/compressed-token/src/actions/compress.ts

@@ -6,7 +6,6 @@ import {
     ComputeBudgetProgram,
 } from '@solana/web3.js';
 import {
-    bn,
     sendAndConfirmTx,
     buildAndSignTx,
     Rpc,
@@ -39,15 +38,13 @@ export async function compress(
     rpc: Rpc,
     payer: Signer,
     mint: PublicKey,
-    amount: number | BN,
+    amount: number | BN | number[] | BN[],
     owner: Signer,
     sourceTokenAccount: PublicKey,
-    toAddress: PublicKey,
+    toAddress: PublicKey | Array<PublicKey>,
     merkleTree?: PublicKey,
     confirmOptions?: ConfirmOptions,
 ): Promise<TransactionSignature> {
-    amount = bn(amount);
-
     const compressIx = await CompressedTokenProgram.compress({
         payer: payer.publicKey,
         owner: owner.publicKey,

js/compressed-token/src/actions/create-token-program-lookup-table.ts

@@ -0,0 +1,73 @@
+import { PublicKey, Signer, TransactionSignature } from '@solana/web3.js';
+import {
+    sendAndConfirmTx,
+    buildAndSignTx,
+    Rpc,
+    dedupeSigner,
+} from '@lightprotocol/stateless.js';
+
+import { CompressedTokenProgram } from '../program';
+
+/**
+ * Create a lookup table for the token program's default accounts
+ *
+ * @param rpc                   Rpc connection to use
+ * @param payer                 Payer of the transaction fees
+ * @param authority             Authority of the lookup table
+ * @param mints                 Optional array of mint public keys to include in
+ *                              the lookup table
+ * @param additionalAccounts    Optional array of additional account public keys
+ *                              to include in the lookup table
+ *
+ * @return Transaction signatures and the address of the created lookup table
+ */
+export async function createTokenProgramLookupTable(
+    rpc: Rpc,
+    payer: Signer,
+    authority: Signer,
+    mints?: PublicKey[],
+    additionalAccounts?: PublicKey[],
+): Promise<{ txIds: TransactionSignature[]; address: PublicKey }> {
+    const recentSlot = await rpc.getSlot('finalized');
+    const { instructions, address } =
+        await CompressedTokenProgram.createTokenProgramLookupTable({
+            payer: payer.publicKey,
+            authority: authority.publicKey,
+            mints,
+            remainingAccounts: additionalAccounts,
+            recentSlot,
+        });
+
+    const additionalSigners = dedupeSigner(payer, [authority]);
+    const blockhashCtx = await rpc.getLatestBlockhash();
+    const signedTx = buildAndSignTx(
+        [instructions[0]],
+        payer,
+        blockhashCtx.blockhash,
+        additionalSigners,
+    );
+
+    /// Must wait for the first instruction to be finalized.
+    const txId = await sendAndConfirmTx(
+        rpc,
+        signedTx,
+        { commitment: 'finalized' },
+        blockhashCtx,
+    );
+
+    const blockhashCtx2 = await rpc.getLatestBlockhash();
+    const signedTx2 = buildAndSignTx(
+        [instructions[1]],
+        payer,
+        blockhashCtx2.blockhash,
+        additionalSigners,
+    );
+    const txId2 = await sendAndConfirmTx(
+        rpc,
+        signedTx2,
+        { commitment: 'finalized' },
+        blockhashCtx2,
+    );
+
+    return { txIds: [txId, txId2], address };
+}

js/compressed-token/src/actions/index.ts

@@ -5,3 +5,4 @@ export * from './create-mint';
 export * from './mint-to';
 export * from './register-mint';
 export * from './transfer';
+export * from './create-token-program-lookup-table';

js/compressed-token/src/actions/mint-to.ts

@@ -20,9 +20,11 @@ import { CompressedTokenProgram } from '../program';
  * @param rpc            Rpc to use
  * @param payer          Payer of the transaction fees
  * @param mint           Mint for the account
- * @param destination    Address of the account to mint to
+ * @param destination    Address of the account to mint to. Can be an array of
+ *                       addresses if the amount is an array of amounts.
  * @param authority      Minting authority
- * @param amount         Amount to mint
+ * @param amount         Amount to mint. Can be an array of amounts if the
+ *                       destination is an array of addresses.
  * @param merkleTree     State tree account that the compressed tokens should be
  *                       part of. Defaults to the default state tree account.
  * @param confirmOptions Options for confirming the transaction
@@ -33,9 +35,9 @@ export async function mintTo(
     rpc: Rpc,
     payer: Signer,
     mint: PublicKey,
-    destination: PublicKey,
+    destination: PublicKey | PublicKey[],
     authority: Signer,
-    amount: number | BN,
+    amount: number | BN | number[] | BN[],
     merkleTree?: PublicKey,
     confirmOptions?: ConfirmOptions,
 ): Promise<TransactionSignature> {

js/compressed-token/src/program.ts

@@ -4,6 +4,7 @@ import {
     TransactionInstruction,
     SystemProgram,
     Connection,
+    AddressLookupTableProgram,
 } from '@solana/web3.js';
 import { BN, Program, AnchorProvider, setProvider } from '@coral-xyz/anchor';
 import { IDL, LightCompressedToken } from './idl/light_compressed_token';
@@ -47,16 +48,17 @@ type CompressParams = {
     source: PublicKey;
     /**
      * owner of the compressed token account.
+     * To compress to a batch of recipients, pass an array of PublicKeys.
      */
-    toAddress: PublicKey;
+    toAddress: PublicKey | PublicKey[];
     /**
      * Mint address of the token to compress.
      */
     mint: PublicKey;
     /**
      * amount of tokens to compress.
      */
-    amount: number | BN;
+    amount: number | BN | number[] | BN[];
     /**
      * The state tree that the tx output should be inserted into. Defaults to a
      * public state tree if unspecified.
@@ -241,6 +243,29 @@ export type ApproveAndMintToParams = {
     merkleTree?: PublicKey;
 };
 
+export type CreateTokenProgramLookupTableParams = {
+    /**
+     * The payer of the transaction.
+     */
+    payer: PublicKey;
+    /**
+     * The authority of the transaction.
+     */
+    authority: PublicKey;
+    /**
+     *  Recently finalized Solana slot.
+     */
+    recentSlot: number;
+    /**
+     * Optional Mint addresses to store in the lookup table.
+     */
+    mints?: PublicKey[];
+    /**
+     * Optional additional addresses to store in the lookup table.
+     */
+    remainingAccounts?: PublicKey[];
+};
+
 /**
  * Sum up the token amounts of the compressed token accounts
  */
@@ -513,6 +538,13 @@ export class CompressedTokenProgram {
         const amounts = toArray<BN | number>(amount).map(amount => bn(amount));
 
         const toPubkeys = toArray(toPubkey);
+
+        if (amounts.length !== toPubkeys.length) {
+            throw new Error(
+                'Amount and toPubkey arrays must have the same length',
+            );
+        }
+
         const instruction = await this.program.methods
             .mintTo(toPubkeys, amounts, null)
             .accounts({
@@ -658,24 +690,102 @@ export class CompressedTokenProgram {
     }
 
     /**
-     * Construct compress instruction
+     * Create lookup table instructions for the token program's default accounts.
+     */
+    static async createTokenProgramLookupTable(
+        params: CreateTokenProgramLookupTableParams,
+    ) {
+        const { authority, mints, recentSlot, payer, remainingAccounts } =
+            params;
+
+        const [createInstruction, lookupTableAddress] =
+            AddressLookupTableProgram.createLookupTable({
+                authority,
+                payer: authority,
+                recentSlot,
+            });
+
+        let optionalMintKeys: PublicKey[] = [];
+        if (mints) {
+            optionalMintKeys = [
+                ...mints,
+                ...mints.map(mint => this.deriveTokenPoolPda(mint)),
+            ];
+        }
+
+        const extendInstruction = AddressLookupTableProgram.extendLookupTable({
+            payer,
+            authority,
+            lookupTable: lookupTableAddress,
+            addresses: [
+                this.deriveCpiAuthorityPda,
+                LightSystemProgram.programId,
+                defaultStaticAccountsStruct().registeredProgramPda,
+                defaultStaticAccountsStruct().noopProgram,
+                defaultStaticAccountsStruct().accountCompressionAuthority,
+                defaultStaticAccountsStruct().accountCompressionProgram,
+                defaultTestStateTreeAccounts().merkleTree,
+                defaultTestStateTreeAccounts().nullifierQueue,
+                defaultTestStateTreeAccounts().addressTree,
+                defaultTestStateTreeAccounts().addressQueue,
+                this.programId,
+                TOKEN_PROGRAM_ID,
+                authority,
+                ...optionalMintKeys,
+                ...(remainingAccounts ?? []),
+            ],
+        });
+
+        return {
+            instructions: [createInstruction, extendInstruction],
+            address: lookupTableAddress,
+        };
+    }
+
+    /**
+     * Create compress instruction
      * @returns compressInstruction
      */
     static async compress(
         params: CompressParams,
     ): Promise<TransactionInstruction> {
         const { payer, owner, source, toAddress, mint, outputStateTree } =
             params;
-        const amount = bn(params.amount);
 
-        const tokenTransferOutputs: TokenTransferOutputData[] = [
-            {
-                owner: toAddress,
-                amount,
-                lamports: bn(0),
-                tlv: null,
-            },
-        ];
+        if (Array.isArray(params.amount) !== Array.isArray(params.toAddress)) {
+            throw new Error(
+                'Both amount and toAddress must be arrays or both must be single values',
+            );
+        }
+
+        let tokenTransferOutputs: TokenTransferOutputData[];
+
+        if (Array.isArray(params.amount) && Array.isArray(params.toAddress)) {
+            if (params.amount.length !== params.toAddress.length) {
+                throw new Error(
+                    'Amount and toAddress arrays must have the same length',
+                );
+            }
+            tokenTransferOutputs = params.amount.map((amt, index) => {
+                const amount = bn(amt);
+                return {
+                    owner: (params.toAddress as PublicKey[])[index],
+                    amount,
+                    lamports: bn(0),
+                    tlv: null,
+                };
+            });
+        } else {
+            tokenTransferOutputs = [
+                {
+                    owner: toAddress as PublicKey,
+                    amount: bn(params.amount as number | BN),
+                    lamports: bn(0),
+                    tlv: null,
+                },
+            ];
+        }
+
         const {
             inputTokenDataWithContext,
             packedOutputTokenData,
@@ -693,7 +803,11 @@ export class CompressedTokenProgram {
             delegatedTransfer: null, // TODO: implement
             inputTokenDataWithContext,
             outputCompressedAccounts: packedOutputTokenData,
-            compressOrDecompressAmount: amount,
+            compressOrDecompressAmount: Array.isArray(params.amount)
+                ? params.amount
+                      .map(amt => new BN(amt))
+                      .reduce((sum, amt) => sum.add(amt), new BN(0))
+                : new BN(params.amount),
             isCompress: true,
             cpiContext: null,
             lamportsChangeAccountMerkleTreeIndex: null,
@@ -704,24 +818,20 @@ export class CompressedTokenProgram {
             data,
         );
 
-        const {
-            accountCompressionAuthority,
-            noopProgram,
-            registeredProgramPda,
-            accountCompressionProgram,
-        } = defaultStaticAccountsStruct();
-
         const instruction = await this.program.methods
             .transfer(encodedData)
             .accounts({
                 feePayer: payer,
                 authority: owner,
                 cpiAuthorityPda: this.deriveCpiAuthorityPda,
                 lightSystemProgram: LightSystemProgram.programId,
-                registeredProgramPda: registeredProgramPda,
-                noopProgram: noopProgram,
-                accountCompressionAuthority: accountCompressionAuthority,
-                accountCompressionProgram: accountCompressionProgram,
+                registeredProgramPda:
+                    defaultStaticAccountsStruct().registeredProgramPda,
+                noopProgram: defaultStaticAccountsStruct().noopProgram,
+                accountCompressionAuthority:
+                    defaultStaticAccountsStruct().accountCompressionAuthority,
+                accountCompressionProgram:
+                    defaultStaticAccountsStruct().accountCompressionProgram,
                 selfProgram: this.programId,
                 tokenPoolPda: this.deriveTokenPoolPda(mint),
                 compressOrDecompressTokenAccount: source, // token