refactor: adapt isContext cheatcode

crates/edr_napi/src/solidity_tests/config.rs

@@ -2,8 +2,10 @@ use std::{collections::HashMap, fmt::Debug, path::PathBuf};
 
 use alloy_primitives::hex;
 use edr_solidity_tests::{
-    executors::invariant::InvariantConfig, fuzz::FuzzConfig,
-    inspectors::cheatcodes::CheatsConfigOptions, SolidityTestRunnerConfig,
+    executors::invariant::InvariantConfig,
+    fuzz::FuzzConfig,
+    inspectors::cheatcodes::{CheatsConfigOptions, ExecutionContextConfig},
+    SolidityTestRunnerConfig,
 };
 use foundry_cheatcodes::{FsPermissions, RpcEndpoint, RpcEndpoints};
 use napi::{
@@ -190,6 +192,9 @@ impl TryFrom<SolidityTestRunnerConfigArgs> for SolidityTestRunnerConfig {
         let fuzz: FuzzConfig = fuzz.map(TryFrom::try_from).transpose()?.unwrap_or_default();
 
         let cheats_config_options = CheatsConfigOptions {
+            // TODO https://github.com/NomicFoundation/edr/issues/657
+            // If gas reporting or coverage is supported, take that into account here.
+            execution_context: ExecutionContextConfig::Test,
             rpc_endpoints: rpc_endpoints
                 .map(|endpoints| {
                     RpcEndpoints::new(

crates/edr_solidity_tests/tests/it/core.rs

@@ -82,6 +82,10 @@ async fn test_core() {
                     None,
                 )],
             ),
+            (
+                "default/core/ExecutionContext.t.sol:ExecutionContextTest",
+                vec![("testContext()", true, None, None, None)],
+            ),
         ]),
     );
 }

crates/edr_solidity_tests/tests/it/helpers.rs

@@ -18,7 +18,7 @@ use edr_solidity_tests::{
     MultiContractRunner, SolidityTestRunnerConfig,
 };
 use edr_test_utils::new_fd_lock;
