agglayer_types/
local_network_state.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3use agglayer_interop_types::{ImportedBridgeExit, TokenInfo};
4use agglayer_primitives::{ruint::UintTryFrom, FromBool, Hashable};
5use agglayer_tries::{roots::LocalExitRoot, smt::Smt};
6use pessimistic_proof::{
7    core::commitment::{PessimisticRootCommitmentValues, PessimisticRootCommitmentVersion},
8    local_balance_tree::{LocalBalancePath, LocalBalanceTree, LOCAL_BALANCE_TREE_DEPTH},
9    local_state::StateCommitment,
10    multi_batch_header::MultiBatchHeader,
11    nullifier_tree::{NullifierKey, NullifierPath, NullifierTree, NULLIFIER_TREE_DEPTH},
12    LocalNetworkState,
13};
14use unified_bridge::LocalExitTree;
15
16use crate::{
17    aggchain_data::{CertificateAggchainDataCtx, CertificateAggchainDataWithCtx},
18    Certificate, Digest, Error, U256, U512,
19};
20
21/// Local state data of one network.
22/// The AggLayer tracks the [`LocalNetworkStateData`] for all networks.
23#[derive(Clone, Debug, Default)]
24pub struct LocalNetworkStateData {
25    /// The local exit tree without leaves.
26    pub exit_tree: LocalExitTree,
27    /// The full local balance tree.
28    pub balance_tree: Smt<LOCAL_BALANCE_TREE_DEPTH>,
29    /// The full nullifier tree.
30    pub nullifier_tree: Smt<NULLIFIER_TREE_DEPTH>,
31}
32
33impl From<LocalNetworkStateData> for LocalNetworkState {
34    fn from(state: LocalNetworkStateData) -> Self {
35        LocalNetworkState {
36            exit_tree: state.exit_tree,
37            balance_tree: LocalBalanceTree::new_with_root(state.balance_tree.root),
38            nullifier_tree: NullifierTree::new_with_root(state.nullifier_tree.root),
39        }
40    }
41}
42
43impl From<LocalNetworkStateData> for pessimistic_proof::NetworkState {
44    fn from(state: LocalNetworkStateData) -> Self {
45        LocalNetworkState::from(state).into()
46    }
47}
48
49/// The last pessimistic root can be either fetched from L1 or recomputed for a
50/// given version.
51#[derive(Debug, Clone)]
52pub enum PessimisticRootInput {
53    /// Computed from the given version.
54    Computed(PessimisticRootCommitmentVersion),
55    /// Fetched from the L1.
56    Fetched(Digest),
57}
58
59/// Represents all context data fetched from the L1 for the witness generation.
60#[derive(Debug, Clone)]
61pub struct L1WitnessCtx {
62    pub l1_info_root: Digest,
63    pub prev_pessimistic_root: PessimisticRootInput,
64    pub aggchain_data_ctx: CertificateAggchainDataCtx,
65}
66
67impl LocalNetworkStateData {
68    /// Prune the SMTs
69    pub fn prune_stale_nodes(&mut self) -> Result<(), Error> {
70        self.balance_tree.traverse_and_prune()?;
71        self.nullifier_tree.traverse_and_prune()?;
72
73        Ok(())
74    }
75
76    /// Apply the [`Certificate`] on the current state and returns the
77    /// [`MultiBatchHeader`] associated to the state transition.
78    pub fn apply_certificate(
79        &mut self,
80        certificate: &Certificate,
81        ctx_from_l1: L1WitnessCtx,
82    ) -> Result<MultiBatchHeader, Error> {
83        let L1WitnessCtx {
84            prev_pessimistic_root,
85            l1_info_root,
86            aggchain_data_ctx,
87        } = ctx_from_l1;
88
89        let gers_are_consistent = certificate
90            .imported_bridge_exits
91            .iter()
92            .all(|ib| ib.valid_claim());
93
94        if !gers_are_consistent {
95            return Err(Error::InconsistentGlobalExitRoot);
96        }
97
98        // Retrieve the pp root
99        let prev_pessimistic_root = match prev_pessimistic_root {
100            PessimisticRootInput::Fetched(settled_from_l1) => settled_from_l1,
101            PessimisticRootInput::Computed(version) => PessimisticRootCommitmentValues {
102                balance_root: self.balance_tree.root.into(),
103                nullifier_root: self.nullifier_tree.root.into(),
104                ler_leaf_count: self.exit_tree.leaf_count(),
105                height: certificate.height.as_u64(),
106                origin_network: certificate.network_id,
107            }
108            .compute_pp_root(version),
109        };
110
111        let prev_local_exit_root = self.exit_tree.get_root().into();
112        if certificate.prev_local_exit_root != prev_local_exit_root {
113            return Err(Error::MismatchPrevLocalExitRoot {
114                computed: prev_local_exit_root,
115                declared: certificate.prev_local_exit_root,
116            });
117        }
118
119        for e in certificate.bridge_exits.iter() {
120            self.exit_tree.add_leaf(e.hash())?;
121        }
122
123        let balances_proofs: BTreeMap<TokenInfo, (U256, LocalBalancePath)> = {
124            // Consider all the imported bridge exits except for the native token
125            let imported_bridge_exits = certificate.imported_bridge_exits.iter().filter(|b| {
126                b.bridge_exit.amount_token_info().origin_network != certificate.network_id
127            });
128
129            // Consider all the bridge exits except for the native token
130            let bridge_exits = certificate
131                .bridge_exits
132                .iter()
133                .filter(|b| b.amount_token_info().origin_network != certificate.network_id);
134
135            // Set of dedup tokens mutated in the transition
136            let mutated_tokens: BTreeSet<TokenInfo> = {
137                let imported_tokens = imported_bridge_exits
138                    .clone()
139                    .map(|exit| exit.bridge_exit.amount_token_info());
140                let exported_tokens = bridge_exits.clone().map(|exit| exit.amount_token_info());
141                imported_tokens.chain(exported_tokens).collect()
142            };
143
144            let initial_balances: BTreeMap<_, _> = mutated_tokens
145                .iter()
146                .map(|&token| {
147                    let balance =
148                        U256::from_be_bytes(*self.balance_tree.get(token).unwrap_or_default());
149                    (token, balance)
150                })
151                .collect();
152
153            let mut new_balances: BTreeMap<_, _> = initial_balances
154                .iter()
155                .map(|(&token, &balance)| (token, U512::from(balance)))
156                .collect();
157
158            for imported_bridge_exit in imported_bridge_exits {
159                let token = imported_bridge_exit.bridge_exit.amount_token_info();
160                new_balances.insert(
161                    token,
162                    new_balances[&token]
163                        .checked_add(U512::from(imported_bridge_exit.bridge_exit.amount))
164                        .ok_or(Error::BalanceOverflow(token))?,
165                );
166            }
167
168            for bridge_exit in bridge_exits {
169                let token = bridge_exit.amount_token_info();
170                new_balances.insert(
171                    token,
172                    new_balances[&token]
173                        .checked_sub(U512::from(bridge_exit.amount))
174                        .ok_or(Error::BalanceUnderflow(token))?,
175                );
176            }
177
178            // Get the proof against the initial balance for each token
179            mutated_tokens
180                .into_iter()
181                .map(|token| {
182                    let initial_balance = initial_balances[&token];
183
184                    let new_balance = U256::uint_try_from(new_balances[&token])
185                        .map_err(|_| Error::BalanceOverflow(token))?;
186
187                    let balance_proof_error =
188                        |source| Error::BalanceProofGenerationFailed { source, token };
189
190                    let path = if initial_balance.is_zero() {
191                        self.balance_tree
192                            .get_inclusion_proof_zero(token)
193                            .map_err(balance_proof_error)?
194                    } else {
195                        self.balance_tree
196                            .get_inclusion_proof(token)
197                            .map_err(balance_proof_error)?
198                    };
199
200                    self.balance_tree
201                        .update(token, new_balance.to_be_bytes().into())
202                        .map_err(balance_proof_error)?;
203
204                    Ok((token, (initial_balance, path)))
205                })
206                .collect::<Result<BTreeMap<_, _>, Error>>()?
207        };
208
209        let imported_bridge_exits: Vec<(ImportedBridgeExit, NullifierPath)> = certificate
210            .imported_bridge_exits
211            .iter()
212            .map(|exit| {
213                let nullifier_key: NullifierKey = exit.global_index.into();
214                let nullifier_error = |source| Error::NullifierPathGenerationFailed {
215                    source,
216                    global_index: exit.global_index,
217                };
218                let nullifier_path = self
219                    .nullifier_tree
220                    .get_non_inclusion_proof(nullifier_key)
221                    .map_err(nullifier_error)?;
222                self.nullifier_tree
223                    .insert(nullifier_key, Digest::from_bool(true))
224                    .map_err(nullifier_error)?;
225                Ok((exit.clone(), nullifier_path))
226            })
227            .collect::<Result<Vec<_>, Error>>()?;
228
229        // Check that the certificate referred to the right target
230        let computed = LocalExitRoot::from(self.exit_tree.get_root());
231        if computed != certificate.new_local_exit_root {
232            return Err(Error::MismatchNewLocalExitRoot {
233                declared: certificate.new_local_exit_root,
234                computed,
235            });
236        }
237
238        let chain_payload = certificate
239            .aggchain_data
240            .clone()
241            .try_into()
242            .map_err(Error::InvalidChainData)?;
243
244        let aggchain_data = CertificateAggchainDataWithCtx(chain_payload, aggchain_data_ctx)
245            .try_into()
246            .map_err(Error::InvalidChainData)?;
247
248        Ok(MultiBatchHeader {
249            origin_network: certificate.network_id,
250            bridge_exits: certificate.bridge_exits.clone(),
251            imported_bridge_exits,
252            balances_proofs,
253            l1_info_root,
254            height: certificate.height.as_u64(),
255            prev_pessimistic_root,
256            aggchain_data,
257            certificate_id: certificate.hash().into(),
258        })
259    }
260
261    /// Generates the [`MultiBatchHeader`] from the state and a [`Certificate`].
262    /// Does not mutate the current state.
263    pub fn make_multi_batch_header(
264        &self,
265        certificate: &Certificate,
266        witness_fetched_from_l1: L1WitnessCtx,
267    ) -> Result<MultiBatchHeader, Error> {
268        self.clone()
269            .apply_certificate(certificate, witness_fetched_from_l1)
270    }
271
272    pub fn get_roots(&self) -> StateCommitment {
273        StateCommitment {
274            exit_root: self.exit_tree.get_root(),
275            ler_leaf_count: self.exit_tree.leaf_count(),
276            balance_root: self.balance_tree.root,
277            nullifier_root: self.nullifier_tree.root,
278        }
279    }
280}
agglayer_types/local_network_state.rs

agglayer_types/
local_network_state.rs
