Account Compression: Add Close Tree Instruction (#3641)

* ac: add close tree instruction

* ac: update sdk with close tree instruction

* ac: test close tree instruction

Author: ngundotra

account-compression/programs/account-compression/src/lib.rs

@@ -57,7 +57,7 @@ pub struct Initialize<'info> {
     /// CHECK: This account will be zeroed out, and the size will be validated
     pub merkle_tree: UncheckedAccount<'info>,
 
-    /// Authority that validates the content of the trees.
+    /// Authority that controls write-access to the tree
     /// Typically a program, e.g., the Bubblegum contract validates that leaves are valid NFTs.
     pub authority: Signer<'info>,
 
@@ -76,7 +76,7 @@ pub struct Modify<'info> {
     /// CHECK: This account is validated in the instruction
     pub merkle_tree: UncheckedAccount<'info>,
 
-    /// Authority that validates the content of the trees.
+    /// Authority that controls write-access to the tree
     /// Typically a program, e.g., the Bubblegum contract validates that leaves are valid NFTs.
     pub authority: Signer<'info>,
 
@@ -100,11 +100,26 @@ pub struct TransferAuthority<'info> {
     /// CHECK: This account is validated in the instruction
     pub merkle_tree: UncheckedAccount<'info>,
 
-    /// Authority that validates the content of the trees.
+    /// Authority that controls write-access to the tree
     /// Typically a program, e.g., the Bubblegum contract validates that leaves are valid NFTs.
     pub authority: Signer<'info>,
 }
 
+/// Context for closing a tree
+#[derive(Accounts)]
+pub struct CloseTree<'info> {
+    #[account(mut)]
+    /// CHECK: This account is validated in the instruction
+    pub merkle_tree: AccountInfo<'info>,
+
+    /// Authority that controls write-access to the tree
+    pub authority: Signer<'info>,
+
+    /// CHECK: Recipient of funds after
+    #[account(mut)]
+    pub recipient: AccountInfo<'info>,
+}
+
 /// This macro applies functions on a ConcurrentMerkleT:ee and emits leaf information
 /// needed to sync the merkle tree state with off-chain indexers.
 macro_rules! merkle_tree_depth_size_apply_fn {
@@ -512,4 +527,39 @@ pub mod spl_account_compression {
             &ctx.accounts.log_wrapper,
         )
     }
+
+    pub fn close_empty_tree(ctx: Context<CloseTree>) -> Result<()> {
+        require_eq!(
+            *ctx.accounts.merkle_tree.owner,
+            crate::id(),
+            AccountCompressionError::IncorrectAccountOwner
+        );
+        let mut merkle_tree_bytes = ctx.accounts.merkle_tree.try_borrow_mut_data()?;
+        let (header_bytes, rest) =
+            merkle_tree_bytes.split_at_mut(CONCURRENT_MERKLE_TREE_HEADER_SIZE_V1);
+
+        let header = ConcurrentMerkleTreeHeader::try_from_slice(header_bytes)?;
+        header.assert_valid_authority(&ctx.accounts.authority.key())?;
+
+        let merkle_tree_size = merkle_tree_get_size(&header)?;
+        let (tree_bytes, canopy_bytes) = rest.split_at_mut(merkle_tree_size);
+
+        let id = ctx.accounts.merkle_tree.key();
+        merkle_tree_apply_fn!(header, id, tree_bytes, prove_tree_is_empty,)?;
+
+        // Close merkle tree account
+        // 1. Move lamports
+        let dest_starting_lamports = ctx.accounts.recipient.lamports();
+        **ctx.accounts.recipient.lamports.borrow_mut() = dest_starting_lamports
+            .checked_add(ctx.accounts.merkle_tree.lamports())
+            .unwrap();
+        **ctx.accounts.merkle_tree.lamports.borrow_mut() = 0;
+
+        // 2. Set all CMT account bytes to 0
+        header_bytes.fill(0);
+        tree_bytes.fill(0);
+        canopy_bytes.fill(0);
+
+        Ok(())
+    }
 }

account-compression/sdk/idl/spl_account_compression.json

@@ -29,7 +29,7 @@
           "isMut": false,
           "isSigner": true,
           "docs": [
-            "Authority that validates the content of the trees.",
+            "Authority that controls write-access to the tree",
             "Typically a program, e.g., the Bubblegum contract validates that leaves are valid NFTs."
           ]
         },
@@ -81,7 +81,7 @@
           "isMut": false,
           "isSigner": true,
           "docs": [
-            "Authority that validates the content of the trees.",
+            "Authority that controls write-access to the tree",
             "Typically a program, e.g., the Bubblegum contract validates that leaves are valid NFTs."
           ]
         },
