LS: Add batching files for diagnostics and plug it into refreshing logic

This commit loses separation between open and other files in trace
profiles, but with recent speedups these spans stopped being useful.

commit-id:bd1161e7

crates/cairo-lang-language-server/src/lang/diagnostics/file_batches.rs

@@ -0,0 +1,50 @@
+use std::collections::HashSet;
+use std::num::NonZero;
+
+use cairo_lang_defs::db::DefsGroup;
+use cairo_lang_filesystem::db::FilesGroup;
+use cairo_lang_filesystem::ids::FileId;
+use lsp_types::Url;
+
+use crate::lang::db::AnalysisDatabase;
+use crate::lang::lsp::LsProtoGroup;
+
+/// Finds all analyzable files in `db` that are open and need to be analysed ASAP, thus _primary_.
+#[tracing::instrument(skip_all)]
+pub fn find_primary_files(db: &AnalysisDatabase, open_files: &HashSet<Url>) -> HashSet<FileId> {
+    open_files.iter().filter_map(|uri| db.file_for_url(uri)).collect()
+}
+
+/// Finds all analyzable files in `db` that are **not** primary.
+#[tracing::instrument(skip_all)]
+pub fn find_secondary_files(
+    db: &AnalysisDatabase,
+    primary_files: &HashSet<FileId>,
+) -> HashSet<FileId> {
+    let mut result = HashSet::new();
+    for crate_id in db.crates() {
+        for module_id in db.crate_modules(crate_id).iter() {
+            // Schedule only module main files for refreshing.
+            // All other related files will be refreshed along with it in a single job.
+            if let Ok(file) = db.module_main_file(*module_id) {
+                if !primary_files.contains(&file) {
+                    result.insert(file);
+                }
+            }
+        }
+    }
+    result
+}
+
+/// Returns `n` optimally distributed batches of the input.
+pub fn batches<'a>(
+    input: impl IntoIterator<Item = &'a FileId> + Clone + 'a,
+    n: NonZero<usize>,
+) -> Vec<Vec<FileId>> {
+    let n = n.get();
+    let batches = (1..=n)
+        .map(|offset| input.clone().into_iter().copied().skip(offset - 1).step_by(n).collect())
+        .collect::<Vec<_>>();
+    debug_assert_eq!(batches.len(), n);
+    batches
+}

crates/cairo-lang-language-server/src/lang/diagnostics/mod.rs

@@ -5,13 +5,15 @@ use lsp_types::Url;
 use tracing::{error, trace};
 
 use self::file_diagnostics::FileDiagnostics;
+use self::refresh::refresh_diagnostics;
 use self::trigger::trigger;
-use crate::lang::diagnostics::refresh::refresh_diagnostics;
+use crate::lang::diagnostics::file_batches::{batches, find_primary_files, find_secondary_files};
 use crate::server::client::Notifier;
 use crate::server::panic::cancelled_anyhow;
 use crate::server::schedule::thread::{self, JoinHandle, ThreadPriority};
 use crate::state::StateSnapshot;
 
+mod file_batches;
 mod file_diagnostics;
 mod lsp;
 mod refresh;
