induce mpt state updating

prover/src/zkevm/circuit/builder.rs

@@ -167,7 +167,7 @@ pub fn finalize_builder(builder: &mut CircuitInputBuilder) -> Result<Block> {
     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?

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) {

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(

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<Item = ZkTrie> {
+        std::iter::once(self.trie).chain(self.storages_cache.into_values())
+    }
+
     /// dump inner data for debugging
     pub fn dump<'a>(&self, addrs: impl Iterator<Item = &'a Address>) {
         for addr in addrs {

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<Item = ZkTrie>) {
+        let updated_data = tries.map(|tr| tr.updated_db()).collect::<Vec<_>>();
+        // 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<Item = &'a Address> + '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<Item = (&'a Address, &'a H256)> + 'd,