[Feat] Optimize zktrieState with flatten proofs

bus-mapping/src/circuit_input_builder/l2.rs

@@ -124,7 +124,16 @@ impl CircuitInputBuilder {
             hex::encode(old_root),
         );
 
-        let mpt_init_state = if !light_mode {
+        let mpt_init_state = if !l2_trace.storage_trace.flatten_proofs.is_empty() {
+            log::info!("always init mpt state with flatten proofs");
+            let state = ZktrieState::construct(old_root);
+            let mut zk_db = state.zk_db.borrow_mut();
+            for (k, bytes) in &l2_trace.storage_trace.flatten_proofs {
+                zk_db.add_node_bytes(bytes, Some(k.as_bytes())).unwrap();
+            }
+            drop(zk_db);
+            Some(state)
+        } else if !light_mode {
             let mpt_init_state = ZktrieState::from_trace_with_additional(
                 old_root,
                 Self::collect_account_proofs(&l2_trace.storage_trace),
@@ -148,21 +157,57 @@ impl CircuitInputBuilder {
         };
 
         let mut sdb = StateDB::new();
-        for parsed in ZktrieState::parse_account_from_proofs(Self::collect_account_proofs(
-            &l2_trace.storage_trace,
-        )) {
-            let (addr, acc) = parsed.map_err(Error::IoError)?;
-            log::trace!("sdb trace {:?} {:?}", addr, acc);
-            sdb.set_account(&addr, state_db::Account::from(&acc));
-        }
+        if let Some(zk_state) = &mpt_init_state {
+            for (addr, acc) in zk_state.query_accounts(
+                Self::collect_account_proofs(&l2_trace.storage_trace).map(|(addr, _)| addr),
+            ) {
+                if let Some(acc) = acc {
+                    log::trace!("sdb trace[query mode] {:?} {:?}", addr, acc);
+                    sdb.set_account(&addr, state_db::Account::from(&acc));
+                } else {
+                    log::warn!("can not query account with addr {:?}", addr);
+                }
+            }
+
+            for ((addr, key), val) in zk_state.query_storages(
+                Self::collect_storage_proofs(&l2_trace.storage_trace)
+                    .map(|(addr, key, _)| (addr, key)),
+            ) {
+                let key = key.to_word();
+                if let Some(val) = val {
+                    log::trace!(
+                        "sdb trace storage[query mode] {:?} {:?} {:?}",
+                        addr,
+                        key,
+                        val
+                    );
+                    *sdb.get_storage_mut(&addr, &key).1 = val.into();
+                } else {
+                    log::trace!(
+                        "sdb trace storage[query mode] {:?} {:?} for zero",
+                        addr,
+                        key
+                    );
+                    *sdb.get_storage_mut(&addr, &key).1 = Default::default();
+                }
+            }
+        } else {
+            for parsed in ZktrieState::parse_account_from_proofs(Self::collect_account_proofs(
+                &l2_trace.storage_trace,
+            )) {
+                let (addr, acc) = parsed.map_err(Error::IoError)?;
+                log::trace!("sdb trace {:?} {:?}", addr, acc);
+                sdb.set_account(&addr, state_db::Account::from(&acc));
+            }
 
-        for parsed in ZktrieState::parse_storage_from_proofs(Self::collect_storage_proofs(
-            &l2_trace.storage_trace,
-        )) {
-            let ((addr, key), val) = parsed.map_err(Error::IoError)?;
-            let key = key.to_word();
-            log::trace!("sdb trace storage {:?} {:?} {:?}", addr, key, val);
-            *sdb.get_storage_mut(&addr, &key).1 = val.into();
+            for parsed in ZktrieState::parse_storage_from_proofs(Self::collect_storage_proofs(
+                &l2_trace.storage_trace,
+            )) {
+                let ((addr, key), val) = parsed.map_err(Error::IoError)?;
+                let key = key.to_word();
+                log::trace!("sdb trace storage {:?} {:?} {:?}", addr, key, val);
+                *sdb.get_storage_mut(&addr, &key).1 = val.into();
+            }
         }
 
         let mut code_db = CodeDB::new();
@@ -191,7 +236,17 @@ impl CircuitInputBuilder {
     /// Apply more l2 traces
     pub fn add_more_l2_trace(&mut self, l2_trace: BlockTrace) -> Result<(), Error> {
         // update init state new data from storage
-        if let Some(mpt_init_state) = &mut self.mpt_init_state {
+        if !l2_trace.storage_trace.flatten_proofs.is_empty() {
+            let mpt_state = self
+                .mpt_init_state
+                .as_ref()
+                .expect("should have inited with flatten proof");
+            log::info!("add more flatten proofs to mpt state");
+            let mut zk_db = mpt_state.zk_db.borrow_mut();
+            for (k, bytes) in &l2_trace.storage_trace.flatten_proofs {
+                zk_db.add_node_bytes(bytes, Some(k.as_bytes())).unwrap();
+            }
+        } else if let Some(mpt_init_state) = &mut self.mpt_init_state {
             mpt_init_state.update_from_trace(
                 Self::collect_account_proofs(&l2_trace.storage_trace),
                 Self::collect_storage_proofs(&l2_trace.storage_trace),
@@ -203,40 +258,62 @@ impl CircuitInputBuilder {
             );
         }
 
-        let new_accounts = ZktrieState::parse_account_from_proofs(
+        let filtered_accounts =
             Self::collect_account_proofs(&l2_trace.storage_trace).filter(|(addr, _)| {
                 let (existed, _) = self.sdb.get_account(addr);
                 !existed
-            }),
-        )
-        .try_fold(
-            HashMap::new(),
-            |mut m, parsed| -> Result<HashMap<_, _>, Error> {
-                let (addr, acc) = parsed.map_err(Error::IoError)?;
-                m.insert(addr, acc);
-                Ok(m)
-            },
-        )?;
+            });
+
+        let new_accounts = if let Some(zk_state) = &self.mpt_init_state {
+            zk_state
+                .query_accounts(filtered_accounts.map(|(addr, _)| addr))
+                .fold(HashMap::new(), |mut m, (addr, acc)| {
+                    if let Some(acc) = acc {
+                        m.insert(addr, acc);
+                    }
+                    m
+                })
+        } else {
+            ZktrieState::parse_account_from_proofs(filtered_accounts).try_fold(
+                HashMap::new(),
+                |mut m, parsed| -> Result<HashMap<_, _>, Error> {
+                    let (addr, acc) = parsed.map_err(Error::IoError)?;
+                    m.insert(addr, acc);
+                    Ok(m)
+                },
+            )?
+        };
 
         for (addr, acc) in new_accounts {
             self.sdb.set_account(&addr, state_db::Account::from(&acc));
         }
 
-        let new_storages = ZktrieState::parse_storage_from_proofs(
+        let filtered_storages =
             Self::collect_storage_proofs(&l2_trace.storage_trace).filter(|(addr, key, _)| {
                 let key = key.to_word();
                 let (existed, _) = self.sdb.get_committed_storage(addr, &key);
                 !existed
-            }),
-        )
-        .try_fold(
-            HashMap::new(),
-            |mut m, parsed| -> Result<HashMap<(Address, Word), Word>, Error> {
-                let ((addr, key), val) = parsed.map_err(Error::IoError)?;
-                m.insert((addr, key.to_word()), val.into());
-                Ok(m)
-            },
-        )?;
+            });
+
+        let new_storages = if let Some(zk_state) = &self.mpt_init_state {
+            zk_state
+                .query_storages(filtered_storages.map(|(addr, key, _)| (addr, key)))
+                .fold(HashMap::new(), |mut m, ((addr, key), val)| {
+                    if let Some(val) = val {
+                        m.insert((addr, key.to_word()), val.into());
+                    }
+                    m
+                })
+        } else {
+            ZktrieState::parse_storage_from_proofs(filtered_storages).try_fold(
+                HashMap::new(),
+                |mut m, parsed| -> Result<HashMap<(Address, Word), Word>, Error> {
+                    let ((addr, key), val) = parsed.map_err(Error::IoError)?;
+                    m.insert((addr, key.to_word()), val.into());
+                    Ok(m)
+                },
+            )?
+        };
 
         for ((addr, key), val) in new_storages {
             *self.sdb.get_storage_mut(&addr, &key).1 = val;

eth-types/src/l2_types.rs

@@ -451,6 +451,9 @@ pub struct StorageTrace {
     #[serde(rename = "deletionProofs", default)]
     /// additional deletion proofs
     pub deletion_proofs: Vec<Bytes>,
+    #[serde(rename = "flattenProofs", default)]
+    ///
+    pub flatten_proofs: HashMap<H256, Bytes>,
 }
 
 /// extension of `GethExecTrace`, with compatible serialize form

integration-tests/tests/l2_trace.rs

@@ -0,0 +1,51 @@
+#![feature(lazy_cell)]
+#![cfg(feature = "scroll")]
+
+use bus_mapping::{
+    circuit_input_builder::{CircuitInputBuilder, CircuitsParams},
+    util::read_env_var,
+};
+use eth_types::l2_types::BlockTrace;
+use integration_tests::log_init;
+use std::fs::File;
+use zkevm_circuits::witness;
+
+fn test_circuit_input_builder_l2block(block_trace: BlockTrace) {
+    let params = CircuitsParams {
+        max_rws: 4_000_000,
+        max_copy_rows: 0, // dynamic
+        max_txs: read_env_var("MAX_TXS", 128),
+        max_calldata: 2_000_000,
+        max_inner_blocks: 64,
+        max_bytecode: 3_000_000,
+        max_mpt_rows: 2_000_000,
+        max_poseidon_rows: 4_000_000,
+        max_keccak_rows: 0,
+        max_exp_steps: 100_000,
+        max_evm_rows: 0,
+        max_rlp_rows: 2_070_000,
+        ..Default::default()
+    };
+
+    let mut builder = CircuitInputBuilder::new_from_l2_trace(params, block_trace, false)
+        .expect("could not handle block tx");
+
+    builder
+        .finalize_building()
+        .expect("could not finalize building block");
+
+    log::trace!("CircuitInputBuilder: {:#?}", builder);
+
+    let mut block = witness::block_convert(&builder.block, &builder.code_db).unwrap();
+    block.apply_mpt_updates(&builder.mpt_init_state.unwrap());
+}
+
+#[test]
+fn local_l2_trace() {
+    log_init();
+    let file_path = read_env_var("TRACE_FILE", "dump.json".to_string());
+    let fd = File::open(file_path).unwrap();
+    let trace: BlockTrace = serde_json::from_reader(fd).unwrap();
+
+    test_circuit_input_builder_l2block(trace);
+}

zktrie/Cargo.toml

@@ -8,7 +8,7 @@ license.workspace = true
 
 [dependencies]
 halo2curves.workspace = true
-zktrie = { git = "https://github.com/scroll-tech/zktrie.git", branch = "main", features= ["rs_zktrie"] }
+zktrie = { git = "https://github.com/scroll-tech/zktrie.git", branch = "feat/node_key_importing", features= ["rs_zktrie"] }
 poseidon-base.workspace = true
 eth-types = { path = "../eth-types" }
 num-bigint.workspace = true

zktrie/src/state.rs

@@ -80,6 +80,44 @@ impl ZktrieState {
         true
     }
 
+    ///
+    pub fn query_accounts<'d: 'a, 'a>(
+        &self,
+        accounts: impl Iterator<Item = &'a Address> + 'd,
+    ) -> impl Iterator<Item = (Address, Option<AccountData>)> + 'a {
+        let trie = self.zk_db.borrow_mut().new_trie(&self.trie_root).unwrap();
+        accounts.map(move |&addr| {
+            let account = trie.get_account(addr.as_bytes()).map(AccountData::from);
+            (addr, account)
+        })
+    }
+
+    ///
+    pub fn query_storages<'d: 'a, 'a>(
+        &self,
+        storages: impl Iterator<Item = (&'a Address, &'a H256)> + 'd,
+    ) -> impl Iterator<Item = ((Address, H256), Option<StorageData>)> + 'a {
+        use std::collections::{hash_map::Entry::*, HashMap};
+        let zk_db = self.zk_db.borrow().clone();
+        let account_trie = zk_db.new_trie(&self.trie_root).unwrap();
+        let mut trie_cache: HashMap<Address, ZkTrie> = HashMap::new();
+        storages.map(move |(&addr, &key)| {
+            let store_val = match trie_cache.entry(addr) {
+                Occupied(entry) => Some(entry.into_mut()),
+                Vacant(entry) => account_trie
+                    .get_account(addr.as_bytes())
+                    .map(AccountData::from)
+                    .and_then(|account| {
+                        zk_db
+                            .new_trie(&account.storage_root.0)
+                            .map(|tr| entry.insert(tr))
+                    }),
+            }
+            .and_then(|tr| tr.get_store(key.as_bytes()).map(StorageData::from));
+            ((addr, key), store_val)
+        })
+    }
+
     /// Helper for parsing account data from external data (mainly storage trace)
     pub fn parse_account_from_proofs<'d: 'a, 'a, BYTES>(
         account_proofs: impl Iterator<Item = (&'a Address, BYTES)> + 'd,
@@ -146,7 +184,7 @@ impl ZktrieState {
             .chain(additional_proofs);
         let mut zk_db = self.zk_db.borrow_mut();
         for bytes in proofs {
-            zk_db.add_node_bytes(bytes).unwrap();
+            zk_db.add_node_data(bytes).unwrap();
         }
     }