Skip to main content

light_client/rpc/
merkle_tree.rs

1use std::mem;
2
3use async_trait::async_trait;
4use light_concurrent_merkle_tree::{
5    copy::ConcurrentMerkleTreeCopy, errors::ConcurrentMerkleTreeError, light_hasher::Poseidon,
6};
7use light_indexed_merkle_tree::{copy::IndexedMerkleTreeCopy, errors::IndexedMerkleTreeError};
8use solana_pubkey::Pubkey;
9use thiserror::Error;
10
11use super::{state::MerkleTreeMetadata, Rpc, RpcError};
12
13#[derive(Error, Debug)]
14pub enum MerkleTreeExtError {
15    #[error(transparent)]
16    Rpc(#[from] RpcError),
17
18    #[error(transparent)]
19    ConcurrentMerkleTree(#[from] ConcurrentMerkleTreeError),
20
21    #[error(transparent)]
22    IndexedMerkleTree(#[from] IndexedMerkleTreeError),
23}
24
25// TODO: hide behind feature to make tree and poseidon deps optional
26/// Extension to the RPC connection which provides convenience utilities for
27/// fetching Merkle trees.
28#[async_trait]
29pub trait MerkleTreeExt: Rpc {
30    // TODO: add v2 state tree
31    async fn get_state_merkle_tree_account(
32        &mut self,
33        pubkey: Pubkey,
34    ) -> Result<ConcurrentMerkleTreeCopy<Poseidon, 26>, MerkleTreeExtError> {
35        let account = self
36            .get_account(pubkey)
37            .await?
38            .ok_or_else(|| RpcError::AccountDoesNotExist(pubkey.to_string()))?;
39        let tree = ConcurrentMerkleTreeCopy::from_bytes_copy(
40            &account.data[8 + mem::size_of::<MerkleTreeMetadata>()..],
41        )?;
42
43        Ok(tree)
44    }
45
46    // TODO: add v2 state tree
47    async fn get_address_merkle_tree_account(
48        &mut self,
49        pubkey: Pubkey,
50    ) -> Result<IndexedMerkleTreeCopy<Poseidon, usize, 26, 16>, MerkleTreeExtError> {
51        let account = self
52            .get_account(pubkey)
53            .await?
54            .ok_or_else(|| RpcError::AccountDoesNotExist(pubkey.to_string()))?;
55        let tree = IndexedMerkleTreeCopy::from_bytes_copy(
56            &account.data[8 + mem::size_of::<MerkleTreeMetadata>()..],
57        )?;
58
59        Ok(tree)
60    }
61}