From 4a3ffcb1e89b3f20b12a37b55595682f4bc866b7 Mon Sep 17 00:00:00 2001
From: Daniel Wagner-Hall <dwagnerhall@apple.com>
Date: Fri, 10 May 2024 17:41:10 +0100
Subject: [PATCH] Add target for running rustfmt directly (#2648)

Previously you either needed to go via our bazel-aware wrapper, or write
a rule to grab the file from the toolchain.
---
 BUILD.bazel                                   |  2 +-
 tools/rustfmt/BUILD.bazel                     | 35 ++++++++++++++++-
 tools/rustfmt/src/upstream_rustfmt_wrapper.rs | 38 +++++++++++++++++++
 3 files changed, 72 insertions(+), 3 deletions(-)
 create mode 100644 tools/rustfmt/src/upstream_rustfmt_wrapper.rs

diff --git a/BUILD.bazel b/BUILD.bazel
index fabaae73c9..27a04d131d 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -110,7 +110,7 @@ label_flag(
 
 alias(
     name = "rustfmt",
-    actual = "//tools/rustfmt",
+    actual = "//tools/rustfmt:target_aware_rustfmt",
     visibility = ["//visibility:public"],
 )
 
diff --git a/tools/rustfmt/BUILD.bazel b/tools/rustfmt/BUILD.bazel
index e7d39615ab..1a74c9971e 100644
--- a/tools/rustfmt/BUILD.bazel
+++ b/tools/rustfmt/BUILD.bazel
@@ -32,8 +32,18 @@ rust_library(
     ],
 )
 
-rust_binary(
+# Deprecated but present for compatibility.
+alias(
     name = "rustfmt",
+    actual = ":target_aware_rustfmt",
+    deprecation = "Prefer //tools/rustfmt:target_aware_rustfmt",
+    visibility = ["//visibility:public"],
+)
+
+# This is a wrapper around the upstream rustfmt binary which is aware of targets,
+# and will try to do things like set the correct edition for files when formatting them based on their owning targets.
+rust_binary(
+    name = "target_aware_rustfmt",
     srcs = [
         "src/main.rs",
     ],
@@ -74,6 +84,27 @@ rust_clippy(
     testonly = True,
     visibility = ["//visibility:private"],
     deps = [
-        ":rustfmt",
+        ":target_aware_rustfmt",
+    ],
+)
+
+# This is a small wrapper around the raw upstream rustfmt binary which can be `bazel run`.
+rust_binary(
+    name = "upstream_rustfmt",
+    srcs = [
+        "src/upstream_rustfmt_wrapper.rs",
+    ],
+    data = ["//rust/toolchain:current_rustfmt_toolchain"],
+    edition = "2018",
+    rustc_env = {
+        "RUSTFMT": "$(rlocationpath //rust/toolchain:current_rustfmt_toolchain)",
+    },
+    toolchains = [
+        "@rules_rust//rust/toolchain:current_rust_toolchain",
+        "@rules_rust//rust/toolchain:current_rustfmt_toolchain",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tools/runfiles",
     ],
 )
diff --git a/tools/rustfmt/src/upstream_rustfmt_wrapper.rs b/tools/rustfmt/src/upstream_rustfmt_wrapper.rs
new file mode 100644
index 0000000000..33c091aa13
--- /dev/null
+++ b/tools/rustfmt/src/upstream_rustfmt_wrapper.rs
@@ -0,0 +1,38 @@
+use std::path::PathBuf;
+use std::process::{exit, Command};
+
+fn main() {
+    let runfiles = runfiles::Runfiles::create().unwrap();
+
+    let rustfmt = runfiles::rlocation!(runfiles, env!("RUSTFMT"));
+    if !rustfmt.exists() {
+        panic!("rustfmt does not exist at: {}", rustfmt.display());
+    }
+
+    let working_directory = std::env::var_os("BUILD_WORKING_DIRECTORY")
+        .map(PathBuf::from)
+        .unwrap_or_else(|| std::env::current_dir().expect("Failed to get working directory"));
+
+    let status = Command::new(rustfmt)
+        .current_dir(&working_directory)
+        .args(std::env::args_os().skip(1))
+        .status()
+        .expect("Failed to run rustfmt");
+    if let Some(exit_code) = status.code() {
+        exit(exit_code);
+    }
+    exit_for_signal(&status);
+    panic!("Child rustfmt process didn't exit or get a signal - don't know how to handle it");
+}
+
+#[cfg(target_family = "unix")]
+fn exit_for_signal(status: &std::process::ExitStatus) {
+    use std::os::unix::process::ExitStatusExt;
+    if let Some(signal) = status.signal() {
+        exit(signal);
+    }
+}
+
+#[cfg(not(target_family = "unix"))]
+#[allow(unused)]
+fn exit_for_signal(status: &std::process::ExitStatus) {}