@@ -146,7 +146,7 @@
           "isMut": false,
           "isSigner": true,
           "docs": [
-            "Authority that validates the content of the trees.",
+            "Authority that controls write-access to the tree",
             "Typically a program, e.g., the Bubblegum contract validates that leaves are valid NFTs."
           ]
         }
@@ -217,7 +217,7 @@
           "isMut": false,
           "isSigner": true,
           "docs": [
-            "Authority that validates the content of the trees.",
+            "Authority that controls write-access to the tree",
             "Typically a program, e.g., the Bubblegum contract validates that leaves are valid NFTs."
           ]
         },
@@ -262,7 +262,7 @@
           "isMut": false,
           "isSigner": true,
           "docs": [
-            "Authority that validates the content of the trees.",
+            "Authority that controls write-access to the tree",
             "Typically a program, e.g., the Bubblegum contract validates that leaves are valid NFTs."
           ]
         },
@@ -300,6 +300,30 @@
           "type": "u32"
         }
       ]
+    },
+    {
+      "name": "closeEmptyTree",
+      "accounts": [
+        {
+          "name": "merkleTree",
+          "isMut": true,
+          "isSigner": false
+        },
+        {
+          "name": "authority",
+          "isMut": false,
+          "isSigner": true,
+          "docs": [
+            "Authority that controls write-access to the tree"
+          ]
+        },
+        {
+          "name": "recipient",
+          "isMut": true,
+          "isSigner": false
+        }
+      ],
+      "args": []
     }
   ],
   "types": [

account-compression/sdk/src/generated/instructions/closeEmptyTree.ts

@@ -0,0 +1,88 @@
+/**
+ * This code was GENERATED using the solita package.
+ * Please DO NOT EDIT THIS FILE, instead rerun solita to update it or write a wrapper to add functionality.
+ *
+ * See: https://github.com/metaplex-foundation/solita
+ */
+
+import * as beet from '@metaplex-foundation/beet'
+import * as web3 from '@solana/web3.js'
+
+/**
+ * @category Instructions
+ * @category CloseEmptyTree
+ * @category generated
+ */
+export const closeEmptyTreeStruct = new beet.BeetArgsStruct<{
+  instructionDiscriminator: number[] /* size: 8 */
+}>(
+  [['instructionDiscriminator', beet.uniformFixedSizeArray(beet.u8, 8)]],
+  'CloseEmptyTreeInstructionArgs'
+)
+/**
+ * Accounts required by the _closeEmptyTree_ instruction
+ *
+ * @property [_writable_] merkleTree
+ * @property [**signer**] authority
+ * @property [_writable_] recipient
+ * @category Instructions
+ * @category CloseEmptyTree
+ * @category generated
+ */
+export type CloseEmptyTreeInstructionAccounts = {
+  merkleTree: web3.PublicKey
+  authority: web3.PublicKey
+  recipient: web3.PublicKey
+  anchorRemainingAccounts?: web3.AccountMeta[]
+}
+
+export const closeEmptyTreeInstructionDiscriminator = [
+  50, 14, 219, 107, 78, 103, 16, 103,
+]
+
+/**
+ * Creates a _CloseEmptyTree_ instruction.
+ *
+ * @param accounts that will be accessed while the instruction is processed
+ * @category Instructions
+ * @category CloseEmptyTree
+ * @category generated
+ */
+export function createCloseEmptyTreeInstruction(
+  accounts: CloseEmptyTreeInstructionAccounts,
+  programId = new web3.PublicKey('cmtDvXumGCrqC1Age74AVPhSRVXJMd8PJS91L8KbNCK')
+) {
+  const [data] = closeEmptyTreeStruct.serialize({
+    instructionDiscriminator: closeEmptyTreeInstructionDiscriminator,
+  })
+  const keys: web3.AccountMeta[] = [
+    {
+      pubkey: accounts.merkleTree,
+      isWritable: true,
+      isSigner: false,
+    },
+    {
+      pubkey: accounts.authority,
+      isWritable: false,
+      isSigner: true,
+    },
+    {
+      pubkey: accounts.recipient,
+      isWritable: true,
+      isSigner: false,
+    },
+  ]
+
+  if (accounts.anchorRemainingAccounts != null) {
+    for (const acc of accounts.anchorRemainingAccounts) {
+      keys.push(acc)
+    }
+  }
+
+  const ix = new web3.TransactionInstruction({
+    programId,
+    keys,
+    data,
+  })
+  return ix
+}

account-compression/sdk/src/generated/instructions/index.ts

@@ -1,4 +1,5 @@
 export * from './append'
+export * from './closeEmptyTree'
 export * from './initEmptyMerkleTree'
 export * from './insertOrAppend'
 export * from './replaceLeaf'

account-compression/sdk/src/instructions/index.ts

@@ -7,7 +7,8 @@ import {
     createTransferAuthorityInstruction,
     createVerifyLeafInstruction,
     PROGRAM_ID,
-    createInitEmptyMerkleTreeInstruction
+    createInitEmptyMerkleTreeInstruction,
+    createCloseEmptyTreeInstruction
 } from "../generated";
 
 /**
@@ -147,4 +148,18 @@ export async function createAllocTreeIx(
         space: requiredSpace,
         programId: PROGRAM_ID
     });
+}
+
+export function createCloseEmptyTreeIx(
+    authority: PublicKey,
+    merkleTree: PublicKey,
+    recipient: PublicKey,
+): TransactionInstruction {
+    return createCloseEmptyTreeInstruction(
+        {
+            merkleTree,
+            authority,
+            recipient
+        },
+    )
 }

account-compression/sdk/tests/accountCompression.test.ts

@@ -30,6 +30,7 @@ import {
   createTransferAuthorityIx,
   createVerifyLeafIx,
   ConcurrentMerkleTreeAccount,
+  createCloseEmptyTreeInstruction,
 } from "../src";
 
 describe("Account Compression", () => {
@@ -324,7 +325,6 @@ describe("Account Compression", () => {
     it(`Replace all of them in a block`, async () => {
       // Replace 64 leaves before syncing off-chain tree with on-chain tree
 
-      // Cache all proofs so we can execute in single block
       let ixArray: TransactionInstruction[] = [];
       let txList: Promise<string>[] = [];
 
@@ -351,7 +351,7 @@ describe("Account Compression", () => {
         ixArray.push(replaceIx);
       }
 
-      // Execute all replaces in a "single block"
+      // Execute all replaces 
       ixArray.map((ix) => {
         txList.push(
           execute(provider, [ix], [payer])
@@ -372,6 +372,70 @@ describe("Account Compression", () => {
         "Updated on chain root does not match root of updated off chain tree"
       );
     });
+    it("Empty all of the leaves and close the tree", async () => {
+      let ixArray: TransactionInstruction[] = [];
+      let txList: Promise<string>[] = [];
+      const leavesToUpdate: Buffer[] = [];
+      for (let i = 0; i < MAX_SIZE; i++) {
+        const index = i;
+        const newLeaf = hash(
+          payer.publicKey.toBuffer(),
+          Buffer.from(new BN(i).toArray())
+        );
+        leavesToUpdate.push(newLeaf);
+        const proof = getProofOfLeaf(offChainTree, index);
+        const replaceIx = createReplaceIx(
+          payer,
+          cmtKeypair.publicKey,
+          offChainTree.root,
+          offChainTree.leaves[i].node,
+          Buffer.alloc(32), // Empty node
+          index,
+          proof.map((treeNode) => {
+            return treeNode.node;
+          })
+        );
+        ixArray.push(replaceIx);
+      }
+      // Execute all replaces 
+      ixArray.map((ix) => {
+        txList.push(
+          execute(provider, [ix], [payer])
+        );
+      });
+      await Promise.all(txList);
+
+      let payerInfo = await provider.connection.getAccountInfo(payer.publicKey, "confirmed")!;
+      let treeInfo = await provider.connection.getAccountInfo(cmtKeypair.publicKey, "confirmed")!;
+
+      let payerLamports = payerInfo!.lamports;
+      let treeLamports = treeInfo!.lamports;
+
+      const ix = createCloseEmptyTreeInstruction({
+        merkleTree: cmtKeypair.publicKey,
+        authority: payer.publicKey,
+        recipient: payer.publicKey,
+      })
+      await execute(provider, [ix], [payer]);
+
+      payerInfo = await provider.connection.getAccountInfo(payer.publicKey, "confirmed")!;
+      const finalLamports = payerInfo!.lamports;
+      assert(finalLamports === (payerLamports + treeLamports - 5000), "Expected payer to have received the lamports from the closed tree account");
+
+      treeInfo = await provider.connection.getAccountInfo(cmtKeypair.publicKey, "confirmed");
+      assert(treeInfo === null, "Expected the merkle tree account info to be null");
+    })
+    it("It cannot be closed until empty", async () => {
+      const ix = createCloseEmptyTreeInstruction({
+        merkleTree: cmtKeypair.publicKey,
+        authority: payer.publicKey,
+        recipient: payer.publicKey,
+      })
+      try {
+        await execute(provider, [ix], [payer]);
+        assert(false, "Closing a tree account before it is empty should ALWAYS error")
+      } catch (e) { }
+    })
   });
 
   describe(`Having created a tree with depth 3`, () => {