From 9f4b54f66929d3b25b334b7d3f1ea63e34519ba8 Mon Sep 17 00:00:00 2001
From: Jonathan LEI <xJonathan@outlook.com>
Date: Sat, 14 Sep 2024 00:56:34 +0800
Subject: [PATCH] refactor: generic `WitnessInput` trait (#52)

---
 crates/executor/client/src/io.rs  | 85 ++++++++++++++++++++++++++-----
 crates/executor/client/src/lib.rs |  3 +-
 2 files changed, 73 insertions(+), 15 deletions(-)

diff --git a/crates/executor/client/src/io.rs b/crates/executor/client/src/io.rs
index c7e6dbc..7d549d7 100644
--- a/crates/executor/client/src/io.rs
+++ b/crates/executor/client/src/io.rs
@@ -31,30 +31,90 @@ pub struct ClientExecutorInput {
 
 impl ClientExecutorInput {
     /// Gets the immediate parent block's header.
+    #[inline(always)]
     pub fn parent_header(&self) -> &Header {
         &self.ancestor_headers[0]
     }
 
-    /// Creates a [WitnessDb] from a [ClientExecutorInput]. To do so, it verifies the state root,
-    /// ancestor headers and account bytecodes, and constructs the account and storage values by
-    /// reading against state tries.
+    /// Creates a [`WitnessDb`].
     pub fn witness_db(&self) -> Result<WitnessDb> {
-        let state_root: B256 = self.parent_header().state_root;
-        if state_root != self.parent_state.state_root() {
+        <Self as WitnessInput>::witness_db(self)
+    }
+}
+
+impl WitnessInput for ClientExecutorInput {
+    #[inline(always)]
+    fn state(&self) -> &EthereumState {
+        &self.parent_state
+    }
+
+    #[inline(always)]
+    fn state_anchor(&self) -> B256 {
+        self.parent_header().state_root
+    }
+
+    #[inline(always)]
+    fn state_requests(&self) -> impl Iterator<Item = (&Address, &Vec<U256>)> {
+        self.state_requests.iter()
+    }
+
+    #[inline(always)]
+    fn bytecodes(&self) -> impl Iterator<Item = &Bytecode> {
+        self.bytecodes.iter()
+    }
+
+    #[inline(always)]
+    fn headers(&self) -> impl Iterator<Item = &Header> {
+        once(&self.current_block.header).chain(self.ancestor_headers.iter())
+    }
+}
+
+/// A trait for constructing [`WitnessDb`].
+pub trait WitnessInput {
+    /// Gets a reference to the state from which account info and storage slots are loaded.
+    fn state(&self) -> &EthereumState;
+
+    /// Gets the state trie root hash that the state referenced by
+    /// [state()](trait.WitnessInput#tymethod.state) must conform to.
+    fn state_anchor(&self) -> B256;
+
+    /// Gets an iterator over address state requests. For each request, the account info and storage
+    /// slots are loaded from the relevant tries in the state returned by
+    /// [state()](trait.WitnessInput#tymethod.state).
+    fn state_requests(&self) -> impl Iterator<Item = (&Address, &Vec<U256>)>;
+
+    /// Gets an iterator over account bytecodes.
+    fn bytecodes(&self) -> impl Iterator<Item = &Bytecode>;
+
+    /// Gets an iterator over references to a consecutive, reverse-chronological block headers
+    /// starting from the current block header.
+    fn headers(&self) -> impl Iterator<Item = &Header>;
+
+    /// Creates a [`WitnessDb`] from a [`WitnessInput`] implementation. To do so, it verifies the
+    /// state root, ancestor headers and account bytecodes, and constructs the account and
+    /// storage values by reading against state tries.
+    ///
+    /// NOTE: For some unknown reasons, calling this trait method directly from outside of the type
+    /// implementing this trait causes a zkVM run to cost over 5M cycles more. To avoid this, define
+    /// a method inside the type that calls this trait method instead.
+    #[inline(always)]
+    fn witness_db(&self) -> Result<WitnessDb> {
+        let state = self.state();
+
+        if self.state_anchor() != state.state_root() {
             eyre::bail!("parent state root mismatch");
         }
 
         let bytecodes_by_hash =
-            self.bytecodes.iter().map(|code| (code.hash_slow(), code)).collect::<HashMap<_, _>>();
+            self.bytecodes().map(|code| (code.hash_slow(), code)).collect::<HashMap<_, _>>();
 
         let mut accounts = HashMap::new();
         let mut storage = HashMap::new();
-        for (&address, slots) in self.state_requests.iter() {
+        for (&address, slots) in self.state_requests() {
             let hashed_address = keccak256(address);
             let hashed_address = hashed_address.as_slice();
 
-            let account_in_trie =
-                self.parent_state.state_trie.get_rlp::<TrieAccount>(hashed_address)?;
+            let account_in_trie = state.state_trie.get_rlp::<TrieAccount>(hashed_address)?;
 
             accounts.insert(
                 address,
@@ -78,8 +138,7 @@ impl ClientExecutorInput {
             if !slots.is_empty() {
                 let mut address_storage = HashMap::new();
 
-                let storage_trie = self
-                    .parent_state
+                let storage_trie = state
                     .storage_tries
                     .get(hashed_address)
                     .ok_or_else(|| eyre::eyre!("parent state does not contain storage trie"))?;
@@ -97,9 +156,7 @@ impl ClientExecutorInput {
 
         // Verify and build block hashes
         let mut block_hashes: HashMap<u64, B256> = HashMap::new();
-        for (child_header, parent_header) in
-            once(&self.current_block.header).chain(self.ancestor_headers.iter()).tuple_windows()
-        {
+        for (child_header, parent_header) in self.headers().tuple_windows() {
             if parent_header.number != child_header.number - 1 {
                 eyre::bail!("non-consecutive blocks");
             }
diff --git a/crates/executor/client/src/lib.rs b/crates/executor/client/src/lib.rs
index 34a1e6c..ed0bb98 100644
--- a/crates/executor/client/src/lib.rs
+++ b/crates/executor/client/src/lib.rs
@@ -1,6 +1,7 @@
+/// Client program input data types.
 pub mod io;
 #[macro_use]
-pub mod utils;
+mod utils;
 
 pub mod custom;