From b1f291d52b7e0cc098cf99439639676cdae8cedf Mon Sep 17 00:00:00 2001
From: grandizzy <grandizzy.the.egg@gmail.com>
Date: Fri, 17 Jan 2025 14:31:36 +0200
Subject: [PATCH] feat(chisel): determine proper path to Vm.sol based on proj
 remappings

---
 Cargo.lock                          |  1 +
 crates/chisel/Cargo.toml            |  1 +
 crates/chisel/src/session_source.rs | 30 ++++++++++++++++++++---------
 3 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6191cc9cfd7d..68c32878c6f5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2129,6 +2129,7 @@ dependencies = [
  "tracing",
  "tracing-subscriber",
  "vergen",
+ "walkdir",
  "yansi",
 ]
 
diff --git a/crates/chisel/Cargo.toml b/crates/chisel/Cargo.toml
index 167329d45036..5975370d49e5 100644
--- a/crates/chisel/Cargo.toml
+++ b/crates/chisel/Cargo.toml
@@ -64,6 +64,7 @@ time = { version = "0.3", features = ["formatting"] }
 tokio = { workspace = true, features = ["full"] }
 yansi.workspace = true
 tracing.workspace = true
+walkdir.workspace = true
 
 [target.'cfg(unix)'.dependencies]
 tikv-jemallocator = { workspace = true, optional = true }
diff --git a/crates/chisel/src/session_source.rs b/crates/chisel/src/session_source.rs
index 128873167a44..f6e36395c507 100644
--- a/crates/chisel/src/session_source.rs
+++ b/crates/chisel/src/session_source.rs
@@ -17,6 +17,7 @@ use semver::Version;
 use serde::{Deserialize, Serialize};
 use solang_parser::{diagnostics::Diagnostic, pt};
 use std::{fs, path::PathBuf};
+use walkdir::WalkDir;
 use yansi::Paint;
 
 /// The minimum Solidity version of the `Vm` interface.
@@ -466,15 +467,26 @@ contract {contract_name} is Script {{
     pub fn to_repl_source(&self) -> String {
         let Version { major, minor, patch, .. } = self.solc.version;
         let Self { contract_name, global_code, top_level_code, run_code, config, .. } = self;
-
-        let (vm_import, vm_constant) = if !config.no_vm {
-            (
-                "import {Vm} from \"forge-std/Vm.sol\";\n",
-                "Vm internal constant vm = Vm(address(uint160(uint256(keccak256(\"hevm cheat code\")))));\n"
-            )
-        } else {
-            ("", "")
-        };
+        let (mut vm_import, mut vm_constant) = (String::new(), String::new());
+        if !config.no_vm {
+            // Check if there's any `forge-std` remapping and determine proper path to it by
+            // searching remapping path.
+            if let Some(remapping) = config
+                .foundry_config
+                .remappings
+                .iter()
+                .find(|remapping| remapping.name == "forge-std/")
+            {
+                if let Some(vm_path) = WalkDir::new(&remapping.path.path)
+                    .into_iter()
+                    .filter_map(|e| e.ok())
+                    .find(|e| e.file_name() == "Vm.sol")
+                {
+                    vm_import = format!("import {{Vm}} from \"{}\";\n", vm_path.path().display());
+                    vm_constant = "Vm internal constant vm = Vm(address(uint160(uint256(keccak256(\"hevm cheat code\")))));\n".to_string();
+                }
+            }
+        }
 
         format!(
             r#"