From 3ab24eac6bb708d4edbc5db7bcaebed322832b38 Mon Sep 17 00:00:00 2001
From: Delweng <delweng@gmail.com>
Date: Wed, 23 Oct 2024 07:26:05 +0800
Subject: [PATCH] feat(rpc-types-trace/prestate): support
 disable_{code,storage} (#1538)

* feat(rpc-types-trace/prestate): support disable_{code,storage}

Signed-off-by: jsvisa <delweng@gmail.com>

* wrong tests

Signed-off-by: jsvisa <delweng@gmail.com>

---------

Signed-off-by: jsvisa <delweng@gmail.com>
---
 crates/rpc-types-trace/src/geth/mux.rs       |  2 +-
 crates/rpc-types-trace/src/geth/pre_state.rs | 45 ++++++++++++++++++--
 2 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/crates/rpc-types-trace/src/geth/mux.rs b/crates/rpc-types-trace/src/geth/mux.rs
index 1714b32c8c3..c7ffd569a17 100644
--- a/crates/rpc-types-trace/src/geth/mux.rs
+++ b/crates/rpc-types-trace/src/geth/mux.rs
@@ -60,7 +60,7 @@ mod tests {
             Some(GethDebugTracerType::BuiltInTracer(GethDebugBuiltInTracerType::MuxTracer));
 
         let call_config = CallConfig { only_top_call: Some(true), with_log: Some(true) };
-        let prestate_config = PreStateConfig { diff_mode: Some(true) };
+        let prestate_config = PreStateConfig { diff_mode: Some(true), ..Default::default() };
 
         opts.tracing_options.tracer_config = MuxConfig(HashMap::from_iter([
             (GethDebugBuiltInTracerType::FourByteTracer, None),
diff --git a/crates/rpc-types-trace/src/geth/pre_state.rs b/crates/rpc-types-trace/src/geth/pre_state.rs
index 4202228b64d..a7339944858 100644
--- a/crates/rpc-types-trace/src/geth/pre_state.rs
+++ b/crates/rpc-types-trace/src/geth/pre_state.rs
@@ -211,6 +211,12 @@ pub struct PreStateConfig {
     /// storage necessary to execute the transaction.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub diff_mode: Option<bool>,
+    /// If `disableCode` is set to true, the response frame will not include code.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub disable_code: Option<bool>,
+    /// If `disableStorage` is set to true, the response frame will not include storage.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub disable_storage: Option<bool>,
 }
 
 impl PreStateConfig {
@@ -225,6 +231,18 @@ impl PreStateConfig {
     pub fn is_default_mode(&self) -> bool {
         !self.is_diff_mode()
     }
+
+    /// Returns true if code is enabled.
+    #[inline]
+    pub fn code_enabled(&self) -> bool {
+        !self.disable_code.unwrap_or_default()
+    }
+
+    /// Returns true if storage is enabled.
+    #[inline]
+    pub fn storage_enabled(&self) -> bool {
+        !self.disable_storage.unwrap_or_default()
+    }
 }
 
 #[cfg(test)]
@@ -245,7 +263,9 @@ mod tests {
         opts.tracing_options.tracer =
             Some(GethDebugTracerType::BuiltInTracer(GethDebugBuiltInTracerType::PreStateTracer));
         opts.tracing_options.tracer_config =
-            serde_json::to_value(PreStateConfig { diff_mode: Some(true) }).unwrap().into();
+            serde_json::to_value(PreStateConfig { diff_mode: Some(true), ..Default::default() })
+                .unwrap()
+                .into();
 
         assert_eq!(
             serde_json::to_string(&opts).unwrap(),
@@ -270,9 +290,26 @@ mod tests {
 
     #[test]
     fn test_is_diff_mode() {
-        assert!(PreStateConfig { diff_mode: Some(true) }.is_diff_mode());
-        assert!(!PreStateConfig { diff_mode: Some(false) }.is_diff_mode());
-        assert!(!PreStateConfig { diff_mode: None }.is_diff_mode());
+        assert!(PreStateConfig { diff_mode: Some(true), ..Default::default() }.is_diff_mode());
+        assert!(!PreStateConfig { diff_mode: Some(false), ..Default::default() }.is_diff_mode());
+        assert!(!PreStateConfig { diff_mode: None, ..Default::default() }.is_diff_mode());
+    }
+
+    #[test]
+    fn test_disable_code() {
+        assert!(PreStateConfig { ..Default::default() }.code_enabled());
+        assert!(PreStateConfig { disable_code: Some(false), ..Default::default() }.code_enabled());
+        assert!(!PreStateConfig { disable_code: Some(true), ..Default::default() }.code_enabled());
+    }
+    #[test]
+    fn test_disable_storage() {
+        assert!(PreStateConfig { ..Default::default() }.storage_enabled());
+        assert!(
+            PreStateConfig { disable_storage: Some(false), ..Default::default() }.storage_enabled()
+        );
+        assert!(
+            !PreStateConfig { disable_storage: Some(true), ..Default::default() }.storage_enabled()
+        );
     }
 
     #[test]