pessimistic_proof/
local_balance_tree.rs

1use agglayer_primitives::{Digest, U256};
2use agglayer_tries::{
3    error::SmtError,
4    smt::{Smt, SmtPath},
5    utils::empty_hash_at_height,
6};
7pub use pessimistic_proof_core::local_balance_tree::{LocalBalancePath, LOCAL_BALANCE_TREE_DEPTH};
8use serde::{Deserialize, Serialize};
9use serde_with::serde_as;
10use unified_bridge::TokenInfo;
11
12/// A commitment to the set of per-network local balance trees maintained by the
13/// local network
14#[serde_as]
15#[derive(Clone, Debug, Serialize, Deserialize)]
16pub struct LocalBalanceTree {
17    /// The Merkle Root of the local balance tree
18    #[serde_as(as = "_")]
19    pub root: Digest,
20}
21
22impl Default for LocalBalanceTree {
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28impl LocalBalanceTree {
29    pub fn new() -> Self {
30        // We add 1 to the depth here because the empty hash at height 0 is
31        // already set to Digest::ZERO.
32        let root = empty_hash_at_height::<{ LOCAL_BALANCE_TREE_DEPTH + 1 }>();
33        LocalBalanceTree { root }
34    }
35
36    pub fn new_with_root(root: Digest) -> Self {
37        LocalBalanceTree { root }
38    }
39}
40
41impl From<LocalBalanceTree> for pessimistic_proof_core::local_balance_tree::LocalBalanceTree {
42    fn from(tree: LocalBalanceTree) -> Self {
43        Self { root: tree.root }
44    }
45}
46
47pub struct BalanceTree(pub Smt<LOCAL_BALANCE_TREE_DEPTH>);
48
49impl BalanceTree {
50    /// Returns all the non-zero token balance contained in this balance tree.
51    pub fn get_all_balances(
52        &self,
53    ) -> Result<impl Iterator<Item = (SmtPath<LOCAL_BALANCE_TREE_DEPTH>, Digest)>, SmtError> {
54        Ok(self.0.entries()?.into_iter())
55    }
56
57    /// Returns the balance for the given [`TokenInfo`].
58    pub fn get_balance(&self, token_info: TokenInfo) -> U256 {
59        self.0
60            .get(token_info)
61            .map(|v| U256::from_be_bytes(*v.as_bytes()))
62            .unwrap_or(U256::ZERO)
63    }
64}
pessimistic_proof/local_balance_tree.rs

pessimistic_proof/
local_balance_tree.rs
