A very basic prototype for concurrent checking

crates/red_knot_workspace/Cargo.toml

@@ -22,12 +22,13 @@ ruff_text_size = { workspace = true }
 anyhow = { workspace = true }
 crossbeam = { workspace = true }
 notify = { workspace = true }
+rayon = { workspace = true }
 rustc-hash = { workspace = true }
 salsa = { workspace = true }
 tracing = { workspace = true }
 
 [dev-dependencies]
-ruff_db = { workspace = true, features = ["testing"]}
+ruff_db = { workspace = true, features = ["testing"] }
 
 [lints]
 workspace = true

crates/red_knot_workspace/src/workspace.rs

@@ -14,7 +14,8 @@ use ruff_db::{
 use ruff_python_ast::{name::Name, PySourceType};
 use ruff_text_size::Ranged;
 
-use crate::workspace::files::{Index, Indexed, PackageFiles};
+use crate::db::RootDatabase;
+use crate::workspace::files::{Index, Indexed, IndexedIter, PackageFiles};
 use crate::{
     db::Db,
     lint::{lint_semantic, lint_syntax},
@@ -190,22 +191,38 @@ impl Workspace {
 
     /// Checks all open files in the workspace and its dependencies.
     #[tracing::instrument(level = "debug", skip_all)]
-    pub fn check(self, db: &dyn Db) -> Vec<String> {
+    pub fn check(self, db: &RootDatabase) -> Vec<String> {
         tracing::debug!("Checking workspace");
 
-        let mut result = Vec::new();
+        let db = db.snapshot();
+        let mut result = Arc::new(std::sync::Mutex::new(Vec::new()));
+        let inner_result = Arc::clone(&result);
 
-        if let Some(open_files) = self.open_files(db) {
-            for file in open_files {
-                result.extend_from_slice(&check_file(db, *file));
-            }
-        } else {
-            for package in self.packages(db) {
-                result.extend(package.check(db));
+        rayon::scope(move |scope| {
+            if let Some(open_files) = self.open_files(&db) {
+                for file in open_files.iter().copied() {
+                    let result = inner_result.clone();
+                    let db = db.snapshot();
+
+                    scope.spawn(move |scope| {
+                        let file_diagnostics = check_file(&db, file);
+                        result.lock().unwrap().extend(file_diagnostics);
+                    });
+                }
+            } else {
+                for package in self.packages(&db) {
+                    let db = db.snapshot();
+                    let result = inner_result.clone();
+
+                    scope.spawn(move |scope| {
+                        let package_diagnostics = package.check(&db);
+                        result.lock().unwrap().extend(package_diagnostics);
+                    });
+                }
             }
-        }
+        });
 
-        result
+        Arc::into_inner(result).unwrap().into_inner().unwrap()
     }
 
     /// Opens a file in the workspace.
@@ -324,20 +341,32 @@ impl Package {
     }
 
     #[tracing::instrument(level = "debug", skip(db))]
-    pub(crate) fn check(self, db: &dyn Db) -> Vec<String> {
+    pub(crate) fn check(self, db: &RootDatabase) -> Vec<String> {
         tracing::debug!("Checking package '{}'", self.root(db));
 
-        let mut result = Vec::new();
-        for file in &self.files(db) {
-            let diagnostics = check_file(db, file);
-            result.extend_from_slice(&diagnostics);
-        }
+        let mut diagnostics = Arc::new(std::sync::Mutex::new(Vec::new()));
+        let files = self.files(db);
+
+        let db = db.snapshot();
+
+        let result = diagnostics.clone();
+        rayon::scope(move |scope| {
+            for file in &files {
+                let db = db.snapshot();
+                let result = result.clone();
+
+                scope.spawn(move |scope| {
+                    let diagnostics = check_file(&db, file);
+                    result.lock().unwrap().extend_from_slice(&diagnostics);
+                })
+            }
+        });
 
-        result
+        Arc::into_inner(diagnostics).unwrap().into_inner().unwrap()
     }
 
     /// Returns the files belonging to this package.
-    pub fn files(self, db: &dyn Db) -> Indexed<'_> {
+    pub fn files<'db>(self, db: &'db dyn Db) -> Indexed<'db> {
         let files = self.file_set(db);
 
         let indexed = match files.get() {
@@ -461,6 +490,39 @@ fn discover_package_files(db: &dyn Db, path: &SystemPath) -> FxHashSet<File> {
 
     files
 }
+//
+// enum WorkspaceFilesIter<'db> {
+//     OpenFiles(std::collections::hash_set::Iter<'db, File>),
+//     PackageFiles {
+//         db: &'db dyn Db,
+//         packages: std::collections::btree_map::Values<'db, SystemPathBuf, Package>,
+//         files: Option<IndexedIter<'db>>,
+//     },
+// }
+//
+// impl Iterator for WorkspaceFilesIter<'_> {
+//     type Item = File;
+//
+//     fn next(&mut self) -> Option<Self::Item> {
+//         match self {
+//             WorkspaceFilesIter::OpenFiles(files) => files.next().copied(),
+//             WorkspaceFilesIter::PackageFiles {
+//                 db,
+//                 packages,
+//                 files,
+//             } => loop {
+//                 if let Some(file) = files.as_mut().and_then(Iterator::next) {
+//                     return Some(file);
+//                 }
+//
+//                 let package = packages.next()?;
+//
+//                 let package_files = package.files(*db);
+//                 *files = Some(package_files.into_iter());
+//             },
+//         }
+//     }
+// }
 
 #[cfg(test)]
 mod tests {

crates/red_knot_workspace/src/workspace/files.rs

@@ -158,15 +158,22 @@ impl Deref for Indexed<'_> {
     }
 }
 
+pub(super) type IndexedIter<'a> = std::iter::Copied<std::collections::hash_set::Iter<'a, File>>;
+
 impl<'a> IntoIterator for &'a Indexed<'_> {
     type Item = File;
-    type IntoIter = std::iter::Copied<std::collections::hash_set::Iter<'a, File>>;
+    type IntoIter = IndexedIter<'a>;
 
     fn into_iter(self) -> Self::IntoIter {
         self.files.iter().copied()
     }
 }
 
+struct IndexedIntoIter<'db> {
+    iter: IndexedIter<'db>,
+    _lifetime: PhantomData<&'db ()>,
+}
+
 /// A Mutable view of a package's indexed files.
 ///
 /// Allows in-place mutation of the files without deep cloning the hash set.