From 01b32934c4fba7fc57198b85d2cf524f269f4078 Mon Sep 17 00:00:00 2001
From: frisitano <francesco.risitano95@gmail.com>
Date: Wed, 2 Oct 2024 14:50:55 +0800
Subject: [PATCH] code size experiment

---
 src/code.rs            | 32 ++++++++++++++++++++++++++++++++
 src/handle_register.rs |  6 ++++--
 src/instruction.rs     | 15 ++++++---------
 src/lib.rs             |  1 +
 src/spec.rs            |  9 ++++++---
 5 files changed, 49 insertions(+), 14 deletions(-)
 create mode 100644 src/code.rs

diff --git a/src/code.rs b/src/code.rs
new file mode 100644
index 0000000..bda9c72
--- /dev/null
+++ b/src/code.rs
@@ -0,0 +1,32 @@
+use primitives::{Address, B256};
+use revm::interpreter::StateLoad;
+
+use crate::ScrollWiring;
+
+type KeccakHash = B256;
+
+pub trait ScrollCodeSizeDatabase {
+    fn code_size(&self, code_hash: KeccakHash) -> Option<usize>;
+}
+
+pub trait ScrollCodeHost {
+    fn code_size(&mut self, code_hash: Address) -> Option<(usize, bool)>;
+}
+
+impl<EvmWiringT> ScrollCodeHost for revm::Context<EvmWiringT>
+where
+    EvmWiringT: ScrollWiring,
+    EvmWiringT::Database: ScrollCodeSizeDatabase,
+{
+    fn code_size(&mut self, address: Address) -> Option<(usize, bool)> {
+        let StateLoad {
+            data: account,
+            is_cold,
+        } = self.evm.load_account(address).ok()?;
+        let code_hash = account.info.code_hash();
+        self.evm
+            .db
+            .code_size(code_hash)
+            .map(|code_size| (code_size, is_cold))
+    }
+}
diff --git a/src/handle_register.rs b/src/handle_register.rs
index 3708a88..ad770a6 100644
--- a/src/handle_register.rs
+++ b/src/handle_register.rs
@@ -16,8 +16,9 @@ use revm::{
 use std::sync::Arc;
 
 use crate::{
-    instruction::make_scroll_instruction_tables, precompile, scroll_spec_to_generic,
-    spec::ScrollSpec, L1BlockInfo, ScrollContext, ScrollTransaction, ScrollWiring,
+    code::ScrollCodeSizeDatabase, instruction::make_scroll_instruction_tables, precompile,
+    scroll_spec_to_generic, spec::ScrollSpec, L1BlockInfo, ScrollContext, ScrollTransaction,
+    ScrollWiring,
 };
 
 /// Configure the handler for the Scroll chain.
@@ -33,6 +34,7 @@ use crate::{
 pub fn scroll_handle_register<EvmWiringT>(handler: &mut EvmHandler<'_, EvmWiringT>)
 where
     EvmWiringT: ScrollWiring,
+    EvmWiringT::Database: ScrollCodeSizeDatabase,
 {
     scroll_spec_to_generic!(handler.spec_id, {
         // Load `L1BlockInfo` from the database and invoke standard `load_accounts` handler.
diff --git a/src/instruction.rs b/src/instruction.rs
index 1ed9a41..695cba3 100644
--- a/src/instruction.rs
+++ b/src/instruction.rs
@@ -12,7 +12,7 @@ use revm::{
     wiring::Block,
 };
 
-use crate::{ScrollSpec, ScrollSpecId};
+use crate::{code::ScrollCodeHost, ScrollSpec, ScrollSpecId};
 
 /// Creates a table of instructions for the Scroll hardfork.
 ///
@@ -24,7 +24,7 @@ use crate::{ScrollSpec, ScrollSpecId};
 /// - `TLOAD`
 /// - `SELFDESTRUCT`
 /// - `MCOPY`
-pub fn make_scroll_instruction_tables<'a, H: Host + ?Sized, SPEC: ScrollSpec>(
+pub fn make_scroll_instruction_tables<'a, H: Host + ?Sized + ScrollCodeHost, SPEC: ScrollSpec>(
 ) -> InstructionTables<'a, H> {
     let mut table = make_instruction_table::<H, SPEC>();
 
@@ -71,18 +71,15 @@ fn blockhash<H: Host + ?Sized>(interpreter: &mut Interpreter, host: &mut H) {
     };
 }
 
-fn extcodesize<H: Host + ?Sized>(interpreter: &mut Interpreter, host: &mut H) {
+fn extcodesize<H: Host + ?Sized + ScrollCodeHost>(interpreter: &mut Interpreter, host: &mut H) {
     pop_address!(interpreter, address);
-    let Some(code) = host.code(address) else {
+    let Some((code_size, is_cold)) = host.code_size(address) else {
         interpreter.instruction_result = InstructionResult::FatalExternalError;
         return;
     };
 
-    let (code, load) = code.into_components();
-
-    gas!(interpreter, warm_cold_cost(load.state_load.is_cold));
-
-    push!(interpreter, U256::from(code.len()));
+    gas!(interpreter, warm_cold_cost(is_cold));
+    push!(interpreter, U256::from(code_size));
 }
 
 fn selfdestruct<H: Host + ?Sized>(interpreter: &mut Interpreter, _host: &mut H) {
diff --git a/src/lib.rs b/src/lib.rs
index 406ae08..75f0dcd 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,5 +1,6 @@
 use revm::{primitives::Bytes, wiring::TransactionValidation};
 
+mod code;
 mod env;
 mod handle_register;
 mod instruction;
diff --git a/src/spec.rs b/src/spec.rs
index be1acf9..955e039 100644
--- a/src/spec.rs
+++ b/src/spec.rs
@@ -1,4 +1,7 @@
-use crate::{env::TxEnv, handle_register::scroll_handle_register, ScrollContext};
+use crate::{
+    code::ScrollCodeSizeDatabase, env::TxEnv, handle_register::scroll_handle_register,
+    ScrollContext,
+};
 use core::marker::PhantomData;
 use revm::{
     database_interface::Database,
@@ -15,7 +18,7 @@ pub struct ScrollEvmWiring<DB: Database, EXT> {
     _phantom: PhantomData<(DB, EXT)>,
 }
 
-impl<DB: Database, EXT> EvmWiring for ScrollEvmWiring<DB, EXT> {
+impl<DB: Database + ScrollCodeSizeDatabase, EXT> EvmWiring for ScrollEvmWiring<DB, EXT> {
     type Block = BlockEnv;
     type Database = DB;
     type ChainContext = Context;
@@ -25,7 +28,7 @@ impl<DB: Database, EXT> EvmWiring for ScrollEvmWiring<DB, EXT> {
     type Transaction = TxEnv;
 }
 
-impl<DB: Database, EXT> revm::EvmWiring for ScrollEvmWiring<DB, EXT> {
+impl<DB: Database + ScrollCodeSizeDatabase, EXT> revm::EvmWiring for ScrollEvmWiring<DB, EXT> {
     fn handler<'evm>(hardfork: Self::Hardfork) -> revm::EvmHandler<'evm, Self> {
         let mut handler = EvmHandler::mainnet_with_spec(hardfork);