[red-knot] Lazy package file discovery (#12452)

Co-authored-by: Carl Meyer <carl@astral.sh>

crates/red_knot/src/db.rs

@@ -12,7 +12,7 @@ use ruff_db::vendored::VendoredFileSystem;
 use ruff_db::{Db as SourceDb, Jar as SourceJar, Upcast};
 
 use crate::lint::{lint_semantic, lint_syntax, unwind_if_cancelled, Diagnostics};
-use crate::workspace::{check_file, Package, Workspace, WorkspaceMetadata};
+use crate::workspace::{check_file, Package, Package_files, Workspace, WorkspaceMetadata};
 
 mod changes;
 
@@ -22,6 +22,7 @@ pub trait Db: DbWithJar<Jar> + SemanticDb + Upcast<dyn SemanticDb> {}
 pub struct Jar(
     Workspace,
     Package,
+    Package_files,
     lint_syntax,
     lint_semantic,
     unwind_if_cancelled,

crates/red_knot/src/workspace.rs

@@ -1,5 +1,5 @@
 // TODO: Fix clippy warnings created by salsa macros
-#![allow(clippy::used_underscore_binding)]
+#![allow(clippy::used_underscore_binding, unreachable_pub)]
 
 use std::{collections::BTreeMap, sync::Arc};
 
@@ -12,11 +12,13 @@ use ruff_db::{
 };
 use ruff_python_ast::{name::Name, PySourceType};
 
+use crate::workspace::files::{Index, IndexedFiles, PackageFiles};
 use crate::{
     db::Db,
     lint::{lint_semantic, lint_syntax, Diagnostics},
 };
 
+mod files;
 mod metadata;
 
 /// The project workspace as a Salsa ingredient.
@@ -93,7 +95,7 @@ pub struct Package {
 
     /// The files that are part of this package.
     #[return_ref]
-    file_set: Arc<FxHashSet<File>>,
+    file_set: PackageFiles,
     // TODO: Add the loaded settings.
 }
 
@@ -240,6 +242,7 @@ impl Workspace {
     }
 }
 
+#[salsa::tracked]
 impl Package {
     pub fn root(self, db: &dyn Db) -> &SystemPath {
         self.root_buf(db)
@@ -247,73 +250,69 @@ impl Package {
 
     /// Returns `true` if `file` is a first-party file part of this package.
     pub fn contains_file(self, db: &dyn Db, file: File) -> bool {
-        self.files(db).contains(&file)
-    }
-
-    pub fn files(self, db: &dyn Db) -> &FxHashSet<File> {
-        self.file_set(db)
+        self.files(db).read().contains(&file)
     }
 
     #[tracing::instrument(level = "debug", skip(db))]
-    pub fn remove_file(self, db: &mut dyn Db, file: File) -> bool {
-        let mut files_arc = self.file_set(db).clone();
+    pub fn remove_file(self, db: &mut dyn Db, file: File) {
+        let Some(mut index) = PackageFiles::indexed_mut(db, self) else {
+            return;
+        };
 
-        // Set a dummy value. Salsa will cancel any pending queries and remove its own reference to `files`
-        // so that the reference counter to `files` now drops to 1.
-        self.set_file_set(db).to(Arc::new(FxHashSet::default()));
-
-        let files = Arc::get_mut(&mut files_arc).unwrap();
-        let removed = files.remove(&file);
-        self.set_file_set(db).to(files_arc);
-
-        removed
+        index.remove(file);
     }
 
     #[tracing::instrument(level = "debug", skip(db))]
-    pub fn add_file(self, db: &mut dyn Db, file: File) -> bool {
-        let mut files_arc = self.file_set(db).clone();
+    pub fn add_file(self, db: &mut dyn Db, file: File) {
+        let Some(mut index) = PackageFiles::indexed_mut(db, self) else {
+            return;
+        };
 
-        // Set a dummy value. Salsa will cancel any pending queries and remove its own reference to `files`
-        // so that the reference counter to `files` now drops to 1.
-        self.set_file_set(db).to(Arc::new(FxHashSet::default()));
-
-        let files = Arc::get_mut(&mut files_arc).unwrap();
-        let added = files.insert(file);
-        self.set_file_set(db).to(files_arc);
-
-        added
+        index.insert(file);
     }
 
     #[tracing::instrument(level = "debug", skip(db))]
     pub(crate) fn check(self, db: &dyn Db) -> Vec<String> {
         let mut result = Vec::new();
-        for file in self.files(db) {
-            let diagnostics = check_file(db, *file);
+        for file in &self.files(db).read() {
+            let diagnostics = check_file(db, file);
             result.extend_from_slice(&diagnostics);
         }
 
         result
     }
 
-    fn from_metadata(db: &dyn Db, metadata: PackageMetadata) -> Self {
-        let files = discover_package_files(db, metadata.root());
+    /// Returns the files belonging to this package.
+    #[salsa::tracked]
+    pub fn files(self, db: &dyn Db) -> IndexedFiles {
+        let files = self.file_set(db);
 
-        Self::new(db, metadata.name, metadata.root, Arc::new(files))
+        let indexed = match files.get() {
+            Index::Lazy(vacant) => {
+                let files = discover_package_files(db, self.root(db));
+                vacant.set(files)
+            }
+            Index::Indexed(indexed) => indexed,
+        };
+
+        indexed
+    }
+
+    fn from_metadata(db: &dyn Db, metadata: PackageMetadata) -> Self {
+        Self::new(db, metadata.name, metadata.root, PackageFiles::default())
     }
 
     fn update(self, db: &mut dyn Db, metadata: PackageMetadata) {
         let root = self.root(db);
         assert_eq!(root, metadata.root());
 
-        self.reload_files(db);
         self.set_name(db).to(metadata.name);
     }
 
     #[tracing::instrument(level = "debug", skip(db))]
     pub fn reload_files(self, db: &mut dyn Db) {
-        let files = discover_package_files(db, self.root(db));
-
-        self.set_file_set(db).to(Arc::new(files));
+        // Force a re-index of the files in the next revision.
+        self.set_file_set(db).to(PackageFiles::lazy());
     }
 }
 

crates/red_knot/src/workspace/files.rs

@@ -0,0 +1,252 @@
+use std::iter::FusedIterator;
+use std::ops::Deref;
+use std::sync::Arc;
+
+use rustc_hash::FxHashSet;
+
+use crate::db::Db;
+use crate::workspace::Package;
+use ruff_db::files::File;
+
+/// The indexed files of a package.
+///
+/// The indexing happens lazily, but the files are then cached for subsequent reads.
+///
+/// ## Implementation
+/// The implementation uses internal mutability to transition between the lazy and indexed state
+/// without triggering a new salsa revision. This is safe because the initial indexing happens on first access,
+/// so no query can be depending on the contents of the indexed files before that. All subsequent mutations to
+/// the indexed files must go through `IndexedFilesMut`, which uses the Salsa setter `package.set_file_set` to
+/// ensure that Salsa always knows when the set of indexed files have changed.
+#[derive(Debug)]
+pub struct PackageFiles {
+    state: std::sync::Mutex<State>,
+}
+
+impl PackageFiles {
+    pub fn lazy() -> Self {
+        Self {
+            state: std::sync::Mutex::new(State::Lazy),
+        }
+    }
+
+    fn indexed(indexed_files: IndexedFiles) -> Self {
+        Self {
+            state: std::sync::Mutex::new(State::Indexed(indexed_files)),
+        }
+    }
+
+    pub fn get(&self) -> Index {
+        let state = self.state.lock().unwrap();
+
+        match &*state {
+            State::Lazy => Index::Lazy(LazyFiles { files: state }),
+            State::Indexed(files) => Index::Indexed(files.clone()),
+        }
+    }
+
+    /// Returns a mutable view on the index that allows cheap in-place mutations.
+    ///
+    /// The changes are automatically written back to the database once the view is dropped.
+    pub fn indexed_mut(db: &mut dyn Db, package: Package) -> Option<IndexedFilesMut> {
+        // Calling `runtime_mut` cancels all pending salsa queries. This ensures that there are no pending
+        // reads to the file set.
+        let _ = db.runtime_mut();
+
+        let files = package.file_set(db);
+
+        let indexed = match &*files.state.lock().unwrap() {
+            State::Lazy => return None,
+            State::Indexed(indexed) => indexed.clone(),
+        };
+
+        Some(IndexedFilesMut {
+            db: Some(db),
+            package,
+            new_revision: indexed.revision,
+            indexed,
+        })
+    }
+}
+
+impl Default for PackageFiles {
+    fn default() -> Self {
+        Self::lazy()
+    }
+}
+
+#[derive(Debug)]
+enum State {
+    /// The files of a package haven't been indexed yet.
+    Lazy,
+
+    /// The files are indexed. Stores the known files of a package.
+    Indexed(IndexedFiles),
+}
+
+pub enum Index<'a> {
+    /// The index has not yet been computed. Allows inserting the files.
+    Lazy(LazyFiles<'a>),
+
+    Indexed(IndexedFiles),
+}
+
+/// Package files that have not been indexed yet.
+pub struct LazyFiles<'a> {
+    files: std::sync::MutexGuard<'a, State>,
+}
+
+impl<'a> LazyFiles<'a> {
+    /// Sets the indexed files of a package to `files`.
+    pub fn set(mut self, files: FxHashSet<File>) -> IndexedFiles {
+        let files = IndexedFiles::new(files);
+        *self.files = State::Indexed(files.clone());
+        files
+    }
+}
+
+/// The indexed files of a package.
+///
+/// # Salsa integration
+/// The type is cheap clonable and allows for in-place mutation of the files. The in-place mutation requires
+/// extra care because the type is used as the result of Salsa queries and Salsa relies on a type's equality
+/// to determine if the output has changed. This is accomplished by using a `revision` that gets incremented
+/// whenever the files are changed. The revision ensures that salsa's comparison of the
+/// previous [`IndexedFiles`] with the next [`IndexedFiles`] returns false even though they both
+/// point to the same underlying hash set.
+///
+/// # Equality
+/// Two [`IndexedFiles`] are only equal if they have the same revision and point to the **same** (identity) hash set.
+#[derive(Debug, Clone)]
+pub struct IndexedFiles {
+    revision: u64,
+    files: Arc<std::sync::Mutex<FxHashSet<File>>>,
+}
+
+impl IndexedFiles {
+    fn new(files: FxHashSet<File>) -> Self {
+        Self {
+            files: Arc::new(std::sync::Mutex::new(files)),
+            revision: 0,
+        }
+    }
+
+    /// Locks the file index for reading.
+    pub fn read(&self) -> IndexedFilesGuard {
+        IndexedFilesGuard {
+            guard: self.files.lock().unwrap(),
+        }
+    }
+}
+
+impl PartialEq for IndexedFiles {
+    fn eq(&self, other: &Self) -> bool {
+        self.revision == other.revision && Arc::ptr_eq(&self.files, &other.files)
+    }
+}
+
+impl Eq for IndexedFiles {}
+
+pub struct IndexedFilesGuard<'a> {
+    guard: std::sync::MutexGuard<'a, FxHashSet<File>>,
+}
+
+impl Deref for IndexedFilesGuard<'_> {
+    type Target = FxHashSet<File>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.guard
+    }
+}
+
+impl<'a> IntoIterator for &'a IndexedFilesGuard<'a> {
+    type Item = File;
+    type IntoIter = IndexedFilesIter<'a>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        IndexedFilesIter {
+            inner: self.guard.iter(),
+        }
+    }
+}
+
+/// Iterator over the indexed files.
+///
+/// # Locks
+/// Holding on to the iterator locks the file index for reading.
+pub struct IndexedFilesIter<'a> {
+    inner: std::collections::hash_set::Iter<'a, File>,
+}
+
+impl<'a> Iterator for IndexedFilesIter<'a> {
+    type Item = File;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.inner.next().copied()
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.inner.size_hint()
+    }
+}
+
+impl FusedIterator for IndexedFilesIter<'_> {}
+
+impl ExactSizeIterator for IndexedFilesIter<'_> {}
+
+/// A Mutable view of a package's indexed files.
+///
+/// Allows in-place mutation of the files without deep cloning the hash set.
+/// The changes are written back when the mutable view is dropped or by calling [`Self::set`] manually.
+pub struct IndexedFilesMut<'db> {
+    db: Option<&'db mut dyn Db>,
+    package: Package,
+    indexed: IndexedFiles,
+    new_revision: u64,
+}
+
+impl IndexedFilesMut<'_> {
+    pub fn insert(&mut self, file: File) -> bool {
+        if self.indexed.files.lock().unwrap().insert(file) {
+            self.new_revision += 1;
+            true
+        } else {
+            false
+        }
+    }
+
+    pub fn remove(&mut self, file: File) -> bool {
+        if self.indexed.files.lock().unwrap().remove(&file) {
+            self.new_revision += 1;
+            true
+        } else {
+            false
+        }
+    }
+
+    /// Writes the changes back to the database.
+    pub fn set(mut self) {
+        self.set_impl();
+    }
+
+    fn set_impl(&mut self) {
+        let Some(db) = self.db.take() else {
+            return;
+        };
+
+        if self.indexed.revision != self.new_revision {
+            self.package
+                .set_file_set(db)
+                .to(PackageFiles::indexed(IndexedFiles {
+                    revision: self.new_revision,
+                    files: self.indexed.files.clone(),
+                }));
+        }
+    }
+}
+
+impl Drop for IndexedFilesMut<'_> {
+    fn drop(&mut self) {
+        self.set_impl();
+    }
+}

