refactor: extract self-contained createNullifierInstruction wrapper

Author: tilo-14

zk/nullifier/ts-tests/nullifier.test.ts

@@ -1,33 +1,73 @@
 import { web3, Program, AnchorProvider, setProvider } from "@coral-xyz/anchor";
-import {
-  bn,
-  createRpc,
-  deriveAddressSeedV2,
-  deriveAddressV2,
-  batchAddressTree,
-  PackedAccounts,
-  Rpc,
-  sleep,
-  SystemAccountMetaConfig,
-  defaultTestStateTreeAccounts,
-  featureFlags,
-  VERSION,
-  confirmTx,
-} from "@lightprotocol/stateless.js";
+import { createRpc, Rpc, sleep, confirmTx } from "@lightprotocol/stateless.js";
 import * as assert from "assert";
 import * as path from "path";
 import * as fs from "fs";
 
-// Force V2 mode
-(featureFlags as any).version = VERSION.V2;
-
-// Load IDL
 const IDL = JSON.parse(
   fs.readFileSync(path.join(process.cwd(), "target/idl/zk_nullifier.json"), "utf8")
 );
 
 const PROGRAM_ID = new web3.PublicKey(IDL.address);
-const NULLIFIER_PREFIX = Buffer.from("nullifier");
+
+async function createNullifierInstruction(
+  rpc: Rpc,
+  program: Program,
+  signer: web3.PublicKey,
+  nullifiers: Uint8Array[]
+): Promise<web3.TransactionInstruction> {
+  const {
+    bn,
+    deriveAddressSeedV2,
+    deriveAddressV2,
+    batchAddressTree,
+    PackedAccounts,
+    SystemAccountMetaConfig,
+    defaultTestStateTreeAccounts,
+    featureFlags,
+    VERSION,
+  } = await import("@lightprotocol/stateless.js");
+
+  (featureFlags as any).version = VERSION.V2;
+
+  const NULLIFIER_PREFIX = Buffer.from("nullifier");
+  const addressTree = new web3.PublicKey(batchAddressTree);
+  const outputStateTree = defaultTestStateTreeAccounts().merkleTree;
+
+  const addressesWithTree = nullifiers.map((nullifier) => {
+    const seed = deriveAddressSeedV2([NULLIFIER_PREFIX, nullifier]);
+    const address = deriveAddressV2(seed, addressTree, program.programId);
+    return { tree: addressTree, queue: addressTree, address: bn(address.toBytes()) };
+  });
+
+  const proofResult = await rpc.getValidityProofV0([], addressesWithTree);
+
+  const remainingAccounts = new PackedAccounts();
+  remainingAccounts.addPreAccountsSigner(signer);
+  remainingAccounts.addSystemAccountsV2(SystemAccountMetaConfig.new(program.programId));
+
+  const addressMerkleTreeIndex = remainingAccounts.insertOrGet(addressTree);
+  const outputStateTreeIndex = remainingAccounts.insertOrGet(outputStateTree);
+
+  const { remainingAccounts: accountMetas, systemStart } = remainingAccounts.toAccountMetas();
+
+  const data = {
+    proof: { 0: proofResult.compressedProof },
+    addressTreeInfo: {
+      addressMerkleTreePubkeyIndex: addressMerkleTreeIndex,
+      addressQueuePubkeyIndex: addressMerkleTreeIndex,
+      rootIndex: proofResult.rootIndices[0],
+    },
+    outputStateTreeIndex,
+    systemAccountsOffset: systemStart,
+  };
+
+  return program.methods
+    .createNullifier(data, nullifiers.map((n) => Array.from(n)))
+    .accounts({ signer })
+    .remainingAccounts(accountMetas)
+    .instruction();
+}
 
 describe("zk-nullifier", () => {
   let rpc: Rpc;
@@ -46,7 +86,6 @@ describe("zk-nullifier", () => {
     await rpc.requestAirdrop(signer.publicKey, web3.LAMPORTS_PER_SOL);
     await sleep(2000);
 
-    // Setup Anchor provider and program
     const connection = new web3.Connection("http://127.0.0.1:8899", "confirmed");
     const wallet = {
       publicKey: signer.publicKey,
@@ -64,71 +103,17 @@ describe("zk-nullifier", () => {
     program = new Program(IDL, provider);
   });
 
-  /** Generate a random 32-byte value */
   function randomBytes32(): Uint8Array {
     return web3.Keypair.generate().publicKey.toBytes();
   }
 
-  async function buildCreateNullifierInstruction(
-    nullifiers: Uint8Array[]
-  ): Promise<web3.TransactionInstruction> {
-    const addressTree = new web3.PublicKey(batchAddressTree);
-    const outputStateTree = defaultTestStateTreeAccounts().merkleTree;
-
-    // Derive addresses for all nullifiers 
-    const addressesWithTree = nullifiers.map((nullifier) => {
-      const seed = deriveAddressSeedV2([NULLIFIER_PREFIX, nullifier]);
-      const address = deriveAddressV2(seed, addressTree, PROGRAM_ID);
-      return { tree: addressTree, queue: addressTree, address: bn(address.toBytes()) };
-    });
-
-    const proofResult = await rpc.getValidityProofV0([], addressesWithTree);
-
-    // Use V2 accounts layout
-    const remainingAccounts = new PackedAccounts();
-    remainingAccounts.addPreAccountsSigner(signer.publicKey);
-    remainingAccounts.addSystemAccountsV2(SystemAccountMetaConfig.new(PROGRAM_ID));
-
-    const addressMerkleTreeIndex = remainingAccounts.insertOrGet(addressTree);
-    const outputStateTreeIndex = remainingAccounts.insertOrGet(outputStateTree);
-
-    const { remainingAccounts: accountMetas, systemStart } = remainingAccounts.toAccountMetas();
-
-    // ValidityProof struct
-    const proof = {
-      0: proofResult.compressedProof,
-    };
-
-    // NullifierInstructionData struct
-    const data = {
-      proof,
-      addressTreeInfo: {
-        addressMerkleTreePubkeyIndex: addressMerkleTreeIndex,
-        addressQueuePubkeyIndex: addressMerkleTreeIndex,
-        rootIndex: proofResult.rootIndices[0],
-      },
-      outputStateTreeIndex,
-      systemAccountsOffset: systemStart,
-    };
-
-    const ix = await program.methods
-      .createNullifier(data, nullifiers.map((n) => Array.from(n)))
-      .accounts({
-        signer: signer.publicKey,
-      })
-      .remainingAccounts(accountMetas)
-      .instruction();
-
-    return ix;
-  }
-
   describe("Single nullifier", () => {
     it("should create a nullifier", async () => {
       const nullifier = randomBytes32();
 
       console.log("Nullifier:", Buffer.from(nullifier).toString("hex").slice(0, 16) + "...");
 
-      const ix = await buildCreateNullifierInstruction([nullifier]);
+      const ix = await createNullifierInstruction(rpc, program, signer.publicKey, [nullifier]);
       const computeIx = web3.ComputeBudgetProgram.setComputeUnitLimit({ units: 200_000 });
 
       const tx = new web3.Transaction().add(computeIx, ix);
@@ -152,7 +137,7 @@ describe("zk-nullifier", () => {
     it("should reject duplicate nullifier", async () => {
       const nullifier = randomBytes32();
 
-      const ix1 = await buildCreateNullifierInstruction([nullifier]);
+      const ix1 = await createNullifierInstruction(rpc, program, signer.publicKey, [nullifier]);
       const computeIx = web3.ComputeBudgetProgram.setComputeUnitLimit({ units: 200_000 });
 
       const tx1 = new web3.Transaction().add(computeIx, ix1);
@@ -163,9 +148,8 @@ describe("zk-nullifier", () => {
       await rpc.sendTransaction(tx1, [signer]);
       await sleep(2000);
 
-      // Attempt to create duplicate - should fail when getting validity proof
       try {
-        await buildCreateNullifierInstruction([nullifier]);
+        await createNullifierInstruction(rpc, program, signer.publicKey, [nullifier]);
         assert.fail("Should have rejected duplicate nullifier");
       } catch (err: any) {
         assert.ok(
@@ -183,7 +167,7 @@ describe("zk-nullifier", () => {
 
       console.log("Creating 3 nullifiers in one transaction...");
 
-      const ix = await buildCreateNullifierInstruction(nullifiers);
+      const ix = await createNullifierInstruction(rpc, program, signer.publicKey, nullifiers);
       const computeIx = web3.ComputeBudgetProgram.setComputeUnitLimit({ units: 300_000 });
 
       const tx = new web3.Transaction().add(computeIx, ix);
@@ -203,4 +187,4 @@ describe("zk-nullifier", () => {
       console.log("Total nullifier accounts:", accounts.items.length);
     });
   });
-});
+});