-use foundry_cheatcodes::{FsPermissions, RpcEndpoint, RpcEndpoints};
+use foundry_cheatcodes::{ExecutionContextConfig, FsPermissions, RpcEndpoint, RpcEndpoints};
 use foundry_compilers::{
     artifacts::{CompactContractBytecode, Libraries},
     Artifact, EvmVersion, Project, ProjectCompileOutput,
@@ -100,7 +100,10 @@ impl ForgeTestProfile {
             trace: true,
             evm_opts: Self::evm_opts(),
             project_root: PROJECT_ROOT.clone(),
-            cheats_config_options: CheatsConfigOptions::default(),
+            cheats_config_options: CheatsConfigOptions {
+                execution_context: ExecutionContextConfig::Test,
+                ..CheatsConfigOptions::default()
+            },
             fuzz: TestFuzzConfig::default().into(),
             invariant: TestInvariantConfig::default().into(),
             coverage: false,

crates/edr_solidity_tests/tests/testdata/cheats/Vm.sol

@@ -8,7 +8,7 @@ pragma experimental ABIEncoderV2;
 interface Vm {
     enum CallerMode { None, Broadcast, RecurrentBroadcast, Prank, RecurrentPrank }
     enum AccountAccessKind { Call, DelegateCall, CallCode, StaticCall, Create, SelfDestruct, Resume, Balance, Extcodesize, Extcodehash, Extcodecopy }
-    enum ForgeContext { TestGroup, Test, Coverage, Snapshot, ScriptGroup, ScriptDryRun, ScriptBroadcast, ScriptResume, Unknown }
+    enum ExecutionContext { TestGroup, Test, Coverage, Snapshot, Unknown }
     struct Log { bytes32[] topics; bytes data; address emitter; }
     struct Rpc { string key; string url; }
     struct EthGetLogs { address emitter; bytes32[] topics; bytes data; bytes32 blockHash; uint64 blockNumber; bytes32 transactionHash; uint64 transactionIndex; uint256 logIndex; bool removed; }
@@ -233,7 +233,7 @@ interface Vm {
     function getNonce(address account) external view returns (uint64 nonce);
     function getRecordedLogs() external returns (Log[] memory logs);
     function indexOf(string calldata input, string calldata key) external pure returns (uint256);
-    function isContext(ForgeContext context) external view returns (bool result);
+    function isContext(ExecutionContext context) external view returns (bool result);
     function isDir(string calldata path) external returns (bool result);
     function isFile(string calldata path) external returns (bool result);
     function isPersistent(address account) external view returns (bool persistent);

crates/edr_solidity_tests/tests/testdata/default/core/ExecutionContext.t.sol

@@ -0,0 +1,17 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity 0.8.18;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+
+contract ExecutionContextTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    function testContext() public {
+        assertEq(vm.isContext(Vm.ExecutionContext.Test), true);
+        assertEq(vm.isContext(Vm.ExecutionContext.TestGroup), true);
+        assertEq(vm.isContext(Vm.ExecutionContext.Coverage), false);
+        assertEq(vm.isContext(Vm.ExecutionContext.Snapshot), false);
+        assertEq(vm.isContext(Vm.ExecutionContext.Unknown), false);
+    }
+}

crates/foundry/cheatcodes/assets/cheatcodes.json

@@ -85,44 +85,28 @@
       ]
     },
     {
-      "name": "ForgeContext",
-      "description": "Forge execution contexts.",
+      "name": "ExecutionContext",
+      "description": "Solidity test execution contexts.",
       "variants": [
         {
           "name": "TestGroup",
-          "description": "Test group execution context (test, coverage or snapshot)."
+          "description": "Test group execution context: any of test, coverage or snapshot."
         },
         {
           "name": "Test",
-          "description": "`forge test` execution context."
+          "description": "Test execution context."
         },
         {
           "name": "Coverage",
-          "description": "`forge coverage` execution context."
+          "description": "Code coverage execution context."
         },
         {
           "name": "Snapshot",
-          "description": "`forge snapshot` execution context."
-        },
-        {
-          "name": "ScriptGroup",
-          "description": "Script group execution context (dry run, broadcast or resume)."
-        },
-        {
-          "name": "ScriptDryRun",
-          "description": "`forge script` execution context."
-        },
-        {
-          "name": "ScriptBroadcast",
-          "description": "`forge script --broadcast` execution context."
-        },
-        {
-          "name": "ScriptResume",
-          "description": "`forge script --resume` execution context."
+          "description": "Gas snapshot execution context."
         },
         {
           "name": "Unknown",
-          "description": "Unknown `forge` execution context."
+          "description": "Unknown execution context."
         }
       ]
     }
@@ -4755,7 +4739,7 @@
       "func": {
         "id": "isContext",
         "description": "Returns true if `forge` command was executed in given context.",
-        "declaration": "function isContext(ForgeContext context) external view returns (bool result);",
+        "declaration": "function isContext(ExecutionContext context) external view returns (bool result);",
         "visibility": "external",
         "mutability": "view",
         "signature": "isContext(uint8)",

crates/foundry/cheatcodes/spec/src/lib.rs

@@ -94,7 +94,7 @@ impl Cheatcodes<'static> {
             enums: Cow::Owned(vec![
                 Vm::CallerMode::ENUM.clone(),
                 Vm::AccountAccessKind::ENUM.clone(),
-                Vm::ForgeContext::ENUM.clone(),
+                Vm::ExecutionContext::ENUM.clone(),
             ]),
             errors: Vm::VM_ERRORS.iter().map(|&x| x.clone()).collect(),
             events: Cow::Borrowed(&[]),

crates/foundry/cheatcodes/spec/src/vm.rs

@@ -10,7 +10,6 @@ use super::{
     Cheatcode, CheatcodeDef, Cow, Enum, EnumVariant, Error, Function, Group, Mutability, Safety,
     Status, Struct, StructField, Visibility,
 };
-use crate::Vm::ForgeContext;
 
 sol! {
 // Cheatcodes are marked as view/pure/none using the following rules:
@@ -69,26 +68,18 @@ interface Vm {
         Extcodecopy,
     }
 
-    /// Forge execution contexts.
-    enum ForgeContext {
-        /// Test group execution context (test, coverage or snapshot).
+    /// Solidity test execution contexts.
+    enum ExecutionContext {
+        /// Test group execution context: any of test, coverage or snapshot.
         TestGroup,
-        /// `forge test` execution context.
+        /// Test execution context.
         Test,
-        /// `forge coverage` execution context.
+        /// Code coverage execution context.
         Coverage,
-        /// `forge snapshot` execution context.
+        /// Gas snapshot execution context.
         Snapshot,
-        /// Script group execution context (dry run, broadcast or resume).
-        ScriptGroup,
-        /// `forge script` execution context.
-        ScriptDryRun,
-        /// `forge script --broadcast` execution context.
-        ScriptBroadcast,
-        /// `forge script --resume` execution context.
-        ScriptResume,
-        /// Unknown `forge` execution context.
-        Unknown,
+        /// Unknown execution context.
+        Unknown
     }
 
     /// An Ethereum log. Returned by `getRecordedLogs`.
@@ -1644,7 +1635,7 @@ interface Vm {
 
     /// Returns true if `forge` command was executed in given context.
     #[cheatcode(group = Environment)]
-    function isContext(ForgeContext context) external view returns (bool result);
+    function isContext(ExecutionContext context) external view returns (bool result);
 
     // ======== Utilities ========
 
@@ -2024,30 +2015,3 @@ interface Vm {
     function ensNamehash(string calldata name) external pure returns (bytes32);
 }
 }
-
-impl PartialEq for ForgeContext {
-    // Handles test group case (any of test, coverage or snapshot)
-    // and script group case (any of dry run, broadcast or resume).
-    fn eq(&self, other: &Self) -> bool {
-        match (self, other) {
-            (_, &ForgeContext::TestGroup) => {
-                self == &ForgeContext::Test
-                    || self == &ForgeContext::Snapshot
-                    || self == &ForgeContext::Coverage
-            }
-            (_, &ForgeContext::ScriptGroup) => {
-                self == &ForgeContext::ScriptDryRun
-                    || self == &ForgeContext::ScriptBroadcast
-                    || self == &ForgeContext::ScriptResume
-            }
-            (&ForgeContext::Test, &ForgeContext::Test)
-            | (&ForgeContext::Snapshot, &ForgeContext::Snapshot)
-            | (&ForgeContext::Coverage, &ForgeContext::Coverage)
-            | (&ForgeContext::ScriptDryRun, &ForgeContext::ScriptDryRun)
-            | (&ForgeContext::ScriptBroadcast, &ForgeContext::ScriptBroadcast)
-            | (&ForgeContext::ScriptResume, &ForgeContext::ScriptResume)
-            | (&ForgeContext::Unknown, &ForgeContext::Unknown) => true,
-            _ => false,
-        }
-    }
-}

crates/foundry/cheatcodes/src/config.rs

@@ -20,6 +20,9 @@ use crate::{cache::StorageCachingConfig, Vm::Rpc};
 /// to know.
 #[derive(Clone, Debug)]
 pub struct CheatsConfig {
+    /// Whether the execution is in the context of a test run, gas snapshot or
+    /// code coverage.
+    pub execution_context: ExecutionContextConfig,
     /// Whether the FFI cheatcode is enabled.
     pub ffi: bool,
     /// Use the create 2 factory in all cases including tests and
@@ -51,9 +54,25 @@ pub struct CheatsConfig {
     pub running_version: Option<Version>,
 }
 
+/// Solidity test execution contexts.
+#[derive(Clone, Debug, Default)]
+pub enum ExecutionContextConfig {
+    /// Test execution context.
+    Test,
+    /// Code coverage execution context.
+    Coverage,
+    /// Gas snapshot execution context.
+    Snapshot,
+    /// Unknown execution context.
+    #[default]
+    Unknown,
+}
+
 /// Configuration options specific to cheat codes.
 #[derive(Clone, Debug, Default)]
 pub struct CheatsConfigOptions {
+    /// Solidity test execution contexts.
+    pub execution_context: ExecutionContextConfig,
     /// Multiple rpc endpoints and their aliases
     pub rpc_endpoints: RpcEndpoints,
     /// Optional RPC cache path. If this is none, then no RPC calls will be
@@ -84,6 +103,7 @@ impl CheatsConfig {
         running_version: Option<Version>,
     ) -> Self {
         let CheatsConfigOptions {
+            execution_context,
             rpc_endpoints,
             rpc_cache_path,
             prompt_timeout,
@@ -95,6 +115,7 @@ impl CheatsConfig {
         let fs_permissions = fs_permissions.joined(&project_root);
 
         Self {
+            execution_context,
             ffi: evm_opts.ffi,
             always_use_create_2_factory: evm_opts.always_use_create_2_factory,
             prompt_timeout: Duration::from_secs(prompt_timeout),
@@ -239,6 +260,7 @@ impl CheatsConfig {
 impl Default for CheatsConfig {
     fn default() -> Self {
         Self {
+            execution_context: ExecutionContextConfig::default(),
             ffi: false,
             always_use_create_2_factory: false,
             prompt_timeout: Duration::from_secs(120),
@@ -262,6 +284,7 @@ mod tests {
 
     fn config(root: &str, fs_permissions: FsPermissions) -> CheatsConfig {
         let cheats_config_options = CheatsConfigOptions {
+            execution_context: ExecutionContextConfig::default(),
             rpc_endpoints: RpcEndpoints::default(),
             rpc_cache_path: None,
             rpc_storage_caching: StorageCachingConfig::default(),

crates/foundry/cheatcodes/src/env.rs

@@ -1,25 +1,23 @@
 //! Implementations of [`Environment`](crate::Group::Environment) cheatcodes.
 
-use std::{env, sync::OnceLock};
+use std::env;
 
 use alloy_dyn_abi::DynSolType;
 use alloy_sol_types::SolValue;
 
 use crate::{
+    config::ExecutionContextConfig,
     string, Cheatcode, Cheatcodes, Error, Result,
     Vm::{
         envAddress_0Call, envAddress_1Call, envBool_0Call, envBool_1Call, envBytes32_0Call,
         envBytes32_1Call, envBytes_0Call, envBytes_1Call, envExistsCall, envInt_0Call,
         envInt_1Call, envOr_0Call, envOr_10Call, envOr_11Call, envOr_12Call, envOr_13Call,
         envOr_1Call, envOr_2Call, envOr_3Call, envOr_4Call, envOr_5Call, envOr_6Call, envOr_7Call,
         envOr_8Call, envOr_9Call, envString_0Call, envString_1Call, envUint_0Call, envUint_1Call,
-        isContextCall, setEnvCall, ForgeContext,
+        isContextCall, setEnvCall, ExecutionContext,
     },
 };
 
-/// Stores the forge execution context for the duration of the program.
-static FORGE_CONTEXT: OnceLock<ForgeContext> = OnceLock::new();
-
 impl Cheatcode for setEnvCall {
     fn apply(&self, _state: &mut Cheatcodes) -> Result {
         let Self { name: key, value } = self;
@@ -291,19 +289,34 @@ impl Cheatcode for envOr_13Call {
 }
 
 impl Cheatcode for isContextCall {
-    fn apply(&self, _state: &mut Cheatcodes) -> Result {
-        let Self { context } = self;
-        Ok((FORGE_CONTEXT.get() == Some(context)).abi_encode())
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
+        let Self {
+            context: context_arg,
+        } = self;
+        let configured_context = &state.config.execution_context;
+
+        let group_match = matches!(
+            (configured_context, context_arg),
+            (
+                &ExecutionContextConfig::Test
+                    | &ExecutionContextConfig::Snapshot
+                    | &ExecutionContextConfig::Coverage,
+                ExecutionContext::TestGroup,
+            )
+        );
+
+        let exact_match = matches!(
+            (configured_context, context_arg),
+            (ExecutionContextConfig::Coverage, ExecutionContext::Coverage)
+                | (ExecutionContextConfig::Snapshot, ExecutionContext::Snapshot)
+                | (ExecutionContextConfig::Test, ExecutionContext::Test)
+                | (ExecutionContextConfig::Unknown, ExecutionContext::Unknown)
+        );
+
+        Ok((group_match || exact_match).abi_encode())
     }
 }
 
-/// Set `forge` command current execution context for the duration of the
-/// program. Execution context is immutable, subsequent calls of this function
-/// won't change the context.
-pub fn set_execution_context(context: ForgeContext) {
-    let _ = FORGE_CONTEXT.set(context);
-}
-
 fn env(key: &str, ty: &DynSolType) -> Result {
     get_env(key).and_then(|val| string::parse(&val, ty).map_err(map_env_err(key, &val)))
 }