feat(lsp): Add Compile code lens for main function and contracts (#2309)

Author: phated

crates/lsp/src/lib.rs

@@ -27,8 +27,11 @@ use noirc_frontend::hir::FunctionNameMatch;
 use serde_json::Value as JsonValue;
 use tower::Service;
 
+const ARROW: &str = "▶\u{fe0e}";
 const TEST_COMMAND: &str = "nargo.test";
-const TEST_CODELENS_TITLE: &str = "▶\u{fe0e} Run Test";
+const TEST_CODELENS_TITLE: &str = "Run Test";
+const COMPILE_COMMAND: &str = "nargo.compile";
+const COMPILE_CODELENS_TITLE: &str = "Compile";
 
 // State for the LSP gets implemented on this struct and is internal to the implementation
 pub struct LspState {
@@ -185,7 +188,7 @@ fn on_code_lens_request(
 
     for package in &workspace {
         let (mut context, crate_id) = prepare_package(package);
-        // We ignore the warnings and errors produced by compilation for producing codelenses
+        // 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);
 
@@ -210,7 +213,7 @@ fn on_code_lens_request(
                 .unwrap_or_default();
 
             let command = Command {
-                title: TEST_CODELENS_TITLE.into(),
+                title: format!("{ARROW} {TEST_CODELENS_TITLE}"),
                 command: TEST_COMMAND.into(),
                 arguments: Some(vec![
                     "--program-dir".into(),
@@ -226,6 +229,73 @@ fn on_code_lens_request(
 
             lenses.push(lens);
         }
+
+        if package.is_binary() {
+            if let Some(main_func_id) = context.get_main_function(&crate_id) {
+                let location = context.function_meta(&main_func_id).name.location;
+                let file_id = location.file;
+
+                // Ignore diagnostics for any file that wasn't the file we saved
+                // TODO: In the future, we could create "related" diagnostics for these files
+                // TODO: This currently just appends the `.nr` file extension that we store as a constant,
+                // but that won't work if we accept other extensions
+                if fm.path(file_id).with_extension(FILE_EXTENSION) != file_path {
+                    continue;
+                }
+
+                let range = byte_span_to_range(files, file_id.as_usize(), location.span.into())
+                    .unwrap_or_default();
+
+                let command = Command {
+                    title: format!("{ARROW} {COMPILE_CODELENS_TITLE}"),
+                    command: COMPILE_COMMAND.into(),
+                    arguments: Some(vec![
+                        "--program-dir".into(),
+                        format!("{}", workspace.root_dir.display()).into(),
+                        "--package".into(),
+                        format!("{}", package.name).into(),
+                    ]),
+                };
+
+                let lens = CodeLens { range, command: command.into(), data: None };
+
+                lenses.push(lens);
+            }
+        }
+
+        if package.is_contract() {
+            // Currently not looking to deduplicate this since we don't have a clear decision on if the Contract stuff is staying
+            for contract in context.get_all_contracts(&crate_id) {
+                let location = contract.location;
+                let file_id = location.file;
+
+                // Ignore diagnostics for any file that wasn't the file we saved
+                // TODO: In the future, we could create "related" diagnostics for these files
+                // TODO: This currently just appends the `.nr` file extension that we store as a constant,
+                // but that won't work if we accept other extensions
+                if fm.path(file_id).with_extension(FILE_EXTENSION) != file_path {
+                    continue;
+                }
+
+                let range = byte_span_to_range(files, file_id.as_usize(), location.span.into())
+                    .unwrap_or_default();
+
+                let command = Command {
+                    title: format!("{ARROW} {COMPILE_CODELENS_TITLE}"),
+                    command: COMPILE_COMMAND.into(),
+                    arguments: Some(vec![
+                        "--program-dir".into(),
+                        format!("{}", workspace.root_dir.display()).into(),
+                        "--package".into(),
+                        format!("{}", package.name).into(),
+                    ]),
+                };
+
+                let lens = CodeLens { range, command: command.into(), data: None };
+
+                lenses.push(lens);
+            }
+        }
     }
 
     let res = if lenses.is_empty() { Ok(None) } else { Ok(Some(lenses)) };
@@ -365,6 +435,9 @@ fn on_did_save_text_document(
         }
     }
 
+    // We need to refresh lenses when we compile since that's the only time they can be accurately reflected
+    let _ = state.client.code_lens_refresh(());
+
     let _ = state.client.publish_diagnostics(PublishDiagnosticsParams {
         uri: params.text_document.uri,
         version: None,

crates/noirc_frontend/src/hir/def_map/mod.rs

@@ -148,7 +148,7 @@ impl CrateDefMap {
                     let functions =
                         module.value_definitions().filter_map(|id| id.as_function()).collect();
                     let name = self.get_module_path(id, module.parent);
-                    Some(Contract { name, functions })
+                    Some(Contract { name, location: module.location, functions })
                 } else {
                     None
                 }
@@ -194,6 +194,7 @@ impl CrateDefMap {
 pub struct Contract {
     /// To keep `name` semi-unique, it is prefixed with the names of parent modules via CrateDefMap::get_module_path
     pub name: String,
+    pub location: Location,
     pub functions: Vec<FuncId>,
 }