Add Xcode toolchain support for an optional binary that can rewrite t…

…he generated header of a Swift module after compilation.

PiperOrigin-RevId: 369323385

swift/internal/compiling.bzl

@@ -48,6 +48,7 @@ load(
     "SWIFT_FEATURE_OPT",
     "SWIFT_FEATURE_OPT_USES_OSIZE",
     "SWIFT_FEATURE_OPT_USES_WMO",
+    "SWIFT_FEATURE_REWRITE_GENERATED_HEADER",
     "SWIFT_FEATURE_STRICT_MODULES",
     "SWIFT_FEATURE_SUPPORTS_LIBRARY_EVOLUTION",
     "SWIFT_FEATURE_SYSTEM_MODULE",
@@ -91,7 +92,8 @@ _WMO_FLAGS = {
 def compile_action_configs(
         *,
         additional_objc_copts = [],
-        additional_swiftc_copts = []):
+        additional_swiftc_copts = [],
+        generated_header_rewriter = None):
     """Returns the list of action configs needed to perform Swift compilation.
 
     Toolchains must add these to their own list of action configs so that
@@ -105,6 +107,9 @@ def compile_action_configs(
         additional_swiftc_copts: An optional list of additional Swift compiler
             flags that should be passed to Swift compile actions only after any
             other toolchain- or user-provided flags.
+        generated_header_rewriter: An executable that will be invoked after
+            compilation to rewrite the generated header, or None if this is not
+            desired.
 
     Returns:
         The list of action configs needed to perform compilation.
@@ -177,6 +182,25 @@ def compile_action_configs(
         ),
     ]
 
+    if generated_header_rewriter:
+        # Only add the generated header rewriter to the command line only if the
+        # toolchain provides one, the relevant feature is requested, and the
+        # particular compilation action is generating a header.
+        def generated_header_rewriter_configurator(prerequisites, args):
+            if prerequisites.generated_header_file:
+                args.add(
+                    generated_header_rewriter,
+                    format = "-Xwrapped-swift=-generated-header-rewriter=%s",
+                )
+
+        action_configs.append(
+            swift_toolchain_config.action_config(
+                actions = [swift_action_names.COMPILE],
+                configurators = [generated_header_rewriter_configurator],
+                features = [SWIFT_FEATURE_REWRITE_GENERATED_HEADER],
+            ),
+        )
+
     #### Compilation-mode-related flags
     #
     # These configs set flags based on the current compilation mode. They mirror

swift/internal/feature_names.bzl

@@ -115,6 +115,11 @@ SWIFT_FEATURE_OPT_USES_WMO = "swift.opt_uses_wmo"
 # the `-Osize` flag instead of `-O`.
 SWIFT_FEATURE_OPT_USES_OSIZE = "swift.opt_uses_osize"
 
+# If enabled, and if the toolchain specifies a generated header rewriting tool,
+# that tool will be invoked after compilation to rewrite the generated header in
+# place.
+SWIFT_FEATURE_REWRITE_GENERATED_HEADER = "swift.rewrite_generated_header"
+
 # If enabled, Swift compiler invocations will use precompiled modules from
 # dependencies instead of module maps and headers, if those dependencies provide
 # them.

swift/internal/xcode_swift_toolchain.bzl

@@ -323,6 +323,7 @@ def _all_action_configs(
         additional_swiftc_copts,
         apple_fragment,
         apple_toolchain,
+        generated_header_rewriter,
         needs_resource_directory,
         target_triple):
     """Returns the action configurations for the Swift toolchain.
@@ -335,6 +336,9 @@ def _all_action_configs(
             the `swift` configuration fragment.
         apple_fragment: The `apple` configuration fragment.
         apple_toolchain: The `apple_common.apple_toolchain()` object.
+        generated_header_rewriter: An executable that will be invoked after
+            compilation to rewrite the generated header, or None if this is not
+            desired.
         needs_resource_directory: If True, the toolchain needs the resource
             directory passed explicitly to the compiler.
         target_triple: The target triple.
@@ -451,13 +455,15 @@ def _all_action_configs(
     action_configs.extend(compile_action_configs(
         additional_objc_copts = additional_objc_copts,
         additional_swiftc_copts = additional_swiftc_copts,
+        generated_header_rewriter = generated_header_rewriter,
     ))
     return action_configs
 
 def _all_tool_configs(
         custom_toolchain,
         env,
         execution_requirements,
+        generated_header_rewriter,
         swift_executable,
         toolchain_root,
         use_param_file,
@@ -469,6 +475,9 @@ def _all_tool_configs(
             one was requested.
         env: The environment variables to set when launching tools.
         execution_requirements: The execution requirements for tools.
+        generated_header_rewriter: An executable that will be invoked after
+            compilation to rewrite the generated header, or None if this is not
+            desired.
         swift_executable: A custom Swift driver executable to be used during the
             build, if provided.
         toolchain_root: The root directory of the toolchain, if provided.
@@ -487,8 +496,13 @@ def _all_tool_configs(
         env = dict(env)
         env["TOOLCHAINS"] = custom_toolchain
 
+    additional_compile_tools = []
+    if generated_header_rewriter:
+        additional_compile_tools.append(generated_header_rewriter)
+
     tool_configs = {
         swift_action_names.COMPILE: swift_toolchain_config.driver_tool_config(
+            additional_tools = additional_compile_tools,
             driver_mode = "swiftc",
             env = env,
             execution_requirements = execution_requirements,
@@ -669,11 +683,13 @@ def _xcode_swift_toolchain_impl(ctx):
 
     env = _xcode_env(platform = platform, xcode_config = xcode_config)
     execution_requirements = xcode_config.execution_info()
+    generated_header_rewriter = ctx.executable.generated_header_rewriter
 
     all_tool_configs = _all_tool_configs(
         custom_toolchain = custom_toolchain,
         env = env,
         execution_requirements = execution_requirements,
+        generated_header_rewriter = generated_header_rewriter,
         swift_executable = swift_executable,
         toolchain_root = toolchain_root,
         use_param_file = use_param_file,
@@ -687,6 +703,7 @@ def _xcode_swift_toolchain_impl(ctx):
         additional_swiftc_copts = ctx.fragments.swift.copts(),
         apple_fragment = apple_fragment,
         apple_toolchain = apple_toolchain,
+        generated_header_rewriter = generated_header_rewriter,
         needs_resource_directory = swift_executable or toolchain_root,
         target_triple = target,
     )
@@ -741,6 +758,20 @@ of a Swift module.
 """,
                 providers = [[SwiftInfo]],
             ),
+            "generated_header_rewriter": attr.label(
+                allow_files = True,
+                cfg = "host",
+                doc = """\
+If present, an executable that will be invoked after compilation to rewrite the
+generated header.
+
+This tool is expected to have a command line interface such that the Swift
+compiler invocation is passed to it following a `"--"` argument, and any
+arguments preceding the `"--"` can be defined by the tool itself (however, at
+this time the worker does not support passing additional flags to the tool).
+""",
+                executable = True,
+            ),
             "implicit_deps": attr.label_list(
                 allow_files = True,
                 doc = """\

tools/worker/swift_runner.cc

@@ -86,6 +86,17 @@ static std::string Unescape(const std::string &arg) {
   return result;
 }
 
+// If `str` starts with `prefix`, `str` is mutated to remove `prefix` and the
+// function returns true. Otherwise, `str` is left unmodified and the function
+// returns `false`.
+static bool StripPrefix(const std::string &prefix, std::string &str) {
+  if (str.find(prefix) != 0) {
+    return false;
+  }
+  str.erase(0, prefix.size());
+  return true;
+}
+
 }  // namespace
 
 SwiftRunner::SwiftRunner(const std::vector<std::string> &args,
@@ -96,6 +107,28 @@ SwiftRunner::SwiftRunner(const std::vector<std::string> &args,
 
 int SwiftRunner::Run(std::ostream *stderr_stream, bool stdout_to_stderr) {
   int exit_code = RunSubProcess(args_, stderr_stream, stdout_to_stderr);
+  if (exit_code != 0) {
+    return exit_code;
+  }
+
+  if (!generated_header_rewriter_path_.empty()) {
+#if __APPLE__
+    // Skip the `xcrun` argument that's added when running on Apple platforms.
+    int initial_args_to_skip = 1;
+#else
+    int initial_args_to_skip = 0;
+#endif
+
+    std::vector<std::string> rewriter_args;
+    rewriter_args.reserve(args_.size() + 2 - initial_args_to_skip);
+    rewriter_args.push_back(generated_header_rewriter_path_);
+    rewriter_args.push_back("--");
+    rewriter_args.insert(rewriter_args.end(),
+                         args_.begin() + initial_args_to_skip, args_.end());
+
+    exit_code = RunSubProcess(rewriter_args, stderr_stream, stdout_to_stderr);
+  }
+
   return exit_code;
 }
 
@@ -155,34 +188,43 @@ bool SwiftRunner::ProcessArgument(
 
   if (arg[0] == '@') {
     changed = ProcessPossibleResponseFile(arg, consumer);
-  } else if (arg == "-Xwrapped-swift=-debug-prefix-pwd-is-dot") {
-    // Get the actual current working directory (the workspace root), which we
-    // didn't know at analysis time.
-    consumer("-debug-prefix-map");
-    consumer(GetCurrentDirectory() + "=.");
-    changed = true;
-  } else if (arg == "-Xwrapped-swift=-ephemeral-module-cache") {
-    // Create a temporary directory to hold the module cache, which will be
-    // deleted after compilation is finished.
-    auto module_cache_dir = TempDirectory::Create("swift_module_cache.XXXXXX");
-    consumer("-module-cache-path");
-    consumer(module_cache_dir->GetPath());
-    temp_directories_.push_back(std::move(module_cache_dir));
-    changed = true;
-  } else if (arg.find("-Xwrapped-swift=") == 0) {
-    // TODO(allevato): Report that an unknown wrapper arg was found and give the
-    // caller a way to exit gracefully.
-    changed = true;
   } else {
-    // Apply any other text substitutions needed in the argument (i.e., for
-    // Apple toolchains).
-    auto new_arg = arg;
-    // Bazel doesn't quote arguments in multi-line params files, so we need to
-    // ensure that our defensive quoting kicks in if an argument contains a
-    // space, even if no other changes would have been made.
-    changed = bazel_placeholder_substitutions_.Apply(new_arg) ||
-              new_arg.find_first_of(' ') != std::string::npos;
-    consumer(new_arg);
+    std::string new_arg = arg;
+    if (StripPrefix("-Xwrapped-swift=", new_arg)) {
+      if (new_arg == "-debug-prefix-pwd-is-dot") {
+        // Get the actual current working directory (the workspace root), which
+        // we didn't know at analysis time.
+        consumer("-debug-prefix-map");
+        consumer(GetCurrentDirectory() + "=.");
+        changed = true;
+      } else if (new_arg == "-ephemeral-module-cache") {
+        // Create a temporary directory to hold the module cache, which will be
+        // deleted after compilation is finished.
+        auto module_cache_dir =
+            TempDirectory::Create("swift_module_cache.XXXXXX");
+        consumer("-module-cache-path");
+        consumer(module_cache_dir->GetPath());
+        temp_directories_.push_back(std::move(module_cache_dir));
+        changed = true;
+      } else if (StripPrefix("-generated-header-rewriter=", new_arg)) {
+        generated_header_rewriter_path_ = new_arg;
+        changed = true;
+      } else {
+        // TODO(allevato): Report that an unknown wrapper arg was found and give
+        // the caller a way to exit gracefully.
+        changed = true;
+      }
+    } else {
+      // Apply any other text substitutions needed in the argument (i.e., for
+      // Apple toolchains).
+      //
+      // Bazel doesn't quote arguments in multi-line params files, so we need to
+      // ensure that our defensive quoting kicks in if an argument contains a
+      // space, even if no other changes would have been made.
+      changed = bazel_placeholder_substitutions_.Apply(new_arg) ||
+                new_arg.find_first_of(' ') != std::string::npos;
+      consumer(new_arg);
+    }
   }
 
   return changed;

tools/worker/swift_runner.h

@@ -124,6 +124,10 @@ class SwiftRunner {
   // Arguments will be unconditionally written into a response file and passed
   // to the tool that way.
   bool force_response_file_;
+
+  // The path to the generated header rewriter tool, if one is being used for
+  // this compilation.
+  std::string generated_header_rewriter_path_;
 };
 
 #endif  // BUILD_BAZEL_RULES_SWIFT_TOOLS_WORKER_SWIFT_RUNNER_H_