feat: Cached LSP parsing

tooling/lsp/Cargo.toml

@@ -25,6 +25,8 @@ async-lsp = { workspace = true, features = ["omni-trait"] }
 serde_with = "3.2.0"
 thiserror.workspace = true
 fm.workspace = true
+rayon = "1.8.0"
+fxhash.workspace = true
 
 [target.'cfg(all(target_arch = "wasm32", not(target_os = "wasi")))'.dependencies]
 wasm-bindgen.workspace = true

tooling/lsp/src/lib.rs

@@ -17,16 +17,20 @@ use async_lsp::{
     router::Router, AnyEvent, AnyNotification, AnyRequest, ClientSocket, Error, LspService,
     ResponseError,
 };
-use fm::codespan_files as files;
+use fm::{codespan_files as files, FileManager};
+use fxhash::FxHashSet;
 use lsp_types::CodeLens;
 use nargo::{parse_all, workspace::Workspace};
 use nargo_toml::{find_file_manifest, resolve_workspace_from_toml, PackageSelection};
 use noirc_driver::{file_manager_with_stdlib, prepare_crate, NOIR_ARTIFACT_VERSION_STRING};
 use noirc_frontend::{
     graph::{CrateId, CrateName},
-    hir::{Context, FunctionNameMatch},
+    hir::{def_map::parse_file, Context, FunctionNameMatch, ParsedFiles},
     node_interner::NodeInterner,
+    parser::ParserError,
+    ParsedModule,
 };
+use rayon::prelude::*;
 
 use notifications::{
     on_did_change_configuration, on_did_change_text_document, on_did_close_text_document,
@@ -65,6 +69,8 @@ pub struct LspState {
     input_files: HashMap<String, String>,
     cached_lenses: HashMap<String, Vec<CodeLens>>,
     cached_definitions: HashMap<String, NodeInterner>,
+    cached_parsed_files: HashMap<PathBuf, (usize, (ParsedModule, Vec<ParserError>))>,
+    parsing_cache_enabled: bool,
 }
 
 impl LspState {
@@ -77,6 +83,8 @@ impl LspState {
             cached_lenses: HashMap::new(),
             cached_definitions: HashMap::new(),
             open_documents_count: 0,
+            cached_parsed_files: HashMap::new(),
+            parsing_cache_enabled: true,
         }
     }
 }
@@ -227,21 +235,78 @@ pub(crate) fn resolve_workspace_for_source_path(file_path: &Path) -> Result<Work
 /// Use case for this is the LSP server and code lenses
 /// which operate on single file and need to understand this file
 /// in order to offer code lenses to the user
