feat: provide completions for labels (#317)

README.md

@@ -85,6 +85,55 @@ def _impl(ctx):
 
 then you'll get autocomplete suggestions for the attributes on `ctx`, like `ctx.actions`, `ctx.attr`, and so on!
 
+## Experimental features
+
+Starpls has a number of experimental features that can be enabled via command-line arguments:
+
+### `--experimental_infer_ctx_attributes`
+
+Infer attributes on a rule implementation function's `ctx` parameter.
+
+```python
+def _foo_impl(ctx):
+    ctx.attr.bar # type: int
+
+foo = rule(
+    implementation = _foo_impl,
+    attrs = {
+        "bar": attr.int(),
+    },
+)
+```
+
+### `--experimental_use_code_flow_analysis`
+
+Use code flow analysis to determine additional information about types.
+
+```python
+if cond:
+    x = 1
+else:
+    x = "abc"
+
+x # type: int | string
+```
+
+### `--experimental_enable_label_completions`
+
+Enables completions for labels within Bazel files. For example, given the following `BUILD.bazel` file at the repository root:
+
+```python
+my_rule(
+    name = "foo"
+)
+
+my_rule(
+    name = "bar",
+    srcs = ["//:"],
+              # ^ ... If the cursor is here, "foo" will be suggested.
+)
+```
+
 ## Roadmap
 
 - Parsing

crates/starpls/src/document.rs

@@ -652,6 +652,16 @@ impl FileLoader for DefaultFileLoader {
             }
         }
     }
+
+    fn resolve_build_file(&self, file_id: FileId) -> Option<String> {
+        let path = self.interner.lookup_by_file_id(file_id);
+        let path = path.strip_prefix(&self.workspace).ok()?;
+        if path.file_name()?.to_string_lossy() == "BUILD.bazel" {
+            Some(path.parent()?.to_string_lossy().to_string())
+        } else {
+            None
+        }
+    }
 }
 
 fn read_dir_packages_and_targets(

crates/starpls/src/event_loop.rs

@@ -41,6 +41,12 @@ pub(crate) struct FetchExternalRepoRequest {
     pub(crate) repo: String,
 }
 
+#[derive(Debug)]
+pub(crate) enum RefreshAllWorkspaceTargetsProgress {
+    Begin,
+    End(Option<Vec<String>>),
+}
+
 #[derive(Debug)]
 pub(crate) enum Task {
     AnalysisRequested(Vec<FileId>),
@@ -54,6 +60,8 @@ pub(crate) enum Task {
     FetchExternalRepos(FetchExternalReposProgress),
     /// A request to fetch an external repository.
     FetchExternalRepoRequest(FetchExternalRepoRequest),
+    /// Events from refreshing targets for the current workspace.
+    RefreshAllWorkspaceTargets(RefreshAllWorkspaceTargetsProgress),
 }
 
 #[derive(Debug)]
@@ -292,6 +300,40 @@ impl Server {
                     self.pending_files.insert(file_id);
                 }
             }
+            Task::RefreshAllWorkspaceTargets(progress) => {
+                let token = "RefreshAllWorkspaceTargets";
+                let work_done = match progress {
+                    RefreshAllWorkspaceTargetsProgress::Begin => {
+                        self.send_request::<lsp_types::request::WorkDoneProgressCreate>(
+                            WorkDoneProgressCreateParams {
+                                token: lsp_types::NumberOrString::String(token.to_string()),
+                            },
+                        );
+
+                        lsp_types::WorkDoneProgress::Begin(lsp_types::WorkDoneProgressBegin {
+                            title: "Refreshing all workspace targets".to_string(),
+                            ..Default::default()
+                        })
+                    }
+                    RefreshAllWorkspaceTargetsProgress::End(targets) => {
+                        self.is_refreshing_all_workspace_targets = false;
+                        if let Some(targets) = targets {
+                            self.analysis.set_all_workspace_targets(targets);
+                        }
+
+                        lsp_types::WorkDoneProgress::End(lsp_types::WorkDoneProgressEnd {
+                            message: None,
+                        })
+                    }
+                };
+
+                self.send_notification::<lsp_types::notification::Progress>(
+                    lsp_types::ProgressParams {
+                        token: lsp_types::NumberOrString::String(token.to_string()),
+                        value: lsp_types::ProgressParamsValue::WorkDone(work_done),
+                    },
+                );
+            }
         }
     }
 

crates/starpls/src/handlers/notifications.rs

