extended the concurrent merkle tree header with a is_batch_initialized flag + comments fixed

Author: StanChe

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

@@ -56,6 +56,10 @@ pub enum AccountCompressionError {
     #[msg("Tree was already initialized")]
     TreeAlreadyInitialized,
 
+    /// The tree header was not initialized for batch processing
+    #[msg("Tree header was not initialized for batch processing")]
+    BatchNotInitialized,
+
     /// The canopy root doesn't match the root of the tree
     #[msg("Canopy root does not match the root of the tree")]
     CanopyRootMismatch,

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

@@ -193,9 +193,9 @@ pub mod spl_account_compression {
     /// expected flow is `init_empty_merkle_tree`. For the latter case, the canopy should be
     /// filled with the necessary nodes to render the tree usable. Thus we need to prefill the
     /// canopy with the necessary nodes. The expected flow for a tree without canopy is
-    /// `prepare_tree` -> `init_merkle_tree_with_root`. The expected flow for a tree with canopy
+    /// `prepare_tree` -> `finalize_merkle_tree_with_root`. The expected flow for a tree with canopy
     /// is `prepare_tree` -> `append_canopy_nodes` (multiple times until all of the canopy is
-    /// filled) -> `init_merkle_tree_with_root`. This instruction initializes the tree header
+    /// filled) -> `finalize_merkle_tree_with_root`. This instruction initializes the tree header
     /// while leaving the tree itself uninitialized. This allows distinguishing between an empty
     /// tree and a tree prepare to be initialized with a root.
     pub fn prepare_tree(
@@ -214,7 +214,7 @@ pub mod spl_account_compression {
             merkle_tree_bytes.split_at_mut(CONCURRENT_MERKLE_TREE_HEADER_SIZE_V1);
 
         let mut header = ConcurrentMerkleTreeHeader::try_from_slice(header_bytes)?;
-        header.initialize(
+        header.initialize_batched(
             max_depth,
             max_buffer_size,
             &ctx.accounts.authority.key(),
@@ -228,7 +228,7 @@ 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
-    /// `init_merkle_tree_with_root` instruction that'll finalize the tree initialization.
+    /// `finalize_merkle_tree_with_root` instruction that'll finalize the tree initialization.
     pub fn append_canopy_nodes(
         ctx: Context<Modify>,
         start_index: u32,
@@ -246,6 +246,7 @@ pub mod spl_account_compression {
 
         let header = ConcurrentMerkleTreeHeader::try_from_slice(header_bytes)?;
         header.assert_valid_authority(&ctx.accounts.authority.key())?;
+        header.assert_is_batch_initialized()?;
         // assert the tree is not initialized yet, we don't want to overwrite the canopy of an
         // initialized tree
         let merkle_tree_size = merkle_tree_get_size(&header)?;
@@ -283,6 +284,7 @@ pub mod spl_account_compression {
         // the header should already be initialized with prepare_tree
         let header = ConcurrentMerkleTreeHeader::try_from_slice(header_bytes)?;
         header.assert_valid_authority(&ctx.accounts.authority.key())?;
+        header.assert_is_batch_initialized()?;
         let merkle_tree_size = merkle_tree_get_size(&header)?;
         let (tree_bytes, canopy_bytes) = rest.split_at_mut(merkle_tree_size);
         // check the canopy root matches the tree root

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

@@ -65,9 +65,14 @@ pub struct ConcurrentMerkleTreeHeaderDataV1 {
     /// Provides a lower-bound on what slot to start (re-)building a tree from.
     creation_slot: u64,
 
+    /// A flag indicating whether the tree has been initialized with a root.
+    /// This field was added together with the `finalize_tree_with_root` instruction.
+    /// It takes 1 byte of space taken from the previous padding for existing accounts.
+    is_batch_initialized: bool,
+
     /// Needs padding for the account to be 8-byte aligned
     /// 8-byte alignment is necessary to zero-copy the SPL ConcurrentMerkleTree
-    _padding: [u8; 6],
+    _padding: [u8; 5],
 }
 
 #[repr(C)]
@@ -95,6 +100,24 @@ impl ConcurrentMerkleTreeHeader {
                 header.max_depth = max_depth;
                 header.authority = *authority;
                 header.creation_slot = creation_slot;
+                // is_batch_initialized is left false by default
+            }
+        }
+    }
+
+    /// Initializes the header with the given parameters and sets the `is_batch_initialized` flag to
+    /// true.
+    pub fn initialize_batched(
+        &mut self,
+        max_depth: u32,
+        max_buffer_size: u32,
+        authority: &Pubkey,
+        creation_slot: u64,
+    ) {
+        self.initialize(max_depth, max_buffer_size, authority, creation_slot);
+        match self.header {
+            ConcurrentMerkleTreeHeaderData::V1(ref mut header) => {
+                header.is_batch_initialized = true;
             }
         }
     }
@@ -155,6 +178,18 @@ impl ConcurrentMerkleTreeHeader {
         }
         Ok(())
     }
+
+    pub fn assert_is_batch_initialized(&self) -> Result<()> {
+        match &self.header {
+            ConcurrentMerkleTreeHeaderData::V1(header) => {
+                require!(
+                    header.is_batch_initialized,
+                    AccountCompressionError::BatchNotInitialized
+                );
+            }
+        }
+        Ok(())
+    }
 }
 
 pub fn merkle_tree_get_size(header: &ConcurrentMerkleTreeHeader) -> Result<usize> {
@@ -166,10 +201,10 @@ pub fn merkle_tree_get_size(header: &ConcurrentMerkleTreeHeader) -> Result<usize
         (7, 16) => Ok(size_of::<ConcurrentMerkleTree<7, 16>>()),
         (8, 16) => Ok(size_of::<ConcurrentMerkleTree<8, 16>>()),
         (9, 16) => Ok(size_of::<ConcurrentMerkleTree<9, 16>>()),
-        (10, 32)=> Ok(size_of::<ConcurrentMerkleTree<10, 32>>()),
-        (11, 32)=> Ok(size_of::<ConcurrentMerkleTree<11, 32>>()),
-        (12, 32)=> Ok(size_of::<ConcurrentMerkleTree<12, 32>>()),
-        (13, 32)=> Ok(size_of::<ConcurrentMerkleTree<13, 32>>()),
+        (10, 32) => Ok(size_of::<ConcurrentMerkleTree<10, 32>>()),
+        (11, 32) => Ok(size_of::<ConcurrentMerkleTree<11, 32>>()),
+        (12, 32) => Ok(size_of::<ConcurrentMerkleTree<12, 32>>()),
+        (13, 32) => Ok(size_of::<ConcurrentMerkleTree<13, 32>>()),
         (14, 64) => Ok(size_of::<ConcurrentMerkleTree<14, 64>>()),
         (14, 256) => Ok(size_of::<ConcurrentMerkleTree<14, 256>>()),
         (14, 1024) => Ok(size_of::<ConcurrentMerkleTree<14, 1024>>()),

account-compression/sdk/idl/spl_account_compression.json

@@ -65,9 +65,9 @@
         "expected flow is `init_empty_merkle_tree`. For the latter case, the canopy should be",
         "filled with the necessary nodes to render the tree usable. Thus we need to prefill the",
         "canopy with the necessary nodes. The expected flow for a tree without canopy is",
-        "`prepare_tree` -> `init_merkle_tree_with_root`. The expected flow for a tree with canopy",
+        "`prepare_tree` -> `finalize_merkle_tree_with_root`. The expected flow for a tree with canopy",
         "is `prepare_tree` -> `append_canopy_nodes` (multiple times until all of the canopy is",
-        "filled) -> `init_merkle_tree_with_root`. This instruction initializes the tree header",
+        "filled) -> `finalize_merkle_tree_with_root`. This instruction initializes the tree header",
         "while leaving the tree itself uninitialized. This allows distinguishing between an empty",
         "tree and a tree prepare to be initialized with a root."
       ],
@@ -109,7 +109,7 @@
       "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",
-        "`init_merkle_tree_with_root` instruction that'll finalize the tree initialization."
+        "`finalize_merkle_tree_with_root` instruction that'll finalize the tree initialization."
       ],
       "accounts": [
         {
@@ -542,14 +542,23 @@
             ],
             "type": "u64"
           },
+          {
+            "name": "isBatchInitialized",
+            "docs": [
+              "A flag indicating whether the tree has been initialized with a root.",
+              "This field was added together with the `finalize_tree_with_root` instruction.",
+              "It takes 1 byte of space taken from the previous padding for existing accounts."
+            ],
+            "type": "bool"
+          },
           {
             "name": "padding",
             "docs": [
               "Needs padding for the account to be 8-byte aligned",
               "8-byte alignment is necessary to zero-copy the SPL ConcurrentMerkleTree"
             ],
             "type": {
-              "array": ["u8", 6]
+              "array": ["u8", 5]
             }
           }
         ]

account-compression/sdk/package.json

@@ -42,7 +42,7 @@
     "lint:fix": "npm run pretty:fix && eslint . --fix --ext .js,.ts",
     "docs": "rm -rf docs/ && typedoc --out docs",
     "deploy:docs": "npm run docs && gh-pages --dest account-compression/sdk --dist docs --dotfiles",
-    "start-validator": "solana-test-validator --reset --quiet --bpf-program cmtDvXumGCrqC1Age74AVPhSRVXJMd8PJS91L8KbNCK ../target/deploy/spl_account_compression.so --bpf-program noopb9bkMVfRPU8AsbpTUg8AQkHtKwMYZiFUjNRtMmV ../target/deploy/spl_noop.so",
+    "start-validator": "solana-test-validator --reset --quiet --bpf-program cmtDvXumGCrqC1Age74AVPhSRVXJMd8PJS91L8KbNCK ../target/deploy/spl_account_compression.so --bpf-program noopb9bkMVfRPU8AsbpTUg8AQkHtKwMYZiFUjNRtMmV ../target/deploy/spl_noop.so --account 27QMkDMpBoAhmWj6xxQNYdqXZL5nnC8tkZcEtkNxCqeX pre-batch-init-tree-account.json",
     "run-tests": "jest tests --detectOpenHandles",
     "run-tests:events": "jest tests/events --detectOpenHandles",
     "run-tests:accounts": "jest tests/accounts --detectOpenHandles",

account-compression/sdk/src/accounts/ConcurrentMerkleTreeAccount.ts

@@ -121,6 +121,14 @@ export class ConcurrentMerkleTreeAccount {
     getCanopyDepth(): number {
         return getCanopyDepth(this.canopy.canopyBytes.length);
     }
+
+    /**
+     * Returns the flag that indicates if the tree has been batch initialized
+     * @returns the flag
+     */
+    getIsBatchInitialized(): boolean {
+        return this.getHeaderV1().isBatchInitialized;
+    }
 }
 
 /**

account-compression/sdk/src/generated/types/ConcurrentMerkleTreeHeaderDataV1.ts

@@ -11,9 +11,10 @@ import * as web3 from '@solana/web3.js';
 export type ConcurrentMerkleTreeHeaderDataV1 = {
     authority: web3.PublicKey;
     creationSlot: beet.bignum;
+    isBatchInitialized: boolean;
     maxBufferSize: number;
     maxDepth: number;
-    padding: number[] /* size: 6 */;
+    padding: number[] /* size: 5 */;
 };
 
 /**
@@ -26,7 +27,8 @@ export const concurrentMerkleTreeHeaderDataV1Beet = new beet.BeetArgsStruct<Conc
         ['maxDepth', beet.u32],
         ['authority', beetSolana.publicKey],
         ['creationSlot', beet.u64],
-        ['padding', beet.uniformFixedSizeArray(beet.u8, 6)],
+        ['isBatchInitialized', beet.bool],
+        ['padding', beet.uniformFixedSizeArray(beet.u8, 5)],
     ],
     'ConcurrentMerkleTreeHeaderDataV1',
 );

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

@@ -19,6 +19,7 @@ import {
     ValidDepthSizePair,
 } from '../src';
 import { hash, MerkleTree } from '../src/merkle-tree';
+import { assertCMTProperties } from './accounts/concurrentMerkleTreeAccount.test';
 import { createTreeOnChain, execute, prepareTree } from './utils';
 
 // eslint-disable-next-line no-empty
@@ -96,7 +97,7 @@ describe('Account Compression', () => {
             const merkleTreeRaw = new MerkleTree(leaves);
             const root = merkleTreeRaw.root;
             const leaf = leaves[leaves.length - 1];
-
+            const canopyDepth = 0;
             const finalize = createFinalizeMerkleTreeWithRootIx(
                 cmt,
                 payer,
@@ -109,11 +110,7 @@ describe('Account Compression', () => {
             await execute(provider, [finalize], [payerKeypair]);
 
             const splCMT = await ConcurrentMerkleTreeAccount.fromAccountAddress(connection, cmt);
-            assert(splCMT.getMaxBufferSize() === size, 'Buffer size does not match');
-            assert(
-                splCMT.getCanopyDepth() === canopyDepth,
-                'Canopy depth does not match: expected ' + canopyDepth + ' but got ' + splCMT.getCanopyDepth(),
-            );
+            assertCMTProperties(splCMT, depth, size, payer, root, canopyDepth, true);
             assert(splCMT.getBufferSize() == 1, 'Buffer size does not match');
         });
         it('Should fail to append canopy node for a tree without canopy', async () => {
@@ -327,6 +324,8 @@ describe('Account Compression', () => {
                 merkleTreeRaw.getProof(leaves.length - 1).proof,
             );
             await execute(provider, [finalize], [payerKeypair]);
+            const splCMT = await ConcurrentMerkleTreeAccount.fromAccountAddress(connection, cmt);
+            assertCMTProperties(splCMT, depth, size, payer, root, canopyDepth, true);
         });
         it('Should be able to setup canopy with several transactions', async () => {
             const merkleTreeRaw = new MerkleTree(leaves);

account-compression/sdk/tests/accounts/concurrentMerkleTreeAccount.test.ts

@@ -8,13 +8,14 @@ import { ALL_DEPTH_SIZE_PAIRS, ConcurrentMerkleTreeAccount, getConcurrentMerkleT
 import { emptyNode, MerkleTree } from '../../src/merkle-tree';
 import { createEmptyTreeOnChain, createTreeOnChain } from '../utils';
 
-function assertCMTProperties(
+export function assertCMTProperties(
     onChainCMT: ConcurrentMerkleTreeAccount,
     expectedMaxDepth: number,
     expectedMaxBufferSize: number,
     expectedAuthority: PublicKey,
     expectedRoot: Buffer,
-    expectedCanopyDepth?: number
+    expectedCanopyDepth?: number,
+    expectedIsBatchInitialized = false,
 ) {
     assert(
         onChainCMT.getMaxDepth() === expectedMaxDepth,
@@ -32,6 +33,10 @@ function assertCMTProperties(
             'On chain canopy depth does not match expected canopy depth'
         );
     }
+    assert(
+        onChainCMT.getIsBatchInitialized() === expectedIsBatchInitialized,
+        'On chain isBatchInitialized does not match expected value'
+    );
 }
 
 describe('ConcurrentMerkleTreeAccount tests', () => {
@@ -142,4 +147,26 @@ describe('ConcurrentMerkleTreeAccount tests', () => {
             }
         });
     });
+
+    describe('Can deserialize an existing CMTAccount from a real on-chain CMT created before the is_batch_initialized field was introduced inplace of the first byte of _padding', () => {
+        it('Interpreted on-chain fields correctly', async () => {
+            // The account data was generated by running:
+            //      $ solana account 27QMkDMpBoAhmWj6xxQNYdqXZL5nnC8tkZcEtkNxCqeX \
+            //                       --output-file tests/fixtures/pre-batch-init-tree-account.json \
+            //                       --output json     
+            const deployedAccount = new PublicKey('27QMkDMpBoAhmWj6xxQNYdqXZL5nnC8tkZcEtkNxCqeX');
+            const cmt = await ConcurrentMerkleTreeAccount.fromAccountAddress(
+                connection,
+                deployedAccount,
+                'confirmed'
+                );
+            const expectedMaxDepth = 10;
+            const expectedMaxBufferSize = 32;
+            const expectedCanopyDepth = 0;
+            const expectedAuthority = new PublicKey('BFNT941iRwYPe2Js64dTJSoksGCptWAwrkKMaSN73XK2');
+            const expectedRoot = new PublicKey('83UjseEuEgxyVyDTmrJCQ9QbeksdRZ7KPDZGQYc5cAgF').toBuffer();
+            const expectedIsBatchInitialized = false;
+            await assertCMTProperties(cmt, expectedMaxDepth, expectedMaxBufferSize, expectedAuthority, expectedRoot, expectedCanopyDepth, expectedIsBatchInitialized);
+        });
+    });
 });