Allow the Swift toolchain to be associated with an execution group and pass an optional execution group name down to actions that use it

This updates the following APIs:

-   `swift_common.get_toolchain` now takes an `exec_group` argument, which is will look up the toolchain under `ctx.exec_groups[NAME]` instead of `ctx`.
-   `swift_common.{compile,compile_module_interface,precompile_clang_module}` now take an `exec_group` argument, which is passed down to the underlying `ctx.actions.run` call for actions using the toolchain.

In all cases the default value of the argument is `None`, so omitting it retains the current behavior.

Most importantly, this lets us declare the Swift toolchain somewhere other than the rule/aspect's `ctx`, which ensures that in a multiple toolchain scenario, the Swift toolchain doesn't fully decide the execution platform for all actions (even non-Swift actions).

Since execution groups are somewhat in flux right now, this may not be the final form—it would be nice to have the coupling between the execution group and the toolchain automated. I considered just having each action be its own execution group (like C++ link actions use `cpp_link`), but that seems worse; Swift actions should almost never be split across different toolchains/platforms, and it would force the user to have to declare multiple exec groups if they wanted to use the feature with multiple actions. This lets them share something like a single "swift" execution group across multiple actions, with the only caveat being that they have to pass that name explicitly.

PiperOrigin-RevId: 502638394
(cherry picked from commit aeee2bb)
Signed-off-by: Brentley Jones <[email protected]>

Author: allevato

swift/internal/actions.bzl

