Account Compression: Change VerifyLeaf to avoid modifying canopy (#3977)

* Make verifyLeaf work for accounts with canopy

* Add verifyLeaf tests

* remove duplicated test

* Make internal macros private, add docs, use enum instead of hard-coded bools

* change macro names

Author: ngundotra

account-compression/Cargo.lock

@@ -19,12 +19,12 @@
 use crate::error::AccountCompressionError;
 use crate::events::ChangeLogEvent;
 use anchor_lang::prelude::*;
-use bytemuck::cast_slice_mut;
+use bytemuck::{cast_slice, cast_slice_mut};
 use spl_concurrent_merkle_tree::node::{empty_node_cached, Node, EMPTY};
 use std::mem::size_of;
 
 #[inline(always)]
-pub fn check_canopy_bytes(canopy_bytes: &mut [u8]) -> Result<()> {
+pub fn check_canopy_bytes(canopy_bytes: &[u8]) -> Result<()> {
     if canopy_bytes.len() % size_of::<Node>() != 0 {
         msg!(
             "Canopy byte length {} is not a multiple of {}",
@@ -38,7 +38,7 @@ pub fn check_canopy_bytes(canopy_bytes: &mut [u8]) -> Result<()> {
 }
 
 #[inline(always)]
-fn get_cached_path_length(canopy: &mut [Node], max_depth: u32) -> Result<u32> {
+fn get_cached_path_length(canopy: &[Node], max_depth: u32) -> Result<u32> {
     // The offset of 2 is applied because the canopy is a full binary tree without the root node
     // Size: (2^n - 2) -> Size + 2 must be a power of 2
     let closest_power_of_2 = (canopy.len() + 2) as u32;
@@ -89,15 +89,15 @@ pub fn update_canopy(
 }
 
 pub fn fill_in_proof_from_canopy(
-    canopy_bytes: &mut [u8],
+    canopy_bytes: &[u8],
     max_depth: u32,
     index: u32,
     proof: &mut Vec<Node>,
 ) -> Result<()> {
     // 30 is hard coded as it is the current max depth that SPL Compression supports
     let mut empty_node_cache = Box::new([EMPTY; 30]);
     check_canopy_bytes(canopy_bytes)?;
-    let canopy = cast_slice_mut::<u8, Node>(canopy_bytes);
+    let canopy = cast_slice::<u8, Node>(canopy_bytes);
     let path_len = get_cached_path_length(canopy, max_depth)?;
 
     // We want to compute the node index (w.r.t. the canopy) where the current path
@@ -115,7 +115,6 @@ pub fn fill_in_proof_from_canopy(
         if canopy[cached_idx] == EMPTY {
             let level = max_depth - (31 - node_idx.leading_zeros());
             let empty_node = empty_node_cached::<30>(level, &mut empty_node_cache);
-            canopy[cached_idx] = empty_node;
             inferred_nodes.push(empty_node);
         } else {
             inferred_nodes.push(canopy[cached_idx]);

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

@@ -27,11 +27,12 @@ use anchor_lang::{
     solana_program::sysvar::{clock::Clock, rent::Rent},
 };
 use borsh::{BorshDeserialize, BorshSerialize};
-use std::mem::size_of;
 
 pub mod canopy;
 pub mod error;
 pub mod events;
+#[macro_use]
+pub mod macros;
 mod noop;
 pub mod state;
 pub mod zero_copy;
@@ -42,7 +43,9 @@ use crate::canopy::{fill_in_proof_from_canopy, update_canopy};
 use crate::error::AccountCompressionError;
 use crate::events::{AccountCompressionEvent, ChangeLogEvent};
 use crate::noop::wrap_event;
-use crate::state::{ConcurrentMerkleTreeHeader, CONCURRENT_MERKLE_TREE_HEADER_SIZE_V1};
+use crate::state::{
+    merkle_tree_get_size, ConcurrentMerkleTreeHeader, CONCURRENT_MERKLE_TREE_HEADER_SIZE_V1,
+};
 use crate::zero_copy::ZeroCopy;
 
 /// Exported for Anchor / Solita
@@ -120,111 +123,6 @@ pub struct CloseTree<'info> {
     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 {
-    ($max_depth:literal, $max_size:literal, $id:ident, $bytes:ident, $func:ident, $($arg:tt)*) => {
-        match ConcurrentMerkleTree::<$max_depth, $max_size>::load_mut_bytes($bytes) {
-            Ok(merkle_tree) => {
-                match merkle_tree.$func($($arg)*) {
-                    Ok(_) => {
-                        Ok(Box::<ChangeLogEvent>::from((merkle_tree.get_change_log(), $id, merkle_tree.sequence_number)))
-                    }
-                    Err(err) => {
-                        msg!("Error using concurrent merkle tree: {}", err);
-                        err!(AccountCompressionError::ConcurrentMerkleTreeError)
-                    }
-                }
-            }
-            Err(err) => {
-                msg!("Error zero copying concurrent merkle tree: {}", err);
-                err!(AccountCompressionError::ZeroCopyError)
-            }
-        }
-    }
-}
-
-fn merkle_tree_get_size(header: &ConcurrentMerkleTreeHeader) -> Result<usize> {
-    // Note: max_buffer_size MUST be a power of 2
-    match (header.get_max_depth(), header.get_max_buffer_size()) {
-        (3, 8) => Ok(size_of::<ConcurrentMerkleTree<3, 8>>()),
-        (5, 8) => Ok(size_of::<ConcurrentMerkleTree<5, 8>>()),
-        (14, 64) => Ok(size_of::<ConcurrentMerkleTree<14, 64>>()),
-        (14, 256) => Ok(size_of::<ConcurrentMerkleTree<14, 256>>()),
-        (14, 1024) => Ok(size_of::<ConcurrentMerkleTree<14, 1024>>()),
-        (14, 2048) => Ok(size_of::<ConcurrentMerkleTree<14, 2048>>()),
-        (15, 64) => Ok(size_of::<ConcurrentMerkleTree<15, 64>>()),
-        (16, 64) => Ok(size_of::<ConcurrentMerkleTree<16, 64>>()),
-        (17, 64) => Ok(size_of::<ConcurrentMerkleTree<17, 64>>()),
-        (18, 64) => Ok(size_of::<ConcurrentMerkleTree<18, 64>>()),
-        (19, 64) => Ok(size_of::<ConcurrentMerkleTree<19, 64>>()),
-        (20, 64) => Ok(size_of::<ConcurrentMerkleTree<20, 64>>()),
-        (20, 256) => Ok(size_of::<ConcurrentMerkleTree<20, 256>>()),
-        (20, 1024) => Ok(size_of::<ConcurrentMerkleTree<20, 1024>>()),
-        (20, 2048) => Ok(size_of::<ConcurrentMerkleTree<20, 2048>>()),
-        (24, 64) => Ok(size_of::<ConcurrentMerkleTree<24, 64>>()),
-        (24, 256) => Ok(size_of::<ConcurrentMerkleTree<24, 256>>()),
-        (24, 512) => Ok(size_of::<ConcurrentMerkleTree<24, 512>>()),
-        (24, 1024) => Ok(size_of::<ConcurrentMerkleTree<24, 1024>>()),
-        (24, 2048) => Ok(size_of::<ConcurrentMerkleTree<24, 2048>>()),
-        (26, 512) => Ok(size_of::<ConcurrentMerkleTree<26, 512>>()),
-        (26, 1024) => Ok(size_of::<ConcurrentMerkleTree<26, 1024>>()),
-        (26, 2048) => Ok(size_of::<ConcurrentMerkleTree<26, 2048>>()),
-        (30, 512) => Ok(size_of::<ConcurrentMerkleTree<30, 512>>()),
-        (30, 1024) => Ok(size_of::<ConcurrentMerkleTree<30, 1024>>()),
-        (30, 2048) => Ok(size_of::<ConcurrentMerkleTree<30, 2048>>()),
-        _ => {
-            msg!(
-                "Failed to get size of max depth {} and max buffer size {}",
-                header.get_max_depth(),
-                header.get_max_buffer_size()
-            );
-            err!(AccountCompressionError::ConcurrentMerkleTreeConstantsError)
-        }
-    }
-}
-
-/// This applies a given function on a ConcurrentMerkleTree by
-/// allowing the compiler to infer the size of the tree based
-/// upon the header information stored on-chain
-macro_rules! merkle_tree_apply_fn {
-    ($header:ident, $id:ident, $bytes:ident, $func:ident, $($arg:tt)*) => {
-        // Note: max_buffer_size MUST be a power of 2
-        match ($header.get_max_depth(), $header.get_max_buffer_size()) {
-            (3, 8) => merkle_tree_depth_size_apply_fn!(3, 8, $id, $bytes, $func, $($arg)*),
-            (5, 8) => merkle_tree_depth_size_apply_fn!(5, 8, $id, $bytes, $func, $($arg)*),
-            (14, 64) => merkle_tree_depth_size_apply_fn!(14, 64, $id, $bytes, $func, $($arg)*),
-            (14, 256) => merkle_tree_depth_size_apply_fn!(14, 256, $id, $bytes, $func, $($arg)*),
-            (14, 1024) => merkle_tree_depth_size_apply_fn!(14, 1024, $id, $bytes, $func, $($arg)*),
-            (14, 2048) => merkle_tree_depth_size_apply_fn!(14, 2048, $id, $bytes, $func, $($arg)*),
-            (15, 64) => merkle_tree_depth_size_apply_fn!(15, 64, $id, $bytes, $func, $($arg)*),
-            (16, 64) => merkle_tree_depth_size_apply_fn!(16, 64, $id, $bytes, $func, $($arg)*),
-            (17, 64) => merkle_tree_depth_size_apply_fn!(17, 64, $id, $bytes, $func, $($arg)*),
-            (18, 64) => merkle_tree_depth_size_apply_fn!(18, 64, $id, $bytes, $func, $($arg)*),
-            (19, 64) => merkle_tree_depth_size_apply_fn!(19, 64, $id, $bytes, $func, $($arg)*),
-            (20, 64) => merkle_tree_depth_size_apply_fn!(20, 64, $id, $bytes, $func, $($arg)*),
-            (20, 256) => merkle_tree_depth_size_apply_fn!(20, 256, $id, $bytes, $func, $($arg)*),
-            (20, 1024) => merkle_tree_depth_size_apply_fn!(20, 1024, $id, $bytes, $func, $($arg)*),
-            (20, 2048) => merkle_tree_depth_size_apply_fn!(20, 2048, $id, $bytes, $func, $($arg)*),
-            (24, 64) => merkle_tree_depth_size_apply_fn!(24, 64, $id, $bytes, $func, $($arg)*),
-            (24, 256) => merkle_tree_depth_size_apply_fn!(24, 256, $id, $bytes, $func, $($arg)*),
-            (24, 512) => merkle_tree_depth_size_apply_fn!(24, 512, $id, $bytes, $func, $($arg)*),
-            (24, 1024) => merkle_tree_depth_size_apply_fn!(24, 1024, $id, $bytes, $func, $($arg)*),
-            (24, 2048) => merkle_tree_depth_size_apply_fn!(24, 2048, $id, $bytes, $func, $($arg)*),
-            (26, 512) => merkle_tree_depth_size_apply_fn!(26, 512, $id, $bytes, $func, $($arg)*),
-            (26, 1024) => merkle_tree_depth_size_apply_fn!(26, 1024, $id, $bytes, $func, $($arg)*),
-            (26, 2048) => merkle_tree_depth_size_apply_fn!(26, 2048, $id, $bytes, $func, $($arg)*),
-            (30, 512) => merkle_tree_depth_size_apply_fn!(30, 512, $id, $bytes, $func, $($arg)*),
-            (30, 1024) => merkle_tree_depth_size_apply_fn!(30, 1024, $id, $bytes, $func, $($arg)*),
-            (30, 2048) => merkle_tree_depth_size_apply_fn!(30, 2048, $id, $bytes, $func, $($arg)*),
-            _ => {
-                msg!("Failed to apply {} on concurrent merkle tree with max depth {} and max buffer size {}", stringify!($func), $header.get_max_depth(), $header.get_max_buffer_size());
-                err!(AccountCompressionError::ConcurrentMerkleTreeConstantsError)
-            }
-        }
-    };
-}
-
 #[program]
 pub mod spl_account_compression {
     use super::*;
@@ -267,7 +165,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!(header, id, tree_bytes, initialize,)?;
+        let change_log_event = merkle_tree_apply_fn_mut!(header, id, tree_bytes, initialize,)?;
         wrap_event(
             &AccountCompressionEvent::ChangeLog(*change_log_event),
             &ctx.accounts.noop,
@@ -372,7 +270,7 @@ 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!(
+        let change_log_event = merkle_tree_apply_fn_mut!(
             header,
             id,
             tree_bytes,
@@ -431,16 +329,16 @@ pub mod spl_account_compression {
             crate::id(),
             AccountCompressionError::IncorrectAccountOwner
         );
-        let mut merkle_tree_bytes = ctx.accounts.merkle_tree.try_borrow_mut_data()?;
+        let merkle_tree_bytes = ctx.accounts.merkle_tree.try_borrow_data()?;
         let (header_bytes, rest) =
-            merkle_tree_bytes.split_at_mut(CONCURRENT_MERKLE_TREE_HEADER_SIZE_V1);
+            merkle_tree_bytes.split_at(CONCURRENT_MERKLE_TREE_HEADER_SIZE_V1);
 
         let header = ConcurrentMerkleTreeHeader::try_from_slice(header_bytes)?;
         header.assert_valid()?;
         header.assert_valid_leaf_index(index)?;
 
         let merkle_tree_size = merkle_tree_get_size(&header)?;
-        let (tree_bytes, canopy_bytes) = rest.split_at_mut(merkle_tree_size);
+        let (tree_bytes, canopy_bytes) = rest.split_at(merkle_tree_size);
 
         let mut proof = vec![];
         for node in ctx.remaining_accounts.iter() {
@@ -475,7 +373,7 @@ pub mod spl_account_compression {
         let id = ctx.accounts.merkle_tree.key();
         let merkle_tree_size = merkle_tree_get_size(&header)?;
         let (tree_bytes, canopy_bytes) = rest.split_at_mut(merkle_tree_size);
-        let change_log_event = merkle_tree_apply_fn!(header, id, tree_bytes, append, leaf)?;
+        let change_log_event = merkle_tree_apply_fn_mut!(header, id, tree_bytes, append, leaf)?;
         update_canopy(
             canopy_bytes,
             header.get_max_depth(),
@@ -520,7 +418,7 @@ 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!(
+        let change_log_event = merkle_tree_apply_fn_mut!(
             header,
             id,
             tree_bytes,
@@ -558,7 +456,7 @@ pub mod spl_account_compression {
         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,)?;
+        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/lib.rs

@@ -0,0 +1,110 @@
+#[allow(dead_code)]
+enum TreeLoad {
+    Immutable,
+    Mutable,
+}
+
+/// 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 {
+    ($max_depth:literal, $max_size:literal, $id:ident, $bytes:ident, $func:ident, TreeLoad::Mutable, $($arg:tt)*)
+     => {
+        match ConcurrentMerkleTree::<$max_depth, $max_size>::load_mut_bytes($bytes) {
+            Ok(merkle_tree) => {
+                match merkle_tree.$func($($arg)*) {
+                    Ok(_) => {
+                        Ok(Box::<ChangeLogEvent>::from((merkle_tree.get_change_log(), $id, merkle_tree.sequence_number)))
+                    }
+                    Err(err) => {
+                        msg!("Error using concurrent merkle tree: {}", err);
+                        err!(AccountCompressionError::ConcurrentMerkleTreeError)
+                    }
+                }
+            }
+            Err(err) => {
+                msg!("Error zero copying concurrent merkle tree: {}", err);
+                err!(AccountCompressionError::ZeroCopyError)
+            }
+        }
+    };
+    ($max_depth:literal, $max_size:literal, $id:ident, $bytes:ident, $func:ident, TreeLoad::Immutable, $($arg:tt)*) => {
+        match ConcurrentMerkleTree::<$max_depth, $max_size>::load_bytes($bytes) {
+            Ok(merkle_tree) => {
+                match merkle_tree.$func($($arg)*) {
+                    Ok(_) => {
+                        Ok(Box::<ChangeLogEvent>::from((merkle_tree.get_change_log(), $id, merkle_tree.sequence_number)))
+                    }
+                    Err(err) => {
+                        msg!("Error using concurrent merkle tree: {}", err);
+                        err!(AccountCompressionError::ConcurrentMerkleTreeError)
+                    }
+                }
+            }
+            Err(err) => {
+                msg!("Error zero copying concurrent merkle tree: {}", err);
+                err!(AccountCompressionError::ZeroCopyError)
+            }
+        }
+    };
+}
+
+/// This applies a given function on a ConcurrentMerkleTree by
+/// allowing the compiler to infer the size of the tree based
+/// upon the header information stored on-chain
+macro_rules! _merkle_tree_apply_fn {
+    ($header:ident, $($arg:tt)*) => {
+        // Note: max_buffer_size MUST be a power of 2
+        match ($header.get_max_depth(), $header.get_max_buffer_size()) {
+            (3, 8) => _merkle_tree_depth_size_apply_fn!(3, 8, $($arg)*),
+            (5, 8) => _merkle_tree_depth_size_apply_fn!(5, 8, $($arg)*),
+            (14, 64) => _merkle_tree_depth_size_apply_fn!(14, 64, $($arg)*),
+            (14, 256) => _merkle_tree_depth_size_apply_fn!(14, 256, $($arg)*),
+            (14, 1024) => _merkle_tree_depth_size_apply_fn!(14, 1024, $($arg)*),
+            (14, 2048) => _merkle_tree_depth_size_apply_fn!(14, 2048, $($arg)*),
+            (15, 64) => _merkle_tree_depth_size_apply_fn!(15, 64, $($arg)*),
+            (16, 64) => _merkle_tree_depth_size_apply_fn!(16, 64, $($arg)*),
+            (17, 64) => _merkle_tree_depth_size_apply_fn!(17, 64, $($arg)*),
+            (18, 64) => _merkle_tree_depth_size_apply_fn!(18, 64, $($arg)*),
+            (19, 64) => _merkle_tree_depth_size_apply_fn!(19, 64, $($arg)*),
+            (20, 64) => _merkle_tree_depth_size_apply_fn!(20, 64, $($arg)*),
+            (20, 256) => _merkle_tree_depth_size_apply_fn!(20, 256, $($arg)*),
+            (20, 1024) => _merkle_tree_depth_size_apply_fn!(20, 1024, $($arg)*),
+            (20, 2048) => _merkle_tree_depth_size_apply_fn!(20, 2048, $($arg)*),
+            (24, 64) => _merkle_tree_depth_size_apply_fn!(24, 64, $($arg)*),
+            (24, 256) => _merkle_tree_depth_size_apply_fn!(24, 256, $($arg)*),
+            (24, 512) => _merkle_tree_depth_size_apply_fn!(24, 512, $($arg)*),
+            (24, 1024) => _merkle_tree_depth_size_apply_fn!(24, 1024, $($arg)*),
+            (24, 2048) => _merkle_tree_depth_size_apply_fn!(24, 2048, $($arg)*),
+            (26, 512) => _merkle_tree_depth_size_apply_fn!(26, 512, $($arg)*),
+            (26, 1024) => _merkle_tree_depth_size_apply_fn!(26, 1024, $($arg)*),
+            (26, 2048) => _merkle_tree_depth_size_apply_fn!(26, 2048, $($arg)*),
+            (30, 512) => _merkle_tree_depth_size_apply_fn!(30, 512, $($arg)*),
+            (30, 1024) => _merkle_tree_depth_size_apply_fn!(30, 1024, $($arg)*),
+            (30, 2048) => _merkle_tree_depth_size_apply_fn!(30, 2048, $($arg)*),
+            _ => {
+                msg!("Failed to apply {} on concurrent merkle tree with max depth {} and max buffer size {}",
+                    stringify!($func),
+                    $header.get_max_depth(),
+                    $header.get_max_buffer_size()
+                );
+                err!(AccountCompressionError::ConcurrentMerkleTreeConstantsError)
+            }
+        }
+    };
+}
+
+/// This applies a given function on a mutable ConcurrentMerkleTree
+#[macro_export]
+macro_rules! merkle_tree_apply_fn_mut {
+    ($header:ident, $id:ident, $bytes:ident, $func:ident, $($arg:tt)*) => {
+        _merkle_tree_apply_fn!($header, $id, $bytes, $func, TreeLoad::Mutable, $($arg)*)
+    };
+}
+
+/// This applies a given function on a read-only ConcurrentMerkleTree
+#[macro_export]
+macro_rules! merkle_tree_apply_fn {
+    ($header:ident, $id:ident, $bytes:ident, $func:ident, $($arg:tt)*) => {
+        _merkle_tree_apply_fn!($header, $id, $bytes, $func, TreeLoad::Immutable, $($arg)*)
+    };
+}