updated close account to allow closing prepared and not finalized accounts + updated comments on append_canopy_nodes to reflect the possibility to replace those

Author: StanChe

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

@@ -229,6 +229,17 @@ pub mod spl_account_compression {
     /// This instruction pre-initializes the canopy with the specified leaf nodes of the canopy.
     /// This is intended to be used after `prepare_tree` and in conjunction with the
     /// `finalize_merkle_tree_with_root` instruction that'll finalize the tree initialization.
+    /// The canopy is used to cache the uppermost nodes of the tree, which allows for a smaller
+    /// proof size when updating the tree. The canopy should be filled with the necessary nodes
+    /// before calling `finalize_merkle_tree_with_root`. You may call this instruction multiple
+    /// times to fill the canopy with the necessary nodes. The canopy may be filled with the
+    /// nodes in any order. The already filled nodes may be replaced with new nodes before calling
+    /// `finalize_merkle_tree_with_root` if the step was done in error.
+    /// The canopy should be filled with all the nodes that are to the left of the rightmost
+    /// leaf of the tree before calling `finalize_merkle_tree_with_root`. The canopy should not
+    /// contain any nodes to the right of the rightmost leaf of the tree.
+    /// This instruction calculates and filles in all the canopy nodes "above" the provided ones.
+    /// The validation of the canopy is done in the `finalize_merkle_tree_with_root` instruction.
     pub fn append_canopy_nodes(
         ctx: Context<Modify>,
         start_index: u32,
@@ -264,8 +275,15 @@ pub mod spl_account_compression {
         )
     }
 
-    /// Initializes a prepared tree with a root and a rightmost leaf. The rightmost leaf is used to verify the canopy if the tree has it. Before calling this instruction, the tree should be prepared with `prepare_tree` and the canopy should be filled with the necessary nodes with `append_canopy_nodes` (if the canopy is used).
-    /// This method should be used for rolluped creation of trees. The indexing of such rollups should be done off-chain. The programs calling this instruction should take care of ensuring the indexing is possible. For example, staking may be required to ensure the tree creator has some responsibility for what is being indexed. If indexing is not possible, there should be a mechanism to penalize the tree creator.
+    /// Initializes a prepared tree with a root and a rightmost leaf. The rightmost leaf is used to
+    /// verify the canopy if the tree has it. Before calling this instruction, the tree should be
+    /// prepared with `prepare_tree` and the canopy should be filled with the necessary nodes with
+    /// `append_canopy_nodes` (if the canopy is used). This method should be used for rolluped
+    /// creation of trees. The indexing of such rollups should be done off-chain. The programs
+    /// calling this instruction should take care of ensuring the indexing is possible. For example,
+    /// staking may be required to ensure the tree creator has some responsibility for what is being
+    /// indexed. If indexing is not possible, there should be a mechanism to penalize the tree
+    /// creator.
     pub fn finalize_merkle_tree_with_root(
         ctx: Context<Modify>,
         root: [u8; 32],
@@ -546,8 +564,11 @@ pub mod spl_account_compression {
         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_mut!(header, id, tree_bytes, prove_tree_is_empty,)?;
+        // Check if the tree is either empty or is batch initialized and not finalized yet.
+        if !header.get_is_batch_initialized() || !tree_bytes.iter().all(|&x| x == 0) {
+            let id = ctx.accounts.merkle_tree.key();
+            merkle_tree_apply_fn_mut!(header, id, tree_bytes, prove_tree_is_empty,)?;
+        }
 
         // Close merkle tree account
         // 1. Move lamports

account-compression/programs/account-compression/src/state/concurrent_merkle_tree_header.rs

@@ -140,6 +140,12 @@ impl ConcurrentMerkleTreeHeader {
         }
     }
 
+    pub fn get_is_batch_initialized(&self) -> bool {
+        match &self.header {
+            ConcurrentMerkleTreeHeaderData::V1(header) => header.is_batch_initialized,
+        }
+    }
+
     pub fn set_new_authority(&mut self, new_authority: &Pubkey) {
         match self.header {
             ConcurrentMerkleTreeHeaderData::V1(ref mut header) => {

account-compression/sdk/idl/spl_account_compression.json

@@ -109,7 +109,18 @@
       "docs": [
         "This instruction pre-initializes the canopy with the specified leaf nodes of the canopy.",
         "This is intended to be used after `prepare_tree` and in conjunction with the",
-        "`finalize_merkle_tree_with_root` instruction that'll finalize the tree initialization."
+        "`finalize_merkle_tree_with_root` instruction that'll finalize the tree initialization.",
+        "The canopy is used to cache the uppermost nodes of the tree, which allows for a smaller",
+        "proof size when updating the tree. The canopy should be filled with the necessary nodes",
+        "before calling `finalize_merkle_tree_with_root`. You may call this instruction multiple",
+        "times to fill the canopy with the necessary nodes. The canopy may be filled with the",
+        "nodes in any order. The already filled nodes may be replaced with new nodes before calling",
+        "`finalize_merkle_tree_with_root` if the step was done in error.",
+        "The canopy should be filled with all the nodes that are to the left of the rightmost",
+        "leaf of the tree before calling `finalize_merkle_tree_with_root`. The canopy should not",
+        "contain any nodes to the right of the rightmost leaf of the tree.",
+        "This instruction calculates and filles in all the canopy nodes \"above\" the provided ones.",
+        "The validation of the canopy is done in the `finalize_merkle_tree_with_root` instruction."
       ],
       "accounts": [
         {
@@ -151,8 +162,15 @@
     {
       "name": "finalizeMerkleTreeWithRoot",
       "docs": [
-        "Initializes a prepared tree with a root and a rightmost leaf. The rightmost leaf is used to verify the canopy if the tree has it. Before calling this instruction, the tree should be prepared with `prepare_tree` and the canopy should be filled with the necessary nodes with `append_canopy_nodes` (if the canopy is used).",
-        "This method should be used for rolluped creation of trees. The indexing of such rollups should be done off-chain. The programs calling this instruction should take care of ensuring the indexing is possible. For example, staking may be required to ensure the tree creator has some responsibility for what is being indexed. If indexing is not possible, there should be a mechanism to penalize the tree creator."
+        "Initializes a prepared tree with a root and a rightmost leaf. The rightmost leaf is used to",
+        "verify the canopy if the tree has it. Before calling this instruction, the tree should be",
+        "prepared with `prepare_tree` and the canopy should be filled with the necessary nodes with",
+        "`append_canopy_nodes` (if the canopy is used). This method should be used for rolluped",
+        "creation of trees. The indexing of such rollups should be done off-chain. The programs",
+        "calling this instruction should take care of ensuring the indexing is possible. For example,",
+        "staking may be required to ensure the tree creator has some responsibility for what is being",
+        "indexed. If indexing is not possible, there should be a mechanism to penalize the tree",
+        "creator."
       ],
       "accounts": [
         {
@@ -727,11 +745,16 @@
     },
     {
       "code": 6011,
+      "name": "BatchNotInitialized",
+      "msg": "Tree header was not initialized for batch processing"
+    },
+    {
+      "code": 6012,
       "name": "CanopyRootMismatch",
       "msg": "Canopy root does not match the root of the tree"
     },
     {
-      "code": 6012,
+      "code": 6013,
       "name": "CanopyRightmostLeafMismatch",
       "msg": "Canopy contains nodes to the right of the rightmost leaf of the tree"
     }

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

@@ -234,14 +234,34 @@ export class TreeAlreadyInitializedError extends Error {
 createErrorFromCodeLookup.set(0x177a, () => new TreeAlreadyInitializedError());
 createErrorFromNameLookup.set('TreeAlreadyInitialized', () => new TreeAlreadyInitializedError());
 
+/**
+ * BatchNotInitialized: 'Tree header was not initialized for batch processing'
+ *
+ * @category Errors
+ * @category generated
+ */
+export class BatchNotInitializedError extends Error {
+    readonly code: number = 0x177b;
+    readonly name: string = 'BatchNotInitialized';
+    constructor() {
+        super('Tree header was not initialized for batch processing');
+        if (typeof Error.captureStackTrace === 'function') {
+            Error.captureStackTrace(this, BatchNotInitializedError);
+        }
+    }
+}
+
+createErrorFromCodeLookup.set(0x177b, () => new BatchNotInitializedError());
+createErrorFromNameLookup.set('BatchNotInitialized', () => new BatchNotInitializedError());
+
 /**
  * CanopyRootMismatch: 'Canopy root does not match the root of the tree'
  *
  * @category Errors
  * @category generated
  */
 export class CanopyRootMismatchError extends Error {
-    readonly code: number = 0x177b;
+    readonly code: number = 0x177c;
     readonly name: string = 'CanopyRootMismatch';
     constructor() {
         super('Canopy root does not match the root of the tree');
@@ -251,7 +271,7 @@ export class CanopyRootMismatchError extends Error {
     }
 }
 
-createErrorFromCodeLookup.set(0x177b, () => new CanopyRootMismatchError());
+createErrorFromCodeLookup.set(0x177c, () => new CanopyRootMismatchError());
 createErrorFromNameLookup.set('CanopyRootMismatch', () => new CanopyRootMismatchError());
 
 /**
@@ -261,7 +281,7 @@ createErrorFromNameLookup.set('CanopyRootMismatch', () => new CanopyRootMismatch
  * @category generated
  */
 export class CanopyRightmostLeafMismatchError extends Error {
-    readonly code: number = 0x177c;
+    readonly code: number = 0x177d;
     readonly name: string = 'CanopyRightmostLeafMismatch';
     constructor() {
         super('Canopy contains nodes to the right of the rightmost leaf of the tree');
@@ -271,7 +291,7 @@ export class CanopyRightmostLeafMismatchError extends Error {
     }
 }
 
-createErrorFromCodeLookup.set(0x177c, () => new CanopyRightmostLeafMismatchError());
+createErrorFromCodeLookup.set(0x177d, () => new CanopyRightmostLeafMismatchError());
 createErrorFromNameLookup.set('CanopyRightmostLeafMismatch', () => new CanopyRightmostLeafMismatchError());
 
 /**

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

@@ -11,6 +11,7 @@ import {
     createAppendCanopyNodesIx,
     createAppendIx,
     createCloseEmptyTreeInstruction,
+    createCloseEmptyTreeIx,
     createFinalizeMerkleTreeWithRootIx,
     createInitEmptyMerkleTreeIx,
     createReplaceIx,
@@ -177,6 +178,27 @@ describe('Account Compression', () => {
                 assert(false, 'Double finalizing should have failed');
             } catch {}
         });
+        
+        it('Should be able to close a prepared tree', async () => {
+            let payerInfo = await provider.connection.getAccountInfo(payer, 'confirmed')!;
+            let treeInfo = await provider.connection.getAccountInfo(cmt, 'confirmed')!;
+
+            const payerLamports = payerInfo!.lamports;
+            const treeLamports = treeInfo!.lamports;
+
+            const closeIx = createCloseEmptyTreeIx(cmt, payer, payer);
+            await execute(provider, [closeIx], [payerKeypair]);
+            
+            payerInfo = await provider.connection.getAccountInfo(payer, '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(cmt, 'confirmed');
+            assert(treeInfo === null, 'Expected the merkle tree account info to be null');
+        });
     });
     describe('Having prepared a tree with canopy', () => {
         const depth = 3;
@@ -474,6 +496,108 @@ describe('Account Compression', () => {
                 assert(false, 'Initializing an empty tree after preparing a tree should have failed');
             } catch {}
         });
+        it('Should be able to close a prepared tree after setting the canopy', async () => {
+            const merkleTreeRaw = new MerkleTree(leaves);
+            
+            const appendIx = createAppendCanopyNodesIx(
+                cmt,
+                payer,
+                merkleTreeRaw.leaves
+                    .slice(0, leaves.length / 2)
+                    .filter((_, i) => i % 2 === 0)
+                    .map(leaf => leaf.parent!.node!),
+                0,
+            );
+            await execute(provider, [appendIx], [payerKeypair]);
+            let payerInfo = await provider.connection.getAccountInfo(payer, 'confirmed')!;
+            let treeInfo = await provider.connection.getAccountInfo(cmt, 'confirmed')!;
+
+            const payerLamports = payerInfo!.lamports;
+            const treeLamports = treeInfo!.lamports;
+
+            const closeIx = createCloseEmptyTreeIx(cmt, payer, payer);
+            await execute(provider, [closeIx], [payerKeypair]);
+            
+            payerInfo = await provider.connection.getAccountInfo(payer, '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(cmt, 'confirmed');
+            assert(treeInfo === null, 'Expected the merkle tree account info to be null');
+        });
+    });
+    describe('Having prepared an empty tree with canopy', () => {
+        const depth = 3;
+        const size = 8;
+        const canopyDepth = 2;
+        // empty leaves represent the empty tree
+        const leaves = [
+            Buffer.alloc(32),
+            Buffer.alloc(32),
+            Buffer.alloc(32),
+            Buffer.alloc(32),
+            Buffer.alloc(32),
+            Buffer.alloc(32),
+            Buffer.alloc(32),
+            Buffer.alloc(32),
+        ];
+        let anotherKeyPair: Keypair;
+        let another: PublicKey;
+        beforeEach(async () => {
+            const cmtKeypair = await prepareTree({
+                canopyDepth,
+                depthSizePair: {
+                    maxBufferSize: size,
+                    maxDepth: depth,
+                },
+                payer: payerKeypair,
+                provider,
+            });
+            cmt = cmtKeypair.publicKey;
+            anotherKeyPair = Keypair.generate();
+            another = anotherKeyPair.publicKey;
+            await provider.connection.confirmTransaction(
+                await provider.connection.requestAirdrop(another, 1e10),
+                'confirmed',
+            );
+        });
+        
+        it('Should be able to finalize an empty tree with empty canopy and close it afterwards', async () => {
+            const merkleTreeRaw = new MerkleTree(leaves);
+            const root = merkleTreeRaw.root;
+            const leaf = leaves[leaves.length - 1];
+
+            const finalize = createFinalizeMerkleTreeWithRootIx(
+                cmt,
+                payer,
+                root,
+                leaf,
+                leaves.length - 1,
+                merkleTreeRaw.getProof(leaves.length - 1).proof,
+            );
+            await execute(provider, [finalize], [payerKeypair]);
+            let payerInfo = await provider.connection.getAccountInfo(payer, 'confirmed')!;
+            let treeInfo = await provider.connection.getAccountInfo(cmt, 'confirmed')!;
+
+            const payerLamports = payerInfo!.lamports;
+            const treeLamports = treeInfo!.lamports;
+
+            const closeIx = createCloseEmptyTreeIx(cmt, payer, payer);
+            await execute(provider, [closeIx], [payerKeypair]);
+            
+            payerInfo = await provider.connection.getAccountInfo(payer, '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(cmt, 'confirmed');
+            assert(treeInfo === null, 'Expected the merkle tree account info to be null');
+        });
     });
 
     describe('Having created a tree with a single leaf', () => {