@@ -120,8 +120,10 @@ def is_action_enabled(action_name, swift_toolchain):
     return bool(tool_config)
 
 def run_toolchain_action(
+        *,
         actions,
         action_name,
+        exec_group = None,
         feature_configuration,
         prerequisites,
         swift_toolchain,
@@ -133,6 +135,8 @@ def run_toolchain_action(
         actions: The rule context's `Actions` object, which will be used to
             create `Args` objects.
         action_name: The name of the action that should be run.
+        exec_group: Runs the Swift compilation action under the given execution
+            group's context. If `None`, the default execution group is used.
         feature_configuration: A feature configuration obtained from
             `swift_common.configure_features`.
         mnemonic: The mnemonic to associate with the action. If not provided,
@@ -206,6 +210,7 @@ def run_toolchain_action(
     actions.run(
         arguments = [tool_executable_args, args],
         env = tool_config.env,
+        exec_group = exec_group,
         executable = executable,
         execution_requirements = execution_requirements,
         inputs = depset(

swift/internal/compiling.bzl

@@ -204,6 +204,7 @@ def compile_module_interface(
         *,
         actions,
         compilation_contexts,
+        exec_group = None,
         feature_configuration,
         module_name,
         swiftinterface_file,
@@ -219,6 +220,8 @@ def compile_module_interface(
             Swift-compatible preprocessor defines, header search paths, and so
             forth. These are typically retrieved from the `CcInfo` providers of
             a target's dependencies.
+        exec_group: Runs the Swift compilation action under the given execution
+            group's context. If `None`, the default execution group is used.
         feature_configuration: A feature configuration obtained from
             `swift_common.configure_features`.
         module_name: The name of the Swift module being compiled. This must be
@@ -327,6 +330,7 @@ def compile_module_interface(
     run_toolchain_action(
         actions = actions,
         action_name = SWIFT_ACTION_COMPILE_MODULE_INTERFACE,
+        exec_group = exec_group,
         feature_configuration = feature_configuration,
         outputs = [swiftmodule_file],
         prerequisites = prerequisites,
@@ -360,6 +364,7 @@ def compile(
         cc_infos,
         copts = [],
         defines = [],
+        exec_group = None,
         extra_swift_infos = [],
         feature_configuration,
         generated_header_name = None,
@@ -392,6 +397,8 @@ def compile(
             determine whether whole module optimization is being requested,
             which affects the nature of the output files.
         defines: Symbols that should be defined by passing `-D` to the compiler.
+        exec_group: Runs the Swift compilation action under the given execution
+            group's context. If `None`, the default execution group is used.
         extra_swift_infos: Extra `SwiftInfo` providers that aren't contained
             by the `deps` of the target being compiled but are required for
             compilation.
@@ -708,6 +715,7 @@ to use swift_common.compile(include_dev_srch_paths = ...) instead.\
     run_toolchain_action(
         actions = actions,
         action_name = SWIFT_ACTION_COMPILE,
+        exec_group = exec_group,
         feature_configuration = feature_configuration,
         outputs = all_compile_outputs,
         prerequisites = prerequisites,
@@ -753,6 +761,7 @@ to use swift_common.compile(include_dev_srch_paths = ...) instead.\
         precompiled_module = _precompile_clang_module(
             actions = actions,
             cc_compilation_context = compilation_context_to_compile,
+            exec_group = exec_group,
             feature_configuration = feature_configuration,
             is_swift_generated_header = True,
             module_map_file = compile_outputs.generated_module_map_file,
@@ -836,6 +845,7 @@ def precompile_clang_module(
         *,
         actions,
         cc_compilation_context,
+        exec_group = None,
         feature_configuration,
         module_map_file,
         module_name,
@@ -855,6 +865,8 @@ def precompile_clang_module(
             of headers of the direct dependencies of the module being compiled,
             which Clang needs to be physically present before it detects that
             they belong to one of the precompiled module dependencies.
+        exec_group: Runs the Swift compilation action under the given execution
+            group's context. If `None`, the default execution group is used.
         feature_configuration: A feature configuration obtained from
             `swift_common.configure_features`.
         module_map_file: A textual module map file that defines the Clang module
@@ -875,6 +887,7 @@ def precompile_clang_module(
     return _precompile_clang_module(
         actions = actions,
         cc_compilation_context = cc_compilation_context,
+        exec_group = exec_group,
         feature_configuration = feature_configuration,
         is_swift_generated_header = False,
         module_map_file = module_map_file,
@@ -888,6 +901,7 @@ def _precompile_clang_module(
         *,
         actions,
         cc_compilation_context,
+        exec_group = None,
         feature_configuration,
         is_swift_generated_header,
         module_map_file,
@@ -908,6 +922,8 @@ def _precompile_clang_module(
             of headers of the direct dependencies of the module being compiled,
             which Clang needs to be physically present before it detects that
             they belong to one of the precompiled module dependencies.
+        exec_group: Runs the Swift compilation action under the given execution
+            group's context. If `None`, the default execution group is used.
         feature_configuration: A feature configuration obtained from
             `swift_common.configure_features`.
         is_swift_generated_header: If True, the action is compiling the
@@ -989,6 +1005,7 @@ def _precompile_clang_module(
     run_toolchain_action(
         actions = actions,
         action_name = SWIFT_ACTION_PRECOMPILE_C_MODULE,
+        exec_group = exec_group,
         feature_configuration = feature_configuration,
         outputs = [precompiled_module],
         prerequisites = prerequisites,

swift/internal/toolchain_utils.bzl

@@ -16,11 +16,16 @@
 
 SWIFT_TOOLCHAIN_TYPE = "@build_bazel_rules_swift//toolchains:toolchain_type"
 
-def get_swift_toolchain(ctx, attr = "_toolchain"):
+def get_swift_toolchain(ctx, *, exec_group = None, attr = "_toolchain"):
     """Gets the Swift toolchain associated with the rule or aspect.
 
     Args:
         ctx: The rule or aspect context.
+        exec_group: The name of the execution group that should contain the
+            toolchain. If this is provided and the toolchain is not declared in
+            that execution group, it will be looked up from `ctx` as a fallback
+            instead. If this argument is `None` (the default), then the
+            toolchain will only be looked up from `ctx.`
         attr: The name of the attribute on the calling rule or aspect that
             should be used to retrieve the toolchain if it is not provided by
             the `toolchains` argument of the rule/aspect. Note that this is only
@@ -30,6 +35,11 @@ def get_swift_toolchain(ctx, attr = "_toolchain"):
     Returns:
         A `SwiftToolchainInfo` provider.
     """
+    if exec_group:
+        group = ctx.exec_groups[exec_group]
+        if group and SWIFT_TOOLCHAIN_TYPE in group.toolchains:
+            return group.toolchains[SWIFT_TOOLCHAIN_TYPE].swift_toolchain
+
     if SWIFT_TOOLCHAIN_TYPE in ctx.toolchains:
         return ctx.toolchains[SWIFT_TOOLCHAIN_TYPE].swift_toolchain
 
@@ -55,8 +65,8 @@ def use_swift_toolchain():
     ```
 
     Returns:
-        A list of toolchain types that should be passed to `rule()` or
-        `aspect()`.
+        A list of toolchain types that should be passed to `rule()`, `aspect()`,
+        or `exec_group`.
     """
 
     # TODO(b/205018581): Intentionally empty for now so that rule definitions