crates/red_knot/tests/file_watching.rs

@@ -9,7 +9,7 @@ use red_knot::watch;
 use red_knot::watch::{directory_watcher, Watcher};
 use red_knot::workspace::WorkspaceMetadata;
 use red_knot_module_resolver::{resolve_module, ModuleName};
-use ruff_db::files::system_path_to_file;
+use ruff_db::files::{system_path_to_file, File};
 use ruff_db::program::{ProgramSettings, SearchPathSettings, TargetVersion};
 use ruff_db::source::source_text;
 use ruff_db::system::{OsSystem, SystemPath, SystemPathBuf};
@@ -54,6 +54,19 @@ impl TestCase {
 
         all_events
     }
+
+    fn collect_package_files(&self, path: &SystemPath) -> Vec<File> {
+        let package = self.db().workspace().package(self.db(), path).unwrap();
+        let files = package.files(self.db());
+        let files = files.read();
+        let mut collected: Vec<_> = files.into_iter().collect();
+        collected.sort_unstable_by_key(|file| file.path(self.db()).as_system_path().unwrap());
+        collected
+    }
+
+    fn system_file(&self, path: impl AsRef<SystemPath>) -> Option<File> {
+        system_path_to_file(self.db(), path.as_ref())
+    }
 }
 
 fn setup<I, P>(workspace_files: I) -> anyhow::Result<TestCase>
