spl_account_compression/
concurrent_tree_wrapper.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
//! 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::*,
};

#[inline(never)]
pub fn merkle_tree_initialize_empty(
    header: &ConcurrentMerkleTreeHeader,
    tree_id: Pubkey,
    tree_bytes: &mut [u8],
) -> Result<Box<ChangeLogEvent>> {
    merkle_tree_apply_fn_mut!(header, tree_id, tree_bytes, initialize,)
}

#[inline(never)]
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)
}

#[inline(never)]
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)
}

#[inline(never)]
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)
}

#[inline(never)]
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)
}

#[inline(never)]
pub fn merkle_tree_append_leaf(
    header: &ConcurrentMerkleTreeHeader,
    tree_id: Pubkey,
    tree_bytes: &mut [u8],
    args: &[u8; 32],
) -> Result<Box<ChangeLogEvent>> {
    merkle_tree_apply_fn_mut!(header, tree_id, tree_bytes, append, *args)
}

/// Checks whether the tree in not initialized yet without doing the deserialization. A rought
/// equivalent to deserializing the tree and calling is_initialized() on it without the heavy
/// lifting with macros. An empty account is a zero'd account. The tree is considered empty if the
/// tree_bytes are all 0. A regular non-batch initialized tree is initialized early on when the
/// init_empty_merkle_tree is called. A batch initialized tree stays uninitialized until the
/// init_prepared_tree_with_root is called.
pub fn tree_bytes_uninitialized(tree_bytes: &[u8]) -> bool {
    tree_bytes.iter().all(|&x| x == 0)
}

#[inline(never)]
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_uninitialized(tree_bytes) {
        return Ok(());
    }
    // check the tree is empty
    merkle_tree_apply_fn_mut!(header, tree_id, tree_bytes, prove_tree_is_empty,)?;
    Ok(())
}