feat: refactor ssload/sstore to JournaledAccount

crates/context/interface/src/journaled_state.rs

@@ -16,6 +16,19 @@ use primitives::{
 use state::{Account, AccountInfo, Bytecode};
 use std::{borrow::Cow, vec::Vec};
 
+/// Result type that contains database error.
+pub type JournaledAccountLoadResult<'a, 'b, JOURNAL> = Result<
+    StateLoad<
+        JournaledAccount<
+            'a,
+            'b,
+            <JOURNAL as JournalTr>::JournalEntry,
+            <JOURNAL as JournalTr>::Database,
+        >,
+    >,
+    <<JOURNAL as JournalTr>::Database as Database>::Error,
+>;
+
 /// 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.
@@ -183,13 +196,7 @@ pub trait JournalTr {
 
     /// Loads the journaled account.
     #[inline]
-    fn load_account_mut(
-        &mut self,
-        address: Address,
-    ) -> Result<
-        StateLoad<JournaledAccount<'_, Self::JournalEntry>>,
-        <Self::Database as Database>::Error,
-    > {
+    fn load_account_mut(&mut self, address: Address) -> JournaledAccountLoadResult<'_, '_, Self> {
         self.load_account_mut_optional_code(address, false)
     }
 
@@ -198,10 +205,7 @@ pub trait JournalTr {
     fn load_account_with_code_mut(
         &mut self,
         address: Address,
-    ) -> Result<
-        StateLoad<JournaledAccount<'_, Self::JournalEntry>>,
-        <Self::Database as Database>::Error,
-    > {
+    ) -> JournaledAccountLoadResult<'_, '_, Self> {
         self.load_account_mut_optional_code(address, true)
     }
 
@@ -210,10 +214,7 @@ pub trait JournalTr {
         &mut self,
         address: Address,
         load_code: bool,
-    ) -> Result<
-        StateLoad<JournaledAccount<'_, Self::JournalEntry>>,
-        <Self::Database as Database>::Error,
-    >;
+    ) -> JournaledAccountLoadResult<'_, '_, Self>;
 
     /// Sets bytecode with hash. Assume that account is warm.
     fn set_code_with_hash(&mut self, address: Address, code: Bytecode, hash: B256);

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

@@ -3,26 +3,143 @@
 //!
 //! Useful to encapsulate account and journal entries together. So when account gets changed, we can add a journal entry for it.
 
+use crate::{
+    context::{SStoreResult, StateLoad},
+    journaled_state::JournalLoadError,
+};
+
 use super::entry::JournalEntryTr;
 use core::ops::Deref;
-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;
 
 /// Journaled account contains both mutable account and journal entries.
 ///
 /// 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> {
+pub struct JournaledAccount<'a, 'b, ENTRY: JournalEntryTr, DB> {
     /// 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: &'b mut DB,
+}
+
+impl<'a, 'b, ENTRY: JournalEntryTr, DB: Database> JournaledAccount<'a, 'b, ENTRY, DB> {
+    /// Creates a new journaled account.
+    #[inline(never)]
+    pub fn sload(
+        &mut self,
+        key: StorageKey,
+        skip_cold_load: bool,
+    ) -> Result<StateLoad<&mut EvmStorageSlot>, JournalLoadError<<DB as Database>::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))
+    }
+
+    /// Warm loads storage slot and stores the new value
+    #[inline]
+    pub fn sstore(
+        &mut self,
+        key: StorageKey,
+        new: StorageValue,
+        skip_cold_load: bool,
+    ) -> Result<StateLoad<SStoreResult>, JournalLoadError<<DB as Database>::Error>> {
+        // assume that acc exists and load the slot.
+        let slot = self.sload(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.
+    #[inline]
+    pub fn load_code(&mut self) -> Result<&Bytecode, JournalLoadError<<DB as Database>::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())
+    }
 }
 
-impl<'a, ENTRY: JournalEntryTr> JournaledAccount<'a, ENTRY> {
+impl<'a, 'b, ENTRY: JournalEntryTr, DB> JournaledAccount<'a, 'b, ENTRY, DB> {
     /// Consumes the journaled account and returns the mutable account.
     #[inline]
     pub fn into_account_ref(self) -> &'a Account {
@@ -35,11 +152,17 @@ impl<'a, ENTRY: JournalEntryTr> JournaledAccount<'a, ENTRY> {
         address: Address,
         account: &'a mut Account,
         journal_entries: &'a mut Vec<ENTRY>,
+        db: &'b mut DB,
+        access_list: &'a HashMap<Address, HashSet<StorageKey>>,
+        transaction_id: usize,
     ) -> Self {
         Self {
             address,
             account,
             journal_entries,
+            db,
+            access_list,
+            transaction_id,
         }
     }
 
@@ -207,7 +330,7 @@ impl<'a, ENTRY: JournalEntryTr> JournaledAccount<'a, ENTRY> {
     }
 }
 
-impl<'a, ENTRY: JournalEntryTr> Deref for JournaledAccount<'a, ENTRY> {
+impl<'a, 'b, ENTRY: JournalEntryTr, DB> Deref for JournaledAccount<'a, 'b, ENTRY, DB> {
     type Target = Account;
 
     fn deref(&self) -> &Self::Target {

crates/context/src/journal.rs

@@ -12,8 +12,8 @@ use bytecode::Bytecode;
 use context_interface::{
     context::{SStoreResult, SelfDestructResult, StateLoad},
     journaled_state::{
-        account::JournaledAccount, AccountInfoLoad, AccountLoad, JournalCheckpoint,
-        JournalLoadError, JournalTr, TransferError,
+        AccountInfoLoad, AccountLoad, JournalCheckpoint, JournalLoadError, JournalTr,
+        JournaledAccountLoadResult, TransferError,
     },
 };
 use core::ops::{Deref, DerefMut};
@@ -257,10 +257,7 @@ impl<DB: Database, ENTRY: JournalEntryTr> JournalTr for Journal<DB, ENTRY> {
         &mut self,
         address: Address,
         load_code: bool,
-    ) -> Result<
-        StateLoad<JournaledAccount<'_, Self::JournalEntry>>,
-        <Self::Database as Database>::Error,
-    > {
+    ) -> JournaledAccountLoadResult<'_, '_, Self> {
         self.inner
             .load_account_mut_optional_code(&mut self.database, address, load_code, false)
             .map_err(JournalLoadError::unwrap_db_error)