chore(cheatcodes): avoid unnecessary lookup in store (#11163)

Author: DaniPopes

crates/cheatcodes/src/error.rs

@@ -1,5 +1,5 @@
 use crate::Vm;
-use alloy_primitives::{Address, Bytes, hex};
+use alloy_primitives::{Bytes, hex};
 use alloy_signer::Error as SignerError;
 use alloy_signer_local::LocalSignerError;
 use alloy_sol_types::SolError;
@@ -65,19 +65,6 @@ macro_rules! ensure {
     };
 }
 
-macro_rules! ensure_not_precompile {
-    ($address:expr, $ctxt:expr) => {
-        if $ctxt.is_precompile($address) {
-            return Err($crate::error::precompile_error($address));
-        }
-    };
-}
-
-#[cold]
-pub(crate) fn precompile_error(address: &Address) -> Error {
-    fmt_err!("cannot use precompile {address} as an argument")
-}
-
 /// Error thrown by cheatcodes.
 // This uses a custom repr to minimize the size of the error.
 // The repr is basically `enum { Cow<'static, str>, Cow<'static, [u8]> }`

crates/cheatcodes/src/evm.rs

@@ -184,7 +184,7 @@ impl Cheatcode for getNonce_1Call {
 impl Cheatcode for loadCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { target, slot } = *self;
-        ensure_not_precompile!(&target, ccx);
+        ccx.ensure_not_precompile(&target)?;
         ccx.ecx.journaled_state.load_account(target)?;
         let mut val = ccx.ecx.journaled_state.sload(target, slot.into())?;
 
@@ -530,7 +530,7 @@ impl Cheatcode for dealCall {
 impl Cheatcode for etchCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { target, newRuntimeBytecode } = self;
-        ensure_not_precompile!(target, ccx);
+        ccx.ensure_not_precompile(target)?;
         ccx.ecx.journaled_state.load_account(*target)?;
         let bytecode = Bytecode::new_raw_checked(newRuntimeBytecode.clone())
             .map_err(|e| fmt_err!("failed to create bytecode: {e}"))?;
@@ -582,9 +582,8 @@ impl Cheatcode for setNonceUnsafeCall {
 impl Cheatcode for storeCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { target, slot, value } = *self;
-        ensure_not_precompile!(&target, ccx);
-        // ensure the account is touched
-        let _ = journaled_account(ccx.ecx, target)?;
+        ccx.ensure_not_precompile(&target)?;
+        ensure_loaded_account(ccx.ecx, target)?;
         ccx.ecx.journaled_state.sstore(target, slot.into(), value.into())?;
         Ok(Default::default())
     }
@@ -1155,9 +1154,14 @@ pub(super) fn journaled_account<'a>(
     ecx: Ecx<'a, '_, '_>,
     addr: Address,
 ) -> Result<&'a mut Account> {
+    ensure_loaded_account(ecx, addr)?;
+    Ok(ecx.journaled_state.state.get_mut(&addr).expect("account is loaded"))
+}
+
+pub(super) fn ensure_loaded_account(ecx: Ecx, addr: Address) -> Result<()> {
     ecx.journaled_state.load_account(addr)?;
     ecx.journaled_state.touch(addr);
-    Ok(ecx.journaled_state.state.get_mut(&addr).expect("account is loaded"))
+    Ok(())
 }
 
 /// Consumes recorded account accesses and returns them as an abi encoded

crates/cheatcodes/src/lib.rs

@@ -162,8 +162,18 @@ impl std::ops::DerefMut for CheatsCtxt<'_, '_, '_, '_> {
 }
 
 impl CheatsCtxt<'_, '_, '_, '_> {
+    #[inline]
+    pub(crate) fn ensure_not_precompile(&self, address: &Address) -> Result<()> {
+        if self.is_precompile(address) { Err(precompile_error(address)) } else { Ok(()) }
+    }
+
     #[inline]
     pub(crate) fn is_precompile(&self, address: &Address) -> bool {
         self.ecx.journaled_state.inner.precompiles.contains(address)
     }
 }
+
+#[cold]
+fn precompile_error(address: &Address) -> Error {
+    fmt_err!("cannot use precompile {address} as an argument")
+}