Allow rules to provide their own rust-analyzer providers (#2487)

This change cleans up the rust-analyzer aspect to support external rules
providing their own crate specs. For now only prost implements behavior
for this and the rust-analyzer interface is still private. In the future
if this proves to be performant and a consistent interface then there
should be no issue making this a public part of the `//rust` package.

This change incorporates
https://github.com/bazelbuild/rules_rust/pull/1875 (special thanks to
@snowp!) and addresses performance issues in the generator tool by
allowing users of `bazelisk` to ensure their `tools/bazel` scripts run
should one be provided and to disable running validation actions when
building crate specs.

7 files changed, 157 insertions, 58 deletions

diff --git a/proto/prost/private/prost.bzl b/proto/prost/private/prost.bzl
index d3c34c71..b7d55c36 100644
--- a/proto/prost/private/prost.bzl
+++ b/proto/prost/private/prost.bzl
@@ -5,9 +5,15 @@ load("//proto/prost:providers.bzl", "ProstProtoInfo")
 load("//rust:defs.bzl", "rust_common")
 
 # buildifier: disable=bzl-visibility
+load("//rust/private:providers.bzl", "RustAnalyzerGroupInfo", "RustAnalyzerInfo")
+
+# buildifier: disable=bzl-visibility
 load("//rust/private:rust.bzl", "RUSTC_ATTRS")
 
 # buildifier: disable=bzl-visibility
+load("//rust/private:rust_analyzer.bzl", "write_rust_analyzer_spec_file")
+
+# buildifier: disable=bzl-visibility
 load("//rust/private:rustc.bzl", "rustc_compile_action")
 
 # buildifier: disable=bzl-visibility
@@ -211,6 +217,7 @@ def _rust_prost_aspect_impl(target, ctx):
 
     direct_deps = []
     transitive_deps = [depset(runtime_deps)]
+    rust_analyzer_deps = []
     for proto_dep in proto_deps:
         proto_info = proto_dep[ProstProtoInfo]
 
@@ -220,6 +227,9 @@ def _rust_prost_aspect_impl(target, ctx):
             transitive = [proto_info.transitive_dep_infos],
         ))
 
+        if RustAnalyzerInfo in proto_dep:
+            rust_analyzer_deps.append(proto_dep[RustAnalyzerInfo])
+
     deps = runtime_deps + direct_deps
 
     crate_name = ctx.label.name.replace("-", "_").replace("/", "_")
@@ -244,12 +254,27 @@ def _rust_prost_aspect_impl(target, ctx):
         edition = RUST_EDITION,
     )
 
