hpl_toolkit/compression/
controlled_merkle_trees.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
use {
    super::{merke_tree_utills::*, merkle_tree_apply_fn, merkle_tree_apply_fn_deep},
    crate::{errors::HplToolkitError, schema::*, utils::VecMap},
    anchor_lang::prelude::*,
};

#[cfg_attr(feature = "debug", derive(Debug))]
#[derive(AnchorSerialize, AnchorDeserialize, Clone, ToSchema, PartialEq)]
pub struct ControlledMerkleTrees {
    pub active: u8,
    pub schema: Schema,
    pub merkle_trees: Vec<Pubkey>,
}

impl ControlledMerkleTrees {
    pub fn assert_append<'info>(
        &mut self,
        merkle_tree_info: AccountInfo<'info>,
    ) -> Result<(u32, u64)> {
        let active_merkle_tree_option = &self.merkle_trees.get(self.active as usize);

        if active_merkle_tree_option.is_none() {
            return Err(HplToolkitError::ActiveTreeNotFound.into());
        }

        let active_merkle_tree = active_merkle_tree_option.unwrap();

        if merkle_tree_info.key != active_merkle_tree {
            return Err(HplToolkitError::TreeNotActive.into());
        }

        let merkle_tree_bytes = unsafe {
            let cloned_tree = merkle_tree_info.data.to_owned().as_ptr();
            cloned_tree.as_ref().unwrap()
        };

        let (tree_capacity, leaf_idx, seq) = merkle_tree_apply_fn!(merkle_tree_bytes get_hpl_info);

        if tree_capacity == 0 {
            return Err(HplToolkitError::TreeFull.into());
        } else if tree_capacity == 1 {
            self.active += 1;
        }

        Ok((leaf_idx, seq + 1))
    }

    pub fn get_size(schema: &Schema, merkle_trees_len: usize) -> usize {
        1 + schema.size() + merkle_trees_len * 32
    }

    pub fn get_size_for_borsh(schema: &Schema, merkle_trees_len: usize) -> usize {
        match (Self {
            active: 0,
            schema: schema.clone(),
            merkle_trees: vec![Pubkey::default(); merkle_trees_len],
        }
        .try_to_vec())
        {
            Ok(vec) => vec.len(),
            Err(_) => 0,
        }
    }

    pub fn size(&self) -> usize {
        1 + self.schema.size() + self.merkle_trees.len() * 32
    }

    pub fn size_for_borsh(&self) -> usize {
        match self.try_to_vec() {
            Ok(vec) => vec.len(),
            Err(_) => 0,
        }
    }
}