pessimistic_proof_core/local_state/
mod.rs

1use std::collections::{btree_map::Entry, BTreeMap};
2
3use agglayer_primitives::{ruint::UintTryFrom, Hashable, U256, U512};
4use agglayer_tries::roots::{LocalBalanceRoot, LocalNullifierRoot};
5use commitment::StateCommitment;
6use serde::{Deserialize, Serialize};
7use unified_bridge::{Error, LocalExitTree, NetworkId, L1_ETH};
8
9use crate::{
10    local_balance_tree::LocalBalanceTree,
11    multi_batch_header::MultiBatchHeader,
12    nullifier_tree::{NullifierKey, NullifierTree},
13    ProofError,
14};
15
16pub mod commitment;
17
18/// State representation of one network without the leaves, taken as input by
19/// the prover.
20#[derive(Clone, Debug, Serialize, Deserialize)]
21pub struct NetworkState {
22    /// Commitment to the [`BridgeExit`](struct@crate::bridge_exit::BridgeExit).
23    pub exit_tree: LocalExitTree,
24    /// Commitment to the balance for each token.
25    pub balance_tree: LocalBalanceTree,
26    /// Commitment to the Nullifier tree for the local network, tracks claimed
27    /// assets on foreign networks
28    pub nullifier_tree: NullifierTree,
29}
30
31impl NetworkState {
32    /// Returns the roots.
33    pub fn get_state_commitment(&self) -> StateCommitment {
34        StateCommitment {
35            exit_root: self.exit_tree.get_root().into(),
36            ler_leaf_count: self.exit_tree.leaf_count,
37            balance_root: LocalBalanceRoot::new(self.balance_tree.root),
38            nullifier_root: LocalNullifierRoot::new(self.nullifier_tree.root),
39        }
40    }
41
42    /// Apply the [`MultiBatchHeader`] on the current [`NetworkState`].
43    /// Checks that the transition reaches the target [`StateCommitment`].
44    /// The state isn't modified on error.
45    pub fn apply_batch_header(
46        &mut self,
47        multi_batch_header: &MultiBatchHeader,
48    ) -> Result<StateCommitment, ProofError> {
49        let mut clone = self.clone();
50        let roots = clone.apply_batch_header_helper(multi_batch_header)?;
51        *self = clone;
52
53        Ok(roots)
54    }
55
56    /// Apply the [`MultiBatchHeader`] on the current [`NetworkState`].
57    /// Returns the resulting [`StateCommitment`] upon success.
58    /// The state can be modified on error.
59    fn apply_batch_header_helper(
60        &mut self,
61        multi_batch_header: &MultiBatchHeader,
62    ) -> Result<StateCommitment, ProofError> {
63        let mut new_balances = BTreeMap::new();
64        for (k, v) in &multi_batch_header.balances_proofs {
65            if new_balances.insert(*k, U512::from(v.0)).is_some() {
66                return Err(ProofError::DuplicateTokenBalanceProof(*k));
67            }
68        }
69
70        // Apply the imported bridge exits
71        for (imported_bridge_exit, nullifier_path) in &multi_batch_header.imported_bridge_exits {
72            if imported_bridge_exit.global_index.network_id() == multi_batch_header.origin_network {
73                // We don't allow a chain to exit to itself
74                return Err(ProofError::CannotExitToSameNetwork);
75            }
76            // Check that the destination network of the bridge exit matches the current
77            // network
78            if imported_bridge_exit.bridge_exit.dest_network != multi_batch_header.origin_network {
79                return Err(ProofError::InvalidImportedBridgeExit {
80                    source: Error::InvalidExitNetwork,
81                    global_index: imported_bridge_exit.global_index,
82                });
83            }
84
85            // Check the inclusion proof
86            imported_bridge_exit
87                .verify_path(multi_batch_header.l1_info_root)
88                .map_err(|source| ProofError::InvalidImportedBridgeExit {
89                    source,
90                    global_index: imported_bridge_exit.global_index,
91                })?;
92
93            // Check the nullifier non-inclusion path and update the nullifier tree
94            let nullifier_key: NullifierKey = imported_bridge_exit.global_index.into();
95            self.nullifier_tree
96                .verify_and_update(nullifier_key, nullifier_path)?;
97
98            // The amount corresponds to L1 ETH if the leaf is a message
99            let token_info = imported_bridge_exit.bridge_exit.amount_token_info();
100
101            if multi_batch_header.origin_network == token_info.origin_network {
102                // When the token is native to the chain, we don't care about the local balance
103                continue;
104            }
105
106            // Update the token balance.
107            let amount = imported_bridge_exit.bridge_exit.amount;
108            let entry = new_balances.entry(token_info);
109            match entry {
110                Entry::Vacant(_) => return Err(ProofError::MissingTokenBalanceProof(token_info)),
111                Entry::Occupied(mut entry) => {
112                    *entry.get_mut() = entry
113                        .get()
114                        .checked_add(U512::from(amount))
115                        .ok_or(ProofError::BalanceOverflowInBridgeExit)?;
116                }
117            }
118        }
119
120        // Apply the bridge exits
121        for bridge_exit in &multi_batch_header.bridge_exits {
122            if bridge_exit.dest_network == multi_batch_header.origin_network {
123                // We don't allow a chain to exit to itself
124                return Err(ProofError::CannotExitToSameNetwork);
125            }
126            self.exit_tree.add_leaf(bridge_exit.hash())?;
127
128            // For message exits, the origin network in token info should be the origin
129            // network of the batch header.
130            if bridge_exit.is_message()
131                && bridge_exit.token_info.origin_network != multi_batch_header.origin_network
132            {
133                return Err(ProofError::InvalidMessageOriginNetwork);
134            }
135
136            // For ETH transfers, we need to check that the origin network is the L1 network
137            if bridge_exit.token_info.origin_token_address == L1_ETH.origin_token_address
138                && bridge_exit.token_info.origin_network != NetworkId::ETH_L1
139            {
140                return Err(ProofError::InvalidL1TokenInfo(bridge_exit.token_info));
141            }
142
143            // The amount corresponds to L1 ETH if the leaf is a message
144            let token_info = bridge_exit.amount_token_info();
145
146            if multi_batch_header.origin_network == token_info.origin_network {
147                // When the token is native to the chain, we don't care about the local balance
148                continue;
149            }
150
151            // Update the token balance.
152            let amount = bridge_exit.amount;
153            let entry = new_balances.entry(token_info);
154            match entry {
155                Entry::Vacant(_) => return Err(ProofError::MissingTokenBalanceProof(token_info)),
156                Entry::Occupied(mut entry) => {
157                    *entry.get_mut() = entry
158                        .get()
159                        .checked_sub(U512::from(amount))
160                        .ok_or(ProofError::BalanceUnderflowInBridgeExit)?;
161                }
162            }
163        }
164
165        // Verify that the original balances were correct and update the local balance
166        // tree with the new balances.
167        for (token, (old_balance, balance_path)) in &multi_batch_header.balances_proofs {
168            let new_balance = new_balances[token];
169            let new_balance = U256::uint_try_from(new_balance)
170                .map_err(|_| ProofError::BalanceOverflowInBridgeExit)?;
171            self.balance_tree
172                .verify_and_update(*token, balance_path, *old_balance, new_balance)?;
173        }
174
175        Ok(self.get_state_commitment())
176    }
177}