feat(lsp): re-add code lens feature with improved performance

tooling/lsp/Cargo.toml

@@ -23,6 +23,7 @@ serde_json.workspace = true
 tower.workspace = true
 async-lsp = { workspace = true, features = ["omni-trait"] }
 serde_with = "3.2.0"
+thiserror.workspace = true
 fm.workspace = true
 
 [target.'cfg(all(target_arch = "wasm32", not(target_os = "wasi")))'.dependencies]

tooling/lsp/src/lib.rs

@@ -7,7 +7,7 @@ use std::{
     collections::HashMap,
     future::Future,
     ops::{self, ControlFlow},
-    path::PathBuf,
+    path::{Path, PathBuf},
     pin::Pin,
     task::{self, Poll},
 };
@@ -18,6 +18,9 @@ use async_lsp::{
     ResponseError,
 };
 use fm::codespan_files as files;
+use lsp_types::CodeLens;
+use nargo_toml::{find_file_manifest, resolve_workspace_from_toml, PackageSelection};
+use noirc_driver::NOIR_ARTIFACT_VERSION_STRING;
 use noirc_frontend::{
     graph::{CrateId, CrateName},
     hir::{Context, FunctionNameMatch},
@@ -27,10 +30,11 @@ use notifications::{
     on_did_open_text_document, on_did_save_text_document, on_exit, on_initialized,
 };
 use requests::{
-    on_formatting, on_goto_definition_request, on_initialize, on_profile_run_request, on_shutdown,
-    on_test_run_request, on_tests_request,
+    on_code_lens_request, on_formatting, on_goto_definition_request, on_initialize,
+    on_profile_run_request, on_shutdown, on_test_run_request, on_tests_request,
 };
 use serde_json::Value as JsonValue;
+use thiserror::Error;
 use tower::Service;
 
 mod notifications;
@@ -41,12 +45,20 @@ mod types;
 use solver::WrapperSolver;
 use types::{notification, request, NargoTest, NargoTestId, Position, Range, Url};
 
+#[derive(Debug, Error)]
+pub enum LspError {
+    /// Error while Resolving Workspace - {0}.
+    #[error("Failed to Resolve Workspace - {0}")]
+    WorkspaceResolutionError(String),
+}
+
 // State for the LSP gets implemented on this struct and is internal to the implementation
 pub struct LspState {
     root_path: Option<PathBuf>,
     client: ClientSocket,
     solver: WrapperSolver,
     input_files: HashMap<String, String>,
+    cached_lenses: HashMap<String, Vec<CodeLens>>,
 }
 
 impl LspState {
@@ -56,6 +68,7 @@ impl LspState {
             root_path: None,
             solver: WrapperSolver(Box::new(solver)),
             input_files: HashMap::new(),
+            cached_lenses: HashMap::new(),
         }
     }
 }
@@ -72,6 +85,7 @@ impl NargoLspService {
             .request::<request::Initialize, _>(on_initialize)
             .request::<request::Formatting, _>(on_formatting)
             .request::<request::Shutdown, _>(on_shutdown)
+            .request::<request::CodeLens, _>(on_code_lens_request)
             .request::<request::NargoTests, _>(on_tests_request)
             .request::<request::NargoTestRun, _>(on_test_run_request)
             .request::<request::NargoProfileRun, _>(on_profile_run_request)
@@ -175,3 +189,29 @@ fn byte_span_to_range<'a, F: files::Files<'a> + ?Sized>(
         None
     }
 }
+
+pub(crate) fn resolve_workspace_for_source_path(
+    file_path: &Path,
+) -> Result<nargo::workspace::Workspace, LspError> {
+    let package_root = find_file_manifest(file_path);
+
+    let toml_path = match package_root {
+        Some(toml_path) => toml_path,
+        None => {
+            return Err(LspError::WorkspaceResolutionError(format!(
+                "Nargo.toml not found for file: {:?}",
+                file_path
+            )))
+        }
+    };
+
+    let workspace = match resolve_workspace_from_toml(
+        &toml_path,
+        PackageSelection::All,
+        Some(NOIR_ARTIFACT_VERSION_STRING.to_string()),
+    ) {
+        Ok(workspace) => workspace,
+        Err(err) => return Err(LspError::WorkspaceResolutionError(format!("{err}"))),
+    };
+    Ok(workspace)
+}

tooling/lsp/src/notifications/mod.rs

@@ -2,18 +2,19 @@ use std::ops::ControlFlow;
 
 use async_lsp::{ErrorCode, LanguageClient, ResponseError};
 use nargo::prepare_package;
