From 3c5a5b14ea66eb6b1ec4982866f4ced0cb616200 Mon Sep 17 00:00:00 2001 From: noelwei Date: Sun, 4 Aug 2024 16:56:41 +0900 Subject: [PATCH 1/8] induce "query" mode for building sdb Signed-off-by: noelwei --- bus-mapping/src/circuit_input_builder/l2.rs | 130 ++++++++++++++------ zktrie/src/state.rs | 38 ++++++ 2 files changed, 132 insertions(+), 36 deletions(-) diff --git a/bus-mapping/src/circuit_input_builder/l2.rs b/bus-mapping/src/circuit_input_builder/l2.rs index 81b9763b50..dbbf0eaf1a 100644 --- a/bus-mapping/src/circuit_input_builder/l2.rs +++ b/bus-mapping/src/circuit_input_builder/l2.rs @@ -148,21 +148,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(); @@ -203,40 +239,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)| { + 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, 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()); + } + 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/zktrie/src/state.rs b/zktrie/src/state.rs index 29da6d554d..5e9d179712 100644 --- a/zktrie/src/state.rs +++ b/zktrie/src/state.rs @@ -80,6 +80,44 @@ impl ZktrieState { true } + /// + pub fn query_accounts<'d: 'a, 'a>( + &self, + accounts: impl Iterator + 'd, + ) -> impl Iterator)> + '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 + 'd, + ) -> impl Iterator)> + '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 = 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 + 'd, From a2656fb7bca229512f56fa28355879a7c6f8b586 Mon Sep 17 00:00:00 2001 From: noelwei Date: Sun, 4 Aug 2024 20:32:10 +0900 Subject: [PATCH 2/8] switch zktrie dep Signed-off-by: noelwei --- Cargo.lock | 4 ++-- zktrie/Cargo.toml | 2 +- zktrie/src/state.rs | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 98598578d4..cc89618bbb 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=feat/node_key_importing#13792d9d2dfbc8f40f9ee9dec8ce254abefdb6e4" 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=feat/node_key_importing#13792d9d2dfbc8f40f9ee9dec8ce254abefdb6e4" dependencies = [ "hex", "lazy_static", diff --git a/zktrie/Cargo.toml b/zktrie/Cargo.toml index 578badba58..c8762a75d3 100644 --- a/zktrie/Cargo.toml +++ b/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 diff --git a/zktrie/src/state.rs b/zktrie/src/state.rs index 5e9d179712..c60edf723a 100644 --- a/zktrie/src/state.rs +++ b/zktrie/src/state.rs @@ -184,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(); } } From fe94845fb3b28698903a0394df7bd8a7a5721a88 Mon Sep 17 00:00:00 2001 From: noelwei Date: Mon, 5 Aug 2024 01:03:04 +0900 Subject: [PATCH 3/8] induce flatten proof and integration test Signed-off-by: noelwei --- bus-mapping/src/circuit_input_builder/l2.rs | 23 +++++++++- eth-types/src/l2_types.rs | 3 ++ integration-tests/tests/l2_trace.rs | 51 +++++++++++++++++++++ 3 files changed, 75 insertions(+), 2 deletions(-) create mode 100644 integration-tests/tests/l2_trace.rs diff --git a/bus-mapping/src/circuit_input_builder/l2.rs b/bus-mapping/src/circuit_input_builder/l2.rs index dbbf0eaf1a..ddb06c07e8 100644 --- a/bus-mapping/src/circuit_input_builder/l2.rs +++ b/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), @@ -227,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), diff --git a/eth-types/src/l2_types.rs b/eth-types/src/l2_types.rs index f912d2652b..c859ed8109 100644 --- a/eth-types/src/l2_types.rs +++ b/eth-types/src/l2_types.rs @@ -451,6 +451,9 @@ pub struct StorageTrace { #[serde(rename = "deletionProofs", default)] /// additional deletion proofs pub deletion_proofs: Vec, + #[serde(rename = "flattenProofs", default)] + /// + pub flatten_proofs: 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); +} From 7677396fc6c3fc8eccd44a5937b917945d34b0dd Mon Sep 17 00:00:00 2001 From: Ho Vei Date: Sun, 11 Aug 2024 22:45:00 +0800 Subject: [PATCH 4/8] update zk_dep use key cache --- Cargo.lock | 4 +-- bus-mapping/src/circuit_input_builder/l2.rs | 38 ++++++++++++++++++--- eth-types/src/l2_types.rs | 6 ++++ zkevm-circuits/src/witness/mpt/witness.rs | 34 +++++++++--------- zktrie/src/lib.rs | 3 +- zktrie/src/state.rs | 32 ++++++++++------- 6 files changed, 80 insertions(+), 37 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index cc89618bbb..9d96052e24 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=feat/node_key_importing#13792d9d2dfbc8f40f9ee9dec8ce254abefdb6e4" +source = "git+https://github.com/scroll-tech/zktrie.git?branch=feat/node_key_importing#966b79453043e3a05a67701dd0db86a1431fa356" 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=feat/node_key_importing#13792d9d2dfbc8f40f9ee9dec8ce254abefdb6e4" +source = "git+https://github.com/scroll-tech/zktrie.git?branch=feat/node_key_importing#966b79453043e3a05a67701dd0db86a1431fa356" 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 ddb06c07e8..a585b4772e 100644 --- a/bus-mapping/src/circuit_input_builder/l2.rs +++ b/bus-mapping/src/circuit_input_builder/l2.rs @@ -126,12 +126,26 @@ impl CircuitInputBuilder { 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(); + 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(); } - drop(zk_db); + 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())), + ); + Some(state) } else if !light_mode { let mpt_init_state = ZktrieState::from_trace_with_additional( @@ -239,13 +253,27 @@ impl CircuitInputBuilder { if !l2_trace.storage_trace.flatten_proofs.is_empty() { let mpt_state = self .mpt_init_state - .as_ref() + .as_mut() .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(); + 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), diff --git a/eth-types/src/l2_types.rs b/eth-types/src/l2_types.rs index c859ed8109..ad22a57459 100644 --- a/eth-types/src/l2_types.rs +++ b/eth-types/src/l2_types.rs @@ -454,6 +454,12 @@ pub struct StorageTrace { #[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/zkevm-circuits/src/witness/mpt/witness.rs b/zkevm-circuits/src/witness/mpt/witness.rs index 498c88b91d..19711c3c3a 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,13 +40,15 @@ 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(), } } } @@ -74,17 +76,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 +98,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 c60edf723a..4120f077d5 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,18 @@ 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") + } + /// 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 @@ -72,7 +80,7 @@ impl ZktrieState { /// 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; } @@ -85,7 +93,7 @@ impl ZktrieState { &self, accounts: impl Iterator + 'd, ) -> impl Iterator)> + 'a { - let trie = self.zk_db.borrow_mut().new_trie(&self.trie_root).unwrap(); + 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) @@ -98,9 +106,9 @@ impl ZktrieState { storages: impl Iterator + 'd, ) -> impl Iterator)> + '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 = HashMap::new(); + 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()), @@ -109,7 +117,7 @@ impl ZktrieState { .map(AccountData::from) .and_then(|account| { zk_db - .new_trie(&account.storage_root.0) + .new_ref_trie(&account.storage_root.0) .map(|tr| entry.insert(tr)) }), } @@ -182,7 +190,7 @@ 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_data(bytes).unwrap(); } @@ -210,6 +218,6 @@ impl ZktrieState { /// get the inner zk memory db pub fn into_inner(self) -> Rc { - self.zk_db.into_inner() + self.zk_db } } From fffc7c31140eb05554d13c01db6309ed24e8465e Mon Sep 17 00:00:00 2001 From: Ho Vei Date: Mon, 12 Aug 2024 12:39:29 +0800 Subject: [PATCH 5/8] fix test issue --- bus-mapping/src/circuit_input_builder/l2.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/bus-mapping/src/circuit_input_builder/l2.rs b/bus-mapping/src/circuit_input_builder/l2.rs index a585b4772e..a12f341916 100644 --- a/bus-mapping/src/circuit_input_builder/l2.rs +++ b/bus-mapping/src/circuit_input_builder/l2.rs @@ -146,6 +146,11 @@ impl CircuitInputBuilder { .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( @@ -179,7 +184,7 @@ impl CircuitInputBuilder { 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); + sdb.set_account(&addr, state_db::Account::zero()); } } From dd22a370dcf4dde7e9e42aeff325798209632d35 Mon Sep 17 00:00:00 2001 From: Ho Vei Date: Tue, 13 Aug 2024 11:21:58 +0800 Subject: [PATCH 6/8] update zktrie dep --- Cargo.lock | 4 ++-- zktrie/Cargo.toml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9d96052e24..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=feat/node_key_importing#966b79453043e3a05a67701dd0db86a1431fa356" +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=feat/node_key_importing#966b79453043e3a05a67701dd0db86a1431fa356" +source = "git+https://github.com/scroll-tech/zktrie.git?branch=main#49e0f027d12abf7f2049c2cf1034128fc14f47dd" dependencies = [ "hex", "lazy_static", diff --git a/zktrie/Cargo.toml b/zktrie/Cargo.toml index c8762a75d3..578badba58 100644 --- a/zktrie/Cargo.toml +++ b/zktrie/Cargo.toml @@ -8,7 +8,7 @@ license.workspace = true [dependencies] halo2curves.workspace = true -zktrie = { git = "https://github.com/scroll-tech/zktrie.git", branch = "feat/node_key_importing", features= ["rs_zktrie"] } +zktrie = { git = "https://github.com/scroll-tech/zktrie.git", branch = "main", features= ["rs_zktrie"] } poseidon-base.workspace = true eth-types = { path = "../eth-types" } num-bigint.workspace = true From 69527c461ffc9186401ad7bf59766785259c3b6d Mon Sep 17 00:00:00 2001 From: Ho Vei Date: Tue, 13 Aug 2024 11:59:45 +0800 Subject: [PATCH 7/8] induce mpt state updating --- prover/src/zkevm/circuit/builder.rs | 2 +- zkevm-circuits/src/witness/block.rs | 11 +++++++++++ zkevm-circuits/src/witness/mpt.rs | 3 ++- zkevm-circuits/src/witness/mpt/witness.rs | 5 +++++ zktrie/src/state.rs | 20 ++++++++++++++++++-- 5 files changed, 37 insertions(+), 4 deletions(-) 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 19711c3c3a..7054997c6e 100644 --- a/zkevm-circuits/src/witness/mpt/witness.rs +++ b/zkevm-circuits/src/witness/mpt/witness.rs @@ -54,6 +54,11 @@ impl From<&ZktrieState> for WitnessGenerator { } 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 { diff --git a/zktrie/src/state.rs b/zktrie/src/state.rs index 4120f077d5..9c0139f155 100644 --- a/zktrie/src/state.rs +++ b/zktrie/src/state.rs @@ -64,6 +64,19 @@ impl ZktrieState { 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 @@ -76,6 +89,9 @@ 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 @@ -88,7 +104,7 @@ impl ZktrieState { true } - /// + /// query a series account data from underlying db pub fn query_accounts<'d: 'a, 'a>( &self, accounts: impl Iterator + 'd, @@ -100,7 +116,7 @@ impl ZktrieState { }) } - /// + /// query a series store data from underlying db pub fn query_storages<'d: 'a, 'a>( &self, storages: impl Iterator + 'd, From 64883536d93e8986760d5e34c8fec1411d51cf91 Mon Sep 17 00:00:00 2001 From: Ho Vei Date: Tue, 13 Aug 2024 16:58:49 +0800 Subject: [PATCH 8/8] fix issue in chunk building --- bus-mapping/src/circuit_input_builder/l2.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/bus-mapping/src/circuit_input_builder/l2.rs b/bus-mapping/src/circuit_input_builder/l2.rs index a12f341916..8d36f89518 100644 --- a/bus-mapping/src/circuit_input_builder/l2.rs +++ b/bus-mapping/src/circuit_input_builder/l2.rs @@ -184,6 +184,7 @@ impl CircuitInputBuilder { 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()); } } @@ -301,9 +302,7 @@ impl CircuitInputBuilder { 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.insert(addr, acc.unwrap_or_default()); m }) } else { @@ -334,6 +333,8 @@ impl CircuitInputBuilder { .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 })