diff --git a/Cargo.lock b/Cargo.lock index 98598578d4..fe058a3e7c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5559,7 +5559,7 @@ dependencies = [ [[package]] name = "zktrie" version = "0.3.0" -source = "git+https://github.com/scroll-tech/zktrie.git?branch=main#2a19ee2155ecb959ba518384e448a638ae9858f2" +source = "git+https://github.com/scroll-tech/zktrie.git?branch=main#49e0f027d12abf7f2049c2cf1034128fc14f47dd" dependencies = [ "gobuild", "zktrie_rust", @@ -5568,7 +5568,7 @@ dependencies = [ [[package]] name = "zktrie_rust" version = "0.3.0" -source = "git+https://github.com/scroll-tech/zktrie.git?branch=main#2a19ee2155ecb959ba518384e448a638ae9858f2" +source = "git+https://github.com/scroll-tech/zktrie.git?branch=main#49e0f027d12abf7f2049c2cf1034128fc14f47dd" dependencies = [ "hex", "lazy_static", diff --git a/bus-mapping/src/circuit_input_builder/l2.rs b/bus-mapping/src/circuit_input_builder/l2.rs index 81b9763b50..8d36f89518 100644 --- a/bus-mapping/src/circuit_input_builder/l2.rs +++ b/bus-mapping/src/circuit_input_builder/l2.rs @@ -124,7 +124,35 @@ 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 mut state = ZktrieState::construct(old_root); + let zk_db = state.expose_db(); + for (k, bytes) in &l2_trace.storage_trace.flatten_proofs { + zk_db.add_node_bytes(bytes, Some(k.as_bytes())).unwrap(); + } + zk_db.with_key_cache( + l2_trace + .storage_trace + .address_hashes + .iter() + .map(|(k, v)| (k.as_bytes(), v.as_bytes())), + ); + zk_db.with_key_cache( + l2_trace + .storage_trace + .store_key_hashes + .iter() + .map(|(k, v)| (k.as_bytes(), v.as_bytes())), + ); + + log::debug!( + "building partial ZktrieState done from new trace, root {}", + hex::encode(state.root()) + ); + + 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 +176,58 @@ 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::trace!("sdb trace[query mode] {:?} for zero account", addr); + sdb.set_account(&addr, state_db::Account::zero()); + } + } + + 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 +256,31 @@ 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_mut() + .expect("should have inited with flatten proof"); + log::info!("add more flatten proofs to mpt state"); + let zk_db = mpt_state.expose_db(); + for (k, bytes) in &l2_trace.storage_trace.flatten_proofs { + zk_db.add_node_bytes(bytes, Some(k.as_bytes())).unwrap(); + } + zk_db.with_key_cache( + l2_trace + .storage_trace + .address_hashes + .iter() + .map(|(k, v)| (k.as_bytes(), v.as_bytes())), + ); + zk_db.with_key_cache( + l2_trace + .storage_trace + .store_key_hashes + .iter() + .map(|(k, v)| (k.as_bytes(), v.as_bytes())), + ); + } 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 +292,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, 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)| { + m.insert(addr, acc.unwrap_or_default()); + m + }) + } else { + ZktrieState::parse_account_from_proofs(filtered_accounts).try_fold( + HashMap::new(), + |mut m, parsed| -> Result, 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, 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()); + } else { + m.insert((addr, key.to_word()), Default::default()); + } + m + }) + } else { + ZktrieState::parse_storage_from_proofs(filtered_storages).try_fold( + HashMap::new(), + |mut m, parsed| -> Result, 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; diff --git a/eth-types/src/l2_types.rs b/eth-types/src/l2_types.rs index f912d2652b..ad22a57459 100644 --- a/eth-types/src/l2_types.rs +++ b/eth-types/src/l2_types.rs @@ -451,6 +451,15 @@ pub struct StorageTrace { #[serde(rename = "deletionProofs", default)] /// additional deletion proofs pub deletion_proofs: Vec, + #[serde(rename = "flattenProofs", default)] + /// + pub flatten_proofs: HashMap, + #[serde(rename = "addressHashes", default)] + /// + pub address_hashes: HashMap, + #[serde(rename = "storeKeyHashes", default)] + /// + pub store_key_hashes: HashMap, } /// extension of `GethExecTrace`, with compatible serialize form diff --git a/integration-tests/tests/l2_trace.rs b/integration-tests/tests/l2_trace.rs new file mode 100644 index 0000000000..a32a88e8ca --- /dev/null +++ b/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); +} diff --git a/prover/src/zkevm/circuit/builder.rs b/prover/src/zkevm/circuit/builder.rs index 8f96a7c59a..69add86a83 100644 --- a/prover/src/zkevm/circuit/builder.rs +++ b/prover/src/zkevm/circuit/builder.rs @@ -167,7 +167,7 @@ pub fn finalize_builder(builder: &mut CircuitInputBuilder) -> Result { if let Some(state) = &mut builder.mpt_init_state { if *state.root() != [0u8; 32] { log::debug!("apply_mpt_updates"); - witness_block.apply_mpt_updates(state); + witness_block.apply_mpt_updates_and_update_mpt_state(state); log::debug!("apply_mpt_updates done"); } else { // Empty state root means circuit capacity checking, or dummy witness block for key gen? diff --git a/zkevm-circuits/src/witness/block.rs b/zkevm-circuits/src/witness/block.rs index a2e2150a93..a2566324cc 100644 --- a/zkevm-circuits/src/witness/block.rs +++ b/zkevm-circuits/src/witness/block.rs @@ -120,6 +120,17 @@ impl Block { pub fn apply_mpt_updates(&mut self, mpt_state: &MptState) { self.mpt_updates.fill_state_roots(mpt_state); } + + /// Replay mpt updates to generate mpt witness, also update the mpt state with + /// calculated mpt updatings + pub fn apply_mpt_updates_and_update_mpt_state(&mut self, mpt_state: &mut MptState) { + let updated_tries = self + .mpt_updates + .fill_state_roots(mpt_state) + .into_updated_trie(); + mpt_state.updated_with_trie(updated_tries); + } + /// For each tx, for each step, print the rwc at the beginning of the step, /// and all the rw operations of the step. pub(crate) fn debug_print_txs_steps_rw_ops(&self) { diff --git a/zkevm-circuits/src/witness/mpt.rs b/zkevm-circuits/src/witness/mpt.rs index d57df98161..0f242e7444 100644 --- a/zkevm-circuits/src/witness/mpt.rs +++ b/zkevm-circuits/src/witness/mpt.rs @@ -173,7 +173,7 @@ impl MptUpdates { self.pretty_print(); } - pub(crate) fn fill_state_roots(&mut self, init_trie: &ZktrieState) { + pub(crate) fn fill_state_roots(&mut self, init_trie: &ZktrieState) -> WitnessGenerator { let root_pair = (self.old_root, self.new_root); self.old_root = init_trie.root().into(); log::trace!("fill_state_roots init {:?}", self.old_root); @@ -223,6 +223,7 @@ impl MptUpdates { } log::debug!("fill_state_roots done"); self.pretty_print(); + wit_gen } fn fill_state_roots_from_generator( diff --git a/zkevm-circuits/src/witness/mpt/witness.rs b/zkevm-circuits/src/witness/mpt/witness.rs index 498c88b91d..7054997c6e 100644 --- a/zkevm-circuits/src/witness/mpt/witness.rs +++ b/zkevm-circuits/src/witness/mpt/witness.rs @@ -9,9 +9,9 @@ use mpt_circuits::{ }; use mpt_zktrie::{ extend_address_to_h256, state::StorageData, AccountData, BytesArray, CanRead, TrieProof, - ZkTrie, ZkTrieNode, ZktrieState, SECURE_HASH_DOMAIN, + ZkMemoryDb, ZkTrie, ZkTrieNode, ZktrieState, SECURE_HASH_DOMAIN, }; -use std::collections::HashMap; +use std::{collections::HashMap, rc::Rc}; use num_bigint::BigUint; use std::{ @@ -40,18 +40,25 @@ fn to_smt_account(acc: AccountData) -> SMTAccount { pub struct WitnessGenerator { trie: ZkTrie, storages_cache: HashMap, + ref_db: Rc, } impl From<&ZktrieState> for WitnessGenerator { fn from(state: &ZktrieState) -> Self { Self { - trie: state.zk_db.borrow_mut().new_trie(&state.trie_root).unwrap(), + trie: state.zk_db.new_trie(&state.trie_root).unwrap(), storages_cache: HashMap::new(), + ref_db: state.zk_db.clone(), } } } impl WitnessGenerator { + /// output all updated ZkTrie + pub fn into_updated_trie(self) -> impl Iterator { + std::iter::once(self.trie).chain(self.storages_cache.into_values()) + } + /// dump inner data for debugging pub fn dump<'a>(&self, addrs: impl Iterator) { for addr in addrs { @@ -74,17 +81,18 @@ impl WitnessGenerator { HexBytes(word_buf) }; - self.storages_cache - .get(&address) - .map(Clone::clone) - .or_else(|| { - self.trie - .get_account(address.as_bytes()) - .map(AccountData::from) - .and_then(|account| self.trie.get_db().new_trie(&account.storage_root.0)) - }) - .and_then(|trie| trie.prove(key.as_ref()).ok()) - .unwrap_or_default() + if let Some(trie) = self.storages_cache.get(&address) { + // trie is under updating + trie.prove(key.as_ref()).ok() + } else { + // create a temporary trie and prove it + self.trie + .get_account(address.as_bytes()) + .map(AccountData::from) + .and_then(|account| self.ref_db.new_trie(&account.storage_root.0)) + .and_then(|trie| trie.prove(key.as_ref()).ok()) + } + .unwrap_or_default() } fn fetch_storage_cache(&mut self, address: Address) -> Option<&mut ZkTrie> { let cache_entry = self.storages_cache.entry(address); @@ -95,8 +103,7 @@ impl WitnessGenerator { .map(AccountData::from) .and_then(|acc_data| { // all trie share the same underlay db, so we can create new trie here - let zk_db = self.trie.get_db(); - zk_db.new_trie(&acc_data.storage_root.0) + self.ref_db.new_trie(&acc_data.storage_root.0) }) .map(|trie| entry.insert(trie)) } diff --git a/zktrie/src/lib.rs b/zktrie/src/lib.rs index 18cafa09a6..40e9e4176c 100644 --- a/zktrie/src/lib.rs +++ b/zktrie/src/lib.rs @@ -5,9 +5,8 @@ /// the state modules include structures represent zktrie and helpers pub mod state; -pub use crate::state::ZktrieState; pub use state::builder::{ self, extend_address_to_h256, AccountData, AccountProof, BytesArray, CanRead, StorageProof, TrieProof, SECURE_HASH_DOMAIN, }; -pub use zktrie::{ZkTrie, ZkTrieNode}; +pub use state::{ZkMemoryDb, ZkTrie, ZkTrieNode, ZktrieState}; diff --git a/zktrie/src/state.rs b/zktrie/src/state.rs index 29da6d554d..9c0139f155 100644 --- a/zktrie/src/state.rs +++ b/zktrie/src/state.rs @@ -2,18 +2,21 @@ use eth_types::{Address, Hash, H256}; use std::{collections::HashSet, io::Error}; -pub use zktrie::{Hash as ZkTrieHash, ZkMemoryDb, ZkTrie, ZkTrieNode}; +pub use zktrie::{Hash as ZkTrieHash, ZkMemoryDb, ZkTrieNode}; +use zktrie::{UpdateDb, ZkTrie as ZktrieT}; +/// export updatable type of zktrie +pub type ZkTrie = ZktrieT; pub mod builder; pub use builder::{AccountData, StorageData}; -use std::{cell::RefCell, fmt, rc::Rc}; +use std::{fmt, rc::Rc}; /// represent a storage state being applied in specified block #[derive(Clone)] pub struct ZktrieState { /// The underlying db - pub zk_db: RefCell>, + pub zk_db: Rc, /// Trie root pub trie_root: ZkTrieHash, addr_cache: HashSet
, @@ -49,13 +52,31 @@ impl ZktrieState { builder::init_hash_scheme(); Self { - zk_db: RefCell::new(ZkMemoryDb::new()), + zk_db: Rc::new(ZkMemoryDb::new()), trie_root: state_root.0, addr_cache: HashSet::new(), storage_cache: HashSet::new(), } } + /// expose writable db, has to be used when no trie is created + pub fn expose_db(&mut self) -> &mut ZkMemoryDb { + Rc::get_mut(&mut self.zk_db).expect("no extra reference") + } + + /// apply the updates in Zktries into underlying db + pub fn updated_with_trie(&mut self, tries: impl Iterator) { + let updated_data = tries.map(|tr| tr.updated_db()).collect::>(); + // note we must first collect the dbs (and drop all reference to current underlying db) + // or `exposed_db` would fail + updated_data + .into_iter() + .fold(self.expose_db(), |db, merged_db| { + db.update(merged_db); + db + }); + } + /// prepare to switch to another root state (trie snapshot) /// it is ok that even the db is not ready for this state /// cache is cleared so user can fill db with new storage traces @@ -68,11 +89,14 @@ impl ZktrieState { /// switch to another root state (trie snapshot) /// return true if the switch success, or false if db have not contain /// corresponding root yet + /// if we have built ZkTrie from the underlying db and updated it + /// (like in mpt witness updating), the updated Zktrie must be merged back + /// to the underlying db by `updated_with_trie` /// notice the cached key would not be clean if we can successfully switch to /// new snapshot since we consider it is not need to send more nodes data /// from storage trace for the updated leaves pub fn switch_to(&mut self, new_root: ZkTrieHash) -> bool { - let test_trie = self.zk_db.borrow_mut().new_trie(&new_root); + let test_trie = self.zk_db.new_trie(&new_root); if test_trie.is_none() { return false; } @@ -80,6 +104,44 @@ impl ZktrieState { true } + /// query a series account data from underlying db + pub fn query_accounts<'d: 'a, 'a>( + &self, + accounts: impl Iterator + 'd, + ) -> impl Iterator)> + 'a { + let trie = self.zk_db.new_ref_trie(&self.trie_root).unwrap(); + accounts.map(move |&addr| { + let account = trie.get_account(addr.as_bytes()).map(AccountData::from); + (addr, account) + }) + } + + /// query a series store data from underlying db + pub fn query_storages<'d: 'a, 'a>( + &self, + storages: impl Iterator + 'd, + ) -> impl Iterator)> + 'a { + use std::collections::{hash_map::Entry::*, HashMap}; + let zk_db = self.zk_db.clone(); + let account_trie = zk_db.new_ref_trie(&self.trie_root).unwrap(); + let mut trie_cache = 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_ref_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 + 'd, @@ -144,9 +206,9 @@ impl ZktrieState { .flat_map(|(_, _, bytes)| bytes), ) .chain(additional_proofs); - let mut zk_db = self.zk_db.borrow_mut(); + let zk_db = Rc::get_mut(&mut self.zk_db).expect("no extra reference"); for bytes in proofs { - zk_db.add_node_bytes(bytes).unwrap(); + zk_db.add_node_data(bytes).unwrap(); } } @@ -172,6 +234,6 @@ impl ZktrieState { /// get the inner zk memory db pub fn into_inner(self) -> Rc { - self.zk_db.into_inner() + self.zk_db } }