feat(state): Block hash cache and overrides (bluealloy#621)

crates/interpreter/src/instructions/host.rs

@@ -9,6 +9,7 @@ use crate::{
 use alloc::boxed::Box;
 use alloc::vec::Vec;
 use core::cmp::min;
+use revm_primitives::BLOCK_HASH_HISTORY;
 
 pub fn balance<SPEC: Spec>(interpreter: &mut Interpreter, host: &mut dyn Host) {
     pop_address!(interpreter, address);
@@ -137,7 +138,7 @@ pub fn blockhash(interpreter: &mut Interpreter, host: &mut dyn Host) {
     if let Some(diff) = host.env().block.number.checked_sub(*number) {
         let diff = as_usize_saturated!(diff);
         // blockhash should push zero if number is same as current block number.
-        if diff <= 256 && diff != 0 {
+        if diff <= BLOCK_HASH_HISTORY && diff != 0 {
             let ret = host.block_hash(*number);
             if ret.is_none() {
                 interpreter.instruction_result = InstructionResult::FatalExternalError;

crates/primitives/src/constants.rs

@@ -9,6 +9,9 @@ pub const CALL_STACK_LIMIT: u64 = 1024;
 /// By default limit is 0x6000 (~25kb)
 pub const MAX_CODE_SIZE: usize = 0x6000;
 
+/// Number of blocks hashes that EVM can access in the past
+pub const BLOCK_HASH_HISTORY: usize = 256;
+
 /// EIP-3860: Limit and meter initcode
 ///
 /// Limit of maximum initcode size is 2 * MAX_CODE_SIZE

crates/revm/src/db/states/state.rs

@@ -2,9 +2,10 @@ use super::{
     cache::CacheState, plain_account::PlainStorage, BundleState, CacheAccount, TransitionState,
 };
 use crate::TransitionAccount;
+use alloc::collections::{btree_map, BTreeMap};
 use revm_interpreter::primitives::{
     db::{Database, DatabaseCommit},
-    hash_map, Account, AccountInfo, Address, Bytecode, HashMap, B256, U256,
+    hash_map, Account, AccountInfo, Bytecode, HashMap, Address, B256, BLOCK_HASH_HISTORY, U256,
 };
 
 /// State of blockchain.
@@ -36,8 +37,12 @@ pub struct State<'a, DBError> {
     /// Bundle is the main output of the state execution and this allows setting previous bundle
     /// and using its values for execution.
     pub use_preloaded_bundle: bool,
-    // if enabled USE Background thread for transitions and bundle
-    //pub use_background_thread: bool,
+    /// If EVM asks for block hash we will first check if they are found here.
+    /// and then ask the database.
+    ///
+    /// This map can be used to give different values for block hashes if in case
+    /// The fork block is different or some blocks are not saved inside database.
+    pub block_hashes: BTreeMap<u64, B256>,
 }
 
 impl<'a, DBError> State<'a, DBError> {
@@ -208,8 +213,25 @@ impl<'a, DBError> Database for State<'a, DBError> {
     }
 
     fn block_hash(&mut self, number: U256) -> Result<B256, Self::Error> {
-        // TODO maybe cache it.
-        self.database.block_hash(number)
+        // block number is never biger then u64::MAX.
+        let u64num: u64 = number.to();
+        match self.block_hashes.entry(u64num) {
+            btree_map::Entry::Occupied(entry) => Ok(*entry.get()),
+            btree_map::Entry::Vacant(entry) => {
+                let ret = *entry.insert(self.database.block_hash(number)?);
+
+                // prune all hashes that are older then BLOCK_HASH_HISTORY
+                while let Some(entry) = self.block_hashes.first_entry() {
+                    if *entry.key() < u64num.saturating_sub(BLOCK_HASH_HISTORY as u64) {
+                        entry.remove();
+                    } else {
+                        break;
+                    }
+                }
+
+                Ok(ret)
+            }
+        }
     }
 }
 
@@ -230,7 +252,31 @@ mod tests {
         },
         StateBuilder,
     };
-    use revm_interpreter::primitives::StorageSlot;
+    use revm_interpreter::primitives::{keccak256, StorageSlot};
+
+    #[test]
+    fn block_hash_cache() {
+        let mut state = StateBuilder::default().build();
+        state.block_hash(U256::from(1)).unwrap();
+        state.block_hash(U256::from(2)).unwrap();
+
+        let test_number = BLOCK_HASH_HISTORY as u64 + 2;
+
+        let block1_hash = keccak256(&U256::from(1).to_be_bytes::<{ U256::BYTES }>());
+        let block2_hash = keccak256(&U256::from(2).to_be_bytes::<{ U256::BYTES }>());
+        let block_test_hash = keccak256(&U256::from(test_number).to_be_bytes::<{ U256::BYTES }>());
+
+        assert_eq!(
+            state.block_hashes,
+            BTreeMap::from([(1, block1_hash), (2, block2_hash)])
+        );
+
+        state.block_hash(U256::from(test_number)).unwrap();
+        assert_eq!(
+            state.block_hashes,
+            BTreeMap::from([(test_number, block_test_hash), (2, block2_hash)])
+        );
+    }
 
     /// Checks that if accounts is touched multiple times in the same block,
     /// then the old values from the first change are preserved and not overwritten.

crates/revm/src/db/states/state_builder.rs

@@ -1,7 +1,8 @@
 use super::{cache::CacheState, BundleState, State, TransitionState};
 use crate::db::EmptyDB;
+use alloc::collections::BTreeMap;
 use core::convert::Infallible;
-use revm_interpreter::primitives::db::Database;
+use revm_interpreter::primitives::{db::Database, B256};
 
 /// Allows building of State and initializing it with different options.
 pub struct StateBuilder<'a, DBError> {
@@ -23,6 +24,8 @@ pub struct StateBuilder<'a, DBError> {
     /// This will allows evm to continue executing.
     /// Default is false.
     pub with_background_transition_merge: bool,
+    /// If we want to set different block hashes
+    pub with_block_hashes: BTreeMap<u64, B256>,
 }
 
 impl Default for StateBuilder<'_, Infallible> {
@@ -34,6 +37,7 @@ impl Default for StateBuilder<'_, Infallible> {
             with_bundle_prestate: None,
             without_bundle_update: false,
             with_background_transition_merge: false,
+            with_block_hashes: BTreeMap::new(),
         }
     }
 }
@@ -57,6 +61,7 @@ impl<'a, DBError> StateBuilder<'a, DBError> {
             with_bundle_prestate: self.with_bundle_prestate,
             without_bundle_update: self.without_bundle_update,
             with_background_transition_merge: self.with_background_transition_merge,
+            with_block_hashes: self.with_block_hashes,
         }
     }
 
@@ -112,6 +117,13 @@ impl<'a, DBError> StateBuilder<'a, DBError> {
         }
     }
 
+    pub fn with_block_hashes(self, block_hashes: BTreeMap<u64, B256>) -> Self {
+        Self {
+            with_block_hashes: block_hashes,
+            ..self
+        }
+    }
+
     pub fn build(mut self) -> State<'a, DBError> {
         let use_preloaded_bundle = if self.with_cache_prestate.is_some() {
             self.with_bundle_prestate = None;
@@ -131,6 +143,7 @@ impl<'a, DBError> StateBuilder<'a, DBError> {
             },
             bundle_state: self.with_bundle_prestate,
             use_preloaded_bundle,
+            block_hashes: self.with_block_hashes,
         }
     }
 }