feat(lsp): cache definitions for goto requests (#3930)

# Description

## Problem\*

Goto Definitions still take take too much time to execute leaving user
with undesirable experience. On a bigger workspace (beyond simple test
cases, eg.
[noir-protocol-circuits](https://github.com/AztecProtocol/aztec-packages/tree/master/yarn-project/noir-protocol-circuits))
it can take as much as a second or more to perform goto request.

Resolves chore(lsp): improve goto performance by caching ast #3942

## Summary\*

Improvement from `~1s` down to `~11ms` on example project
[noir-protocol-circuits](https://github.com/AztecProtocol/aztec-packages/tree/master/yarn-project/noir-protocol-circuits)
is achieved by caching definitions resolved during file save or file
open.

The cost for this optimisation is when user opens noir source file, at
which point parsing and resolving is done. This however is not
noticeable for user as source code can be seen and operated on.

## Additional Context



## Documentation\*

Check one:
- [x] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[Exceptional Case]** Documentation to be submitted in a separate
PR.

# PR Checklist\*

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.

Author: kobyhallx

compiler/noirc_frontend/src/hir/mod.rs

@@ -194,14 +194,6 @@ impl Context<'_> {
             .collect()
     }
 
-    /// Returns the [Location] of the definition of the given Ident found at [Span] of the given [FileId].
-    /// Returns [None] when definition is not found.
-    pub fn get_definition_location_from(&self, location: Location) -> Option<Location> {
-        let interner = &self.def_interner;
-
-        interner.find_location_index(location).and_then(|index| interner.resolve_location(index))
-    }
-
     /// Return a Vec of all `contract` declarations in the source code and the functions they contain
     pub fn get_all_contracts(&self, crate_id: &CrateId) -> Vec<Contract> {
         self.def_map(crate_id)

compiler/noirc_frontend/src/node_interner.rs

@@ -1271,10 +1271,16 @@ impl NodeInterner {
         self.selected_trait_implementations.get(&ident_id).cloned()
     }
 
+    /// Returns the [Location] of the definition of the given Ident found at [Span] of the given [FileId].
+    /// Returns [None] when definition is not found.
+    pub fn get_definition_location_from(&self, location: Location) -> Option<Location> {
+        self.find_location_index(location).and_then(|index| self.resolve_location(index))
+    }
+
     /// For a given [Index] we return [Location] to which we resolved to
     /// We currently return None for features not yet implemented
     /// TODO(#3659): LSP goto def should error when Ident at Location could not resolve
-    pub(crate) fn resolve_location(&self, index: impl Into<Index>) -> Option<Location> {
+    fn resolve_location(&self, index: impl Into<Index>) -> Option<Location> {
         let node = self.nodes.get(index.into())?;
 
         match node {

tooling/lsp/src/lib.rs

@@ -25,6 +25,7 @@ use noirc_driver::{file_manager_with_stdlib, prepare_crate, NOIR_ARTIFACT_VERSIO
 use noirc_frontend::{
     graph::{CrateId, CrateName},
     hir::{Context, FunctionNameMatch},
+    node_interner::NodeInterner,
 };
 
 use notifications::{
@@ -59,8 +60,10 @@ pub struct LspState {
     root_path: Option<PathBuf>,
     client: ClientSocket,
     solver: WrapperSolver,
+    open_documents_count: usize,
     input_files: HashMap<String, String>,
     cached_lenses: HashMap<String, Vec<CodeLens>>,
+    cached_definitions: HashMap<String, NodeInterner>,
 }
 
 impl LspState {
@@ -71,6 +74,8 @@ impl LspState {
             solver: WrapperSolver(Box::new(solver)),
             input_files: HashMap::new(),
             cached_lenses: HashMap::new(),
+            cached_definitions: HashMap::new(),
+            open_documents_count: 0,
         }
     }
 }

tooling/lsp/src/notifications/mod.rs

@@ -36,7 +36,16 @@ pub(super) fn on_did_open_text_document(
     params: DidOpenTextDocumentParams,
 ) -> ControlFlow<Result<(), async_lsp::Error>> {
     state.input_files.insert(params.text_document.uri.to_string(), params.text_document.text);
-    ControlFlow::Continue(())
+
+    let document_uri = params.text_document.uri;
+
+    match process_noir_document(document_uri, state) {
+        Ok(_) => {
+            state.open_documents_count += 1;
+            ControlFlow::Continue(())
+        }
+        Err(err) => ControlFlow::Break(Err(err)),
+    }
 }
 
 pub(super) fn on_did_change_text_document(
@@ -85,34 +94,39 @@ pub(super) fn on_did_close_text_document(
 ) -> ControlFlow<Result<(), async_lsp::Error>> {
     state.input_files.remove(&params.text_document.uri.to_string());
     state.cached_lenses.remove(&params.text_document.uri.to_string());
+
+    state.open_documents_count -= 1;
+
+    if state.open_documents_count == 0 {
+        state.cached_definitions.clear();
+    }
+
     ControlFlow::Continue(())
 }
 
 pub(super) fn on_did_save_text_document(
     state: &mut LspState,
     params: DidSaveTextDocumentParams,
 ) -> ControlFlow<Result<(), async_lsp::Error>> {
-    let file_path = match params.text_document.uri.to_file_path() {
-        Ok(file_path) => file_path,
-        Err(()) => {
-            return ControlFlow::Break(Err(ResponseError::new(
-                ErrorCode::REQUEST_FAILED,
-                "URI is not a valid file path",
-            )
-            .into()))
-        }
-    };
+    let document_uri = params.text_document.uri;
 
-    let workspace = match resolve_workspace_for_source_path(&file_path) {
-        Ok(value) => value,
-        Err(lsp_error) => {
-            return ControlFlow::Break(Err(ResponseError::new(
-                ErrorCode::REQUEST_FAILED,
-                lsp_error.to_string(),
-            )
-            .into()))
-        }
-    };
+    match process_noir_document(document_uri, state) {
+        Ok(_) => ControlFlow::Continue(()),
+        Err(err) => ControlFlow::Break(Err(err)),
+    }
+}
+
+fn process_noir_document(
+    document_uri: lsp_types::Url,
+    state: &mut LspState,
+) -> Result<(), 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")
+    })?;
+
+    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 = file_manager_with_stdlib(&workspace.root_dir);
     insert_all_files_for_workspace_into_file_manager(&workspace, &mut workspace_file_manager);
@@ -127,6 +141,8 @@ pub(super) fn on_did_save_text_document(
                 Err(errors_and_warnings) => errors_and_warnings,
             };
 
+            let package_root_dir: String = package.root_dir.as_os_str().to_string_lossy().into();
+
             // We don't add test headings for a package if it contains no `#[test]` functions
             if let Some(tests) = get_package_tests_in_crate(&context, &crate_id, &package.name) {
                 let _ = state.client.notify::<notification::NargoUpdateTests>(NargoPackageTests {
@@ -142,7 +158,9 @@ pub(super) fn on_did_save_text_document(
                 package,
                 Some(&file_path),
             );
-            state.cached_lenses.insert(params.text_document.uri.to_string(), collected_lenses);
+            state.cached_lenses.insert(document_uri.to_string(), collected_lenses);
+
+            state.cached_definitions.insert(package_root_dir, context.def_interner);
 
             let fm = &context.file_manager;
             let files = fm.as_file_map();
@@ -178,14 +196,13 @@ pub(super) fn on_did_save_text_document(
                 .collect()
         })
         .collect();
-
     let _ = state.client.publish_diagnostics(PublishDiagnosticsParams {
-        uri: params.text_document.uri,
+        uri: document_uri,
         version: None,
         diagnostics,
     });
 
-    ControlFlow::Continue(())
+    Ok(())
 }
 
 pub(super) fn on_exit(

tooling/lsp/src/requests/goto_definition.rs

@@ -29,20 +29,27 @@ fn on_goto_definition_inner(
     let workspace = resolve_workspace_for_source_path(file_path.as_path()).unwrap();
     let package = workspace.members.first().unwrap();
 
+    let package_root_path: String = package.root_dir.as_os_str().to_string_lossy().into();
+
     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 (mut context, crate_id) = nargo::prepare_package(&workspace_file_manager, package);
 
-    // We ignore the warnings and errors produced by compilation while resolving the definition
-    let _ = noirc_driver::check_crate(&mut context, crate_id, false, false);
+    let interner;
+    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
+        let _ = noirc_driver::check_crate(&mut context, crate_id, false, false);
+        interner = &context.def_interner;
+    }
 
     let files = context.file_manager.as_file_map();
     let file_id = context.file_manager.name_to_id(file_path.clone()).ok_or(ResponseError::new(
         ErrorCode::REQUEST_FAILED,
         format!("Could not find file in file manager. File path: {:?}", file_path),
     ))?;
-
     let byte_index =
         position_to_byte_index(files, file_id, &params.text_document_position_params.position)
             .map_err(|err| {
@@ -58,7 +65,7 @@ fn on_goto_definition_inner(
     };
 
     let goto_definition_response =
-        context.get_definition_location_from(search_for_location).and_then(|found_location| {
+        interner.get_definition_location_from(search_for_location).and_then(|found_location| {
             let file_id = found_location.file;
             let definition_position = to_lsp_location(files, file_id, found_location.span)?;
             let response: GotoDefinitionResponse =