Show code lens on entrypoints to invoke commands (#1142)

This adds "Run | Histogram | Estimate | Debug" commands above the
entrypoint in source files.


![image](https://github.com/microsoft/qsharp/assets/16928427/6c07576f-07ce-4f3a-9134-3212caf06327)

language_service/src/code_lens.rs

@@ -0,0 +1,69 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#[cfg(test)]
+mod tests;
+
+use crate::{
+    compilation::{Compilation, CompilationKind},
+    protocol::{CodeLens, CodeLensCommand},
+    qsc_utils::{into_range, span_contains},
+};
+use qsc::{
+    hir::{Attr, ItemKind},
+    line_column::Encoding,
+};
+
+pub(crate) fn get_code_lenses(
+    compilation: &Compilation,
+    source_name: &str,
+    position_encoding: Encoding,
+) -> Vec<CodeLens> {
+    if matches!(compilation.kind, CompilationKind::Notebook) {
+        // entrypoint actions don't work in notebooks
+        return vec![];
+    }
+
+    let user_unit = compilation.user_unit();
+    let source_span = compilation.package_span_of_source(source_name);
+
+    // Get callables in the current source file with the @EntryPoint() attribute.
+    // If there is more than one entrypoint, not our problem, we'll go ahead
+    // and return code lenses for all. The duplicate entrypoint diagnostic
+    // will be reported from elsewhere.
+    let entry_point_decls = user_unit.package.items.values().filter_map(|item| {
+        if span_contains(source_span, item.span.lo) {
+            if let ItemKind::Callable(decl) = &item.kind {
+                if item.attrs.iter().any(|a| a == &Attr::EntryPoint) {
+                    return Some(decl);
+                }
+            }
+        }
+        None
+    });
+
+    entry_point_decls
+        .flat_map(|decl| {
+            let range = into_range(position_encoding, decl.span, &user_unit.sources);
+
+            [
+                CodeLens {
+                    range,
+                    command: CodeLensCommand::Run,
+                },
+                CodeLens {
+                    range,
+                    command: CodeLensCommand::Histogram,
+                },
+                CodeLens {
+                    range,
+                    command: CodeLensCommand::Estimate,
+                },
+                CodeLens {
+                    range,
+                    command: CodeLensCommand::Debug,
+                },
+            ]
+        })
+        .collect()
+}

language_service/src/code_lens/tests.rs

@@ -0,0 +1,132 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#![allow(clippy::needless_raw_string_hashes)]
+
+use super::get_code_lenses;
+use crate::{
+    test_utils::{
+        compile_notebook_with_fake_stdlib, compile_with_fake_stdlib_and_markers_no_cursor,
+    },
+    Encoding,
+};
+use expect_test::{expect, Expect};
+
+fn check(source_with_markers: &str, expect: &Expect) {
+    let (compilation, expected_code_lens_ranges) =
+        compile_with_fake_stdlib_and_markers_no_cursor(source_with_markers);
+    let actual_code_lenses = get_code_lenses(&compilation, "<source>", Encoding::Utf8);
+
+    for expected_range in &expected_code_lens_ranges {
+        assert!(
+            actual_code_lenses
+                .iter()
+                .any(|cl| cl.range == *expected_range),
+            "expected range not found in actual code lenses: {expected_range:?}"
+        );
+    }
+
+    for actual_range in actual_code_lenses.iter().map(|cl| cl.range) {
+        assert!(
+            expected_code_lens_ranges.iter().any(|r| r == &actual_range),
+            "got code lens for unexpected range: {actual_range:?}"
+        );
+    }
+
+    let actual = expected_code_lens_ranges
+        .iter()
+        .enumerate()
+        .map(|(i, r)| {
+            (
+                i,
+                actual_code_lenses
+                    .iter()
+                    .filter(|cl| cl.range == *r)
+                    .map(|cl| cl.command)
+                    .collect::<Vec<_>>(),
+            )
+        })
+        .collect::<Vec<_>>();
+    expect.assert_debug_eq(&actual);
+}
+
+#[test]
+fn one_entrypoint() {
+    check(
+        r#"
+        namespace Test {
+            @EntryPoint()
+            ◉operation Main() : Unit{
+            }◉
+        }"#,
+        &expect![[r#"
+            [
+                (
+                    0,
+                    [
+                        Run,
+                        Histogram,
+                        Estimate,
+                        Debug,
+                    ],
+                ),
+            ]
+        "#]],
+    );
+}
+
+#[test]
+fn two_entrypoints() {
+    check(
+        r#"
+        namespace Test {
+            @EntryPoint()
+            ◉operation Main() : Unit{
+            }◉
+
+            @EntryPoint()
+            ◉operation Foo() : Unit{
+            }◉
+        }"#,
+        &expect![[r#"
+            [
+                (
+                    0,
+                    [
+                        Run,
+                        Histogram,
+                        Estimate,
+                        Debug,
+                    ],
+                ),
+                (
+                    1,
+                    [
+                        Run,
+                        Histogram,
+                        Estimate,
+                        Debug,
+                    ],
+                ),
+            ]
+        "#]],
+    );
+}
+
+#[test]
+fn no_entrypoint_code_lens_in_notebook() {
+    let compilation = compile_notebook_with_fake_stdlib(
+        [(
+            "cell1",
+            "@EntryPoint()
+            operation Main() : Unit {}",
+        )]
+        .into_iter(),
+    );
+
+    let lenses = get_code_lenses(&compilation, "cell1", Encoding::Utf8);
+    assert!(
+        lenses.is_empty(),
+        "entrypoint code lenses should not be present in notebooks"
+    );
+}

language_service/src/compilation.rs

@@ -11,7 +11,7 @@ use qsc::{
     line_column::{Encoding, Position},
     resolve,
     target::Profile,
-    CompileUnit, PackageStore, PackageType, SourceMap,
+    CompileUnit, PackageStore, PackageType, SourceMap, Span,
 };
 use std::sync::Arc;
 
@@ -154,6 +154,23 @@ impl Compilation {
         source.offset + offset
     }
 
+    /// Gets the span of the whole source file.
+    pub(crate) fn package_span_of_source(&self, source_name: &str) -> Span {
+        let unit = self.user_unit();
+
+        let source = unit
+            .sources
+            .find_by_name(source_name)
+            .expect("source should exist in the user source map");
+
+        let len = u32::try_from(source.contents.len()).expect("source length should fit into u32");
+
+        Span {
+            lo: source.offset,
+            hi: source.offset + len,
+        }
+    }
+
     /// Regenerates the compilation with the same sources but the passed in workspace configuration options.
     pub fn recompile(&mut self, package_type: PackageType, target_profile: Profile) {
         let sources = self

language_service/src/lib.rs

@@ -4,6 +4,7 @@
 #![warn(clippy::mod_module_files, clippy::pedantic, clippy::unwrap_used)]
 #![allow(clippy::missing_errors_doc, clippy::missing_panics_doc)]
 
+pub mod code_lens;
 mod compilation;
 pub mod completion;
 pub mod definition;
@@ -26,7 +27,7 @@ use futures::channel::mpsc::{unbounded, UnboundedReceiver, UnboundedSender};
 use futures_util::StreamExt;
 use log::{trace, warn};
 use protocol::{
-    CompletionList, DiagnosticUpdate, Hover, NotebookMetadata, SignatureHelp,
+    CodeLens, CompletionList, DiagnosticUpdate, Hover, NotebookMetadata, SignatureHelp,
     WorkspaceConfigurationUpdate,
 };
 use qsc::{
@@ -205,12 +206,12 @@ impl LanguageService {
         include_declaration: bool,
     ) -> Vec<Location> {
         self.document_op(
-            |position_encoding, compilation, uri, position| {
+            |compilation, uri, position, position_encoding| {
                 references::get_references(
-                    position_encoding,
                     compilation,
                     uri,
                     position,
+                    position_encoding,
                     include_declaration,
                 )
             },
@@ -249,25 +250,39 @@ impl LanguageService {
         self.document_op(rename::prepare_rename, "prepare_rename", uri, position)
     }
 
-    /// Executes an operation that takes a document uri and offset, using the current compilation for that document.
+    /// LSP: textDocument/codeLens
+    #[must_use]
+    pub fn get_code_lenses(&self, uri: &str) -> Vec<CodeLens> {
+        self.document_op(
+            |compilation, uri, (), position_encoding| {
+                code_lens::get_code_lenses(compilation, uri, position_encoding)
+            },
+            "get_code_lenses",
+            uri,
+            (),
+        )
+    }
+
+    /// Executes an operation that takes a document uri, using the current compilation for that document.
     /// All "read" operations should go through this method. This method will borrow the current
     /// compilation state to perform the request.
     ///
     /// If there are outstanding updates to the compilation in the update message queue,
     /// this method will still just return the current compilation state.
-    fn document_op<F, T>(&self, op: F, op_name: &str, uri: &str, position: Position) -> T
+    fn document_op<F, T, A>(&self, op: F, op_name: &str, uri: &str, arg: A) -> T
     where
-        F: Fn(&Compilation, &str, Position, Encoding) -> T,
+        F: Fn(&Compilation, &str, A, Encoding) -> T,
         T: Debug + Default,
+        A: Debug,
     {
-        trace!("{op_name}: uri: {uri}, position: {position:?}");
+        trace!("{op_name}: uri: {uri}, arg: {arg:?}");
 
         // Borrow must succeed here. If it doesn't succeed, a writer
         // (i.e. [`state::CompilationStateUpdater`]) must be holding a mutable reference across
         // an `await` point. Which it shouldn't be doing.
         let compilation_state = self.state.borrow();
         if let Some(compilation) = compilation_state.get_compilation(uri) {
-            let res = op(compilation, uri, position, self.position_encoding);
+            let res = op(compilation, uri, arg, self.position_encoding);
             trace!("{op_name} result: {res:?}");
             res
         } else {

language_service/src/protocol.rs

@@ -119,3 +119,17 @@ pub struct ParameterInformation {
 pub struct NotebookMetadata {
     pub target_profile: Option<Profile>,
 }
+
+#[derive(Debug)]
+pub struct CodeLens {
+    pub range: Range,
+    pub command: CodeLensCommand,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub enum CodeLensCommand {
+    Histogram,
+    Debug,
+    Run,
+    Estimate,
+}