refactor ZktrieState [dup #903]

bus-mapping/src/circuit_input_builder.rs

@@ -179,7 +179,7 @@ pub struct CircuitInputBuilder {
     pub block_ctx: BlockContext,
     #[cfg(feature = "scroll")]
     /// Initial Zktrie Status for a incremental updating
-    pub mpt_init_state: ZktrieState,
+    pub mpt_init_state: Option<ZktrieState>,
 }
 
 impl<'a> CircuitInputBuilder {

bus-mapping/src/circuit_input_builder/l2.rs

@@ -8,11 +8,11 @@ use eth_types::{
     self,
     evm_types::OpcodeId,
     l2_types::{BlockTrace, EthBlock, ExecStep, StorageTrace},
-    Address, ToAddress, Word, H256,
+    Address, ToAddress, ToWord, Word, H256,
 };
-use ethers_core::types::{Bytes, U256};
+use ethers_core::types::Bytes;
 use mpt_zktrie::state::{AccountData, ZktrieState};
-use std::collections::hash_map::Entry;
+use std::collections::hash_map::{Entry, HashMap};
 
 impl From<&AccountData> for state_db::Account {
     fn from(acc_data: &AccountData) -> Self {
@@ -31,29 +31,6 @@ impl From<&AccountData> for state_db::Account {
     }
 }
 
-impl From<&ZktrieState> for StateDB {
-    fn from(mpt_state: &ZktrieState) -> Self {
-        let mut sdb = StateDB::new();
-
-        for (addr, acc) in mpt_state.state() {
-            sdb.set_account(addr, acc.into())
-        }
-
-        for (storage_key, data) in mpt_state.storage() {
-            // Since the StateDB is a partical db, 0 means we know it is zero instead of "unknown".
-            log::trace!(
-                "trace sdb: addr {:?} key {:?} value {:?}",
-                storage_key.0,
-                storage_key.1,
-                *data.as_ref()
-            );
-            let (_, acc) = sdb.get_account_mut(&storage_key.0);
-            acc.storage.insert(storage_key.1, *data.as_ref());
-        }
-        sdb
-    }
-}
-
 fn decode_bytecode(bytecode: &str) -> Result<Vec<u8>, Error> {
     let mut stripped = if let Some(stripped) = bytecode.strip_prefix("0x") {
         stripped.to_string()
@@ -302,7 +279,7 @@ impl CircuitInputBuilder {
             code_db,
             block: block.clone(),
             block_ctx: BlockContext::new(),
-            mpt_init_state,
+            mpt_init_state: Some(mpt_init_state),
         }
     }
 
@@ -322,25 +299,43 @@ impl CircuitInputBuilder {
             hex::encode(old_root),
         );
 
-        let mpt_init_state = ZktrieState::from_trace_with_additional(
-            old_root,
-            Self::collect_account_proofs(&l2_trace.storage_trace),
-            Self::collect_storage_proofs(&l2_trace.storage_trace),
-            l2_trace
-                .storage_trace
-                .deletion_proofs
-                .iter()
-                .map(Bytes::as_ref),
-            light_mode,
-        )
-        .unwrap();
+        let mpt_init_state = if !light_mode {
+            let mpt_init_state = ZktrieState::from_trace_with_additional(
+                old_root,
+                Self::collect_account_proofs(&l2_trace.storage_trace),
+                Self::collect_storage_proofs(&l2_trace.storage_trace),
+                l2_trace
+                    .storage_trace
+                    .deletion_proofs
+                    .iter()
+                    .map(Bytes::as_ref),
+            )
+            .map_err(Error::IoError)?;
 
-        log::debug!(
-            "building partial statedb done, root {}",
-            hex::encode(mpt_init_state.root())
-        );
+            log::debug!(
+                "building partial statedb done, root {}",
+                hex::encode(mpt_init_state.root())
+            );
 
-        let sdb = StateDB::from(&mpt_init_state);
+            Some(mpt_init_state)
+        } else {
+            None
+        };
+
+        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)?;
+            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)?;
+            *sdb.get_storage_mut(&addr, &key).1 = val.into();
+        }
 
         /*
         let (zero_coinbase_exist, _) = sdb.get_account(&Default::default());
@@ -355,7 +350,7 @@ impl CircuitInputBuilder {
 
         let mut builder_block = circuit_input_builder::Block::from_headers(&[], circuits_params);
         builder_block.chain_id = chain_id;
-        builder_block.prev_state_root = U256::from(mpt_init_state.root());
+        builder_block.prev_state_root = old_root.to_word();
         builder_block.start_l1_queue_index = l2_trace.start_l1_queue_index;
         let mut builder = Self {
             sdb,
@@ -370,15 +365,10 @@ impl CircuitInputBuilder {
     }
 
     /// ...
-    pub fn add_more_l2_trace(
-        &mut self,
-        l2_trace: &BlockTrace,
-        more: bool,
-        light_mode: bool,
-    ) -> Result<(), Error> {
-        // update sdb for new data from storage
-        if !light_mode {
-            self.mpt_init_state.update_nodes_from_proofs(
+    pub fn add_more_l2_trace(&mut self, l2_trace: &BlockTrace, more: bool) -> Result<(), Error> {
+        // update init state new data from storage
+        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),
                 l2_trace
@@ -389,25 +379,45 @@ impl CircuitInputBuilder {
             );
         }
 
-        self.mpt_init_state
-            .update_account_from_proofs(
-                Self::collect_account_proofs(&l2_trace.storage_trace),
-                |addr, acc_data| {
-                    self.sdb.set_account(addr, acc_data.into());
-                    Ok(())
-                },
-            )
-            .map_err(Error::IoError)?;
+        let new_accounts = ZktrieState::parse_account_from_proofs(
+            Self::collect_account_proofs(&l2_trace.storage_trace).filter(|(addr, _)| {
+                let (existed, _) = self.sdb.get_account(addr);
+                !existed
+            }),
+        )
+        .fold(
+            Ok(HashMap::new()),
+            |m, parsed| -> Result<HashMap<_, _>, Error> {
+                let mut m = m?;
+                let (addr, acc) = parsed.map_err(Error::IoError)?;
+                m.insert(addr, acc);
+                Ok(m)
+            },
+        )?;
 
-        self.mpt_init_state
-            .update_storage_from_proofs(
-                Self::collect_storage_proofs(&l2_trace.storage_trace),
-                |storage_key, value| {
-                    *self.sdb.get_storage_mut(&storage_key.0, &storage_key.1).1 = *value.as_ref();
-                    Ok(())
-                },
-            )
-            .map_err(Error::IoError)?;
+        for (addr, acc) in new_accounts {
+            self.sdb.set_account(&addr, state_db::Account::from(&acc));
+        }
+
+        let new_storages = ZktrieState::parse_storage_from_proofs(
+            Self::collect_storage_proofs(&l2_trace.storage_trace).filter(|(addr, key, _)| {
+                let (existed, _) = self.sdb.get_committed_storage(addr, key);
+                !existed
+            }),
+        )
+        .fold(
+            Ok(HashMap::new()),
+            |m, parsed| -> Result<HashMap<(Address, Word), Word>, Error> {
+                let mut m = m?;
+                let ((addr, key), val) = parsed.map_err(Error::IoError)?;
+                m.insert((addr, key), val.into());
+                Ok(m)
+            },
+        )?;
+
+        for ((addr, key), val) in new_storages {
+            *self.sdb.get_storage_mut(&addr, &key).1 = val;
+        }
 
         update_codedb(&mut self.code_db, &self.sdb, l2_trace)?;
 

prover/src/zkevm/capacity_checker.rs

@@ -119,7 +119,7 @@ pub struct CircuitCapacityChecker {
     pub light_mode: bool,
     pub acc_row_usage: RowUsage,
     pub row_usages: Vec<RowUsage>,
-    pub builder_ctx: Option<(CodeDB, StateDB, ZktrieState)>,
+    pub builder_ctx: Option<(CodeDB, StateDB, Option<ZktrieState>)>,
 }
 
 // Currently TxTrace is same as BlockTrace, with "transactions" and "executionResults" should be of
@@ -177,17 +177,26 @@ impl CircuitCapacityChecker {
                     circuit_input_builder::Block::from_headers(&[], get_super_circuit_params());
                 builder_block.chain_id = txs[0].chain_id;
                 builder_block.start_l1_queue_index = txs[0].start_l1_queue_index;
-                builder_block.prev_state_root = H256(*mpt_state.root()).to_word();
+                builder_block.prev_state_root = mpt_state
+                    .as_ref()
+                    .map(|state| state.root())
+                    .map(|root| H256(*root))
+                    .unwrap_or(txs[0].header.state_root)
+                    .to_word();
                 // notice the trace has included all code required for builidng witness block,
                 // so we do not need to pick them from previous one, but we still keep the
                 // old codedb in previous run for some dedup work
-                let mut builder = CircuitInputBuilder::new_with_trie_state(
-                    sdb,
-                    CodeDB::new(),
-                    mpt_state,
-                    &builder_block,
-                );
-                builder.add_more_l2_trace(&txs[0], txs.len() > 1, self.light_mode)?;
+                let mut builder = if let Some(mpt_state) = mpt_state {
+                    CircuitInputBuilder::new_with_trie_state(
+                        sdb,
+                        CodeDB::new(),
+                        mpt_state,
+                        &builder_block,
+                    )
+                } else {
+                    CircuitInputBuilder::new(sdb, CodeDB::new(), &builder_block)
+                };
+                builder.add_more_l2_trace(&txs[0], txs.len() > 1)?;
                 (builder, Some(code_db))
             } else {
                 (
@@ -201,11 +210,8 @@ impl CircuitCapacityChecker {
                 )
             };
         let traces = &txs[1..];
-        let witness_block = block_traces_to_witness_block_with_updated_state(
-            traces,
-            &mut estimate_builder,
-            self.light_mode,
-        )?;
+        let witness_block =
+            block_traces_to_witness_block_with_updated_state(traces, &mut estimate_builder)?;
         let mut rows = calculate_row_usage_of_witness_block(&witness_block)?;
 
         let mut code_db = codedb_prev.unwrap_or_else(CodeDB::new);

prover/src/zkevm/circuit/l2_builder.rs

@@ -8,7 +8,7 @@ use bus_mapping::{
 use eth_types::{l2_types::BlockTrace, ToWord, H256};
 use halo2_proofs::halo2curves::bn256::Fr;
 use itertools::Itertools;
-use mpt_zktrie::state::ZktrieState;
+use mpt_zktrie::state::{ZkTrieHash, ZktrieState};
 use once_cell::sync::Lazy;
 use std::time::Instant;
 use zkevm_circuits::{
@@ -198,19 +198,17 @@ fn prepare_default_builder(
     builder_block.prev_state_root = old_root.to_word();
     let code_db = CodeDB::new();
 
-    if let Some(mpt_state) = initial_mpt_state {
+    if let Some(mpt_state) = &initial_mpt_state {
         assert_eq!(
             H256::from_slice(mpt_state.root()),
             old_root,
             "the provided zktrie state must be the prev state"
         );
-        let state_db = StateDB::from(&mpt_state);
-        let mut builder = CircuitInputBuilder::new(state_db, code_db, &builder_block);
-        builder.mpt_init_state = mpt_state;
-        builder
-    } else {
-        CircuitInputBuilder::new(StateDB::new(), code_db, &builder_block)
     }
+
+    let mut builder = CircuitInputBuilder::new(StateDB::new(), code_db, &builder_block);
+    builder.mpt_init_state = initial_mpt_state;
+    builder
 }
 
 /// check if block traces match preset parameters
@@ -257,15 +255,15 @@ pub fn block_traces_to_witness_block(block_traces: &[BlockTrace]) -> Result<Bloc
     // etc, so the generated block maybe invalid without any message
     if block_traces.is_empty() {
         let mut builder = prepare_default_builder(eth_types::Hash::zero(), None);
-        block_traces_to_witness_block_with_updated_state(&[], &mut builder, false)
+        block_traces_to_witness_block_with_updated_state(&[], &mut builder)
     } else {
         let mut builder = CircuitInputBuilder::new_from_l2_trace(
             get_super_circuit_params(),
             &block_traces[0],
             block_traces.len() > 1,
             false,
         )?;
-        block_traces_to_witness_block_with_updated_state(&block_traces[1..], &mut builder, false)
+        block_traces_to_witness_block_with_updated_state(&block_traces[1..], &mut builder)
     }
 }
 
@@ -276,7 +274,6 @@ pub fn block_traces_to_witness_block(block_traces: &[BlockTrace]) -> Result<Bloc
 pub fn block_traces_to_witness_block_with_updated_state(
     block_traces: &[BlockTrace],
     builder: &mut CircuitInputBuilder,
-    light_mode: bool,
 ) -> Result<Block<Fr>> {
     let metric = |builder: &CircuitInputBuilder, idx: usize| -> Result<(), bus_mapping::Error> {
         let t = Instant::now();
@@ -321,7 +318,7 @@ pub fn block_traces_to_witness_block_with_updated_state(
             "add_more_l2_trace idx {idx}, block num {:?}",
             block_trace.header.number
         );
-        builder.add_more_l2_trace(block_trace, !is_last, false)?;
+        builder.add_more_l2_trace(block_trace, !is_last)?;
         if per_block_metric {
             metric(builder, idx + initial_blk_index)?;
         }
@@ -338,14 +335,24 @@ pub fn block_traces_to_witness_block_with_updated_state(
         witness_block.circuits_params
     );
 
-    if !light_mode && *builder.mpt_init_state.root() != [0u8; 32] {
-        log::debug!("block_apply_mpt_state");
-        block_apply_mpt_state(&mut witness_block, &builder.mpt_init_state);
-        log::debug!("block_apply_mpt_state done");
+    if let Some(state) = &mut builder.mpt_init_state {
+        if *state.root() != [0u8; 32] {
+            log::debug!("block_apply_mpt_state");
+            block_apply_mpt_state(&mut witness_block, state);
+            log::debug!("block_apply_mpt_state done");
+        };
+        let root_after = witness_block.state_root.unwrap_or_default();
+
+        log::debug!(
+            "finish replay trie updates, root {}, root after {:#x?}",
+            hex::encode(state.root()),
+            root_after,
+        );
+        // switch state to new root
+        let mut new_root_hash = ZkTrieHash::default();
+        root_after.to_big_endian(&mut new_root_hash);
+        assert!(state.switch_to(new_root_hash));
     }
-    log::debug!(
-        "finish replay trie updates, root {}",
-        hex::encode(builder.mpt_init_state.root())
-    );
+
     Ok(witness_block)
 }

testool/src/statetest/executor.rs

@@ -314,7 +314,10 @@ fn trace_config_to_witness_block_l2(
         .expect("could not finalize building block");
     let mut block =
         zkevm_circuits::witness::block_convert(&builder.block, &builder.code_db).unwrap();
-    zkevm_circuits::witness::block_apply_mpt_state(&mut block, &builder.mpt_init_state);
+    zkevm_circuits::witness::block_apply_mpt_state(
+        &mut block,
+        builder.mpt_init_state.as_ref().unwrap(),
+    );
     // as mentioned above, we cannot fit the trace into circuit
     // stop here
     if exceed_max_steps != 0 {