From 16c88568f2c59c77ed9c23d2642dadae461ba754 Mon Sep 17 00:00:00 2001
From: Ed Schouten <eschouten@apple.com>
Date: Fri, 7 Jun 2024 16:08:49 +0200
Subject: [PATCH] Depend on a copy of rustfmt for the target

The toolchain_files() and current_rust_toolchain() rules return files
for the current 'exec' platform. This is perfecly reasonable if you want
to use these tools as part of ctx.actions.run(). But if you want to use
these as part of 'data = []' (runfiles), they must be built for the
target.

This change adds a helper rule, toolchain_files_for_target(), which can
take the aforementioned rules as an input, transitioning it so that it
yields files for the target platform. I'm not happy with how this rule
is implemented, but there is little we can do about that until
https://github.com/bazelbuild/bazel/issues/19645 is addressed.
---
 rust/private/toolchain_utils.bzl | 31 ++++++++++++++++++++++++++++++-
 tools/rustfmt/BUILD.bazel        | 14 ++++++++++----
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/rust/private/toolchain_utils.bzl b/rust/private/toolchain_utils.bzl
index 4fb2ae83e3..542aad41dd 100644
--- a/rust/private/toolchain_utils.bzl
+++ b/rust/private/toolchain_utils.bzl
@@ -53,7 +53,7 @@ def _toolchain_files_impl(ctx):
     )]
 
 toolchain_files = rule(
-    doc = "A rule for fetching files from a rust toolchain.",
+    doc = "A rule for fetching files from a rust toolchain for the exec platform.",
     implementation = _toolchain_files_impl,
     attrs = {
         "tool": attr.string(
@@ -95,3 +95,32 @@ current_rust_toolchain = rule(
         str(Label("@rules_rust//rust:toolchain_type")),
     ],
 )
+
+def _transition_to_target_impl(settings, attr):
+    return {
+        # String conversion is needed to prevent a crash with Bazel 6.x.
+        "//command_line_option:extra_execution_platforms": [
+            str(platform)
+            for platform in settings["//command_line_option:platforms"]
+        ],
+    }
+
+_transition_to_target = transition(
+    implementation = _transition_to_target_impl,
+    inputs = ["//command_line_option:platforms"],
+    outputs = ["//command_line_option:extra_execution_platforms"],
+)
+
+def _toolchain_files_for_target_impl(ctx):
+    return [ctx.attr.toolchain_files[0][DefaultInfo]]
+
+toolchain_files_for_target = rule(
+    doc = "A rule for fetching files from a rust toolchain for the target platform.",
+    implementation = _toolchain_files_for_target_impl,
+    attrs = {
+        "toolchain_files": attr.label(cfg = _transition_to_target, mandatory = True),
+        "_allowlist_function_transition": attr.label(
+            default = "@bazel_tools//tools/allowlists/function_transition_allowlist",
+        ),
+    },
+)
diff --git a/tools/rustfmt/BUILD.bazel b/tools/rustfmt/BUILD.bazel
index 691d29caac..b35b7ef5c2 100644
--- a/tools/rustfmt/BUILD.bazel
+++ b/tools/rustfmt/BUILD.bazel
@@ -1,4 +1,5 @@
 load("//rust:defs.bzl", "rust_binary", "rust_clippy", "rust_library")
+load("//rust/private:toolchain_utils.bzl", "toolchain_files_for_target")
 load("//tools:tool_utils.bzl", "aspect_repository")
 
 exports_files(
@@ -9,6 +10,11 @@ exports_files(
     visibility = ["//visibility:public"],
 )
 
+toolchain_files_for_target(
+    name = "current_rustfmt_toolchain_for_target",
+    toolchain_files = "//rust/toolchain:current_rustfmt_toolchain",
+)
+
 rust_library(
     name = "rustfmt_lib",
     srcs = glob(
@@ -19,12 +25,12 @@ rust_library(
         ],
     ),
     data = [
+        ":current_rustfmt_toolchain_for_target",
         "//:rustfmt.toml",
-        "//rust/toolchain:current_rustfmt_toolchain",
     ],
     edition = "2018",
     rustc_env = {
-        "RUSTFMT": "$(rlocationpath //rust/toolchain:current_rustfmt_toolchain)",
+        "RUSTFMT": "$(rlocationpath :current_rustfmt_toolchain_for_target)",
         "RUSTFMT_CONFIG": "$(rlocationpath //:rustfmt.toml)",
     },
     deps = [
@@ -91,10 +97,10 @@ rust_binary(
     srcs = [
         "src/upstream_rustfmt_wrapper.rs",
     ],
-    data = ["//rust/toolchain:current_rustfmt_toolchain"],
+    data = [":current_rustfmt_toolchain_for_target"],
     edition = "2018",
     rustc_env = {
-        "RUSTFMT": "$(rlocationpath //rust/toolchain:current_rustfmt_toolchain)",
+        "RUSTFMT": "$(rlocationpath :current_rustfmt_toolchain_for_target)",
     },
     toolchains = ["@rules_rust//rust/toolchain:current_rust_toolchain"],
     visibility = ["//visibility:public"],