spl_account_compression/
canopy.rs

1//! Canopy is way to cache the upper `N` levels of a SPL ConcurrentMerkleTree.
2//!
3//! By caching the upper `N` levels of a depth `D` SPL ConcurrentMerkleTree,
4//! proofs can be truncated to the first `D - N` nodes. This helps reduce the size of account
5//! compression transactions, and makes it possible to
6//! modify trees up to depth 31, which store more than 1 billion leaves.
7//!
8//! Note: this means that creating a tree of depth > 24 without a canopy will be impossible to modify
9//! on-chain until TransactionV2 is launched.
10//!
11//! To initialize a canopy on a ConcurrentMerkleTree account, you must initialize
12//! the ConcurrentMerkleTree account with additional bytes. The number of additional bytes
13//! needed is `(pow(2, N+1)-1) * 32`, where `N` is the number of levels of the merkle tree
14//! you want the canopy to cache.
15//!
16//! The canopy will be updated everytime the concurrent merkle tree is modified. No additional work
17//! needed.
18
19use crate::error::AccountCompressionError;
20use crate::events::ChangeLogEvent;
21use anchor_lang::prelude::*;
22use bytemuck::{cast_slice, cast_slice_mut};
23use spl_concurrent_merkle_tree::node::{empty_node_cached, Node, EMPTY};
24use std::mem::size_of;
25
26#[inline(always)]
27pub fn check_canopy_bytes(canopy_bytes: &[u8]) -> Result<()> {
28    if canopy_bytes.len() % size_of::<Node>() != 0 {
29        msg!(
30            "Canopy byte length {} is not a multiple of {}",
31            canopy_bytes.len(),
32            size_of::<Node>()
33        );
34        err!(AccountCompressionError::CanopyLengthMismatch)
35    } else {
36        Ok(())
37    }
38}
39
40#[inline(always)]
41fn get_cached_path_length(canopy: &[Node], max_depth: u32) -> Result<u32> {
42    // The offset of 2 is applied because the canopy is a full binary tree without the root node
43    // Size: (2^n - 2) -> Size + 2 must be a power of 2
44    let closest_power_of_2 = (canopy.len() + 2) as u32;
45    // This expression will return true if `closest_power_of_2` is actually a power of 2
46    if closest_power_of_2 & (closest_power_of_2 - 1) == 0 {
47        // (1 << max_depth) returns the number of leaves in the full merkle tree
48        // (1 << (max_depth + 1)) - 1 returns the number of nodes in the full tree
49        // The canopy size cannot exceed the size of the tree
50        if closest_power_of_2 > (1 << (max_depth + 1)) {
51            msg!(
52                "Canopy size is too large. Size: {}. Max size: {}",
53                closest_power_of_2 - 2,
54                (1 << (max_depth + 1)) - 2
55            );
56            return err!(AccountCompressionError::CanopyLengthMismatch);
57        }
58    } else {
59        msg!(
60            "Canopy length {} is not 2 less than a power of 2",
61            canopy.len()
62        );
63        return err!(AccountCompressionError::CanopyLengthMismatch);
64    }
65    // 1 is subtracted from the trailing zeros because the root is not stored in the canopy
66    Ok(closest_power_of_2.trailing_zeros() - 1)
67}
68
69pub fn update_canopy(
70    canopy_bytes: &mut [u8],
71    max_depth: u32,
72    change_log: Option<&ChangeLogEvent>,
73) -> Result<()> {
74    check_canopy_bytes(canopy_bytes)?;
75    let canopy = cast_slice_mut::<u8, Node>(canopy_bytes);
76    let path_len = get_cached_path_length(canopy, max_depth)?;
77    if let Some(cl_event) = change_log {
78        match &*cl_event {
79            ChangeLogEvent::V1(cl) => {
80                // Update the canopy from the newest change log
81                for path_node in cl.path.iter().rev().skip(1).take(path_len as usize) {
82                    // node_idx - 2 maps to the canopy index
83                    canopy[(path_node.index - 2) as usize] = path_node.node;
84                }
85            }
86        }
87    }
88    Ok(())
89}
90
91pub fn fill_in_proof_from_canopy(
92    canopy_bytes: &[u8],
93    max_depth: u32,
94    index: u32,
95    proof: &mut Vec<Node>,
96) -> Result<()> {
97    // 30 is hard coded as it is the current max depth that SPL Compression supports
98    let mut empty_node_cache = Box::new([EMPTY; 30]);
99    check_canopy_bytes(canopy_bytes)?;
100    let canopy = cast_slice::<u8, Node>(canopy_bytes);
101    let path_len = get_cached_path_length(canopy, max_depth)?;
102
103    // We want to compute the node index (w.r.t. the canopy) where the current path
104    // intersects the leaves of the canopy
105    let mut node_idx = ((1 << max_depth) + index) >> (max_depth - path_len);
106    let mut inferred_nodes = vec![];
107    while node_idx > 1 {
108        // node_idx - 2 maps to the canopy index
109        let shifted_index = node_idx as usize - 2;
110        let cached_idx = if shifted_index % 2 == 0 {
111            shifted_index + 1
112        } else {
113            shifted_index - 1
114        };
115        if canopy[cached_idx] == EMPTY {
116            let level = max_depth - (31 - node_idx.leading_zeros());
117            let empty_node = empty_node_cached::<30>(level, &mut empty_node_cache);
118            inferred_nodes.push(empty_node);
119        } else {
120            inferred_nodes.push(canopy[cached_idx]);
121        }
122        node_idx >>= 1;
123    }
124    // We only want to add inferred canopy nodes such that the proof length
125    // is equal to the tree depth. If the length of proof + length of canopy nodes is
126    // less than the tree depth, the instruction will fail.
127    let overlap = (proof.len() + inferred_nodes.len()).saturating_sub(max_depth as usize);
128    proof.extend(inferred_nodes.iter().skip(overlap));
129    Ok(())
130}