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

# Description

## Problem\*

Code Lens part of protocol previously supported, attempted to do to many
things at once.

## Summary\*

By introducing statfullnes into LSP client we benefit from not repeating
preparation procedure on each editor code lens requst.

By reducing scope of interest when resolving code lens we gain speed -
only single source file needs to get prepared at a time to provide code
lenses to client.

By introducing `enableCodeLens` initialization option (a2943b1:
feat(lsp): accept option to disable code lens
) we are able to control if code lens feature is enabled. This will
require plugin setting this value, while by default it is turned on.

## 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.

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,19 +18,27 @@ use async_lsp::{
     ResponseError,
 };
 use fm::codespan_files as files;
+use lsp_types::CodeLens;
+use nargo::workspace::Workspace;
+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},
 };
+
+use fm::FileManager;
+
 use notifications::{
     on_did_change_configuration, on_did_change_text_document, on_did_close_text_document,
     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 +49,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.
+    #[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 +72,7 @@ impl LspState {
             root_path: None,
             solver: WrapperSolver(Box::new(solver)),
             input_files: HashMap::new(),
+            cached_lenses: HashMap::new(),
         }
     }
 }
@@ -72,6 +89,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 +193,60 @@ fn byte_span_to_range<'a, F: files::Files<'a> + ?Sized>(
         None
     }
 }
+
+pub(crate) fn resolve_workspace_for_source_path(file_path: &Path) -> Result<Workspace, LspError> {
+    let package_root = find_file_manifest(file_path);
+
+    let toml_path = package_root.ok_or_else(|| {
+        LspError::WorkspaceResolutionError(format!(
+            "Nargo.toml not found for file: {:?}",
+            file_path
+        ))
+    })?;
+
+    let workspace = resolve_workspace_from_toml(
+        &toml_path,
+        PackageSelection::All,
+        Some(NOIR_ARTIFACT_VERSION_STRING.to_string()),
+    )
+    .map_err(|err| LspError::WorkspaceResolutionError(err.to_string()))?;
+
+    Ok(workspace)
+}
+
+/// Prepares a package from a source string
+/// This is useful for situations when we don't need dependencies
+/// and just need to operate on single file.
+///
+/// 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>, CrateId) {
+    let root = Path::new("");
+    let mut file_manager = FileManager::new(root);
+    let root_file_id = file_manager.add_file_with_source(Path::new("main.nr"), source).expect(
+        "Adding source buffer to file manager should never fail when file manager is empty",
+    );
+
+    let mut context = Context::new(file_manager);
+
+    let root_crate_id = context.crate_graph.add_crate_root(root_file_id);
+
+    (context, root_crate_id)
+}
+
+#[test]
+fn prepare_package_from_source_string() {
+    let source = r#"
+    fn main() {
+        let x = 1;
+        let y = 2;
+        let z = x + y;
+    }
+    "#;
+
+    let (mut context, crate_id) = crate::prepare_source(source.to_string());
+    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

@@ -2,18 +2,20 @@ use std::ops::ControlFlow;
 
 use async_lsp::{ErrorCode, LanguageClient, ResponseError};
 use nargo::{insert_all_files_for_workspace_into_file_manager, prepare_package};
-use nargo_toml::{find_file_manifest, resolve_workspace_from_toml, PackageSelection};
-use noirc_driver::{check_crate, file_manager_with_stdlib, NOIR_ARTIFACT_VERSION_STRING};
+use noirc_driver::{check_crate, file_manager_with_stdlib};
 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, prepare_source,
+    resolve_workspace_for_source_path, LspState,
+};
 
 pub(super) fn on_initialized(
     _state: &mut LspState,
@@ -42,7 +44,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) = 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,
+                lsp_error.to_string(),
+            )
+            .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 +84,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 +103,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}"),
+                lsp_error.to_string(),
             )
-            .into()));
+            .into()))
         }
     };
 
@@ -121,6 +135,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();