-fn prepare_source(source: String) -> (Context<'static, 'static>, CrateId) {
+fn prepare_source(source: String, state: &mut LspState) -> (Context<'static, 'static>, CrateId) {
     let root = Path::new("");
     let file_name = Path::new("main.nr");
     let mut file_manager = file_manager_with_stdlib(root);
     file_manager.add_file_with_source(file_name, source).expect(
         "Adding source buffer to file manager should never fail when file manager is empty",
     );
-    let parsed_files = parse_all(&file_manager);
+    let parsed_files = parse_diff(&file_manager, state);
 
     let mut context = Context::new(file_manager, parsed_files);
     let root_crate_id = prepare_crate(&mut context, file_name);
 
     (context, root_crate_id)
 }
 
+fn parse_diff(file_manager: &FileManager, state: &mut LspState) -> ParsedFiles {
+    if state.parsing_cache_enabled {
+        let noir_file_hashes: Vec<_> = file_manager
+            .as_file_map()
+            .all_file_ids()
+            .par_bridge()
+            .filter_map(|&file_id| {
+                let file_path = file_manager.path(file_id).expect("expected file to exist");
+                let file_extension =
+                    file_path.extension().expect("expected all file paths to have an extension");
+                if file_extension == "nr" {
+                    Some((
+                        file_id,
+                        file_path.to_path_buf(),
+                        fxhash::hash(file_manager.fetch_file(file_id).expect("file must exist")),
+                    ))
+                } else {
+                    None
+                }
+            })
+            .collect();
+
+        let cache_hits: Vec<_> = noir_file_hashes
+            .par_iter()
+            .filter_map(|(file_id, file_path, current_hash)| {
+                let cached_version = state.cached_parsed_files.get(file_path);
+                if let Some((hash, cached_parsing)) = cached_version {
+                    if hash == current_hash {
+                        return Some((*file_id, cached_parsing.clone()));
+                    }
+                }
+                None
+            })
+            .collect();
+
+        let cache_hits_ids: FxHashSet<_> = cache_hits.iter().map(|(file_id, _)| *file_id).collect();
+
+        let cache_misses: Vec<_> = noir_file_hashes
+            .into_par_iter()
+            .filter(|(id, _, _)| !cache_hits_ids.contains(id))
+            .map(|(file_id, path, hash)| (file_id, path, hash, parse_file(file_manager, file_id)))
+            .collect();
+
+        cache_misses.iter().for_each(|(_, path, hash, parse_results)| {
+            state.cached_parsed_files.insert(path.clone(), (*hash, parse_results.clone()));
+        });
+
+        cache_misses
+            .into_iter()
+            .map(|(id, _, _, parse_results)| (id, parse_results))
+            .chain(cache_hits.into_iter())
+            .collect()
+    } else {
+        parse_all(file_manager)
+    }
+}
+
 #[test]
 fn prepare_package_from_source_string() {
     let source = r#"
@@ -252,7 +317,10 @@ fn prepare_package_from_source_string() {
     }
     "#;
 
-    let (mut context, crate_id) = crate::prepare_source(source.to_string());
+    let client = ClientSocket::new_closed();
+    let mut state = LspState::new(&client, acvm::blackbox_solver::StubbedBlackBoxSolver);
+
+    let (mut context, crate_id) = crate::prepare_source(source.to_string(), &mut state);
     let _check_result = noirc_driver::check_crate(&mut context, crate_id, false, false);
     let main_func_id = context.get_main_function(&crate_id);
     assert!(main_func_id.is_some());

tooling/lsp/src/notifications/mod.rs

@@ -1,7 +1,7 @@
 use std::ops::ControlFlow;
 
 use async_lsp::{ErrorCode, LanguageClient, ResponseError};
-use nargo::{insert_all_files_for_workspace_into_file_manager, parse_all, prepare_package};
+use nargo::{insert_all_files_for_workspace_into_file_manager, prepare_package};
 use noirc_driver::{check_crate, file_manager_with_stdlib};
 use noirc_errors::{DiagnosticKind, FileDiagnostic};
 
@@ -13,7 +13,7 @@ use crate::types::{
 };
 
 use crate::{
-    byte_span_to_range, get_package_tests_in_crate, prepare_source,
+    byte_span_to_range, get_package_tests_in_crate, parse_diff, prepare_source,
     resolve_workspace_for_source_path, LspState,
 };
 
@@ -55,7 +55,7 @@ pub(super) fn on_did_change_text_document(
     let text = params.content_changes.into_iter().next().unwrap().text;
     state.input_files.insert(params.text_document.uri.to_string(), text.clone());
 
-    let (mut context, crate_id) = prepare_source(text);
+    let (mut context, crate_id) = prepare_source(text, state);
     let _ = check_crate(&mut context, crate_id, false, false);
 
     let workspace = match resolve_workspace_for_source_path(
@@ -130,7 +130,8 @@ fn process_noir_document(
 
     let mut workspace_file_manager = file_manager_with_stdlib(&workspace.root_dir);
     insert_all_files_for_workspace_into_file_manager(&workspace, &mut workspace_file_manager);
-    let parsed_files = parse_all(&workspace_file_manager);
+
+    let parsed_files = parse_diff(&workspace_file_manager, state);
 
     let diagnostics: Vec<_> = workspace
         .into_iter()

tooling/lsp/src/requests/code_lens_request.rs

@@ -64,7 +64,7 @@ fn on_code_lens_request_inner(
     let workspace = resolve_workspace_for_source_path(file_path.as_path()).unwrap();
     let package = workspace.members.first().unwrap();
 
-    let (mut context, crate_id) = prepare_source(source_string);
+    let (mut context, crate_id) = prepare_source(source_string, state);
     // We ignore the warnings and errors produced by compilation for producing code lenses
     // because we can still get the test functions even if compilation fails
     let _ = check_crate(&mut context, crate_id, false, false);

tooling/lsp/src/requests/goto_declaration.rs

@@ -1,13 +1,13 @@
 use std::future::{self, Future};
 
-use crate::resolve_workspace_for_source_path;
 use crate::types::GotoDeclarationResult;
 use crate::LspState;
+use crate::{parse_diff, resolve_workspace_for_source_path};
 use async_lsp::{ErrorCode, ResponseError};
 
 use lsp_types::request::{GotoDeclarationParams, GotoDeclarationResponse};
 
-use nargo::{insert_all_files_for_workspace_into_file_manager, parse_all};
+use nargo::insert_all_files_for_workspace_into_file_manager;
 use noirc_driver::file_manager_with_stdlib;
 
 use super::{position_to_byte_index, to_lsp_location};
@@ -21,7 +21,7 @@ pub(crate) fn on_goto_declaration_request(
 }
 
 fn on_goto_definition_inner(
-    _state: &mut LspState,
+    state: &mut LspState,
     params: GotoDeclarationParams,
 ) -> Result<GotoDeclarationResult, ResponseError> {
     let file_path =
@@ -36,13 +36,13 @@ fn on_goto_definition_inner(
 
     let mut workspace_file_manager = file_manager_with_stdlib(&workspace.root_dir);
     insert_all_files_for_workspace_into_file_manager(&workspace, &mut workspace_file_manager);
-    let parsed_files = parse_all(&workspace_file_manager);
+    let parsed_files = parse_diff(&workspace_file_manager, state);
 
     let (mut context, crate_id) =
         nargo::prepare_package(&workspace_file_manager, &parsed_files, package);
 
     let interner;
-    if let Some(def_interner) = _state.cached_definitions.get(&package_root_path) {
+    if let Some(def_interner) = state.cached_definitions.get(&package_root_path) {
         interner = def_interner;
     } else {
         // We ignore the warnings and errors produced by compilation while resolving the definition

tooling/lsp/src/requests/goto_definition.rs

@@ -1,12 +1,12 @@
 use std::future::{self, Future};
 
-use crate::resolve_workspace_for_source_path;
+use crate::{parse_diff, resolve_workspace_for_source_path};
 use crate::{types::GotoDefinitionResult, LspState};
 use async_lsp::{ErrorCode, ResponseError};
 
 use lsp_types::request::GotoTypeDefinitionParams;
 use lsp_types::{GotoDefinitionParams, GotoDefinitionResponse};
-use nargo::{insert_all_files_for_workspace_into_file_manager, parse_all};
+use nargo::insert_all_files_for_workspace_into_file_manager;
 use noirc_driver::file_manager_with_stdlib;
 
 use super::{position_to_byte_index, to_lsp_location};
@@ -28,7 +28,7 @@ pub(crate) fn on_goto_type_definition_request(
 }
 
 fn on_goto_definition_inner(
-    _state: &mut LspState,
+    state: &mut LspState,
     params: GotoDefinitionParams,
     return_type_location_instead: bool,
 ) -> Result<GotoDefinitionResult, ResponseError> {
@@ -44,13 +44,13 @@ fn on_goto_definition_inner(
 
     let mut workspace_file_manager = file_manager_with_stdlib(&workspace.root_dir);
     insert_all_files_for_workspace_into_file_manager(&workspace, &mut workspace_file_manager);
-    let parsed_files = parse_all(&workspace_file_manager);
+    let parsed_files = parse_diff(&workspace_file_manager, state);
 
     let (mut context, crate_id) =
         nargo::prepare_package(&workspace_file_manager, &parsed_files, package);
 
     let interner;
-    if let Some(def_interner) = _state.cached_definitions.get(&package_root_path) {
+    if let Some(def_interner) = state.cached_definitions.get(&package_root_path) {
         interner = def_interner;
     } else {
         // We ignore the warnings and errors produced by compilation while resolving the definition

tooling/lsp/src/requests/mod.rs

@@ -47,15 +47,25 @@ struct LspInitializationOptions {
     /// By default this will be set to true (enabled).
     #[serde(rename = "enableCodeLens", default = "default_enable_code_lens")]
     enable_code_lens: bool,
+
+    #[serde(rename = "enableParsingCache", default = "default_enable_parsing_cache")]
+    enable_parsing_cache: bool,
 }
 
 fn default_enable_code_lens() -> bool {
     true
 }
 
+fn default_enable_parsing_cache() -> bool {
+    true
+}
+
 impl Default for LspInitializationOptions {
     fn default() -> Self {
-        Self { enable_code_lens: default_enable_code_lens() }
+        Self {
+            enable_code_lens: default_enable_code_lens(),
+            enable_parsing_cache: default_enable_parsing_cache(),
+        }
     }
 }
 
@@ -64,11 +74,11 @@ pub(crate) fn on_initialize(
     params: InitializeParams,
 ) -> impl Future<Output = Result<InitializeResult, ResponseError>> {
     state.root_path = params.root_uri.and_then(|root_uri| root_uri.to_file_path().ok());
-
     let initialization_options: LspInitializationOptions = params
         .initialization_options
         .and_then(|value| serde_json::from_value(value).ok())
         .unwrap_or_default();
+    state.parsing_cache_enabled = initialization_options.enable_parsing_cache;
 
     async move {
         let text_document_sync = TextDocumentSyncCapability::Kind(TextDocumentSyncKind::FULL);

tooling/lsp/src/requests/profile_run.rs

@@ -5,16 +5,15 @@ use std::{
 
 use acvm::ExpressionWidth;
 use async_lsp::{ErrorCode, ResponseError};
-use nargo::{
-    artifacts::debug::DebugArtifact, insert_all_files_for_workspace_into_file_manager, parse_all,
-};
+use nargo::{artifacts::debug::DebugArtifact, insert_all_files_for_workspace_into_file_manager};
 use nargo_toml::{find_package_manifest, resolve_workspace_from_toml, PackageSelection};
 use noirc_driver::{
     file_manager_with_stdlib, CompileOptions, DebugFile, NOIR_ARTIFACT_VERSION_STRING,
 };
 use noirc_errors::{debug_info::OpCodesCount, Location};
 
 use crate::{
+    parse_diff,
     types::{NargoProfileRunParams, NargoProfileRunResult},
     LspState,
 };
@@ -28,7 +27,7 @@ pub(crate) fn on_profile_run_request(
 }
 
 fn on_profile_run_request_inner(
-    state: &LspState,
+    state: &mut LspState,
     params: NargoProfileRunParams,
 ) -> Result<NargoProfileRunResult, ResponseError> {
     let root_path = state.root_path.as_deref().ok_or_else(|| {
@@ -54,7 +53,7 @@ fn on_profile_run_request_inner(
 
     let mut workspace_file_manager = file_manager_with_stdlib(&workspace.root_dir);
     insert_all_files_for_workspace_into_file_manager(&workspace, &mut workspace_file_manager);
-    let parsed_files = parse_all(&workspace_file_manager);
+    let parsed_files = parse_diff(&workspace_file_manager, state);
 
     // Since we filtered on crate name, this should be the only item in the iterator
     match workspace.into_iter().next() {

tooling/lsp/src/requests/test_run.rs

@@ -4,7 +4,7 @@ use async_lsp::{ErrorCode, ResponseError};
 use nargo::{
     insert_all_files_for_workspace_into_file_manager,
     ops::{run_test, TestStatus},
-    parse_all, prepare_package,
+    prepare_package,
 };
 use nargo_toml::{find_package_manifest, resolve_workspace_from_toml, PackageSelection};
 use noirc_driver::{
@@ -13,6 +13,7 @@ use noirc_driver::{
 use noirc_frontend::hir::FunctionNameMatch;
 
 use crate::{
+    parse_diff,
     types::{NargoTestRunParams, NargoTestRunResult},
     LspState,
 };
@@ -25,7 +26,7 @@ pub(crate) fn on_test_run_request(
 }
 
 fn on_test_run_request_inner(
-    state: &LspState,
+    state: &mut LspState,
     params: NargoTestRunParams,
 ) -> Result<NargoTestRunResult, ResponseError> {
     let root_path = state.root_path.as_deref().ok_or_else(|| {
@@ -52,7 +53,7 @@ fn on_test_run_request_inner(
 
     let mut workspace_file_manager = file_manager_with_stdlib(&workspace.root_dir);
     insert_all_files_for_workspace_into_file_manager(&workspace, &mut workspace_file_manager);
-    let parsed_files = parse_all(&workspace_file_manager);
+    let parsed_files = parse_diff(&workspace_file_manager, state);
 
     // Since we filtered on crate name, this should be the only item in the iterator
     match workspace.into_iter().next() {