Codelens support for the LSP

starlark/src/analysis/find_call_name.rs

@@ -27,6 +27,15 @@ use crate::codemap::Span;
 use crate::codemap::Spanned;
 use crate::syntax::AstModule;
 
+#[derive(Debug, PartialEq, Eq)]
+/// Function calls that have a name attribute
+pub struct NamedFunctionCall {
+    /// The value of the name attribute passed to the function call
+    pub name: String,
+    /// The span of the function call
+    pub span: Span,
+}
+
 /// Find the location of a top level function call that has a kwarg "name", and a string value
 /// matching `name`.
 pub trait AstModuleFindCallName {
@@ -36,17 +45,23 @@ pub trait AstModuleFindCallName {
     /// NOTE: If the AST is exposed in the future, this function may be removed and implemented
     ///       by specific programs instead.
     fn find_function_call_with_name(&self, name: &str) -> Option<Span>;
+
+    /// Find all top level function calls that have a kwarg "name"
+    fn find_named_function_calls(&self) -> Vec<NamedFunctionCall>;
 }
 
 impl AstModuleFindCallName for AstModule {
     fn find_function_call_with_name(&self, name: &str) -> Option<Span> {
-        let mut ret = None;
+        self.find_named_function_calls()
+            .iter()
+            .find(|call| call.name == name)
+            .map(|call| call.span)
+    }
 
-        fn visit_expr(ret: &mut Option<Span>, name: &str, node: &AstExpr) {
-            if ret.is_some() {
-                return;
-            }
+    fn find_named_function_calls<'a>(&'a self) -> Vec<NamedFunctionCall> {
+        let mut ret = Vec::new();
 
+        fn visit_expr(ret: &mut Vec<NamedFunctionCall>, node: &AstExpr) {
             match node {
                 Spanned {
                     node: Expr::Call(identifier, arguments),
@@ -60,20 +75,22 @@ impl AstModuleFindCallName for AstModule {
                                     node: Expr::Literal(AstLiteral::String(s)),
                                     ..
                                 },
-                            ) if arg_name.node == "name" && s.node == name => Some(identifier.span),
+                            ) if arg_name.node == "name" => Some(NamedFunctionCall {
+                                name: s.node.clone(),
+                                span: identifier.span,
+                            }),
                             _ => None,
                         });
-                        if found.is_some() {
-                            *ret = found;
+                        if let Some(found) = found {
+                            ret.push(found);
                         }
                     }
                 }
-                _ => node.visit_expr(|x| visit_expr(ret, name, x)),
+                _ => node.visit_expr(|x| visit_expr(ret, x)),
             }
         }
 
-        self.statement()
-            .visit_expr(|x| visit_expr(&mut ret, name, x));
+        self.statement().visit_expr(|x| visit_expr(&mut ret, x));
         ret
     }
 }
@@ -113,4 +130,44 @@ def x(name = "foo_name"):
         assert_eq!(None, module.find_function_call_with_name("bar_name"));
         Ok(())
     }
+
+    #[test]
+    fn finds_all_named_function_calls() -> anyhow::Result<()> {
+        let contents = r#"
+foo(name = "foo_name")
+bar("bar_name")
+baz(name = "baz_name")
+
+def x(name = "foo_name"):
+    pass
+"#;
+
+        let module = AstModule::parse("foo.star", contents.to_owned(), &Dialect::Extended).unwrap();
+
+        let calls = module.find_named_function_calls();
+
+        assert_eq!(
+            calls.iter().map(|call| &call.name).collect::<Vec<_>>(),
+            &["foo_name", "baz_name"]
+        );
+
+        assert_eq!(
+            calls
+                .iter()
+                .map(|call| module.codemap().resolve_span(call.span))
+                .collect::<Vec<_>>(),
+            &[
+                ResolvedSpan {
+                    begin: ResolvedPos { line: 1, column: 0 },
+                    end: ResolvedPos { line: 1, column: 3 }
+                },
+                ResolvedSpan {
+                    begin: ResolvedPos { line: 3, column: 0 },
+                    end: ResolvedPos { line: 3, column: 3 }
+                }
+            ]
+        );
+
+        Ok(())
+    }
 }

