chore(lsp): avoid redundant type checking

tooling/lsp/src/lib.rs

@@ -29,6 +29,7 @@ use nargo::{
 };
 use nargo_toml::{PackageSelection, find_file_manifest, resolve_workspace_from_toml};
 use noirc_driver::NOIR_ARTIFACT_VERSION_STRING;
+use noirc_errors::CustomDiagnostic;
 use noirc_frontend::{
     ParsedModule,
     graph::{CrateGraph, CrateId, CrateName},
@@ -90,12 +91,15 @@ pub struct LspState {
     root_path: Option<PathBuf>,
     client: ClientSocket,
     solver: WrapperSolver,
-    open_documents_count: usize,
     input_files: HashMap<String, String>,
     cached_parsed_files: HashMap<PathBuf, (usize, (ParsedModule, Vec<ParserError>))>,
     workspace_cache: HashMap<PathBuf, WorkspaceCacheData>,
     package_cache: HashMap<PathBuf, PackageCacheData>,
     workspace_symbol_cache: WorkspaceSymbolCache,
+
+    /// Whenever a file in a workspace is changed we'll add it to this set.
+    workspaces_to_process: HashSet<PathBuf>,
+
     options: LspInitializationOptions,
 
     // Tracks files that currently have errors, by package root.
@@ -112,6 +116,8 @@ struct PackageCacheData {
     node_interner: NodeInterner,
     def_maps: BTreeMap<CrateId, CrateDefMap>,
     usage_tracker: UsageTracker,
+    diagnostics: Vec<CustomDiagnostic>,
+    diagnostics_just_published: bool,
 }
 
 impl LspState {
@@ -128,7 +134,7 @@ impl LspState {
             workspace_cache: HashMap::new(),
             package_cache: HashMap::new(),
             workspace_symbol_cache: WorkspaceSymbolCache::default(),
-            open_documents_count: 0,
+            workspaces_to_process: HashSet::new(),
             options: Default::default(),
             files_with_errors: HashMap::new(),
         }

tooling/lsp/src/notifications/mod.rs

@@ -1,6 +1,6 @@
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::ops::ControlFlow;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 
 use crate::{
     PackageCacheData, WorkspaceCacheData, insert_all_files_for_workspace_into_file_manager,
@@ -9,7 +9,7 @@ use async_lsp::lsp_types;
 use async_lsp::lsp_types::{DiagnosticRelatedInformation, DiagnosticTag, Url};
 use async_lsp::{ErrorCode, LanguageClient, ResponseError};
 use fm::{FileManager, FileMap};
-use fxhash::FxHashMap as HashMap;
+use nargo::workspace::Workspace;
 use noirc_driver::check_crate;
 use noirc_errors::reporter::CustomLabel;
 use noirc_errors::{CustomDiagnostic, DiagnosticKind, Location};
@@ -46,13 +46,10 @@ pub(crate) fn on_did_open_text_document(
     state.input_files.insert(params.text_document.uri.to_string(), params.text_document.text);
 
     let document_uri = params.text_document.uri;
-    let output_diagnostics = true;
+    let change = false;
 
-    match process_workspace_for_noir_document(state, document_uri, output_diagnostics) {
-        Ok(_) => {
-            state.open_documents_count += 1;
-            ControlFlow::Continue(())
-        }
+    match handle_text_document_notification(state, document_uri, change) {
+        Ok(_) => ControlFlow::Continue(()),
         Err(err) => ControlFlow::Break(Err(err)),
     }
 }
@@ -66,9 +63,9 @@ pub(super) fn on_did_change_text_document(
     state.workspace_symbol_cache.reprocess_uri(&params.text_document.uri);
 
     let document_uri = params.text_document.uri;
-    let output_diagnostics = false;
+    let change = true;
 
-    match process_workspace_for_noir_document(state, document_uri, output_diagnostics) {
+    match handle_text_document_notification(state, document_uri, change) {
         Ok(_) => ControlFlow::Continue(()),
         Err(err) => ControlFlow::Break(Err(err)),
     }
@@ -81,17 +78,10 @@ pub(super) fn on_did_close_text_document(
     state.input_files.remove(&params.text_document.uri.to_string());
     state.workspace_symbol_cache.reprocess_uri(&params.text_document.uri);
 
-    state.open_documents_count -= 1;
-
-    if state.open_documents_count == 0 {
-        state.package_cache.clear();
-        state.workspace_cache.clear();
-    }
-
     let document_uri = params.text_document.uri;
-    let output_diagnostics = false;
+    let change = false;
 
-    match process_workspace_for_noir_document(state, document_uri, output_diagnostics) {
+    match handle_text_document_notification(state, document_uri, change) {
         Ok(_) => ControlFlow::Continue(()),
         Err(err) => ControlFlow::Break(Err(err)),
     }
@@ -101,23 +91,84 @@ pub(super) fn on_did_save_text_document(
     state: &mut LspState,
     params: DidSaveTextDocumentParams,
 ) -> ControlFlow<Result<(), async_lsp::Error>> {
-    let document_uri = params.text_document.uri;
-    let output_diagnostics = true;
+    let workspace = match workspace_from_document_uri(params.text_document.uri) {
+        Ok(workspace) => workspace,
+        Err(err) => return ControlFlow::Break(Err(err)),
+    };
 
-    match process_workspace_for_noir_document(state, document_uri, output_diagnostics) {
-        Ok(_) => ControlFlow::Continue(()),
-        Err(err) => ControlFlow::Break(Err(err)),
+    // Process any pending changes
+    if state.workspaces_to_process.remove(&workspace.root_dir) {
+        let _ = process_workspace_for_noir_document(state, &workspace.root_dir, false);
     }
+
+    // Cached data should be here but, if it doesn't, we'll just type-check and output diagnostics
+    let (Some(workspace_cache), Some(package_cache)) = (
+        state.workspace_cache.get(&workspace.root_dir),
+        state.package_cache.get(&workspace.root_dir),
+    ) else {
+        let output_diagnostics = true;
+        return match process_workspace_for_noir_document(
+            state,
+            &workspace.root_dir,
+            output_diagnostics,
+        ) {
+            Ok(_) => ControlFlow::Continue(()),
+            Err(err) => return ControlFlow::Break(Err(err)),
+        };
+    };
+
+    // If the last thing the user did was to save a file in the workspace, it could be that
+    // the underlying files in the filesystem have changed (for example a `git checkout`),
+    // so here we force a type-check just in case.
+    if package_cache.diagnostics_just_published {
+        let output_diagnostics = true;
+        return match process_workspace_for_noir_document(
+            state,
+            &workspace.root_dir,
+            output_diagnostics,
+        ) {
+            Ok(_) => ControlFlow::Continue(()),
+            Err(err) => return ControlFlow::Break(Err(err)),
+        };
+    }
+
+    // Otherwise, we can publish the diagnostics we computed in the last type-check
+    publish_diagnostics(
+        state,
+        &workspace.root_dir,
+        &workspace_cache.file_manager.clone(),
+        package_cache.diagnostics.clone(),
+    );
+
+    if let Some(package_cache) = state.package_cache.get_mut(&workspace.root_dir) {
+        package_cache.diagnostics_just_published = true;
+    }
+
+    ControlFlow::Continue(())
 }
 
-// Given a Noir document, find the workspace it's contained in (an assumed workspace is created if
-// it's only contained in a package), then type-checks the workspace's packages,
-// caching code lenses and type definitions, and notifying about compilation errors.
-pub(crate) fn process_workspace_for_noir_document(
+fn handle_text_document_notification(
     state: &mut LspState,
     document_uri: Url,
-    output_diagnostics: bool,
+    change: bool,
 ) -> Result<(), async_lsp::Error> {
+    let workspace = workspace_from_document_uri(document_uri.clone())?;
+
+    if state.package_cache.contains_key(&workspace.root_dir) {
+        // If we have cached data but the file didn't change there's nothing to do
+        if change {
+            state.workspaces_to_process.insert(workspace.root_dir.clone());
+        }
+        Ok(())
+    } else {
+        // If it's the first time we see this package, show diagnostics.
+        // This can happen for example when a user opens a Noir file in a package for the first time.
+        let output_diagnostics = true;
+        process_workspace_for_noir_document(state, &workspace.root_dir, output_diagnostics)
+    }
+}
+
+fn workspace_from_document_uri(document_uri: Url) -> Result<Workspace, async_lsp::Error> {
     let file_path = document_uri.to_file_path().map_err(|_| {
         ResponseError::new(ErrorCode::REQUEST_FAILED, "URI is not a valid file path")
     })?;
@@ -126,6 +177,21 @@ pub(crate) fn process_workspace_for_noir_document(
         ResponseError::new(ErrorCode::REQUEST_FAILED, lsp_error.to_string())
     })?;
 
+    Ok(workspace)
+}
+
+// Given a Noir document, find the workspace it's contained in (an assumed workspace is created if
+// it's only contained in a package), then type-checks the workspace's packages,
+// caching code lenses and type definitions, and notifying about compilation errors.
+pub(crate) fn process_workspace_for_noir_document(
+    state: &mut LspState,
+    file_path: &Path,
+    output_diagnostics: bool,
+) -> Result<(), async_lsp::Error> {
+    let workspace = resolve_workspace_for_source_path(file_path).map_err(|lsp_error| {
+        ResponseError::new(ErrorCode::REQUEST_FAILED, lsp_error.to_string())
+    })?;
+
     let mut workspace_file_manager = workspace.new_file_manager();
 
     insert_all_files_for_workspace_into_file_manager(
@@ -161,14 +227,15 @@ pub(crate) fn process_workspace_for_noir_document(
                 node_interner: context.def_interner,
                 def_maps: context.def_maps,
                 usage_tracker: context.usage_tracker,
+                diagnostics: file_diagnostics.clone(),
+                diagnostics_just_published: output_diagnostics,
             },
         );
 
         let fm = &context.file_manager;
-        let files = fm.as_file_map();
 
         if output_diagnostics {
-            publish_diagnostics(state, &package.root_dir, files, fm, file_diagnostics);
+            publish_diagnostics(state, &package.root_dir, fm, file_diagnostics);
         }
     }
 
@@ -183,10 +250,10 @@ pub(crate) fn process_workspace_for_noir_document(
 fn publish_diagnostics(
     state: &mut LspState,
     package_root_dir: &PathBuf,
-    files: &FileMap,
     fm: &FileManager,
     custom_diagnostics: Vec<CustomDiagnostic>,
 ) {
+    let files = fm.as_file_map();
     let mut diagnostics_per_url: HashMap<Url, Vec<Diagnostic>> = HashMap::default();
 
     for custom_diagnostic in custom_diagnostics.into_iter() {

tooling/lsp/src/requests/mod.rs

@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
 use std::path::{Path, PathBuf};
 use std::{collections::HashMap, future::Future};
 
+use crate::notifications::process_workspace_for_noir_document;
 use crate::{PackageCacheData, insert_all_files_for_workspace_into_file_manager, parse_diff};
 use crate::{
     resolve_workspace_for_source_path,
@@ -488,6 +489,12 @@ where
         })?;
 
     let workspace = resolve_workspace_for_source_path(file_path.as_path()).unwrap();
+
+    // First type-check the workspace if needed
+    if state.workspaces_to_process.remove(&workspace.root_dir) {
+        let _ = process_workspace_for_noir_document(state, &workspace.root_dir, false);
+    }
+
     let package = crate::workspace_package_for_file(&workspace, &file_path).ok_or_else(|| {
         ResponseError::new(ErrorCode::REQUEST_FAILED, "Could not find package for file")
     })?;

tooling/lsp/src/requests/references.rs

@@ -116,7 +116,7 @@ mod references_tests {
 
         notifications::process_workspace_for_noir_document(
             &mut state,
-            one_lib.clone(),
+            &one_lib.to_file_path().unwrap(),
             output_diagnostics,
         )
         .unwrap();