fix(forge): preprocess mocks declared in test file (#12516)

* fix(forge): preprocess mockes declared in same test file

* Bump compilers version

Author: grandizzy

Cargo.lock

@@ -223,7 +223,7 @@ foundry-linking = { path = "crates/linking" }
 
 # solc & compilation utilities
 foundry-block-explorers = { version = "0.22.0", default-features = false }
-foundry-compilers = { version = "0.19.5", default-features = false, features = [
+foundry-compilers = { version = "0.19.6", default-features = false, features = [
     "rustls",
     "svm-solc",
 ] }

Cargo.toml

@@ -33,41 +33,54 @@ impl PreprocessorDependencies {
     ) -> Self {
         let mut preprocessed_contracts = BTreeMap::new();
         let mut referenced_contracts = HashSet::new();
-        for contract_id in gcx.hir.contract_ids() {
-            let contract = gcx.hir.contract(contract_id);
-            let source = gcx.hir.source(contract.source);
+        let mut current_mocks = HashSet::new();
+
+        // Helper closure for iterating candidate contracts to preprocess (tests and scripts).
+        let candidate_contracts = || {
+            gcx.hir.contract_ids().filter_map(|id| {
+                let contract = gcx.hir.contract(id);
+                let source = gcx.hir.source(contract.source);
+                let FileName::Real(path) = &source.file.name else {
+                    return None;
+                };
 
-            let FileName::Real(path) = &source.file.name else {
-                continue;
-            };
+                if !paths.contains(path) {
+                    trace!("{} is not test or script", path.display());
+                    return None;
+                }
 
-            // Collect dependencies only for tests and scripts.
-            if !paths.contains(path) {
-                let path = path.display();
-                trace!("{path} is not test or script");
-                continue;
-            }
+                Some((id, contract, source, path))
+            })
+        };
 
-            // Do not collect dependencies for mock contracts. Walk through base contracts and
-            // check if they're from src dir.
-            if contract.linearized_bases.iter().any(|base_contract_id| {
-                let base_contract = gcx.hir.contract(*base_contract_id);
-                let FileName::Real(path) = &gcx.hir.source(base_contract.source).file.name else {
-                    return false;
-                };
-                path.starts_with(src_dir)
+        // Collect current mocks.
+        for (_, contract, _, path) in candidate_contracts() {
+            if contract.linearized_bases.iter().any(|base_id| {
+                let base = gcx.hir.contract(*base_id);
+                matches!(
+                    &gcx.hir.source(base.source).file.name,
+                    FileName::Real(base_path) if base_path.starts_with(src_dir)
+                )
             }) {
-                // Record mock contracts to be evicted from preprocessed cache.
-                mocks.insert(root_dir.join(path));
-                let path = path.display();
-                trace!("found mock contract {path}");
+                let mock_path = root_dir.join(path);
+                trace!("found mock contract {}", mock_path.display());
+                current_mocks.insert(mock_path);
+            }
+        }
+
+        // Collect dependencies for non-mock test/script contracts.
+        for (contract_id, contract, source, path) in candidate_contracts() {
+            let full_path = root_dir.join(path);
+
+            if current_mocks.contains(&full_path) {
+                trace!("{} is a mock, skipping", path.display());
                 continue;
-            } else {
-                // Make sure current contract is not in list of mocks (could happen when a contract
-                // which used to be a mock is refactored to a non-mock implementation).
-                mocks.remove(&root_dir.join(path));
             }
 
+            // Make sure current contract is not in list of mocks (could happen when a contract
+            // which used to be a mock is refactored to a non-mock implementation).
+            mocks.remove(&full_path);
+
             let mut deps_collector =
                 BytecodeDependencyCollector::new(gcx, source.file.src.as_str(), src_dir);
             // Analyze current contract.
@@ -76,9 +89,14 @@ impl PreprocessorDependencies {
             if !deps_collector.dependencies.is_empty() {
                 preprocessed_contracts.insert(contract_id, deps_collector.dependencies);
             }
+
             // Record collected referenced contract ids.
             referenced_contracts.extend(deps_collector.referenced_contracts);
         }
+
+        // Add current mocks.
+        mocks.extend(current_mocks);
+
         Self { preprocessed_contracts, referenced_contracts }
     }
 }

crates/common/src/preprocessor/deps.rs

@@ -716,6 +716,82 @@ Compiling 2 files with [..]
 "#]]);
 });
 
+// <https://github.com/foundry-rs/foundry/issues/12452>
+// - CounterMock contract is Counter contract
+// - CounterMock declared in CounterTest
+//
+// ├── src
+// │ └── Counter.sol
+// └── test
+//    ├── Counter.t.sol
+forgetest_init!(preprocess_mock_declared_in_test_contract, |prj, cmd| {
+    prj.update_config(|config| {
+        config.dynamic_test_linking = true;
+    });
+
+    prj.add_source(
+        "Counter.sol",
+        r#"
+contract Counter {
+    function add(uint256 x, uint256 y) public pure returns (uint256) {
+        return x + y;
+    }
+}
+    "#,
+    );
+
+    prj.add_test(
+        "Counter.t.sol",
+        r#"
+import {Test} from "forge-std/Test.sol";
+import {Counter} from "src/Counter.sol";
+
+contract CounterMock is Counter {}
+
+contract CounterTest is Test {
+    function test_add() public {
+        CounterMock impl = new CounterMock();
+        assertEq(impl.add(2, 2), 4);
+    }
+}
+    "#,
+    );
+    // 20 files plus one mock file are compiled on first run.
+    cmd.args(["test"]).with_no_redact().assert_success().stdout_eq(str![[r#"
+...
+Compiling 21 files with [..]
+...
+
+"#]]);
+    cmd.with_no_redact().assert_success().stdout_eq(str![[r#"
+...
+No files changed, compilation skipped
+...
+
+"#]]);
+
+    // Change Counter implementation to fail tests.
+    prj.add_source(
+        "Counter.sol",
+        r#"
+contract Counter {
+    function add(uint256 x, uint256 y) public pure returns (uint256) {
+        return x + y + 1;
+    }
+}
+    "#,
+    );
+    // Assert that Counter and CounterTest files are compiled and tests fail.
+    cmd.with_no_redact().assert_failure().stdout_eq(str![[r#"
+...
+Compiling 2 files with [..]
+...
+[FAIL: assertion failed: 5 != 4] test_add() (gas: [..])
+...
+
+"#]]);
+});
+
 // ├── src
 // │ ├── CounterA.sol
 // │ ├── CounterB.sol