@@ -55,6 +55,10 @@ pub(crate) fn did_save_text_document(
         match path.file_name().and_then(|file_name| file_name.to_str()) {
             Some("MODULE.bazel" | "WORKSPACE" | "WORKSPACE.bazel" | "WORKSPACE.bzlmod") => {}
             Some(file_name) if file_name.ends_with(".MODULE.bazel") => {}
+            Some("BUILD.bazel") => {
+                server.refresh_all_workspace_targets();
+                return Ok(());
+            }
             _ => return Ok(()),
         }
         server.bazel_client.clear_repo_mappings();

crates/starpls/src/handlers/requests.rs

@@ -84,7 +84,7 @@ pub(crate) fn completion(
     Ok(Some(
         snapshot
             .analysis_snapshot
-            .completion(
+            .completions(
                 FilePosition { file_id, pos },
                 params.context.and_then(|cx| cx.trigger_character),
             )?

crates/starpls/src/main.rs

@@ -56,12 +56,21 @@ pub(crate) struct ServerArgs {
     /// Path to the Bazel binary.
     #[clap(long = "bazel_path")]
     bazel_path: Option<String>,
+
     /// Infer attributes on a rule implementation function's context parameter.
     #[clap(long = "experimental_infer_ctx_attributes", default_value_t = false)]
     infer_ctx_attributes: bool,
-    /// Enable code-flow analysis during typechecking.
+
+    /// Use code-flow analysis during typechecking.
     #[clap(long = "experimental_use_code_flow_analysis", default_value_t = false)]
     use_code_flow_analysis: bool,
+
+    /// Enable completions for labels for targets in the current workspace.
+    #[clap(
+        long = "experimental_enable_label_completions",
+        default_value_t = false
+    )]
+    enable_label_completions: bool,
 }
 
 fn main() -> anyhow::Result<()> {

crates/starpls/src/server.rs

@@ -32,6 +32,7 @@ use crate::document::DocumentChangeKind;
 use crate::document::DocumentManager;
 use crate::document::PathInterner;
 use crate::event_loop::FetchExternalReposProgress;
+use crate::event_loop::RefreshAllWorkspaceTargetsProgress;
 use crate::event_loop::Task;
 use crate::task_pool::TaskPool;
 use crate::task_pool::TaskPoolHandle;
@@ -56,6 +57,7 @@ pub(crate) struct Server {
     pub(crate) force_analysis_for_files: FxHashSet<FileId>,
     pub(crate) fetched_repos: FxHashSet<String>,
     pub(crate) is_fetching_repos: bool,
+    pub(crate) is_refreshing_all_workspace_targets: bool,
 }
 
 pub(crate) struct ServerSnapshot {
@@ -171,6 +173,24 @@ impl Server {
             }
         };
 
+        // Query for all targets in the current workspace, to use for label completion.
+        let targets = if config.args.enable_label_completions {
+            eprintln!("server: querying for all targets in the current workspace");
+            match bazel_client.query_all_workspace_targets() {
+                Ok(targets) => {
+                    eprintln!("server: successfully queried for all targets");
+                    targets
+                }
+                Err(err) => {
+                    eprintln!("server: failed to query all workspace targets: {}", err);
+                    has_bazel_init_err = true;
+                    Default::default()
+                }
+            }
+        } else {
+            Default::default()
+        };
+
         let path_interner = Arc::new(PathInterner::default());
         let loader = DefaultFileLoader::new(
             bazel_client.clone(),
@@ -189,6 +209,7 @@ impl Server {
             },
         );
 
+        analysis.set_all_workspace_targets(targets);
         analysis.set_builtin_defs(builtins, rules);
 
         // Check for a prelude file. We skip verifying that `//tools/build_tools` is actually a package (i.e.
@@ -229,6 +250,7 @@ impl Server {
             force_analysis_for_files: Default::default(),
             fetched_repos: Default::default(),
             is_fetching_repos: false,
+            is_refreshing_all_workspace_targets: false,
         };
 
         if has_bazel_init_err {
@@ -336,6 +358,7 @@ impl Server {
         let repos = mem::take(&mut self.pending_repos);
         let files = mem::take(&mut self.pending_files);
         let bazel_client = self.bazel_client.clone();
+
         self.is_fetching_repos = true;
         self.fetched_repos.extend(repos.clone());
         self.task_pool_handle.spawn_with_sender(move |sender| {
@@ -366,6 +389,37 @@ impl Server {
                 .unwrap();
         });
     }
+
+    pub(crate) fn refresh_all_workspace_targets(&mut self) {
+        if self.is_refreshing_all_workspace_targets || !self.config.args.enable_label_completions {
+            return;
+        }
+
+        let bazel_client = self.bazel_client.clone();
+
+        self.is_refreshing_all_workspace_targets = true;
+        self.task_pool_handle.spawn_with_sender(move |sender| {
+            sender
+                .send(Task::RefreshAllWorkspaceTargets(
+                    RefreshAllWorkspaceTargetsProgress::Begin,
+                ))
+                .unwrap();
+
+            let targets = match bazel_client.query_all_workspace_targets() {
+                Ok(targets) => Some(targets),
+                Err(err) => {
+                    eprintln!("server: failed to query all workspace targets: {}", err);
+                    None
+                }
+            };
+
+            sender
+                .send(Task::RefreshAllWorkspaceTargets(
+                    RefreshAllWorkspaceTargetsProgress::End(targets),
+                ))
+                .unwrap();
+        });
+    }
 }
 
 impl panic::RefUnwindSafe for ServerSnapshot {}

crates/starpls_bazel/src/client.rs

@@ -31,6 +31,7 @@ pub trait BazelClient: Send + Sync + 'static {
     fn clear_repo_mappings(&self);
     fn null_query_external_repo_targets(&self, repo: &str) -> anyhow::Result<()>;
     fn repo_mapping_keys(&self, from_repo: &str) -> anyhow::Result<Vec<String>>;
+    fn query_all_workspace_targets(&self) -> anyhow::Result<Vec<String>>;
 }
 
 pub struct BazelCLI {
@@ -174,6 +175,15 @@ impl BazelClient for BazelCLI {
             .insert(from_repo.to_string(), mapping);
         Ok(keys)
     }
+
+    fn query_all_workspace_targets(&self) -> anyhow::Result<Vec<String>> {
+        let output = self.run_command(&["query", "kind('.* rule', ...)"])?;
+        let targets = str::from_utf8(&output)?
+            .lines()
+            .map(|line| line.to_string())
+            .collect();
+        Ok(targets)
+    }
 }
 
 impl Default for BazelCLI {

crates/starpls_common/src/lib.rs

@@ -96,6 +96,8 @@ pub trait Db: salsa::DbWithJar<Jar> {
         dialect: Dialect,
         from: FileId,
     ) -> anyhow::Result<Option<ResolvedPath>>;
+
+    fn resolve_build_file(&self, file_id: FileId) -> Option<String>;
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]

crates/starpls_hir/src/lib.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use starpls_bazel::Builtins;
 use starpls_common::parse;
 use starpls_common::Dialect;
@@ -77,6 +79,8 @@ pub trait Db: salsa::DbWithJar<Jar> + starpls_common::Db {
     fn get_builtin_defs(&self, dialect: &Dialect) -> BuiltinDefs;
     fn set_bazel_prelude_file(&mut self, file_id: FileId);
     fn get_bazel_prelude_file(&self) -> Option<FileId>;
+    fn set_all_workspace_targets(&mut self, targets: Vec<String>);
+    fn get_all_workspace_targets(&self) -> Arc<Vec<String>>;
 }
 
 #[salsa::tracked]

crates/starpls_hir/src/test_database.rs

@@ -24,6 +24,7 @@ pub(crate) struct TestDatabase {
     storage: salsa::Storage<Self>,
     files: Arc<DashMap<FileId, File>>,
     prelude_file: Option<FileId>,
+    all_workspace_targets: Arc<Vec<String>>,
     pub(crate) gcx: Arc<GlobalContext>,
 }
 
@@ -83,6 +84,10 @@ impl starpls_common::Db for TestDatabase {
     ) -> anyhow::Result<Option<ResolvedPath>> {
         Ok(None)
     }
+
+    fn resolve_build_file(&self, _file_id: FileId) -> Option<String> {
+        None
+    }
 }
 
 impl crate::Db for TestDatabase {
@@ -119,6 +124,14 @@ impl crate::Db for TestDatabase {
     fn gcx(&self) -> &GlobalContext {
         &self.gcx
     }
+
+    fn set_all_workspace_targets(&mut self, targets: Vec<String>) {
+        self.all_workspace_targets = Arc::new(targets)
+    }
+
+    fn get_all_workspace_targets(&self) -> Arc<Vec<String>> {
+        Arc::clone(&self.all_workspace_targets)
+    }
 }
 
 #[allow(unused)]