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,