[feat] add feature switch of memory/stack/storage (scroll-tech#1038)

* add feature switch

* add missing

* fix merge

bus-mapping/Cargo.toml

@@ -41,8 +41,12 @@ mock = { path = "../mock" }
 rand.workspace = true
 
 [features]
-default = ["test"]
+default = ["test", "enable-stack", "enable-storage"]
 test = ["mock", "rand"]
 scroll = ["eth-types/scroll", "mock?/scroll"]
 # Enable shanghai feature of mock only if mock is enabled (by test).
 shanghai = ["eth-types/shanghai", "mock?/shanghai"]
+tracer-tests = ["enable-memory"]
+enable-stack = ["eth-types/enable-stack", "mock?/enable-stack"]
+enable-memory = ["eth-types/enable-memory", "mock?/enable-memory"]
+enable-storage = ["eth-types/enable-storage", "mock?/enable-storage"]

bus-mapping/src/circuit_input_builder.rs

@@ -8,7 +8,7 @@ mod execution;
 mod input_state_ref;
 #[cfg(feature = "scroll")]
 mod l2;
-#[cfg(test)]
+#[cfg(all(feature = "tracer-tests", feature = "enable-memory", test))]
 mod tracer_tests;
 mod transaction;
 

bus-mapping/src/evm/opcodes.rs

@@ -12,7 +12,8 @@ use crate::{
 use core::fmt::Debug;
 use eth_types::{evm_unimplemented, GethExecStep, ToAddress, ToWord, Word};
 
-use crate::util::CHECK_MEM_STRICT;
+#[cfg(feature = "enable-memory")]
+use crate::util::GETH_TRACE_CHECK_LEVEL;
 
 #[cfg(any(feature = "test", test))]
 pub use self::sha3::sha3_tests::{gen_sha3_code, MemoryKind};
@@ -70,7 +71,7 @@ mod error_precompile_failed;
 mod error_return_data_outofbound;
 mod error_write_protection;
 
-#[cfg(test)]
+#[cfg(all(feature = "enable-memory", test))]
 mod memory_expansion_test;
 #[cfg(feature = "test")]
 pub use callop::tests::PrecompileCallArgs;
@@ -383,8 +384,8 @@ pub fn gen_associated_ops(
     state: &mut CircuitInputStateRef,
     geth_steps: &[GethExecStep],
 ) -> Result<Vec<ExecStep>, Error> {
-    let check_level = if *CHECK_MEM_STRICT { 2 } else { 0 }; // 0: no check, 1: check and log error and fix, 2: check and assert_eq
-    if check_level >= 1 {
+    #[cfg(feature = "enable-memory")]
+    if GETH_TRACE_CHECK_LEVEL.should_check() {
         let memory_enabled = !geth_steps.iter().all(|s| s.memory.is_empty());
         assert!(memory_enabled);
         if memory_enabled {
@@ -414,7 +415,7 @@ pub fn gen_associated_ops(
                         );
                     }
                 }
-                if check_level >= 2 {
+                if GETH_TRACE_CHECK_LEVEL.should_panic() {
                     panic!("mem wrong");
                 }
                 state.call_ctx_mut()?.memory = geth_steps[0].memory.clone();

bus-mapping/src/evm/opcodes/memory_expansion_test.rs

@@ -15,14 +15,13 @@ fn might_neg_index(index: isize, len: usize) -> usize {
     }
 }
 
-fn assert_expanded(_traces: &[GethExecStep], _before: isize, _after: isize) {
-    // FIXME: memory is removed
-    // let traces_len = traces.len();
-    // let before = might_neg_index(before, traces_len);
-    // let after = might_neg_index(after, traces_len);
-    // let size_before = traces[before].memory.borrow().len();
-    // let size_after = traces[after].memory.borrow().len();
-    // assert_ne!(size_before, size_after);
+fn assert_expanded(traces: &[GethExecStep], before: isize, after: isize) {
+    let traces_len = traces.len();
+    let before = might_neg_index(before, traces_len);
+    let after = might_neg_index(after, traces_len);
+    let size_before = traces[before].memory.0.len();
+    let size_after = traces[after].memory.0.len();
+    assert_ne!(size_before, size_after);
 }
 
 fn trace_and_assert<FN>(code: Bytecode, before: isize, after: isize, assert_fn: FN)

bus-mapping/src/rpc.rs

@@ -11,7 +11,7 @@ use ethers_providers::JsonRpcClient;
 use serde::Serialize;
 use std::collections::HashMap;
 
-use crate::util::CHECK_MEM_STRICT;
+use crate::util::GETH_TRACE_CHECK_LEVEL;
 
 /// Serialize a type.
 ///
@@ -45,9 +45,9 @@ pub(crate) struct GethLoggerConfig {
 impl Default for GethLoggerConfig {
     fn default() -> Self {
         Self {
-            enable_memory: false,
-            disable_stack: false,
-            disable_storage: false,
+            enable_memory: cfg!(feature = "enable-memory"),
+            disable_stack: !cfg!(feature = "enable-stack"),
+            disable_storage: !cfg!(feature = "enable-storage"),
             enable_return_data: true,
             timeout: None,
         }
@@ -156,7 +156,7 @@ impl<P: JsonRpcClient> GethClient<P> {
     pub async fn trace_tx_by_hash(&self, hash: H256) -> Result<GethExecTrace, Error> {
         let hash = serialize(&hash);
         let cfg = GethLoggerConfig {
-            enable_memory: *CHECK_MEM_STRICT,
+            enable_memory: cfg!(feature = "enable-memory") && GETH_TRACE_CHECK_LEVEL.should_check(),
             ..Default::default()
         };
         let cfg = serialize(&cfg);

bus-mapping/src/util.rs

@@ -3,6 +3,7 @@ use eth_types::{Hash, U256};
 pub use eth_types::{KECCAK_CODE_HASH_EMPTY, POSEIDON_CODE_HASH_EMPTY};
 use halo2_proofs::halo2curves::{bn256::Fr, group::ff::PrimeField};
 use once_cell::sync::Lazy;
+use std::convert::Infallible;
 
 use std::str::FromStr;
 
@@ -12,8 +13,44 @@ pub fn read_env_var<T: Clone + FromStr>(var_name: &'static str, default: T) -> T
         .map(|s| s.parse::<T>().unwrap_or_else(|_| default.clone()))
         .unwrap_or(default)
 }
-/// ..
-pub static CHECK_MEM_STRICT: Lazy<bool> = Lazy::new(|| read_env_var("CHECK_MEM_STRICT", false));
+/// env var for Geth trace sanity check level
+pub static GETH_TRACE_CHECK_LEVEL: Lazy<GethTraceSanityCheckLevel> =
+    Lazy::new(|| read_env_var("GETH_TRACE_CHECK_LEVEL", GethTraceSanityCheckLevel::None));
+
+/// Geth trace sanity check level
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum GethTraceSanityCheckLevel {
+    /// No sanity check
+    None,
+    /// Check sanity and log error
+    Check,
+    /// Panic on error
+    Strict,
+}
+
+impl GethTraceSanityCheckLevel {
+    /// Whether to do sanity check
+    pub fn should_check(&self) -> bool {
+        *self != GethTraceSanityCheckLevel::None
+    }
+
+    /// Whether to panic on error
+    pub fn should_panic(&self) -> bool {
+        *self == GethTraceSanityCheckLevel::Strict
+    }
+}
+
+impl FromStr for GethTraceSanityCheckLevel {
+    type Err = Infallible;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "strict" => Ok(GethTraceSanityCheckLevel::Strict),
+            _ if !s.is_empty() => Ok(GethTraceSanityCheckLevel::Check),
+            _ => Ok(GethTraceSanityCheckLevel::None),
+        }
+    }
+}
 
 /// Default number of bytes to pack into a field element.
 pub const POSEIDON_HASH_BYTES_IN_FIELD: usize = 31;

eth-types/Cargo.toml

@@ -30,3 +30,8 @@ default = ["warn-unimplemented"]
 warn-unimplemented = []
 shanghai = []
 scroll = []
+
+# trace heap allocation related feature switches
+enable-stack = []
+enable-memory = []
+enable-storage = []

eth-types/src/l2_types.rs

@@ -1,13 +1,20 @@
 //! L2 types used to deserialize traces for l2geth.
 
 use crate::{
-    evm_types::{Gas, GasCost, Memory, OpcodeId, ProgramCounter, Stack, Storage},
+    evm_types::{Gas, GasCost, OpcodeId, ProgramCounter},
     Block, GethExecError, GethExecStep, GethExecTrace, Hash, Transaction, Word, H256,
 };
 use ethers_core::types::{Address, Bytes, U256, U64};
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 
+#[cfg(feature = "enable-memory")]
+use crate::evm_types::Memory;
+#[cfg(feature = "enable-stack")]
+use crate::evm_types::Stack;
+#[cfg(feature = "enable-storage")]
+use crate::evm_types::Storage;
+
 /// l2 block full trace
 #[derive(Deserialize, Serialize, Default, Debug, Clone)]
 pub struct BlockTrace {
@@ -226,19 +233,18 @@ pub struct ExecStep {
     pub refund: u64,
     pub depth: isize,
     pub error: Option<GethExecError>,
+    #[cfg(feature = "enable-stack")]
     pub stack: Option<Vec<Word>>,
+    #[cfg(feature = "enable-memory")]
     pub memory: Option<Vec<Word>>,
+    #[cfg(feature = "enable-storage")]
     pub storage: Option<HashMap<Word, Word>>,
     #[serde(rename = "extraData")]
     pub extra_data: Option<ExtraData>,
 }
 
 impl From<ExecStep> for GethExecStep {
     fn from(e: ExecStep) -> Self {
-        let stack = e.stack.map_or_else(Stack::new, Stack::from);
-        let storage = e.storage.map_or_else(Storage::empty, Storage::from);
-        let memory = e.memory.map_or_else(Memory::default, Memory::from);
-
         GethExecStep {
             pc: ProgramCounter(e.pc as usize),
             // FIXME
@@ -248,9 +254,12 @@ impl From<ExecStep> for GethExecStep {
             refund: Gas(e.refund),
             depth: e.depth as u16,
             error: e.error,
-            stack,
-            memory,
-            storage,
+            #[cfg(feature = "enable-stack")]
+            stack: e.stack.map_or_else(Stack::new, Stack::from),
+            #[cfg(feature = "enable-memory")]
+            memory: e.memory.map_or_else(Memory::default, Memory::from),
+            #[cfg(feature = "enable-storage")]
+            storage: e.storage.map_or_else(Storage::empty, Storage::from),
         }
     }
 }