feat: JournaledAccount sload/sstore

crates/context/interface/src/journaled_state.rs

@@ -7,6 +7,7 @@ use crate::{
     context::{SStoreResult, SelfDestructResult},
     host::LoadError,
     journaled_state::account::JournaledAccountTr,
+    ErasedError,
 };
 use core::ops::{Deref, DerefMut};
 use database_interface::Database;
@@ -15,7 +16,6 @@ use primitives::{
 };
 use state::{Account, AccountInfo, Bytecode};
 use std::{borrow::Cow, vec::Vec};
-
 /// Trait that contains database and journal of all changes that were made to the state.
 pub trait JournalTr {
     /// Database type that is used in the journal.
@@ -304,6 +304,9 @@ pub enum JournalLoadError<E> {
     ColdLoadSkipped,
 }
 
+/// Journal error on loading of storage or account with Boxed Database error.
+pub type JournalLoadErasedError = JournalLoadError<ErasedError>;
+
 impl<E> JournalLoadError<E> {
     /// Returns true if the error is a database error.
     #[inline]
@@ -462,6 +465,7 @@ pub struct AccountInfoLoad<'a> {
 
 impl<'a> AccountInfoLoad<'a> {
     /// Creates new [`AccountInfoLoad`] with the given account info, cold load status and empty status.
+    #[inline]
     pub fn new(account: &'a AccountInfo, is_cold: bool, is_empty: bool) -> Self {
         Self {
             account: Cow::Borrowed(account),
@@ -473,6 +477,7 @@ impl<'a> AccountInfoLoad<'a> {
     /// Maps the account info of the [`AccountInfoLoad`] to a new [`StateLoad`].
     ///
     /// Useful for transforming the account info of the [`AccountInfoLoad`] and preserving the cold load status.
+    #[inline]
     pub fn into_state_load<F, O>(self, f: F) -> StateLoad<O>
     where
         F: FnOnce(Cow<'a, AccountInfo>) -> O,

crates/context/interface/src/journaled_state/account.rs

@@ -3,12 +3,19 @@
 //!
 //! Useful to encapsulate account and journal entries together. So when account gets changed, we can add a journal entry for it.
 
-use crate::journaled_state::entry::JournalEntry;
+use crate::{
+    context::{SStoreResult, StateLoad},
+    journaled_state::{entry::JournalEntry, JournalLoadErasedError, JournalLoadError},
+    ErasedError,
+};
 
 use super::entry::JournalEntryTr;
 use auto_impl::auto_impl;
-use primitives::{Address, B256, KECCAK_EMPTY, U256};
-use state::{Account, Bytecode};
+use database_interface::Database;
+use primitives::{
+    hash_map::Entry, Address, HashMap, HashSet, StorageKey, StorageValue, B256, KECCAK_EMPTY, U256,
+};
+use state::{Account, Bytecode, EvmStorageSlot};
 use std::vec::Vec;
 
 /// Trait that contains database and journal of all changes that were made to the account.
@@ -17,6 +24,24 @@ pub trait JournaledAccountTr {
     /// Returns the account.
     fn account(&self) -> &Account;
 
+    /// Sloads the storage slot and returns its mutable reference
+    fn sload(
+        &mut self,
+        key: StorageKey,
+        skip_cold_load: bool,
+    ) -> Result<StateLoad<&mut EvmStorageSlot>, JournalLoadErasedError>;
+
+    /// Loads the storage slot and stores the new value
+    fn sstore(
+        &mut self,
+        key: StorageKey,
+        new: StorageValue,
+        skip_cold_load: bool,
+    ) -> Result<StateLoad<SStoreResult>, JournalLoadErasedError>;
+
+    /// Loads the code of the account. and returns it as reference.
+    fn load_code(&mut self) -> Result<&Bytecode, JournalLoadErasedError>;
+
     /// Returns the balance of the account.
     fn balance(&self) -> &U256;
 
@@ -96,38 +121,164 @@ pub trait JournaledAccountTr {
 ///
 /// Useful to encapsulate account and journal entries together. So when account gets changed, we can add a journal entry for it.
 #[derive(Debug, PartialEq, Eq)]
-pub struct JournaledAccount<'a, ENTRY: JournalEntryTr = JournalEntry> {
+pub struct JournaledAccount<'a, DB, ENTRY: JournalEntryTr = JournalEntry> {
     /// Address of the account.
     address: Address,
     /// Mutable account.
     account: &'a mut Account,
     /// Journal entries.
     journal_entries: &'a mut Vec<ENTRY>,
+    /// Access list.
+    access_list: &'a HashMap<Address, HashSet<StorageKey>>,
+    /// Transaction ID.
+    transaction_id: usize,
+    /// Database used to load storage.
+    db: &'a mut DB,
 }
 
-impl<'a, ENTRY: JournalEntryTr> JournaledAccount<'a, ENTRY> {
-    /// Consumes the journaled account and returns the mutable account.
-    #[inline]
-    pub fn into_account_ref(self) -> &'a Account {
-        self.account
-    }
-
-    /// Creates a new journaled account.
+impl<'a, DB: Database, ENTRY: JournalEntryTr> JournaledAccount<'a, DB, ENTRY> {
+    /// Creates new JournaledAccount
     #[inline]
     pub fn new(
         address: Address,
         account: &'a mut Account,
         journal_entries: &'a mut Vec<ENTRY>,
+        db: &'a mut DB,
+        access_list: &'a HashMap<Address, HashSet<StorageKey>>,
+        transaction_id: usize,
     ) -> Self {
         Self {
             address,
             account,
             journal_entries,
+            access_list,
+            transaction_id,
+            db,
+        }
+    }
+
+    /// Creates a new journaled account.
+    ///
+    /// Does not erase the db error.
+    #[inline(never)]
+    pub fn sload_preserve_error(
+        &mut self,
+        key: StorageKey,
+        skip_cold_load: bool,
+    ) -> Result<StateLoad<&mut EvmStorageSlot>, JournalLoadError<DB::Error>> {
+        let is_newly_created = self.account.is_created();
+        let (slot, is_cold) = match self.account.storage.entry(key) {
+            Entry::Occupied(occ) => {
+                let slot = occ.into_mut();
+                // skip load if account is cold.
+                let is_cold = slot.is_cold_transaction_id(self.transaction_id);
+                if is_cold && skip_cold_load {
+                    return Err(JournalLoadError::ColdLoadSkipped);
+                }
+                slot.mark_warm_with_transaction_id(self.transaction_id);
+                (slot, is_cold)
+            }
+            Entry::Vacant(vac) => {
+                // is storage cold
+                let is_cold = self
+                    .access_list
+                    .get(&self.address)
+                    .and_then(|v| v.get(&key))
+                    .is_none();
+
+                if is_cold && skip_cold_load {
+                    return Err(JournalLoadError::ColdLoadSkipped);
+                }
+                // if storage was cleared, we don't need to ping db.
+                let value = if is_newly_created {
+                    StorageValue::ZERO
+                } else {
+                    self.db.storage(self.address, key)?
+                };
+
+                let slot = vac.insert(EvmStorageSlot::new(value, self.transaction_id));
+                (slot, is_cold)
+            }
+        };
+
+        if is_cold {
+            // add it to journal as cold loaded.
+            self.journal_entries
+                .push(ENTRY::storage_warmed(self.address, key));
+        }
+
+        Ok(StateLoad::new(slot, is_cold))
+    }
+
+    /// Stores the storage slot.
+    ///
+    /// If account is not present it will load from database.
+    ///
+    /// Does not erase the db error.
+    #[inline]
+    pub fn sstore_preserve_error(
+        &mut self,
+        key: StorageKey,
+        new: StorageValue,
+        skip_cold_load: bool,
+    ) -> Result<StateLoad<SStoreResult>, JournalLoadError<DB::Error>> {
+        // touch the account so changes are tracked.
+        self.touch();
+
+        // assume that acc exists and load the slot.
+        let slot = self.sload_preserve_error(key, skip_cold_load)?;
+
+        let ret = Ok(StateLoad::new(
+            SStoreResult {
+                original_value: slot.original_value(),
+                present_value: slot.present_value(),
+                new_value: new,
+            },
+            slot.is_cold,
+        ));
+
+        // when new value is different from present, we need to add a journal entry and make a change.
+        if slot.present_value != new {
+            let previous_value = slot.present_value;
+            // insert value into present state.
+            slot.data.present_value = new;
+
+            // add journal entry.
+            self.journal_entries
+                .push(ENTRY::storage_changed(self.address, key, previous_value));
+        }
+
+        ret
+    }
+
+    /// Loads the code of the account. and returns it as reference.
+    ///
+    /// Does not erase the db error.
+    #[inline]
+    pub fn load_code_preserve_error(&mut self) -> Result<&Bytecode, JournalLoadError<DB::Error>> {
+        if self.account.info.code.is_none() {
+            let hash = *self.code_hash();
+            let code = if hash == KECCAK_EMPTY {
+                Bytecode::default()
+            } else {
+                self.db.code_by_hash(hash)?
+            };
+            self.account.info.code = Some(code);
         }
+
+        Ok(self.account.info.code.as_ref().unwrap())
+    }
+
+    /// Consumes the journaled account and returns the account.
+    #[inline]
+    pub fn into_account(self) -> &'a Account {
+        self.account
     }
 }
 
-impl<'a, ENTRY: JournalEntryTr> JournaledAccountTr for JournaledAccount<'a, ENTRY> {
+impl<'a, DB: Database, ENTRY: JournalEntryTr> JournaledAccountTr
+    for JournaledAccount<'a, DB, ENTRY>
+{
     /// Returns the account.
     fn account(&self) -> &Account {
         self.account
@@ -295,4 +446,34 @@ impl<'a, ENTRY: JournalEntryTr> JournaledAccountTr for JournaledAccount<'a, ENTR
         self.set_code(hash, bytecode);
         self.bump_nonce();
     }
+
+    /// Loads the storage slot.
+    #[inline]
+    fn sload(
+        &mut self,
+        key: StorageKey,
+        skip_cold_load: bool,
+    ) -> Result<StateLoad<&mut EvmStorageSlot>, JournalLoadErasedError> {
+        self.sload_preserve_error(key, skip_cold_load)
+            .map_err(|i| i.map(ErasedError::new))
+    }
+
+    /// Stores the storage slot.
+    #[inline]
+    fn sstore(
+        &mut self,
+        key: StorageKey,
+        new: StorageValue,
+        skip_cold_load: bool,
+    ) -> Result<StateLoad<SStoreResult>, JournalLoadErasedError> {
+        self.sstore_preserve_error(key, new, skip_cold_load)
+            .map_err(|i| i.map(ErasedError::new))
+    }
+
+    /// Loads the code of the account. and returns it as reference.
+    #[inline]
+    fn load_code(&mut self) -> Result<&Bytecode, JournalLoadErasedError> {
+        self.load_code_preserve_error()
+            .map_err(|i| i.map(ErasedError::new))
+    }
 }

crates/context/interface/src/lib.rs

@@ -17,8 +17,7 @@ pub mod transaction;
 pub use block::Block;
 pub use cfg::{Cfg, CreateScheme, TransactTo};
 pub use context::{ContextError, ContextSetters, ContextTr};
-pub use database_interface::erased_error::ErasedError;
-pub use database_interface::{DBErrorMarker, Database};
+pub use database_interface::{erased_error::ErasedError, DBErrorMarker, Database};
 pub use either;
 pub use host::{DummyHost, Host};
 pub use journaled_state::JournalTr;

crates/context/src/context.rs

@@ -8,7 +8,9 @@ use context_interface::{
 };
 use database_interface::{Database, DatabaseRef, EmptyDB, WrapDatabaseRef};
 use derive_where::derive_where;
-use primitives::{hardfork::SpecId, Address, Log, StorageKey, StorageValue, B256, U256};
+use primitives::{
+    hardfork::SpecId, hints_util::cold_path, Address, Log, StorageKey, StorageValue, B256, U256,
+};
 
 /// EVM context contains data that EVM needs for execution.
 #[derive_where(Clone, Debug; BLOCK, CFG, CHAIN, TX, DB, JOURNAL, <DB as Database>::Error, LOCAL)]
@@ -499,6 +501,7 @@ impl<
         self.db_mut()
             .block_hash(requested_number)
             .map_err(|e| {
+                cold_path();
                 *self.error() = Err(e.into());
             })
             .ok()
@@ -532,6 +535,7 @@ impl<
         self.journal_mut()
             .selfdestruct(address, target, skip_cold_load)
             .map_err(|e| {
+                cold_path();
                 let (ret, err) = e.into_parts();
                 if let Some(err) = err {
                     *self.error() = Err(err.into());
@@ -551,6 +555,7 @@ impl<
         self.journal_mut()
             .sstore_skip_cold_load(address, key, value, skip_cold_load)
             .map_err(|e| {
+                cold_path();
                 let (ret, err) = e.into_parts();
                 if let Some(err) = err {
                     *self.error() = Err(err.into());
@@ -569,6 +574,7 @@ impl<
         self.journal_mut()
             .sload_skip_cold_load(address, key, skip_cold_load)
             .map_err(|e| {
+                cold_path();
                 let (ret, err) = e.into_parts();
                 if let Some(err) = err {
                     *self.error() = Err(err.into());
@@ -584,14 +590,17 @@ impl<
         load_code: bool,
         skip_cold_load: bool,
     ) -> Result<AccountInfoLoad<'_>, LoadError> {
-        let error = &mut self.error;
-        let journal = &mut self.journaled_state;
-        match journal.load_account_info_skip_cold_load(address, load_code, skip_cold_load) {
+        match self.journaled_state.load_account_info_skip_cold_load(
+            address,
+            load_code,
+            skip_cold_load,
+        ) {
             Ok(a) => Ok(a),
             Err(e) => {
+                cold_path();
                 let (ret, err) = e.into_parts();
                 if let Some(err) = err {
-                    *error = Err(err.into());
+                    self.error = Err(err.into());
                 }
                 Err(ret)
             }