Speed up mdtests (#13832)

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>

Cargo.lock

@@ -2978,7 +2978,7 @@ checksum = "e86697c916019a8588c99b5fac3cead74ec0b4b819707a682fd4d23fa0ce1ba1"
 [[package]]
 name = "salsa"
 version = "0.18.0"
-source = "git+https://github.com/salsa-rs/salsa.git?rev=b14be5c0392f4c55eca60b92e457a35549372382#b14be5c0392f4c55eca60b92e457a35549372382"
+source = "git+https://github.com/salsa-rs/salsa.git?rev=254c749b02cde2fd29852a7463a33e800b771758#254c749b02cde2fd29852a7463a33e800b771758"
 dependencies = [
  "append-only-vec",
  "arc-swap",
@@ -2998,12 +2998,12 @@ dependencies = [
 [[package]]
 name = "salsa-macro-rules"
 version = "0.1.0"
-source = "git+https://github.com/salsa-rs/salsa.git?rev=b14be5c0392f4c55eca60b92e457a35549372382#b14be5c0392f4c55eca60b92e457a35549372382"
+source = "git+https://github.com/salsa-rs/salsa.git?rev=254c749b02cde2fd29852a7463a33e800b771758#254c749b02cde2fd29852a7463a33e800b771758"
 
 [[package]]
 name = "salsa-macros"
 version = "0.18.0"
-source = "git+https://github.com/salsa-rs/salsa.git?rev=b14be5c0392f4c55eca60b92e457a35549372382#b14be5c0392f4c55eca60b92e457a35549372382"
+source = "git+https://github.com/salsa-rs/salsa.git?rev=254c749b02cde2fd29852a7463a33e800b771758#254c749b02cde2fd29852a7463a33e800b771758"
 dependencies = [
  "heck",
  "proc-macro2",

Cargo.toml

@@ -116,7 +116,7 @@ rand = { version = "0.8.5" }
 rayon = { version = "1.10.0" }
 regex = { version = "1.10.2" }
 rustc-hash = { version = "2.0.0" }
-salsa = { git = "https://github.com/salsa-rs/salsa.git", rev = "b14be5c0392f4c55eca60b92e457a35549372382" }
+salsa = { git = "https://github.com/salsa-rs/salsa.git", rev = "254c749b02cde2fd29852a7463a33e800b771758" }
 schemars = { version = "0.8.16" }
 seahash = { version = "4.1.0" }
 serde = { version = "1.0.197", features = ["derive"] }

crates/red_knot_test/src/db.rs

@@ -2,13 +2,13 @@ use red_knot_python_semantic::{
     Db as SemanticDb, Program, ProgramSettings, PythonVersion, SearchPathSettings,
 };
 use ruff_db::files::{File, Files};
-use ruff_db::system::SystemPathBuf;
-use ruff_db::system::{DbWithTestSystem, System, TestSystem};
+use ruff_db::system::{DbWithTestSystem, System, SystemPath, SystemPathBuf, TestSystem};
 use ruff_db::vendored::VendoredFileSystem;
 use ruff_db::{Db as SourceDb, Upcast};
 
 #[salsa::db]
 pub(crate) struct Db {
+    workspace_root: SystemPathBuf,
     storage: salsa::Storage<Self>,
     files: Files,
     system: TestSystem,
@@ -18,6 +18,7 @@ pub(crate) struct Db {
 impl Db {
     pub(crate) fn setup(workspace_root: SystemPathBuf) -> Self {
         let db = Self {
+            workspace_root,
             storage: salsa::Storage::default(),
             system: TestSystem::default(),
             vendored: red_knot_vendored::file_system().clone(),
@@ -25,20 +26,24 @@ impl Db {
         };
 
         db.memory_file_system()
-            .create_directory_all(&workspace_root)
+            .create_directory_all(&db.workspace_root)
             .unwrap();
 
         Program::from_settings(
             &db,
             &ProgramSettings {
                 target_version: PythonVersion::default(),
-                search_paths: SearchPathSettings::new(workspace_root),
+                search_paths: SearchPathSettings::new(db.workspace_root.clone()),
             },
         )
         .expect("Invalid search path settings");
 
         db
     }
+
+    pub(crate) fn workspace_root(&self) -> &SystemPath {
+        &self.workspace_root
+    }
 }
 
 impl DbWithTestSystem for Db {

crates/red_knot_test/src/lib.rs

@@ -1,7 +1,7 @@
 use colored::Colorize;
 use parser as test_parser;
 use red_knot_python_semantic::types::check_types;
-use ruff_db::files::system_path_to_file;
+use ruff_db::files::{system_path_to_file, Files};
 use ruff_db::parsed::parsed_module;
 use ruff_db::system::{DbWithTestSystem, SystemPathBuf};
 use std::collections::BTreeMap;
@@ -28,9 +28,15 @@ pub fn run(path: &Path, title: &str) {
         }
     };
 
+    let mut db = db::Db::setup(SystemPathBuf::from("/src"));
+
     let mut any_failures = false;
     for test in suite.tests() {
-        if let Err(failures) = run_test(&test) {
+        // Remove all files so that the db is in a "fresh" state.
+        db.memory_file_system().remove_all();
+        Files::sync_all(&mut db);
+
+        if let Err(failures) = run_test(&mut db, &test) {
             any_failures = true;
             println!("\n{}\n", test.name().bold().underline());
 
@@ -52,9 +58,8 @@ pub fn run(path: &Path, title: &str) {
     assert!(!any_failures, "Some tests failed.");
 }
 
-fn run_test(test: &parser::MarkdownTest) -> Result<(), Failures> {
-    let workspace_root = SystemPathBuf::from("/src");
-    let mut db = db::Db::setup(workspace_root.clone());
+fn run_test(db: &mut db::Db, test: &parser::MarkdownTest) -> Result<(), Failures> {
+    let workspace_root = db.workspace_root().to_path_buf();
 
     let mut system_paths = vec![];
 
@@ -71,8 +76,8 @@ fn run_test(test: &parser::MarkdownTest) -> Result<(), Failures> {
     let mut failures = BTreeMap::default();
 
     for path in system_paths {
-        let file = system_path_to_file(&db, path.clone()).unwrap();
-        let parsed = parsed_module(&db, file);
+        let file = system_path_to_file(db, path.clone()).unwrap();
+        let parsed = parsed_module(db, file);
 
         // TODO allow testing against code with syntax errors
         assert!(
@@ -83,7 +88,7 @@ fn run_test(test: &parser::MarkdownTest) -> Result<(), Failures> {
             parsed.errors()
         );
 
-        matcher::match_file(&db, file, check_types(&db, file)).unwrap_or_else(|line_failures| {
+        matcher::match_file(db, file, check_types(db, file)).unwrap_or_else(|line_failures| {
             failures.insert(path, line_failures);
         });
     }

crates/ruff_db/src/system/memory_fs.rs

@@ -351,6 +351,17 @@ impl MemoryFileSystem {
 
         Ok(ReadDirectory::new(collected))
     }
+
+    /// Removes all files and directories except the current working directory.
+    pub fn remove_all(&self) {
+        self.inner.virtual_files.write().unwrap().clear();
+
+        self.inner
+            .by_path
+            .write()
+            .unwrap()
+            .retain(|key, _| key == self.inner.cwd.as_utf8_path());
+    }
 }
 
 impl Default for MemoryFileSystem {