+    # Always add `test` & `debug_assertions`. See rust-analyzer source code:
+    # https://github.com/rust-analyzer/rust-analyzer/blob/2021-11-15/crates/project_model/src/workspace.rs#L529-L531
+    cfgs = ["test", "debug_assertions"]
+
+    rust_analyzer_info = write_rust_analyzer_spec_file(ctx, ctx.rule.attr, ctx.label, RustAnalyzerInfo(
+        crate = dep_variant_info.crate_info,
+        cfgs = cfgs,
+        env = dep_variant_info.crate_info.rustc_env,
+        deps = rust_analyzer_deps,
+        crate_specs = depset(transitive = [dep.crate_specs for dep in rust_analyzer_deps]),
+        proc_macro_dylib_path = None,
+        build_info = dep_variant_info.build_info,
+    ))
+
     return [
         ProstProtoInfo(
             dep_variant_info = dep_variant_info,
             transitive_dep_infos = depset(transitive = transitive_deps),
             package_info = package_info_file,
         ),
+        rust_analyzer_info,
     ]
 
 rust_prost_aspect = aspect(
@@ -290,13 +315,13 @@ def _rust_prost_library_impl(ctx):
 
     return [
         DefaultInfo(files = depset([dep_variant_info.crate_info.output])),
-        rust_proto_info,
         rust_common.crate_group_info(
             dep_variant_infos = depset(
                 [dep_variant_info],
                 transitive = [rust_proto_info.transitive_dep_infos],
             ),
         ),
+        RustAnalyzerGroupInfo(deps = [proto_dep[RustAnalyzerInfo]]),
     ]
 
 rust_prost_library = rule(
diff --git a/rust/private/providers.bzl b/rust/private/providers.bzl
index a1ab2fe6..2bb81ef7 100644
--- a/rust/private/providers.bzl
+++ b/rust/private/providers.bzl
@@ -151,3 +151,23 @@ TestCrateInfo = provider(
         "crate": "CrateInfo: The underlying CrateInfo of the dependency",
     },
 )
+
+RustAnalyzerInfo = provider(
+    doc = "RustAnalyzerInfo holds rust crate metadata for targets",
+    fields = {
+        "build_info": "BuildInfo: build info for this crate if present",
+        "cfgs": "List[String]: features or other compilation `--cfg` settings",
+        "crate": "CrateInfo: Crate information.",
+        "crate_specs": "Depset[File]: transitive closure of OutputGroupInfo files",
+        "deps": "List[RustAnalyzerInfo]: direct dependencies",
+        "env": "Dict[String: String]: Environment variables, used for the `env!` macro",
+        "proc_macro_dylib_path": "File: compiled shared library output of proc-macro rule",
+    },
+)
+
+RustAnalyzerGroupInfo = provider(
+    doc = "RustAnalyzerGroupInfo holds multiple RustAnalyzerInfos",
+    fields = {
+        "deps": "List[RustAnalyzerInfo]: direct dependencies",
+    },
+)
diff --git a/rust/private/rust_analyzer.bzl b/rust/private/rust_analyzer.bzl
index 74b46648..2682d966 100644
--- a/rust/private/rust_analyzer.bzl
+++ b/rust/private/rust_analyzer.bzl
@@ -20,9 +20,9 @@ given targets. This file can be consumed by rust-analyzer as an alternative
 to Cargo.toml files.
 """
 
-load("//proto/prost:providers.bzl", "ProstProtoInfo")
 load("//rust/platform:triple_mappings.bzl", "system_to_dylib_ext", "triple_to_system")
 load("//rust/private:common.bzl", "rust_common")
+load("//rust/private:providers.bzl", "RustAnalyzerGroupInfo", "RustAnalyzerInfo")
 load("//rust/private:rustc.bzl", "BuildInfo")
 load(
     "//rust/private:utils.bzl",
@@ -32,25 +32,43 @@ load(
     "find_toolchain",
 )
 
-RustAnalyzerInfo = provider(
-    doc = "RustAnalyzerInfo holds rust crate metadata for targets",
-    fields = {
-        "build_info": "BuildInfo: build info for this crate if present",
-        "cfgs": "List[String]: features or other compilation --cfg settings",
-        "crate": "rust_common.crate_info",
-        "crate_specs": "Depset[File]: transitive closure of OutputGroupInfo files",
-        "deps": "List[RustAnalyzerInfo]: direct dependencies",
-        "env": "Dict{String: String}: Environment variables, used for the `env!` macro",
-        "proc_macro_dylib_path": "File: compiled shared library output of proc-macro rule",
-    },
-)
+def write_rust_analyzer_spec_file(ctx, attrs, owner, base_info):
+    """Write a rust-analyzer spec info file.
 
-RustAnalyzerGroupInfo = provider(
-    doc = "RustAnalyzerGroupInfo holds multiple RustAnalyzerInfos",
-    fields = {
-        "deps": "List[RustAnalyzerInfo]: direct dependencies",
-    },
-)
+    Args:
+        ctx (ctx): The current rule's context object.
+        attrs (dict): A mapping of attributes.
+        owner (Label): The label of the owner of the spec info.
+        base_info (RustAnalyzerInfo): The data the resulting RustAnalyzerInfo is based on.
+
+    Returns:
+        RustAnalyzerInfo: Info with the embedded spec file.
+    """
+    crate_spec = ctx.actions.declare_file("{}.rust_analyzer_crate_spec.json".format(owner.name))
+
+    rust_analyzer_info = RustAnalyzerInfo(
+        crate = base_info.crate,
+        cfgs = base_info.cfgs,
+        env = base_info.env,
+        deps = base_info.deps,
+        crate_specs = depset(direct = [crate_spec], transitive = [base_info.crate_specs]),
+        proc_macro_dylib_path = base_info.proc_macro_dylib_path,
+        build_info = base_info.build_info,
+    )
+
+    ctx.actions.write(
+        output = crate_spec,
+        content = json.encode_indent(
+            _create_single_crate(
+                ctx,
+                attrs,
+                rust_analyzer_info,
+            ),
+            indent = " " * 4,
+        ),
+    )
+
+    return rust_analyzer_info
 
 def _rust_analyzer_aspect_impl(target, ctx):
     if (rust_common.crate_info not in target and
@@ -58,6 +76,9 @@ def _rust_analyzer_aspect_impl(target, ctx):
         rust_common.crate_group_info not in target):
         return []
 
+    if RustAnalyzerInfo in target or RustAnalyzerGroupInfo in target:
+        return []
+
     toolchain = find_toolchain(ctx)
 
     # Always add `test` & `debug_assertions`. See rust-analyzer source code:
@@ -102,28 +123,7 @@ def _rust_analyzer_aspect_impl(target, ctx):
         if RustAnalyzerGroupInfo in ctx.rule.attr.actual:
             dep_infos.extend(ctx.rule.attr.actual[RustAnalyzerGroupInfo])
 
-    if ProstProtoInfo in target:
-        for info in target[ProstProtoInfo].transitive_dep_infos.to_list():
-            crate_info = info.crate_info
-            crate_spec = ctx.actions.declare_file(crate_info.owner.name + ".rust_analyzer_crate_spec")
-            rust_analyzer_info = RustAnalyzerInfo(
-                crate = crate_info,
-                cfgs = cfgs,
-                env = crate_info.rustc_env,
-                deps = [],
-                crate_specs = depset(direct = [crate_spec]),
-                proc_macro_dylib_path = None,
-                build_info = info.build_info,
-            )
-            ctx.actions.write(
-                output = crate_spec,
-                content = json.encode(_create_single_crate(ctx, rust_analyzer_info)),
-            )
-            dep_infos.append(rust_analyzer_info)
-
-    if ProstProtoInfo in target:
-        crate_info = target[ProstProtoInfo].dep_variant_info.crate_info
-    elif rust_common.crate_group_info in target:
+    if rust_common.crate_group_info in target:
         return [RustAnalyzerGroupInfo(deps = dep_infos)]
     elif rust_common.crate_info in target:
         crate_info = target[rust_common.crate_info]
@@ -132,22 +132,15 @@ def _rust_analyzer_aspect_impl(target, ctx):
     else:
         fail("Unexpected target type: {}".format(target))
 
-    crate_spec = ctx.actions.declare_file(ctx.label.name + ".rust_analyzer_crate_spec")
-
-    rust_analyzer_info = RustAnalyzerInfo(
+    rust_analyzer_info = write_rust_analyzer_spec_file(ctx, ctx.rule.attr, ctx.label, RustAnalyzerInfo(
         crate = crate_info,
         cfgs = cfgs,
         env = crate_info.rustc_env,
         deps = dep_infos,
-        crate_specs = depset(direct = [crate_spec], transitive = [dep.crate_specs for dep in dep_infos]),
+        crate_specs = depset(transitive = [dep.crate_specs for dep in dep_infos]),
         proc_macro_dylib_path = find_proc_macro_dylib_path(toolchain, target),
         build_info = build_info,
-    )
-
-    ctx.actions.write(
-        output = crate_spec,
-        content = json.encode(_create_single_crate(ctx, rust_analyzer_info)),
-    )
+    ))
 
     return [
         rust_analyzer_info,
@@ -201,12 +194,13 @@ def _crate_id(crate_info):
     """
     return "ID-" + crate_info.root.path
 
-def _create_single_crate(ctx, info):
+def _create_single_crate(ctx, attrs, info):
     """Creates a crate in the rust-project.json format.
 
     Args:
-        ctx (ctx): The rule context
-        info (RustAnalyzerInfo): RustAnalyzerInfo for the current crate
+        ctx (ctx): The rule context.
+        attrs (dict): A mapping of attributes.
+        info (RustAnalyzerInfo): RustAnalyzerInfo for the current crate.
 
     Returns:
         (dict) The crate rust-project.json representation
@@ -240,7 +234,7 @@ def _create_single_crate(ctx, info):
 
     # TODO: The only imagined use case is an env var holding a filename in the workspace passed to a
     # macro like include_bytes!. Other use cases might exist that require more complex logic.
-    expand_targets = concat([getattr(ctx.rule.attr, attr, []) for attr in ["data", "compile_data"]])
+    expand_targets = concat([getattr(attrs, attr, []) for attr in ["data", "compile_data"]])
 
     crate["env"].update({k: dedup_expand_location(ctx, v, expand_targets) for k, v in info.env.items()})
 
diff --git a/tools/rust_analyzer/aquery.rs b/tools/rust_analyzer/aquery.rs
index 3697b0e4..7208aa47 100644
--- a/tools/rust_analyzer/aquery.rs
+++ b/tools/rust_analyzer/aquery.rs
@@ -77,6 +77,9 @@ pub fn get_crate_specs(
 
     let aquery_output = Command::new(bazel)
         .current_dir(workspace)
+        .env_remove("BAZELISK_SKIP_WRAPPER")
+        .env_remove("BUILD_WORKING_DIRECTORY")
+        .env_remove("BUILD_WORKSPACE_DIRECTORY")
         .arg("aquery")
         .arg("--include_aspects")
         .arg("--include_artifacts")
@@ -85,7 +88,7 @@ pub fn get_crate_specs(
         ))
         .arg("--output_groups=rust_analyzer_crate_spec")
         .arg(format!(
-            r#"outputs(".*[.]rust_analyzer_crate_spec",{target_pattern})"#
+            r#"outputs(".*\.rust_analyzer_crate_spec\.json",{target_pattern})"#
         ))
         .arg("--output=jsonproto")
         .output()?;
diff --git a/tools/rust_analyzer/lib.rs b/tools/rust_analyzer/lib.rs
index 53cf0bf7..0ab395ad 100644
--- a/tools/rust_analyzer/lib.rs
+++ b/tools/rust_analyzer/lib.rs
@@ -20,7 +20,11 @@ pub fn generate_crate_info(
 
     let output = Command::new(bazel.as_ref())
         .current_dir(workspace.as_ref())
+        .env_remove("BAZELISK_SKIP_WRAPPER")
+        .env_remove("BUILD_WORKING_DIRECTORY")
+        .env_remove("BUILD_WORKSPACE_DIRECTORY")
         .arg("build")
+        .arg("--norun_validations")
         .arg(format!(
             "--aspects={}//rust:defs.bzl%rust_analyzer_aspect",
             rules_rust.as_ref()
diff --git a/tools/rust_analyzer/main.rs b/tools/rust_analyzer/main.rs
index fb7b69e0..df4b2f9b 100644
--- a/tools/rust_analyzer/main.rs
+++ b/tools/rust_analyzer/main.rs
@@ -65,7 +65,11 @@ fn parse_config() -> anyhow::Result<Config> {
 
     // We need some info from `bazel info`. Fetch it now.
     let mut bazel_info_command = Command::new(&config.bazel);
-    bazel_info_command.arg("info");
+    bazel_info_command
+        .env_remove("BAZELISK_SKIP_WRAPPER")
+        .env_remove("BUILD_WORKING_DIRECTORY")
+        .env_remove("BUILD_WORKSPACE_DIRECTORY")
+        .arg("info");
     if let Some(workspace) = &config.workspace {
         bazel_info_command.current_dir(workspace);
     }
diff --git a/tools/rust_analyzer/rust_project.rs b/tools/rust_analyzer/rust_project.rs
index e4314a23..ba87c581 100644
--- a/tools/rust_analyzer/rust_project.rs
+++ b/tools/rust_analyzer/rust_project.rs
@@ -168,6 +168,23 @@ pub fn generate_rust_project(
                 skipped_crates.len(),
                 skipped_crates
             );
+            let crate_map: BTreeMap<String, &CrateSpec> = unmerged_crates
+                .iter()
+                .map(|c| (c.crate_id.to_string(), *c))
+                .collect();
+
+            for unmerged_crate in &unmerged_crates {
+                let mut path = vec![];
+                if let Some(cycle) = detect_cycle(unmerged_crate, &crate_map, &mut path) {
+                    log::warn!(
+                        "Cycle detected: {:?}",
+                        cycle
+                            .iter()
+                            .map(|c| c.crate_id.to_string())
+                            .collect::<Vec<String>>()
+                    );
+                }
+            }
             return Err(anyhow!(
                 "Failed to make progress on building crate dependency graph"
             ));
@@ -179,6 +196,38 @@ pub fn generate_rust_project(
     Ok(project)
 }
 
+fn detect_cycle<'a>(
+    current_crate: &'a CrateSpec,
+    all_crates: &'a BTreeMap<String, &'a CrateSpec>,
+    path: &mut Vec<&'a CrateSpec>,
+) -> Option<Vec<&'a CrateSpec>> {
+    if path
+        .iter()
+        .any(|dependent_crate| dependent_crate.crate_id == current_crate.crate_id)
+    {
+        let mut cycle_path = path.clone();
+        cycle_path.push(current_crate);
+        return Some(cycle_path);
+    }
+
+    path.push(current_crate);
+
+    for dep in &current_crate.deps {
+        match all_crates.get(dep) {
+            Some(dep_crate) => {
+                if let Some(cycle) = detect_cycle(dep_crate, all_crates, path) {
+                    return Some(cycle);
+                }
+            }
+            None => log::debug!("dep {dep} not found in unmerged crate map"),
+        }
+    }
+
+    path.pop();
+
+    None
+}
+
 pub fn write_rust_project(
     rust_project_path: &Path,
     execution_root: &Path,