Implement transition.and_then (bazelbuild#29542)

### Description

Implements the proposal https://github.com/bazelbuild/proposals/blob/main/designs/2024-04-16-transition-composition.md.

### Motivation

This change in particular allows non-toolchain dependencies of rules to inherit both the execution platform (via `config.exec()`) as well as the current target platform (via a custom transition that records `//command_line_option:platforms`), which is important for certain cross-compilation scenarios.

### Build API Changes

Yes, as discussed and approved in https://github.com/bazelbuild/proposals/blob/main/designs/2024-04-16-transition-composition.md.

### Checklist

- [x] I have added tests for the new use cases (if any).
- [x] I have updated the documentation (if applicable).

### Release Notes

RELNOTES[NEW]: The `and_then` method on `transition`s can be used to compose transitions. Both Starlark transitions and native transitions (e.g. `config.exec()`) are supported.

Closes bazelbuild#29542.

PiperOrigin-RevId: 934095515
Change-Id: I3f36ec0907c4af5bdf43e38323d166422b47051c

Author: fmeum

docs/extending/config.mdx

@@ -665,6 +665,34 @@ hot_chocolate_transition = transition(
 )
 ```
 
+### Composing transitions {#composing-transitions}
+
+Two transitions can be combined into a single new transition using
+[`transition.and_then`](/rules/lib/builtins/transition#and_then):
+
+```python
+combined_transition = first_transition.and_then(second_transition)
+```
+
+The composition applies `first_transition` to the input configuration and then
+runs `second_transition` against each of its output configurations. The original
+transitions are not modified, and the result is itself a `transition` object
+that can be composed further with another `and_then` call.
+
+A composed transition can be attached to a rule or attribute wherever its
+component transitions could be used (subject to the usual restriction that an
+[incoming edge transition](#incoming-edge-transitions) must be 1:1). At most one
+of the composed transitions may be an exec transition (`"exec"` or
+[`config.exec`](/rules/lib/toplevel/config#exec)).
+
+When two of the composed transitions are 1:2+, the composition produces the
+cross product of their splits. The key for each combined split is the
+comma-separated concatenation of the component keys: if the first transition
+produces keys `{a, b}` and the second produces `{x, y}` for each of those, the
+composition produces the four keys `a,x`, `a,y`, `b,x`, and `b,y`. Note that
+each additional 1:2+ transition multiplies the resulting dependency count, so
+chain them with care.
+
 ### Accessing attributes with transitions {#accessing-attributes-with-transitions}
 
 [End to end example](https://github.com/bazelbuild/examples/tree/HEAD/configurations/read_attr_in_transition)

site/en/extending/config.md

@@ -668,6 +668,34 @@ hot_chocolate_transition = transition(
 )
 ```
 
+### Composing transitions {:#composing-transitions}
+
+Two transitions can be combined into a single new transition using
+[`transition.and_then`](/rules/lib/builtins/transition#and_then):
+
+```python
+combined_transition = first_transition.and_then(second_transition)
+```
+
+The composition applies `first_transition` to the input configuration and then
+runs `second_transition` against each of its output configurations. The original
+transitions are not modified, and the result is itself a `transition` object
+that can be composed further with another `and_then` call.
+
+A composed transition can be attached to a rule or attribute wherever its
+component transitions could be used (subject to the usual restriction that an
+[incoming edge transition](#incoming-edge-transitions) must be 1:1). At most one
+of the composed transitions may be an exec transition (`"exec"` or
+[`config.exec`](/rules/lib/toplevel/config#exec)).
+
+When two of the composed transitions are 1:2+, the composition produces the
+cross product of their splits. The key for each combined split is the
+comma-separated concatenation of the component keys: if the first transition
+produces keys `{a, b}` and the second produces `{x, y}` for each of those, the
+composition produces the four keys `a,x`, `a,y`, `b,x`, and `b,y`. Note that
+each additional 1:2+ transition multiplies the resulting dependency count, so
+chain them with care.
+
 ### Accessing attributes with transitions {:#accessing-attributes-with-transitions}
 
 [End to end example](https://github.com/bazelbuild/examples/tree/HEAD/configurations/read_attr_in_transition){: .external}

src/main/java/com/google/devtools/build/lib/analysis/BUILD

@@ -101,6 +101,7 @@ java_library(
         "extra/ExtraAction.java",
         "extra/ExtraActionMapProvider.java",
         "extra/ExtraActionSpec.java",
+        "starlark/ComposedTransitionMaterializer.java",
         "starlark/StarlarkActionFactory.java",
         "starlark/StarlarkAspectPropagationContext.java",
         "starlark/StarlarkAttrModule.java",

src/main/java/com/google/devtools/build/lib/analysis/ConfiguredRuleClassProvider.java

@@ -392,7 +392,7 @@ public Builder addTrimmingTransitionFactory(TransitionFactory<RuleTransitionData
         trimmingTransitionFactory = factory;
       } else {
         trimmingTransitionFactory =
-            ComposingTransitionFactory.of(trimmingTransitionFactory, factory);
+            ComposingTransitionFactory.ofUnchecked(trimmingTransitionFactory, factory);
       }
       return this;
     }

src/main/java/com/google/devtools/build/lib/analysis/DependencyResolutionHelpers.java

@@ -228,17 +228,19 @@ private static ExecutionPlatformResult getExecutionPlatformLabel(
       return ExecutionPlatformResult.ofNullLabel();
     }
 
-    TransitionFactory<AttributeTransitionData> transitionFactory =
-        kind.getAttribute().getTransitionFactory();
-    if (!(transitionFactory instanceof ExecutionTransitionFactory)) {
+    // The exec transition may be composed with other transitions (via transition.and_then), so
+    // look for it anywhere in the (possibly composed) transition factory.
+    ExecutionTransitionFactory execTransitionFactory =
+        findExecutionTransitionFactory(kind.getAttribute().getTransitionFactory());
+    if (execTransitionFactory == null) {
       return ExecutionPlatformResult.ofLabel(
           toolchainContexts
               .getToolchainContext(DeclaredExecGroup.DEFAULT_EXEC_GROUP_NAME)
               .executionPlatform()
               .label());
     }
 
-    String execGroup = ((ExecutionTransitionFactory) transitionFactory).getExecGroup();
+    String execGroup = execTransitionFactory.getExecGroup();
     if (toolchainContexts.hasToolchainContext(execGroup)) {
       PlatformInfo platform = toolchainContexts.getToolchainContext(execGroup).executionPlatform();
       return platform == null
@@ -252,6 +254,24 @@ private static ExecutionPlatformResult getExecutionPlatformLabel(
             kind.getAttribute().getName(), execGroup));
   }
 
+  /**
+   * Returns the {@link ExecutionTransitionFactory} within the given transition factory, or {@code
+   * null} if there is none. A composition contains at most one native transition, so there is at
+   * most one match.
+   */
+  @Nullable
+  private static ExecutionTransitionFactory findExecutionTransitionFactory(
+      TransitionFactory<AttributeTransitionData> transitionFactory) {
+    var found = new ExecutionTransitionFactory[1];
+    transitionFactory.visit(
+        factory -> {
+          if (factory instanceof ExecutionTransitionFactory execFactory) {
+            found[0] = execFactory;
+          }
+        });
+    return found[0];
+  }
+
   /** True if {@code owningAspect} is the main aspect, the last one in {@code aspectsList}. */
   private static boolean isMainAspect(
       ImmutableList<Aspect> aspectsList, @Nullable AspectClass owningAspect) {

src/main/java/com/google/devtools/build/lib/analysis/config/transitions/ComposingTransitionFactory.java

@@ -48,16 +48,20 @@ public abstract class ComposingTransitionFactory<T extends TransitionFactory.Dat
    * one of the transitions is {@link NoTransition}, and returns an efficiently composed transition.
    */
   public static <T extends TransitionFactory.Data> TransitionFactory<T> of(
-      TransitionFactory<T> transitionFactory1, TransitionFactory<T> transitionFactory2) {
-
+      TransitionFactory<T> transitionFactory1, TransitionFactory<T> transitionFactory2)
+      throws IncompatibleTransitionsException {
     Preconditions.checkNotNull(transitionFactory1);
     Preconditions.checkNotNull(transitionFactory2);
-    Preconditions.checkArgument(
-        transitionFactory1.transitionType().isCompatibleWith(transitionFactory2.transitionType()),
-        "transition factory types must be compatible");
-    Preconditions.checkArgument(
-        !transitionFactory1.isSplit() || !transitionFactory2.isSplit(),
-        "can't compose two split transition factories");
+    if (!transitionFactory1
+        .transitionType()
+        .isCompatibleWith(transitionFactory2.transitionType())) {
+      throw new IncompatibleTransitionsException(
+          "transition types must be compatible, got %s and %s"
+              .formatted(transitionFactory1.transitionType(), transitionFactory2.transitionType()));
+    }
+    if (transitionFactory1.isTool() && transitionFactory2.isTool()) {
+      throw new IncompatibleTransitionsException("can't compose two exec transitions");
+    }
 
     if (NoTransition.isInstance(transitionFactory1)) {
       // Since transitionFactory1 causes no changes, use transitionFactory2 directly.
@@ -70,6 +74,26 @@ public static <T extends TransitionFactory.Data> TransitionFactory<T> of(
     return create(transitionFactory1, transitionFactory2);
   }
 
+  /**
+   * Use {@link #of} instead unless the two transition factories are statically guaranteed to be
+   * compatible.
+   */
+  public static <T extends TransitionFactory.Data> TransitionFactory<T> ofUnchecked(
+      TransitionFactory<T> transitionFactory1, TransitionFactory<T> transitionFactory2) {
+    try {
+      return of(transitionFactory1, transitionFactory2);
+    } catch (IncompatibleTransitionsException e) {
+      throw new IllegalArgumentException(e);
+    }
+  }
+
+  /** Thrown when two {@link TransitionFactory} instances cannot be composed. */
+  public static final class IncompatibleTransitionsException extends Exception {
+    IncompatibleTransitionsException(String message) {
+      super(message);
+    }
+  }
+
   private static <T extends TransitionFactory.Data> TransitionFactory<T> create(
       TransitionFactory<T> transitionFactory1, TransitionFactory<T> transitionFactory2) {
     return new AutoValue_ComposingTransitionFactory<T>(transitionFactory1, transitionFactory2);
@@ -84,8 +108,11 @@ public ConfigurationTransition create(T data) {
 
   @Override
   public TransitionType transitionType() {
-    // Both types must match so this is correct.
-    return transitionFactory1().transitionType();
+    // The two types are compatible, so at most one of them is non-ANY; return the more specific
+    // one.
+    return transitionFactory1().transitionType() == TransitionType.ANY
+        ? transitionFactory2().transitionType()
+        : transitionFactory1().transitionType();
   }
 
   abstract TransitionFactory<T> transitionFactory1();
@@ -186,22 +213,19 @@ public boolean equals(Object other) {
     }
 
     /**
-     * Composes a new key out of two given keys. Composing two split transitions is not allowed at
-     * the moment, so what this essentially does are (1) make sure not both transitions are split
-     * and (2) choose one from a split transition, if there's any, or return {@code
-     * PATCH_TRANSITION_KEY}, if there isn't.
+     * Composes a new key out of two given keys. If either transition is a patch (1:1) its
+     * placeholder {@code PATCH_TRANSITION_KEY} is absorbed; if both transitions split, the keys are
+     * joined by a comma so the composition produces the cross product of their splits (e.g. keys
+     * {@code "a"} and {@code "x"} compose to {@code "a,x"}).
      */
-    private String composeKeys(String key1, String key2) {
-      if (!key1.equals(PATCH_TRANSITION_KEY)) {
-        if (!key2.equals(PATCH_TRANSITION_KEY)) {
-          throw new IllegalStateException(
-              String.format(
-                  "can't compose two split transitions %s and %s",
-                  transition1.getName(), transition2.getName()));
-        }
+    private static String composeKeys(String key1, String key2) {
+      if (key1.equals(PATCH_TRANSITION_KEY)) {
+        return key2;
+      }
+      if (key2.equals(PATCH_TRANSITION_KEY)) {
         return key1;
       }
-      return key2;
+      return key1 + "," + key2;
     }
   }
 }

src/main/java/com/google/devtools/build/lib/analysis/producers/RuleTransitionApplier.java

@@ -254,7 +254,7 @@ public StateMachine computeTransition(Tasks tasks) {
         targetAndConfigurationData.getTrimmingTransitionFactory();
     if (trimmingTransitionFactory != null) {
       transitionFactory =
-          ComposingTransitionFactory.of(transitionFactory, trimmingTransitionFactory);
+          ComposingTransitionFactory.ofUnchecked(transitionFactory, trimmingTransitionFactory);
     }
     ConfiguredTargetKey preRuleTransitionKey = targetAndConfigurationData.getPreRuleTransitionKey();
     var transitionData =

src/main/java/com/google/devtools/build/lib/analysis/starlark/ComposedTransitionMaterializer.java

@@ -0,0 +1,76 @@
+// Copyright 2026 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.analysis.starlark;
+
+import com.google.devtools.build.lib.analysis.config.transitions.ComposingTransitionFactory;
+import com.google.devtools.build.lib.analysis.config.transitions.ComposingTransitionFactory.IncompatibleTransitionsException;
+import com.google.devtools.build.lib.analysis.config.transitions.TransitionFactory;
+import com.google.devtools.build.lib.starlarkbuildapi.config.ComposedConfigurationTransition;
+import com.google.devtools.build.lib.starlarkbuildapi.config.ConfigurationTransitionApi;
+import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.Starlark;
+
+/**
+ * Helpers for turning a {@link ComposedConfigurationTransition} (the deferred result of {@code
+ * transition.and_then}) into a single {@link TransitionFactory} for a specific use context (rule
+ * vs. attribute).
+ */
+final class ComposedTransitionMaterializer {
+
+  private ComposedTransitionMaterializer() {}
+
+  /** Converts one element of a composed transition into a {@link TransitionFactory}. */
+  @FunctionalInterface
+  interface ElementConverter<T extends TransitionFactory.Data> {
+    TransitionFactory<T> convert(ConfigurationTransitionApi element) throws EvalException;
+  }
+
+  /**
+   * Folds {@code composition}'s elements into a single {@link TransitionFactory} by converting each
+   * with {@code converter} and combining them with {@link ComposingTransitionFactory#of}.
+   *
+   * @param incompatibleElementMessage error suffix used when an element can't be converted in this
+   *     context (e.g. a native attribute-only transition used as a rule {@code cfg})
+   */
+  static <T extends TransitionFactory.Data> TransitionFactory<T> fold(
+      ComposedConfigurationTransition composition,
+      ElementConverter<T> converter,
+      String incompatibleElementMessage)
+      throws EvalException {
+    TransitionFactory<T> result = null;
+    for (ConfigurationTransitionApi element : composition.getElements()) {
+      TransitionFactory<T> factory;
+      try {
+        factory = converter.convert(element);
+      } catch (EvalException unused) {
+        throw Starlark.errorf(
+            "invalid composed transition for `cfg`: %s (composed at %s)",
+            incompatibleElementMessage, composition.getLocation());
+      }
+      if (result == null) {
+        result = factory;
+      } else {
+        try {
+          result = ComposingTransitionFactory.of(result, factory);
+        } catch (IncompatibleTransitionsException e) {
+          throw Starlark.errorf(
+              "invalid composed transition for `cfg`: %s (composed at %s)",
+              e.getMessage(), composition.getLocation());
+        }
+      }
+    }
+    return result;
+  }
+}

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkAttrModule.java

@@ -56,6 +56,7 @@
 import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
 import com.google.devtools.build.lib.starlarkbuildapi.NativeComputedDefaultApi;
 import com.google.devtools.build.lib.starlarkbuildapi.StarlarkAttrModuleApi;
+import com.google.devtools.build.lib.starlarkbuildapi.config.ComposedConfigurationTransition;
 import com.google.devtools.build.lib.starlarkbuildapi.config.ConfigurationTransitionApi;
 import com.google.devtools.build.lib.starlarkbuildapi.core.StructApi;
 import com.google.devtools.build.lib.util.FileType;
@@ -535,6 +536,12 @@ private static TransitionFactory<AttributeTransitionData> convertCfg(
     if (trans instanceof StarlarkDefinedConfigTransition starlarkDefinedTransition) {
       return new StarlarkAttributeTransitionProvider(starlarkDefinedTransition);
     }
+    if (trans instanceof ComposedConfigurationTransition composition) {
+      return ComposedTransitionMaterializer.fold(
+          composition,
+          element -> convertCfg(thread, element),
+          "it contains a native transition that can only be used as a rule transition");
+    }
     if (trans instanceof ConfigurationTransitionApi cta) {
       // Every ConfigurationTransitionApi must be a TransitionFactory instance to be usable.
       if (cta instanceof TransitionFactory<?> tf) {

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java

@@ -114,6 +114,7 @@
 import com.google.devtools.build.lib.starlarkbuildapi.MacroFunctionApi;
 import com.google.devtools.build.lib.starlarkbuildapi.StarlarkRuleFunctionsApi;
 import com.google.devtools.build.lib.starlarkbuildapi.StarlarkSubruleApi;
+import com.google.devtools.build.lib.starlarkbuildapi.config.ComposedConfigurationTransition;
 import com.google.devtools.build.lib.starlarkbuildapi.config.ConfigurationTransitionApi;
 import com.google.devtools.build.lib.util.FileTypeSet;
 import com.google.devtools.build.lib.util.Pair;
@@ -1098,7 +1099,7 @@ public static StarlarkRuleFunction createRule(
         });
     if (parent != null) {
       transitionFactory =
-          ComposingTransitionFactory.of(transitionFactory, parent.getTransitionFactory());
+          ComposingTransitionFactory.ofUnchecked(transitionFactory, parent.getTransitionFactory());
     }
     // Check if the transition has any Starlark code.
     StarlarkTransitionCheckingVisitor visitor = new StarlarkTransitionCheckingVisitor();
@@ -1250,6 +1251,13 @@ private static TransitionFactory<RuleTransitionData> convertConfig(@Nullable Obj
       // defined in Starlark via, cfg = transition
       return new StarlarkRuleTransitionProvider(starlarkDefinedConfigTransition);
     }
+    if (cfg instanceof ComposedConfigurationTransition composition) {
+      return ComposedTransitionMaterializer.fold(
+          composition,
+          StarlarkRuleClassFunctions::convertConfig,
+          "it contains a native transition that can only be used as an attribute transition, such"
+              + " as the exec transition");
+    }
     if (cfg instanceof ConfigurationTransitionApi cta) {
       // Every ConfigurationTransitionApi must be a TransitionFactory instance to be usable.
       if (cta instanceof TransitionFactory<?> tf) {