@@ -59,13 +61,7 @@ impl DiagnosticsController {
 
         while let Some(WorkerArgs { state, notifier }) = receiver.wait() {
             if let Err(err) = catch_unwind(AssertUnwindSafe(|| {
-                refresh_diagnostics(
-                    &state.db,
-                    &state.open_files,
-                    state.config.trace_macro_diagnostics,
-                    &mut file_diagnostics,
-                    notifier,
-                );
+                Self::control_tick(state, &mut file_diagnostics, notifier);
             })) {
                 if let Ok(err) = cancelled_anyhow(err, "diagnostics refreshing has been cancelled")
                 {
@@ -76,4 +72,26 @@ impl DiagnosticsController {
             }
         }
     }
+
+    /// Runs a single tick of the diagnostics controller's event loop.
+    #[tracing::instrument(skip_all)]
+    fn control_tick(
+        state: StateSnapshot,
+        file_diagnostics: &mut HashMap<Url, FileDiagnostics>,
+        notifier: Notifier,
+    ) {
+        // TODO(mkaput): Make multiple batches and run them in parallel.
+        let primary_files = find_primary_files(&state.db, &state.open_files);
+        let secondary_files = find_secondary_files(&state.db, &primary_files);
+        let files = primary_files.into_iter().chain(secondary_files).collect::<Vec<_>>();
+        for batch in batches(&files, 1.try_into().unwrap()) {
+            refresh_diagnostics(
+                &state.db,
+                batch,
+                state.config.trace_macro_diagnostics,
+                file_diagnostics,
+                notifier.clone(),
+            );
+        }
+    }
 }

crates/cairo-lang-language-server/src/lang/diagnostics/refresh.rs

@@ -1,8 +1,6 @@
 use std::collections::{HashMap, HashSet};
 
-use cairo_lang_defs::db::DefsGroup;
 use cairo_lang_defs::ids::ModuleId;
-use cairo_lang_filesystem::db::FilesGroup;
 use cairo_lang_filesystem::ids::FileId;
 use cairo_lang_utils::LookupIntern;
 use lsp_types::notification::PublishDiagnostics;
@@ -18,65 +16,26 @@ use crate::server::client::Notifier;
 #[tracing::instrument(skip_all)]
 pub fn refresh_diagnostics(
     db: &AnalysisDatabase,
-    open_files: &HashSet<Url>,
+    batch: Vec<FileId>,
     trace_macro_diagnostics: bool,
     file_diagnostics: &mut HashMap<Url, FileDiagnostics>,
     notifier: Notifier,
 ) {
     let mut files_with_set_diagnostics: HashSet<Url> = HashSet::default();
     let mut processed_modules: HashSet<ModuleId> = HashSet::default();
 
-    let open_files_ids = info_span!("get_open_files_ids").in_scope(|| {
-        open_files.iter().filter_map(|uri| db.file_for_url(uri)).collect::<HashSet<FileId>>()
-    });
-
-    // Refresh open files modules first for better UX
-    info_span!("refresh_open_files_modules").in_scope(|| {
-        for &file in &open_files_ids {
-            refresh_file_diagnostics(
-                db,
-                file,
-                trace_macro_diagnostics,
-                &mut processed_modules,
-                &mut files_with_set_diagnostics,
-                file_diagnostics,
-                &notifier,
-            );
-        }
-    });
-
-    let rest_of_files = info_span!("get_rest_of_files").in_scope(|| {
-        let mut rest_of_files: HashSet<FileId> = HashSet::default();
-        for crate_id in db.crates() {
-            for module_id in db.crate_modules(crate_id).iter() {
-                // Schedule only module main files for refreshing.
-                // All other related files will be refreshed along with it in a single job.
-                if let Ok(file) = db.module_main_file(*module_id) {
-                    if !open_files_ids.contains(&file) {
-                        rest_of_files.insert(file);
-                    }
-                }
-            }
-        }
-        rest_of_files
-    });
-
-    // Refresh the rest of files after, since they are not viewed currently
-    info_span!("refresh_other_files_modules").in_scope(|| {
-        for file in rest_of_files {
-            refresh_file_diagnostics(
-                db,
-                file,
-                trace_macro_diagnostics,
-                &mut processed_modules,
-                &mut files_with_set_diagnostics,
-                file_diagnostics,
-                &notifier,
-            );
-        }
-    });
+    for file in batch {
+        refresh_file_diagnostics(
+            db,
+            file,
+            trace_macro_diagnostics,
+            &mut processed_modules,
+            &mut files_with_set_diagnostics,
+            file_diagnostics,
+            &notifier,
+        );
+    }
 
-    // Clear old diagnostics
     info_span!("clear_old_diagnostics").in_scope(|| {
         let mut removed_files = Vec::new();