refactored tree unitialized check on bytes directly, also moved tree initialization call in a wrapper as it started reporting a stack overflow

Author: StanChe

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

@@ -63,3 +63,31 @@ pub fn merkle_tree_prove_leaf(
 ) -> Result<Box<ChangeLogEvent>> {
     merkle_tree_apply_fn!(header, tree_id, tree_bytes, prove_leaf, args)
 }
+
+pub fn merkle_tree_initialize(
+    header: &ConcurrentMerkleTreeHeader,
+    tree_id: Pubkey,
+    tree_bytes: &mut [u8],
+) -> Result<Box<ChangeLogEvent>> {
+    merkle_tree_apply_fn_mut!(header, tree_id, tree_bytes, initialize,)
+}
+
+pub fn tree_bytes_unititialized(tree_bytes: &[u8]) -> bool {
+    tree_bytes.iter().all(|&x| x == 0)
+}
+
+pub fn assert_tree_is_empty(
+    header: &ConcurrentMerkleTreeHeader,
+    tree_id: Pubkey,
+    tree_bytes: &mut [u8],
+) -> Result<()> {
+    // If the tree is batch initialized and not finalized yet, we can treat it as empty.
+    // Before the tree is finalized, the tree_bytes will be all 0 as only the header will be
+    // initialized at that point, so we may skip the deserialization.
+    if header.get_is_batch_initialized() && tree_bytes_unititialized(tree_bytes) {
+        return Ok(());
+    }
+    // check the tree is empty
+    merkle_tree_apply_fn_mut!(header, tree_id, tree_bytes, prove_tree_is_empty,)?;
+    Ok(())
+}

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

@@ -174,7 +174,7 @@ 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();
-        let change_log_event = merkle_tree_apply_fn_mut!(header, id, tree_bytes, initialize,)?;
+        let change_log_event = merkle_tree_initialize(&header, id, tree_bytes)?;
         wrap_event(
             &AccountCompressionEvent::ChangeLog(*change_log_event),
             &ctx.accounts.noop,
@@ -264,7 +264,7 @@ pub mod spl_account_compression {
         let (tree_bytes, canopy_bytes) = rest.split_at_mut(merkle_tree_size);
         // ensure the tree is not initialized, the hacky way
         require!(
-            tree_bytes.iter().all(|&x| x == 0),
+            tree_bytes_unititialized(tree_bytes),
             AccountCompressionError::TreeAlreadyInitialized
         );
         set_canopy_leaf_nodes(
@@ -564,11 +564,8 @@ 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);
 
-        // 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,)?;
-        }
+        let id = ctx.accounts.merkle_tree.key();
+        assert_tree_is_empty(&header, id, tree_bytes)?;
 
         // Close merkle tree account
         // 1. Move lamports