account-compression: fix for a stack overflow issue on macro usage (#6827)

* account-compression: fix for a stack overflow issue on macro usage

* clippy: removed unnecessary references

* update cmt tests

* fix rustfmt

* bump spl-concurrent-merkle-tree to 0.3.0 for API changes

---------

Co-authored-by: ngundotra <noah@gundotra.org>
Co-authored-by: Noah Gundotra <ngundotra@users.noreply.github.com>

Author: StanChe

Cargo.lock

@@ -21,8 +21,8 @@ default = []
 anchor-lang = "0.29.0"
 bytemuck = "1.13"
 solana-program = ">=1.18.11,<=2"
-spl-concurrent-merkle-tree = { version = "0.2.0", path = "../../../libraries/concurrent-merkle-tree", features = [
-  "sol-log",
+spl-concurrent-merkle-tree = { version = "0.3.0", path = "../../../libraries/concurrent-merkle-tree", features = [
+    "sol-log",
 ] }
 spl-noop = { version = "0.2.0", path = "../noop", features = ["no-entrypoint"] }
 

account-compression/programs/account-compression/Cargo.toml

@@ -0,0 +1,65 @@
+//! This module provides a wrapper around the `ConcurrentMerkleTree` struct from
+//! the `spl_concurrent_merkle_tree` crate. It provides a set of functions that
+//! can be called from the Anchor program to interact with the tree.
+//! The functions are used to initialize the tree, set a leaf, fill empty or
+//! append a leaf, and prove a leaf. As the tree is generic over the depth and
+//! buffer size, the functions are implemented using macros that infer the depth
+//! and buffer size from the header information stored on-chain. Usage of the
+//! macros directly is discouraged, as they have huge match statements with
+//! every case taking it's own stack frame. Instead, use the exported functions
+//! from this module and refenrece or Box the arguments to the functions to
+//! avoid the stack frame explosion.
+
+pub use crate::error::AccountCompressionError;
+/// Exported for Anchor / Solita
+pub use spl_concurrent_merkle_tree::{
+    concurrent_merkle_tree::{
+        ConcurrentMerkleTree, FillEmptyOrAppendArgs, InitializeWithRootArgs, ProveLeafArgs,
+        SetLeafArgs,
+    },
+    error::ConcurrentMerkleTreeError,
+    node::Node,
+    node::EMPTY,
+};
+use {
+    crate::{
+        events::ChangeLogEvent, macros::*, state::ConcurrentMerkleTreeHeader, zero_copy::ZeroCopy,
+    },
+    anchor_lang::prelude::*,
+};
+
+pub fn merkle_tree_initialize_with_root(
+    header: &ConcurrentMerkleTreeHeader,
+    tree_id: Pubkey,
+    tree_bytes: &mut [u8],
+    args: &InitializeWithRootArgs,
+) -> Result<Box<ChangeLogEvent>> {
+    merkle_tree_apply_fn_mut!(header, tree_id, tree_bytes, initialize_with_root, args)
+}
+
+pub fn merkle_tree_set_leaf(
+    header: &ConcurrentMerkleTreeHeader,
+    tree_id: Pubkey,
+    tree_bytes: &mut [u8],
+    args: &SetLeafArgs,
+) -> Result<Box<ChangeLogEvent>> {
+    merkle_tree_apply_fn_mut!(header, tree_id, tree_bytes, set_leaf, args)
+}
+
+pub fn merkle_tree_fill_empty_or_append(
+    header: &ConcurrentMerkleTreeHeader,
+    tree_id: Pubkey,
+    tree_bytes: &mut [u8],
+    args: &FillEmptyOrAppendArgs,
+) -> Result<Box<ChangeLogEvent>> {
+    merkle_tree_apply_fn_mut!(header, tree_id, tree_bytes, fill_empty_or_append, args)
+}
+
+pub fn merkle_tree_prove_leaf(
+    header: &ConcurrentMerkleTreeHeader,
+    tree_id: Pubkey,
+    tree_bytes: &[u8],
+    args: &ProveLeafArgs,
+) -> Result<Box<ChangeLogEvent>> {
+    merkle_tree_apply_fn!(header, tree_id, tree_bytes, prove_leaf, args)
+}

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

@@ -29,6 +29,7 @@ use anchor_lang::{
 use borsh::{BorshDeserialize, BorshSerialize};
 
 pub mod canopy;
+pub mod concurrent_tree_wrapper;
 pub mod error;
 pub mod events;
 #[macro_use]
@@ -40,6 +41,7 @@ pub mod zero_copy;
 pub use crate::noop::{wrap_application_data_v1, Noop};
 
 use crate::canopy::{fill_in_proof_from_canopy, update_canopy};
+use crate::concurrent_tree_wrapper::*;
 pub use crate::error::AccountCompressionError;
 pub use crate::events::{AccountCompressionEvent, ChangeLogEvent};
 use crate::noop::wrap_event;
@@ -50,7 +52,9 @@ use crate::zero_copy::ZeroCopy;
 
 /// Exported for Anchor / Solita
 pub use spl_concurrent_merkle_tree::{
-    concurrent_merkle_tree::ConcurrentMerkleTree, error::ConcurrentMerkleTreeError, node::Node,
+    concurrent_merkle_tree::{ConcurrentMerkleTree, FillEmptyOrAppendArgs},
+    error::ConcurrentMerkleTreeError,
+    node::Node,
 };
 
 declare_id!("cmtDvXumGCrqC1Age74AVPhSRVXJMd8PJS91L8KbNCK");
@@ -270,17 +274,15 @@ pub mod spl_account_compression {
         fill_in_proof_from_canopy(canopy_bytes, header.get_max_depth(), index, &mut proof)?;
         let id = ctx.accounts.merkle_tree.key();
         // A call is made to ConcurrentMerkleTree::set_leaf(root, previous_leaf, new_leaf, proof, index)
-        let change_log_event = merkle_tree_apply_fn_mut!(
-            header,
-            id,
-            tree_bytes,
-            set_leaf,
-            root,
+        let args = &SetLeafArgs {
+            current_root: root,
             previous_leaf,
             new_leaf,
-            &proof,
+            proof_vec: proof,
             index,
-        )?;
+        };
+        let change_log_event = merkle_tree_set_leaf(&header, id, tree_bytes, args)?;
+
         update_canopy(
             canopy_bytes,
             header.get_max_depth(),
@@ -347,7 +349,14 @@ pub mod spl_account_compression {
         fill_in_proof_from_canopy(canopy_bytes, header.get_max_depth(), index, &mut proof)?;
         let id = ctx.accounts.merkle_tree.key();
 
-        merkle_tree_apply_fn!(header, id, tree_bytes, prove_leaf, root, leaf, &proof, index)?;
+        let args = &ProveLeafArgs {
+            current_root: root,
+            leaf,
+            proof_vec: proof,
+            index,
+        };
+        merkle_tree_prove_leaf(&header, id, tree_bytes, args)?;
+
         Ok(())
     }
 
@@ -418,16 +427,14 @@ pub mod spl_account_compression {
         fill_in_proof_from_canopy(canopy_bytes, header.get_max_depth(), index, &mut proof)?;
         // A call is made to ConcurrentMerkleTree::fill_empty_or_append
         let id = ctx.accounts.merkle_tree.key();
-        let change_log_event = merkle_tree_apply_fn_mut!(
-            header,
-            id,
-            tree_bytes,
-            fill_empty_or_append,
-            root,
+        let args = &FillEmptyOrAppendArgs {
+            current_root: root,
             leaf,
-            &proof,
+            proof_vec: proof,
             index,
-        )?;
+        };
+        let change_log_event = merkle_tree_fill_empty_or_append(&header, id, tree_bytes, args)?;
+
         update_canopy(
             canopy_bytes,
             header.get_max_depth(),

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

@@ -4,8 +4,8 @@ enum TreeLoad {
     Mutable,
 }
 
-/// This macro applies functions on a ConcurrentMerkleT:ee and emits leaf information
-/// needed to sync the merkle tree state with off-chain indexers.
+/// This macro applies functions on a ConcurrentMerkleT:ee and emits leaf
+/// information needed to sync the merkle tree state with off-chain indexers.
 #[macro_export]
 macro_rules! _merkle_tree_depth_size_apply_fn {
     ($max_depth:literal, $max_size:literal, $id:ident, $bytes:ident, $func:ident, TreeLoad::Mutable, $($arg:tt)*)
@@ -118,3 +118,8 @@ macro_rules! merkle_tree_apply_fn {
         _merkle_tree_apply_fn!($header, $id, $bytes, $func, TreeLoad::Immutable, $($arg)*)
     };
 }
+
+pub(crate) use {
+    _merkle_tree_apply_fn, _merkle_tree_depth_size_apply_fn, merkle_tree_apply_fn,
+    merkle_tree_apply_fn_mut,
+};

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

@@ -1,6 +1,6 @@
 [package]
 name = "spl-concurrent-merkle-tree"
-version = "0.2.0"
+version = "0.3.0"
 description = "Solana Program Library Concurrent Merkle Tree"
 authors = ["Solana Labs Maintainers <maintainers@solanalabs.com>"]
 repository = "https://github.com/solana-labs/solana-program-library"
@@ -9,7 +9,7 @@ edition = "2021"
 
 [features]
 log = []
-sol-log = [ "log" ]
+sol-log = ["log"]
 
 [dependencies]
 solana-program = "1.16"