@@ -135,9 +148,12 @@ where
 #[test]
 fn new_file() -> anyhow::Result<()> {
     let mut case = setup([("bar.py", "")])?;
+    let bar_path = case.workspace_path("bar.py");
+    let bar_file = case.system_file(&bar_path).unwrap();
     let foo_path = case.workspace_path("foo.py");
 
-    assert_eq!(system_path_to_file(case.db(), &foo_path), None);
+    assert_eq!(case.system_file(&foo_path), None);
+    assert_eq!(&case.collect_package_files(&bar_path), &[bar_file]);
 
     std::fs::write(foo_path.as_std_path(), "print('Hello')")?;
 
@@ -145,15 +161,9 @@ fn new_file() -> anyhow::Result<()> {
 
     case.db_mut().apply_changes(changes);
 
-    let foo = system_path_to_file(case.db(), &foo_path).expect("foo.py to exist.");
+    let foo = case.system_file(&foo_path).expect("foo.py to exist.");
 
-    let package = case
-        .db()
-        .workspace()
-        .package(case.db(), &foo_path)
-        .expect("foo.py to belong to a package.");
-
-    assert!(package.contains_file(case.db(), foo));
+    assert_eq!(&case.collect_package_files(&bar_path), &[bar_file, foo]);
 
     Ok(())
 }
@@ -161,9 +171,12 @@ fn new_file() -> anyhow::Result<()> {
 #[test]
 fn new_ignored_file() -> anyhow::Result<()> {
     let mut case = setup([("bar.py", ""), (".ignore", "foo.py")])?;
+    let bar_path = case.workspace_path("bar.py");
+    let bar_file = case.system_file(&bar_path).unwrap();
     let foo_path = case.workspace_path("foo.py");
 
-    assert_eq!(system_path_to_file(case.db(), &foo_path), None);
+    assert_eq!(case.system_file(&foo_path), None);
+    assert_eq!(&case.collect_package_files(&bar_path), &[bar_file]);
 
     std::fs::write(foo_path.as_std_path(), "print('Hello')")?;
 
@@ -171,15 +184,8 @@ fn new_ignored_file() -> anyhow::Result<()> {
 
     case.db_mut().apply_changes(changes);
 
-    let foo = system_path_to_file(case.db(), &foo_path).expect("foo.py to exist.");
-
-    let package = case
-        .db()
-        .workspace()
-        .package(case.db(), &foo_path)
-        .expect("foo.py to belong to a package.");
-
-    assert!(!package.contains_file(case.db(), foo));
+    assert!(case.system_file(&foo_path).is_some());
+    assert_eq!(&case.collect_package_files(&bar_path), &[bar_file]);
 
     Ok(())
 }
@@ -190,8 +196,11 @@ fn changed_file() -> anyhow::Result<()> {
     let mut case = setup([("foo.py", foo_source)])?;
     let foo_path = case.workspace_path("foo.py");
 
-    let foo = system_path_to_file(case.db(), &foo_path).ok_or_else(|| anyhow!("Foo not found"))?;
+    let foo = case
+        .system_file(&foo_path)
+        .ok_or_else(|| anyhow!("Foo not found"))?;
     assert_eq!(source_text(case.db(), foo).as_str(), foo_source);
+    assert_eq!(&case.collect_package_files(&foo_path), &[foo]);
 
     std::fs::write(foo_path.as_std_path(), "print('Version 2')")?;
 
@@ -200,6 +209,7 @@ fn changed_file() -> anyhow::Result<()> {
     case.db_mut().apply_changes(changes);
 
     assert_eq!(source_text(case.db(), foo).as_str(), "print('Version 2')");
+    assert_eq!(&case.collect_package_files(&foo_path), &[foo]);
 
     Ok(())
 }
@@ -212,7 +222,9 @@ fn changed_metadata() -> anyhow::Result<()> {
     let mut case = setup([("foo.py", "")])?;
     let foo_path = case.workspace_path("foo.py");
 
-    let foo = system_path_to_file(case.db(), &foo_path).ok_or_else(|| anyhow!("Foo not found"))?;
+    let foo = case
+        .system_file(&foo_path)
+        .ok_or_else(|| anyhow!("Foo not found"))?;
     assert_eq!(
         foo.permissions(case.db()),
         Some(
@@ -252,14 +264,12 @@ fn deleted_file() -> anyhow::Result<()> {
     let mut case = setup([("foo.py", foo_source)])?;
     let foo_path = case.workspace_path("foo.py");
 
-    let foo = system_path_to_file(case.db(), &foo_path).ok_or_else(|| anyhow!("Foo not found"))?;
-
-    let Some(package) = case.db().workspace().package(case.db(), &foo_path) else {
-        panic!("Expected foo.py to belong to a package.");
-    };
+    let foo = case
+        .system_file(&foo_path)
+        .ok_or_else(|| anyhow!("Foo not found"))?;
 
     assert!(foo.exists(case.db()));
-    assert!(package.contains_file(case.db(), foo));
+    assert_eq!(&case.collect_package_files(&foo_path), &[foo]);
 
     std::fs::remove_file(foo_path.as_std_path())?;
 
@@ -268,7 +278,7 @@ fn deleted_file() -> anyhow::Result<()> {
     case.db_mut().apply_changes(changes);
 
     assert!(!foo.exists(case.db()));
-    assert!(!package.contains_file(case.db(), foo));
+    assert_eq!(&case.collect_package_files(&foo_path), &[] as &[File]);
 
     Ok(())
 }
@@ -285,14 +295,12 @@ fn move_file_to_trash() -> anyhow::Result<()> {
     let trash_path = case.root_path().join(".trash");
     std::fs::create_dir_all(trash_path.as_std_path())?;
 
-    let foo = system_path_to_file(case.db(), &foo_path).ok_or_else(|| anyhow!("Foo not found"))?;
-
-    let Some(package) = case.db().workspace().package(case.db(), &foo_path) else {
-        panic!("Expected foo.py to belong to a package.");
-    };
+    let foo = case
+        .system_file(&foo_path)
+        .ok_or_else(|| anyhow!("Foo not found"))?;
 
     assert!(foo.exists(case.db()));
-    assert!(package.contains_file(case.db(), foo));
+    assert_eq!(&case.collect_package_files(&foo_path), &[foo]);
 
     std::fs::rename(
         foo_path.as_std_path(),
@@ -304,7 +312,7 @@ fn move_file_to_trash() -> anyhow::Result<()> {
     case.db_mut().apply_changes(changes);
 
     assert!(!foo.exists(case.db()));
-    assert!(!package.contains_file(case.db(), foo));
+    assert_eq!(&case.collect_package_files(&foo_path), &[] as &[File]);
 
     Ok(())
 }
@@ -313,13 +321,16 @@ fn move_file_to_trash() -> anyhow::Result<()> {
 #[test]
 fn move_file_to_workspace() -> anyhow::Result<()> {
     let mut case = setup([("bar.py", "")])?;
+    let bar_path = case.workspace_path("bar.py");
+    let bar = case.system_file(&bar_path).unwrap();
+
     let foo_path = case.root_path().join("foo.py");
     std::fs::write(foo_path.as_std_path(), "")?;
 
     let foo_in_workspace_path = case.workspace_path("foo.py");
 
-    assert!(system_path_to_file(case.db(), &foo_path).is_some());
-
+    assert!(case.system_file(&foo_path).is_some());
+    assert_eq!(&case.collect_package_files(&bar_path), &[bar]);
     assert!(case
         .db()
         .workspace()
@@ -332,19 +343,15 @@ fn move_file_to_workspace() -> anyhow::Result<()> {
 
     case.db_mut().apply_changes(changes);
 
-    let foo_in_workspace = system_path_to_file(case.db(), &foo_in_workspace_path)
+    let foo_in_workspace = case
+        .system_file(&foo_in_workspace_path)
         .ok_or_else(|| anyhow!("Foo not found"))?;
 
-    let Some(package) = case
-        .db()
-        .workspace()
-        .package(case.db(), &foo_in_workspace_path)
-    else {
-        panic!("Expected foo.py to belong to a package.");
-    };
-
     assert!(foo_in_workspace.exists(case.db()));
-    assert!(package.contains_file(case.db(), foo_in_workspace));
+    assert_eq!(
+        &case.collect_package_files(&foo_in_workspace_path),
+        &[bar, foo_in_workspace]
+    );
 
     Ok(())
 }
@@ -356,11 +363,11 @@ fn rename_file() -> anyhow::Result<()> {
     let foo_path = case.workspace_path("foo.py");
     let bar_path = case.workspace_path("bar.py");
 
-    let foo = system_path_to_file(case.db(), &foo_path).ok_or_else(|| anyhow!("Foo not found"))?;
+    let foo = case
+        .system_file(&foo_path)
+        .ok_or_else(|| anyhow!("Foo not found"))?;
 
-    let Some(package) = case.db().workspace().package(case.db(), &foo_path) else {
-        panic!("Expected foo.py to belong to a package.");
-    };
+    assert_eq!(case.collect_package_files(&foo_path), [foo]);
 
     std::fs::rename(foo_path.as_std_path(), bar_path.as_std_path())?;
 
@@ -369,16 +376,13 @@ fn rename_file() -> anyhow::Result<()> {
     case.db_mut().apply_changes(changes);
 
     assert!(!foo.exists(case.db()));
-    assert!(!package.contains_file(case.db(), foo));
 
-    let bar = system_path_to_file(case.db(), &bar_path).ok_or_else(|| anyhow!("Bar not found"))?;
-
-    let Some(package) = case.db().workspace().package(case.db(), &bar_path) else {
-        panic!("Expected bar.py to belong to a package.");
-    };
+    let bar = case
+        .system_file(&bar_path)
+        .ok_or_else(|| anyhow!("Bar not found"))?;
 
     assert!(bar.exists(case.db()));
-    assert!(package.contains_file(case.db(), bar));
+    assert_eq!(case.collect_package_files(&foo_path), [bar]);
 
     Ok(())
 }
@@ -386,6 +390,7 @@ fn rename_file() -> anyhow::Result<()> {
 #[test]
 fn directory_moved_to_workspace() -> anyhow::Result<()> {
     let mut case = setup([("bar.py", "import sub.a")])?;
+    let bar = case.system_file(case.workspace_path("bar.py")).unwrap();
 
     let sub_original_path = case.root_path().join("sub");
     let init_original_path = sub_original_path.join("__init__.py");
@@ -400,6 +405,10 @@ fn directory_moved_to_workspace() -> anyhow::Result<()> {
     let sub_a_module = resolve_module(case.db().upcast(), ModuleName::new_static("sub.a").unwrap());
 
     assert_eq!(sub_a_module, None);
+    assert_eq!(
+        case.collect_package_files(&case.workspace_path("bar.py")),
+        &[bar]
+    );
 
     let sub_new_path = case.workspace_path("sub");
     std::fs::rename(sub_original_path.as_std_path(), sub_new_path.as_std_path())
@@ -409,21 +418,20 @@ fn directory_moved_to_workspace() -> anyhow::Result<()> {
 
     case.db_mut().apply_changes(changes);
 
-    let init_file = system_path_to_file(case.db(), sub_new_path.join("__init__.py"))
+    let init_file = case
+        .system_file(sub_new_path.join("__init__.py"))
         .expect("__init__.py to exist");
-    let a_file = system_path_to_file(case.db(), sub_new_path.join("a.py")).expect("a.py to exist");
+    let a_file = case
+        .system_file(sub_new_path.join("a.py"))
+        .expect("a.py to exist");
 
     // `import sub.a` should now resolve
     assert!(resolve_module(case.db().upcast(), ModuleName::new_static("sub.a").unwrap()).is_some());
 
-    let package = case
-        .db()
-        .workspace()
-        .package(case.db(), &sub_new_path)
-        .expect("sub to belong to a package");
-
-    assert!(package.contains_file(case.db(), init_file));
-    assert!(package.contains_file(case.db(), a_file));
+    assert_eq!(
+        case.collect_package_files(&case.workspace_path("bar.py")),
+        &[bar, init_file, a_file]
+    );
 
     Ok(())
 }
@@ -435,23 +443,22 @@ fn directory_moved_to_trash() -> anyhow::Result<()> {
         ("sub/__init__.py", ""),
         ("sub/a.py", ""),
     ])?;
+    let bar = case.system_file(case.workspace_path("bar.py")).unwrap();
 
     assert!(resolve_module(case.db().upcast(), ModuleName::new_static("sub.a").unwrap()).is_some(),);
 
     let sub_path = case.workspace_path("sub");
+    let init_file = case
+        .system_file(sub_path.join("__init__.py"))
+        .expect("__init__.py to exist");
+    let a_file = case
+        .system_file(sub_path.join("a.py"))
+        .expect("a.py to exist");
 
-    let package = case
-        .db()
-        .workspace()
-        .package(case.db(), &sub_path)
-        .expect("sub to belong to a package");
-
-    let init_file =
-        system_path_to_file(case.db(), sub_path.join("__init__.py")).expect("__init__.py to exist");
-    let a_file = system_path_to_file(case.db(), sub_path.join("a.py")).expect("a.py to exist");
-
-    assert!(package.contains_file(case.db(), init_file));
-    assert!(package.contains_file(case.db(), a_file));
+    assert_eq!(
+        case.collect_package_files(&case.workspace_path("bar.py")),
+        &[bar, init_file, a_file]
+    );
 
     std::fs::create_dir(case.root_path().join(".trash").as_std_path())?;
     let trashed_sub = case.root_path().join(".trash/sub");
@@ -468,8 +475,10 @@ fn directory_moved_to_trash() -> anyhow::Result<()> {
     assert!(!init_file.exists(case.db()));
     assert!(!a_file.exists(case.db()));
 
-    assert!(!package.contains_file(case.db(), init_file));
-    assert!(!package.contains_file(case.db(), a_file));
+    assert_eq!(
+        case.collect_package_files(&case.workspace_path("bar.py")),
+        &[bar]
+    );
 
     Ok(())
 }
@@ -482,6 +491,8 @@ fn directory_renamed() -> anyhow::Result<()> {
         ("sub/a.py", ""),
     ])?;
 
+    let bar = case.system_file(case.workspace_path("bar.py")).unwrap();
+
     assert!(resolve_module(case.db().upcast(), ModuleName::new_static("sub.a").unwrap()).is_some());
     assert!(resolve_module(
         case.db().upcast(),
@@ -490,19 +501,17 @@ fn directory_renamed() -> anyhow::Result<()> {
     .is_none());
 
     let sub_path = case.workspace_path("sub");
+    let sub_init = case
+        .system_file(sub_path.join("__init__.py"))
+        .expect("__init__.py to exist");
+    let sub_a = case
+        .system_file(sub_path.join("a.py"))
+        .expect("a.py to exist");
 
-    let package = case
-        .db()
-        .workspace()
-        .package(case.db(), &sub_path)
-        .expect("sub to belong to a package");
-
-    let sub_init =
-        system_path_to_file(case.db(), sub_path.join("__init__.py")).expect("__init__.py to exist");
-    let sub_a = system_path_to_file(case.db(), sub_path.join("a.py")).expect("a.py to exist");
-
-    assert!(package.contains_file(case.db(), sub_init));
-    assert!(package.contains_file(case.db(), sub_a));
+    assert_eq!(
+        case.collect_package_files(&sub_path),
+        &[bar, sub_init, sub_a]
+    );
 
     let foo_baz = case.workspace_path("foo/baz");
 
@@ -527,20 +536,22 @@ fn directory_renamed() -> anyhow::Result<()> {
     assert!(!sub_init.exists(case.db()));
     assert!(!sub_a.exists(case.db()));
 
-    assert!(!package.contains_file(case.db(), sub_init));
-    assert!(!package.contains_file(case.db(), sub_a));
-
-    let foo_baz_init =
-        system_path_to_file(case.db(), foo_baz.join("__init__.py")).expect("__init__.py to exist");
-    let foo_baz_a = system_path_to_file(case.db(), foo_baz.join("a.py")).expect("a.py to exist");
+    let foo_baz_init = case
+        .system_file(foo_baz.join("__init__.py"))
+        .expect("__init__.py to exist");
+    let foo_baz_a = case
+        .system_file(foo_baz.join("a.py"))
+        .expect("a.py to exist");
 
     // The new paths are synced
 
     assert!(foo_baz_init.exists(case.db()));
     assert!(foo_baz_a.exists(case.db()));
 
-    assert!(package.contains_file(case.db(), foo_baz_init));
-    assert!(package.contains_file(case.db(), foo_baz_a));
+    assert_eq!(
+        case.collect_package_files(&sub_path),
+        &[bar, foo_baz_init, foo_baz_a]
+    );
 
     Ok(())
 }
@@ -553,22 +564,22 @@ fn directory_deleted() -> anyhow::Result<()> {
         ("sub/a.py", ""),
     ])?;
 
+    let bar = case.system_file(case.workspace_path("bar.py")).unwrap();
+
     assert!(resolve_module(case.db().upcast(), ModuleName::new_static("sub.a").unwrap()).is_some(),);
 
     let sub_path = case.workspace_path("sub");
 
-    let package = case
-        .db()
-        .workspace()
-        .package(case.db(), &sub_path)
-        .expect("sub to belong to a package");
-
-    let init_file =
-        system_path_to_file(case.db(), sub_path.join("__init__.py")).expect("__init__.py to exist");
-    let a_file = system_path_to_file(case.db(), sub_path.join("a.py")).expect("a.py to exist");
-
-    assert!(package.contains_file(case.db(), init_file));
-    assert!(package.contains_file(case.db(), a_file));
+    let init_file = case
+        .system_file(sub_path.join("__init__.py"))
+        .expect("__init__.py to exist");
+    let a_file = case
+        .system_file(sub_path.join("a.py"))
+        .expect("a.py to exist");
+    assert_eq!(
+        case.collect_package_files(&sub_path),
+        &[bar, init_file, a_file]
+    );
 
     std::fs::remove_dir_all(sub_path.as_std_path())
         .with_context(|| "Failed to remove the sub directory")?;
@@ -582,9 +593,7 @@ fn directory_deleted() -> anyhow::Result<()> {
 
     assert!(!init_file.exists(case.db()));
     assert!(!a_file.exists(case.db()));
-
-    assert!(!package.contains_file(case.db(), init_file));
-    assert!(!package.contains_file(case.db(), a_file));
+    assert_eq!(case.collect_package_files(&sub_path), &[bar]);
 
     Ok(())
 }

crates/ruff_db/src/lib.rs

@@ -19,7 +19,8 @@ pub mod system;
 pub mod testing;
 pub mod vendored;
 
-pub(crate) type FxDashMap<K, V> = dashmap::DashMap<K, V, BuildHasherDefault<FxHasher>>;
+pub type FxDashMap<K, V> = dashmap::DashMap<K, V, BuildHasherDefault<FxHasher>>;
+pub type FxDashSet<K> = dashmap::DashSet<K, BuildHasherDefault<FxHasher>>;
 
 #[salsa::jar(db=Db)]
 pub struct Jar(File, Program, source_text, line_index, parsed_module);