-use nargo_toml::{find_file_manifest, resolve_workspace_from_toml, PackageSelection};
-use noirc_driver::{check_crate, NOIR_ARTIFACT_VERSION_STRING};
+use noirc_driver::check_crate;
 use noirc_errors::{DiagnosticKind, FileDiagnostic};
 
+use crate::requests::collect_lenses_for_package;
 use crate::types::{
     notification, Diagnostic, DiagnosticSeverity, DidChangeConfigurationParams,
     DidChangeTextDocumentParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams,
-    DidSaveTextDocumentParams, InitializedParams, LogMessageParams, MessageType, NargoPackageTests,
-    PublishDiagnosticsParams,
+    DidSaveTextDocumentParams, InitializedParams, NargoPackageTests, PublishDiagnosticsParams,
 };
 
-use crate::{byte_span_to_range, get_package_tests_in_crate, LspState};
+use crate::{
+    byte_span_to_range, get_package_tests_in_crate, resolve_workspace_for_source_path, LspState,
+};
 
 pub(super) fn on_initialized(
     _state: &mut LspState,
@@ -42,7 +43,38 @@ pub(super) fn on_did_change_text_document(
     params: DidChangeTextDocumentParams,
 ) -> ControlFlow<Result<(), async_lsp::Error>> {
     let text = params.content_changes.into_iter().next().unwrap().text;
-    state.input_files.insert(params.text_document.uri.to_string(), text);
+    state.input_files.insert(params.text_document.uri.to_string(), text.clone());
+
+    let (mut context, crate_id) = nargo::prepare_source(text);
+    let _ = check_crate(&mut context, crate_id, false, false);
+
+    let workspace = match resolve_workspace_for_source_path(
+        params.text_document.uri.to_file_path().unwrap().as_path(),
+    ) {
+        Ok(workspace) => workspace,
+        Err(lsp_error) => {
+            return ControlFlow::Break(Err(ResponseError::new(
+                ErrorCode::REQUEST_FAILED,
+                format!("{}", lsp_error),
+            )
+            .into()))
+        }
+    };
+    let package = match workspace.members.first() {
+        Some(package) => package,
+        None => {
+            return ControlFlow::Break(Err(ResponseError::new(
+                ErrorCode::REQUEST_FAILED,
+                "Selected workspace has no members",
+            )
+            .into()))
+        }
+    };
+
+    let lenses = collect_lenses_for_package(&context, crate_id, &workspace, package, None);
+
+    state.cached_lenses.insert(params.text_document.uri.to_string(), lenses);
+
     ControlFlow::Continue(())
 }
 
@@ -51,6 +83,7 @@ pub(super) fn on_did_close_text_document(
     params: DidCloseTextDocumentParams,
 ) -> 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());
     ControlFlow::Continue(())
 }
 
@@ -69,34 +102,14 @@ pub(super) fn on_did_save_text_document(
         }
     };
 
-    let package_root = find_file_manifest(file_path.as_path());
-
-    let toml_path = match package_root {
-        Some(toml_path) => toml_path,
-        None => {
-            // If we cannot find a manifest, we log a warning but return no diagnostics
-            // We can reconsider this when we can build a file without the need for a Nargo.toml file to resolve deps
-            let _ = state.client.log_message(LogMessageParams {
-                typ: MessageType::WARNING,
-                message: format!("Nargo.toml not found for file: {:}", file_path.display()),
-            });
-            return ControlFlow::Continue(());
-        }
-    };
-
-    let workspace = match resolve_workspace_from_toml(
-        &toml_path,
-        PackageSelection::All,
-        Some(NOIR_ARTIFACT_VERSION_STRING.to_string()),
-    ) {
-        Ok(workspace) => workspace,
-        Err(err) => {
-            // If we found a manifest, but the workspace is invalid, we raise an error about it
+    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,
-                format!("{err}"),
+                format!("{}", lsp_error),
             )
-            .into()));
+            .into()))
         }
     };
 
@@ -118,6 +131,15 @@ pub(super) fn on_did_save_text_document(
                 });
             }
 
+            let collected_lenses = crate::requests::collect_lenses_for_package(
+                &context,
+                crate_id,
+                &workspace,
+                package,
+                Some(&file_path),
+            );
+            state.cached_lenses.insert(params.text_document.uri.to_string(), collected_lenses);
+
             let fm = &context.file_manager;
             let files = fm.as_file_map();