starlark_bin/bin/bazel.rs

@@ -38,6 +38,8 @@ use std::process::Command;
 use either::Either;
 use lsp_types::CompletionItemKind;
 use lsp_types::Url;
+use serde::Deserialize;
+use serde::Serialize;
 use starlark::analysis::find_call_name::AstModuleFindCallName;
 use starlark::analysis::AstModuleLint;
 use starlark::docs::get_registered_starlark_docs;
@@ -53,6 +55,7 @@ use starlark::syntax::AstModule;
 use starlark_lsp::completion::StringCompletionResult;
 use starlark_lsp::completion::StringCompletionType;
 use starlark_lsp::error::eval_message_to_lsp_diagnostic;
+use starlark_lsp::server::Codelens;
 use starlark_lsp::server::LspContext;
 use starlark_lsp::server::LspEvalResult;
 use starlark_lsp::server::LspUrl;
@@ -542,7 +545,7 @@ impl BazelContext {
                         if let Some(targets) = self.query_buildable_targets(
                             &format!(
                                 "{render_base}{}",
-                                if render_base.ends_with(':') { "" } else { ":" }
+                                render_base.strip_suffix(":").unwrap_or(&render_base)
                             ),
                             workspace_root,
                         ) {
@@ -605,7 +608,7 @@ impl BazelContext {
         workspace_dir: Option<&Path>,
     ) -> Option<Vec<String>> {
         let mut raw_command = Command::new("bazel");
-        let mut command = raw_command.arg("query").arg(format!("{module}*"));
+        let mut command = raw_command.arg("query").arg(format!("{module}:*"));
         if let Some(workspace_dir) = workspace_dir {
             command = command.current_dir(workspace_dir);
         }
@@ -619,13 +622,102 @@ impl BazelContext {
         Some(
             output
                 .lines()
-                .filter_map(|line| line.strip_prefix(module).map(|str| str.to_owned()))
+                .filter_map(|line| {
+                    line.strip_prefix(module)
+                        .and_then(|line| line.strip_prefix(":"))
+                        .map(|str| str.to_owned())
+                })
                 .collect(),
         )
     }
+
+    fn query_executable_targets(
+        &self,
+        module: &str,
+        workspace_dir: Option<&Path>,
+    ) -> Option<Vec<String>> {
+        let mut raw_command = Command::new("bazel");
+        let mut command = raw_command
+            .arg("cquery")
+            .arg(format!("{module}:*"))
+            .arg("--output=starlark")
+            .arg("--starlark:expr")
+            .arg("target.label.name if providers(target)['FilesToRunProvider'].executable != None else ''");
+
+        if let Some(workspace_dir) = workspace_dir {
+            command = command.current_dir(workspace_dir);
+        }
+
+        let output = command.output().ok()?;
+        if !output.status.success() {
+            return None;
+        }
+
+        let output = String::from_utf8(output.stdout).ok()?;
+        Some(
+            output
+                .lines()
+                .filter(|line| !line.is_empty())
+                .map(|line| line.to_owned())
+                .collect(),
+        )
+    }
+
+    /// Resolves a label, excluding the target name, given a BUILD file
+    fn resolve_label_from_build_file(
+        &self,
+        build_file_path: &Path,
+        workspace_root: Option<&Path>,
+    ) -> Option<String> {
+        let (repo_name, path) =
+            if let Some((repo_name, path)) = self.get_repository_for_path(build_file_path) {
+                (Some(repo_name), path)
+            } else {
+                (None, build_file_path.strip_prefix(workspace_root?).ok()?)
+            };
+
+        let package_directory = path.parent()?;
+
+        if let Some(repo_name) = repo_name {
+            Some(format!("@{repo_name}//{}", package_directory.display()))
+        } else {
+            Some(format!("//{}", package_directory.display()))
+        }
+    }
+}
+
+#[derive(Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct BazelTargets {
+    workspace_root: PathBuf,
+    targets: Vec<String>,
+}
+
+pub enum BazelCommand {
+    Build(BazelTargets),
+    Run(BazelTargets),
+}
+
+impl starlark_lsp::server::Command for BazelCommand {
+    fn to_lsp(&self) -> lsp_types::Command {
+        match self {
+            BazelCommand::Build(targets) => lsp_types::Command {
+                title: "Build".to_string(),
+                command: "bazel.buildTarget".to_string(),
+                arguments: Some(vec![serde_json::to_value(targets).unwrap()]),
+            },
+            BazelCommand::Run(targets) => lsp_types::Command {
+                title: "Run".to_string(),
+                command: "bazel.runTarget".to_string(),
+                arguments: Some(vec![serde_json::to_value(targets).unwrap()]),
+            },
+        }
+    }
 }
 
 impl LspContext for BazelContext {
+    type Command = BazelCommand;
+
     fn parse_file_with_contents(&self, uri: &LspUrl, content: String) -> LspEvalResult {
         match uri {
             LspUrl::File(uri) => {
@@ -871,4 +963,61 @@ impl LspContext for BazelContext {
 
         Ok(names)
     }
+
+    fn codelens(
+        &self,
+        document_uri: &LspUrl,
+        workspace_root: Option<&Path>,
+        ast: &AstModule,
+    ) -> Vec<Codelens<Self::Command>> {
+        if let (LspUrl::File(build_file_path), Some(workspace_root)) =
+            (document_uri, workspace_root)
+        {
+            if let Some(label_prefix) =
+                self.resolve_label_from_build_file(&build_file_path, Some(workspace_root))
+            {
+                if let (Some(buildable_targets), Some(executable_targets)) = (
+                    self.query_buildable_targets(&label_prefix, Some(workspace_root)),
+                    self.query_executable_targets(&label_prefix, Some(workspace_root)),
+                ) {
+                    let buildable_targets: HashSet<_> = buildable_targets.into_iter().collect();
+                    let executable_targets: HashSet<_> = executable_targets.into_iter().collect();
+
+                    let mut codelenses = Vec::new();
+
+                    for call in ast.find_named_function_calls() {
+                        let label = format!("{label_prefix}:{}", call.name);
+
+                        if buildable_targets.contains(&call.name) {
+                            codelenses.push(Codelens {
+                                command: BazelCommand::Build(BazelTargets {
+                                    workspace_root: workspace_root.to_owned(),
+                                    targets: vec![label.clone()],
+                                }),
+                                span: call.span,
+                            })
+                        }
+
+                        if executable_targets.contains(&call.name) {
+                            codelenses.push(Codelens {
+                                command: BazelCommand::Run(BazelTargets {
+                                    workspace_root: workspace_root.to_owned(),
+                                    targets: vec![label],
+                                }),
+                                span: call.span,
+                            });
+                        }
+                    }
+
+                    codelenses
+                } else {
+                    Vec::new()
+                }
+            } else {
+                Vec::new()
+            }
+        } else {
+            Vec::new()
+        }
+    }
 }

starlark_bin/bin/eval.rs

@@ -39,6 +39,7 @@ use starlark::eval::Evaluator;
 use starlark::syntax::AstModule;
 use starlark::syntax::Dialect;
 use starlark_lsp::error::eval_message_to_lsp_diagnostic;
+use starlark_lsp::server::Command;
 use starlark_lsp::server::LspContext;
 use starlark_lsp::server::LspEvalResult;
 use starlark_lsp::server::LspUrl;
@@ -287,7 +288,17 @@ impl Context {
     }
 }
 
+pub enum EvalCommand {}
+
+impl Command for EvalCommand {
+    fn to_lsp(&self) -> lsp_types::Command {
+        match *self {}
+    }
+}
+
 impl LspContext for Context {
+    type Command = EvalCommand;
+
     fn parse_file_with_contents(&self, uri: &LspUrl, content: String) -> LspEvalResult {
         match uri {
             LspUrl::File(uri) => {
@@ -375,6 +386,10 @@ impl LspContext for Context {
     fn get_environment(&self, _uri: &LspUrl) -> DocModule {
         DocModule::default()
     }
+
+    fn codelens(&self, document_uri: &LspUrl, workspace_root: Option<&Path>, ast: &AstModule) -> Vec<starlark_lsp::server::Codelens<EvalCommand>> {
+        todo!()
+    }
 }
 
 pub(crate) fn globals() -> Globals {