add withdraw_proof to MptUpdates (#678)

* withdraw proof

* Update witness.rs

* Update mpt.rs

Author: lispc

Cargo.lock

@@ -28,6 +28,7 @@ mpt-zktrie = { path = "../zktrie" }
 keccak256 = { path = "../keccak256"}
 log = "0.4"
 env_logger = "0.9"
+serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0.78"
 
 hash-circuit = { package = "poseidon-circuit", git = "https://github.com/scroll-tech/poseidon-circuit.git", branch = "scroll-dev-0619", features=['short']}
@@ -58,8 +59,6 @@ itertools = "0.10.1"
 mock = { path = "../mock" }
 pretty_assertions = "1.0.0"
 cli-table = "0.4"
-serde = { version = "1.0.130", features = ["derive"] }
-serde_json = "1.0.78"
 paste = "1.0"
 
 [features]

zkevm-circuits/Cargo.toml

@@ -12,7 +12,7 @@ mod call;
 pub use call::Call;
 
 mod mpt;
-pub use mpt::{MptUpdate, MptUpdateRow, MptUpdates};
+pub use mpt::{MptUpdate, MptUpdateRow, MptUpdates, WithdrawProof};
 
 mod receipt;
 pub use receipt::Receipt;

zkevm-circuits/src/witness.rs

@@ -10,10 +10,22 @@ use mpt_zktrie::{
     state,
     state::witness::WitnessGenerator,
 };
+use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 
 pub use state::ZktrieState;
 
+/// Used to store withdraw proof
+#[derive(Default, Debug, Clone, Serialize, Deserialize)]
+pub struct WithdrawProof {
+    /// state root after a block
+    pub state_root: U256,
+    /// account proof for withdraw bridge contract
+    pub account_proof: Vec<Vec<u8>>,
+    /// storage proof for withdraw bridge contract, withdraw root storage key
+    pub storage_proof: Vec<Vec<u8>>,
+}
+
 /// An MPT update whose validity is proved by the MptCircuit
 #[derive(Debug, Clone)]
 pub struct MptUpdate {
@@ -56,6 +68,9 @@ pub struct MptUpdates {
     old_root: Word,
     new_root: Word,
     updates: BTreeMap<Key, MptUpdate>,
+    /// TODO: is here the best place for this?
+    /// Withdraw proof after this block
+    pub withdraw_proof: WithdrawProof,
     pub(crate) smt_traces: Vec<SMTTrace>,
     pub(crate) proof_types: Vec<MPTProofType>,
 }
@@ -135,6 +150,20 @@ impl MptUpdates {
             );
             wit_gen.dump();
         }
+
+        // generate withdraw proof
+        let address = *bus_mapping::l2_predeployed::message_queue::ADDRESS;
+        let key = *bus_mapping::l2_predeployed::message_queue::WITHDRAW_TRIE_ROOT_SLOT;
+        let account_proof = wit_gen.account_proof(address);
+        let storage_proof = wit_gen.storage_proof(address, key);
+        // TODO: add withdraw_root to WithdrawProof?
+        let withdraw_proof = WithdrawProof {
+            state_root: self.new_root,
+            account_proof,
+            storage_proof,
+        };
+        log::debug!("withdraw proof {withdraw_proof:?}");
+        self.withdraw_proof = withdraw_proof;
     }
 
     fn fill_state_roots_from_generator(
@@ -176,6 +205,7 @@ impl MptUpdates {
             self.old_root,
             self.new_root
         );
+
         wit_gen
     }
 

zkevm-circuits/src/witness/mpt.rs

@@ -99,7 +99,26 @@ impl WitnessGenerator {
     pub fn dump(&self) {
         log::info!("account data {:#?}", self.accounts);
     }
-
+    /// get account proof
+    pub fn account_proof(&self, address: Address) -> Vec<Vec<u8>> {
+        self.trie.prove(address.as_bytes()).unwrap()
+    }
+    /// get storage proof
+    pub fn storage_proof(&self, address: Address, key: Word) -> Vec<Vec<u8>> {
+        let (_storage_key, key) = {
+            let mut word_buf = [0u8; 32];
+            key.to_big_endian(word_buf.as_mut_slice());
+            (hash_zktrie_key(&word_buf), HexBytes(word_buf))
+        };
+        // TODO: use or_else to optimize
+        let default_trie = &ZktrieState::default()
+            .zk_db
+            .borrow_mut()
+            .new_trie(&ZkTrieHash::default())
+            .unwrap();
+        let trie: &ZkTrie = self.storages.get(&address).unwrap_or(default_trie);
+        trie.prove(key.as_ref()).unwrap()
+    }
     fn trace_storage_update(
         &mut self,
         address: Address,