update to ActiveTreeBundle type

js/stateless.js/src/constants.ts

@@ -1,6 +1,7 @@
 import BN from 'bn.js';
 import { Buffer } from 'buffer';
 import { ConfirmOptions, PublicKey } from '@solana/web3.js';
+import { ActiveTreeBundle, TreeType } from './state/types';
 
 export const FIELD_SIZE = new BN(
     '21888242871839275222246405745257275088548364400416034343698204186575808495617',
@@ -49,36 +50,71 @@ export const defaultStaticAccountsStruct = () => {
     };
 };
 
-export const defaultStateTreeLookupTables = () => {
+export type StateTreeLUTPair = {
+    stateTreeLookupTable: PublicKey;
+    nullifyTable: PublicKey;
+};
+
+/**
+ * Returns the Default Public State Tree LUTs for Devnet and Mainnet-Beta.
+ */
+export const defaultStateTreeLookupTables = (): {
+    mainnet: StateTreeLUTPair[];
+    devnet: StateTreeLUTPair[];
+} => {
     return {
-        mainnet: {
-            stateTreeLookupTable: new PublicKey(stateTreeLookupTableMainnet),
-            nullifyTable: new PublicKey(nullifiedStateTreeLookupTableMainnet),
-        },
-        devnet: {
-            stateTreeLookupTable: new PublicKey(stateTreeLookupTableDevnet),
-            nullifyTable: new PublicKey(nullifiedStateTreeLookupTableDevnet),
-        },
+        mainnet: [
+            {
+                stateTreeLookupTable: new PublicKey(
+                    stateTreeLookupTableMainnet,
+                ),
+                nullifyTable: new PublicKey(
+                    nullifiedStateTreeLookupTableMainnet,
+                ),
+            },
+        ],
+        devnet: [
+            {
+                stateTreeLookupTable: new PublicKey(stateTreeLookupTableDevnet),
+                nullifyTable: new PublicKey(
+                    nullifiedStateTreeLookupTableDevnet,
+                ),
+            },
+        ],
     };
 };
 
+/**
+ * @internal
+ */
 export const isLocalTest = (url: string) => {
     return url.includes('localhost') || url.includes('127.0.0.1');
 };
 
-export const localTestActiveStateTreeInfo = () => {
-    return {
-        activeStateTrees: [
-            new PublicKey(merkletreePubkey),
-            new PublicKey(merkleTree2Pubkey),
-        ],
-        activeQueues: [
-            new PublicKey(nullifierQueuePubkey),
-            new PublicKey(nullifierQueue2Pubkey),
-        ],
-    };
+/**
+ * @internal
+ */
+export const localTestActiveStateTreeInfo = (): ActiveTreeBundle[] => {
+    return [
+        {
+            tree: new PublicKey(merkletreePubkey),
+            queue: new PublicKey(nullifierQueuePubkey),
+            cpiContext: new PublicKey(cpiContextPubkey),
+            treeType: TreeType.State,
+        },
+        {
+            tree: new PublicKey(merkleTree2Pubkey),
+            queue: new PublicKey(nullifierQueue2Pubkey),
+            cpiContext: new PublicKey(cpiContext2Pubkey),
+            treeType: TreeType.State,
+        },
+    ];
 };
 
+/**
+ * Use only with Localnet testing.
+ * For public networks, fetch via {@link defaultStateTreeLookupTables} and {@link getLightStateTreeInfo}.
+ */
 export const defaultTestStateTreeAccounts = () => {
     return {
         nullifierQueue: new PublicKey(nullifierQueuePubkey),
@@ -88,6 +124,10 @@ export const defaultTestStateTreeAccounts = () => {
         addressQueue: new PublicKey(addressQueue),
     };
 };
+
+/**
+ * @internal testing only
+ */
 export const defaultTestStateTreeAccounts2 = () => {
     return {
         nullifierQueue2: new PublicKey(nullifierQueue2Pubkey),
@@ -108,6 +148,7 @@ export const nullifiedStateTreeLookupTableDevnet =
 
 export const nullifierQueuePubkey =
     'nfq1NvQDJ2GEgnS8zt9prAe8rjjpAW1zFkrvZoBR148';
+export const cpiContextPubkey = 'cpi1uHzrEhBG733DoEJNgHCyRS3XmmyVNZx5fonubE4';
 
 export const merkletreePubkey = 'smt1NamzXdq4AMqS2fS2F1i5KTYPZRhoHgWx38d8WsT';
 export const addressTree = 'amt1Ayt45jfbdw5YSo7iz6WZxUmnZsQTYXy82hVwyC2';
@@ -116,6 +157,7 @@ export const addressQueue = 'aq1S9z4reTSQAdgWHGD2zDaS39sjGrAxbR31vxJ2F4F';
 export const merkleTree2Pubkey = 'smt2rJAFdyJJupwMKAqTNAJwvjhmiZ4JYGZmbVRw1Ho';
 export const nullifierQueue2Pubkey =
     'nfq2hgS7NYemXsFaFUCe3EMXSDSfnZnAe27jC6aPP1X';
+export const cpiContext2Pubkey = 'cpi2cdhkH5roePvcudTgUL8ppEBfTay1desGh8G8QxK';
 
 export const confirmConfig: ConfirmOptions = {
     commitment: 'confirmed',

js/stateless.js/src/rpc.ts

@@ -71,10 +71,8 @@ import {
     negateAndCompressProof,
 } from './utils/parse-validity-proof';
 import { LightWasm } from './test-helpers';
-import {
-    ActiveStateTreeAddresses,
-    getLightStateTreeInfo,
-} from './utils/get-light-state-tree-info';
+import { getLightStateTreeInfo } from './utils/get-light-state-tree-info';
+import { ActiveTreeBundle } from './state/types';
 
 /** @internal */
 export function parseAccountData({
@@ -189,7 +187,7 @@ async function getCompressedTokenAccountsByOwnerOrDelegate(
 /** @internal */
 function buildCompressedAccountWithMaybeTokenData(
     accountStructWithOptionalTokenData: any,
-    activeStateTreeInfo: ActiveStateTreeAddresses,
+    activeStateTreeInfo: ActiveTreeBundle[],
 ): {
     account: CompressedAccountWithMerkleContext;
     maybeTokenData: TokenData | null;
@@ -542,17 +540,19 @@ export function getPublicInputHash(
  * @returns The queue for the given tree, or undefined if not found
  */
 export function getQueueForTree(
-    info: ActiveStateTreeAddresses,
+    info: ActiveTreeBundle[],
     tree: PublicKey,
 ): PublicKey {
-    const index = info.activeStateTrees.findIndex(t => t.equals(tree));
-
+    const index = info.findIndex(t => t.tree.equals(tree));
     if (index === -1) {
         throw new Error(
             'No associated queue found for tree. Please set activeStateTreeInfo with latest Tree accounts. If you use custom state trees, set manually.',
         );
     }
-    return info.activeQueues[index];
+    if (!info[index].queue) {
+        throw new Error('Queue must not be null for state tree');
+    }
+    return info[index].queue;
 }
 
 /**
@@ -563,16 +563,19 @@ export function getQueueForTree(
  * @returns The tree for the given queue, or undefined if not found
  */
 export function getTreeForQueue(
-    info: ActiveStateTreeAddresses,
+    info: ActiveTreeBundle[],
     queue: PublicKey,
 ): PublicKey {
-    const index = info.activeQueues.findIndex(q => q.equals(queue));
+    const index = info.findIndex(q => q.queue?.equals(queue));
     if (index === -1) {
         throw new Error(
             'No associated tree found for queue. Please set activeStateTreeInfo with latest Tree accounts. If you use custom state trees, set manually.',
         );
     }
-    return info.activeStateTrees[index];
+    if (!info[index].tree) {
+        throw new Error('Tree must not be null for state tree');
+    }
+    return info[index].tree;
 }
 
 /**
@@ -583,16 +586,19 @@ export function getTreeForQueue(
  * @param info - The active state tree addresses
  * @returns A random tree and queue
  */
-export function pickRandomTreeAndQueue(info: ActiveStateTreeAddresses): {
+export function pickRandomTreeAndQueue(info: ActiveTreeBundle[]): {
     tree: PublicKey;
     queue: PublicKey;
 } {
-    const length = info.activeStateTrees.length;
+    const length = info.length;
     const index = Math.floor(Math.random() * length);
 
+    if (!info[index].queue) {
+        throw new Error('Queue must not be null for state tree');
+    }
     return {
-        tree: info.activeStateTrees[index],
-        queue: info.activeQueues[index],
+        tree: info[index].tree,
+        queue: info[index].queue,
     };
 }
 
@@ -602,8 +608,7 @@ export function pickRandomTreeAndQueue(info: ActiveStateTreeAddresses): {
 export class Rpc extends Connection implements CompressionApiInterface {
     compressionApiEndpoint: string;
     proverEndpoint: string;
-    activeStateTreeInfo: ActiveStateTreeAddresses | null = null;
-    fetchStateTreePromise: Promise<ActiveStateTreeAddresses> | null = null;
+    activeStateTreeInfo: ActiveTreeBundle[] | null = null;
 
     constructor(
         endpoint: string,
@@ -619,34 +624,35 @@ export class Rpc extends Connection implements CompressionApiInterface {
     /**
      * Manually set state tree addresses
      */
-    setStateTreeInfo(info: ActiveStateTreeAddresses): void {
+    setStateTreeInfo(info: ActiveTreeBundle[]): void {
         this.activeStateTreeInfo = info;
     }
 
     /**
      * Get the active state tree addresses from the cluster.
      * If not already cached, fetches from the cluster.
      */
-    async getCachedActiveStateTreeInfo(): Promise<ActiveStateTreeAddresses> {
+    async getCachedActiveStateTreeInfo(): Promise<ActiveTreeBundle[]> {
         if (isLocalTest(this.rpcEndpoint)) {
             return localTestActiveStateTreeInfo();
         }
 
-        let info: ActiveStateTreeAddresses | null = null;
+        let info: ActiveTreeBundle[] | null = null;
         if (!this.activeStateTreeInfo) {
             const { mainnet, devnet } = defaultStateTreeLookupTables();
             try {
                 info = await getLightStateTreeInfo({
                     connection: this,
-                    stateTreeLookupTableAddress: mainnet.stateTreeLookupTable,
-                    nullifyTableAddress: mainnet.nullifyTable,
+                    stateTreeLookupTableAddress:
+                        mainnet[0].stateTreeLookupTable,
+                    nullifyTableAddress: mainnet[0].nullifyTable,
                 });
                 this.activeStateTreeInfo = info;
             } catch {
                 info = await getLightStateTreeInfo({
                     connection: this,
-                    stateTreeLookupTableAddress: devnet.stateTreeLookupTable,
-                    nullifyTableAddress: devnet.nullifyTable,
+                    stateTreeLookupTableAddress: devnet[0].stateTreeLookupTable,
+                    nullifyTableAddress: devnet[0].nullifyTable,
                 });
                 this.activeStateTreeInfo = info;
             }
@@ -665,7 +671,7 @@ export class Rpc extends Connection implements CompressionApiInterface {
     /**
      * Fetch the latest state tree addresses from the cluster.
      */
-    async getLatestActiveStateTreeInfo(): Promise<ActiveStateTreeAddresses> {
+    async getLatestActiveStateTreeInfo(): Promise<ActiveTreeBundle[]> {
         this.activeStateTreeInfo = null;
         return await this.getCachedActiveStateTreeInfo();
     }

js/stateless.js/src/state/types.ts

@@ -3,6 +3,32 @@ import { PublicKey } from '@solana/web3.js';
 import { Buffer } from 'buffer';
 import { NewAddressParamsPacked } from '../utils';
 
+export enum TreeType {
+    /**
+     * v1 state merkle tree
+     */
+    State = 0,
+    /**
+     * v1 address merkle tree
+     */
+    Address = 1,
+    /**
+     * v2 state merkle tree
+     */
+    BatchedState = 2,
+    /**
+     * v2 address merkle tree
+     */
+    BatchedAddress = 3,
+}
+
+export type ActiveTreeBundle = {
+    tree: PublicKey;
+    queue: PublicKey | null;
+    cpiContext: PublicKey | null;
+    treeType: TreeType;
+};
+
 export interface PackedCompressedAccountWithMerkleContext {
     compressedAccount: CompressedAccount;
     merkleContext: PackedMerkleContext;

js/stateless.js/src/test-helpers/test-rpc/test-rpc.ts

@@ -47,15 +47,11 @@ import {
 import { IndexedArray } from '../merkle-tree';
 import {
     MerkleContextWithNewAddressProof,
-    Rpc,
     convertMerkleProofsWithContextToHex,
     convertNonInclusionMerkleProofInputsToHex,
     proverRequest,
 } from '../../rpc';
-import {
-    ActiveStateTreeAddresses,
-    getLightStateTreeInfo,
-} from '../../utils/get-light-state-tree-info';
+import { ActiveTreeBundle } from '../../state/types';
 
 export interface TestRpcConfig {
     /**
@@ -155,8 +151,7 @@ export class TestRpc extends Connection implements CompressionApiInterface {
     lightWasm: LightWasm;
     depth: number;
     log = false;
-    activeStateTreeInfo: ActiveStateTreeAddresses | null = null;
-    fetchStateTreePromise: Promise<ActiveStateTreeAddresses> | null = null;
+    activeStateTreeInfo: ActiveTreeBundle[] | null = null;
 
     /**
      * Establish a Compression-compatible JSON RPC mock-connection
@@ -212,21 +207,21 @@ export class TestRpc extends Connection implements CompressionApiInterface {
     /**
      * Manually set state tree addresses
      */
-    setStateTreeInfo(info: ActiveStateTreeAddresses): void {
+    setStateTreeInfo(info: ActiveTreeBundle[]): void {
         this.activeStateTreeInfo = info;
     }
 
     /**
      * Returns local test state trees.
      */
-    async getCachedActiveStateTreeInfo(): Promise<ActiveStateTreeAddresses> {
+    async getCachedActiveStateTreeInfo(): Promise<ActiveTreeBundle[]> {
         return localTestActiveStateTreeInfo();
     }
 
     /**
      * Returns local test state trees.
      */
-    async getLatestActiveStateTreeInfo(): Promise<ActiveStateTreeAddresses> {
+    async getLatestActiveStateTreeInfo(): Promise<ActiveTreeBundle[]> {
         return localTestActiveStateTreeInfo();
     }