Allow rule writers to create exec_group related exec transitions via config.exec(exec_group) and attach them to starlark dependency attrs.

PiperOrigin-RevId: 308290724
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/BUILD b/src/main/java/com/google/devtools/build/lib/analysis/BUILD
index 92efb5b..4de5d46 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/analysis/BUILD
@@ -1557,9 +1557,11 @@
         ":config/transitions/patch_transition",
         ":config/transitions/transition_factory",
         ":platform_options",
+        ":toolchain_collection",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/events",
         "//src/main/java/com/google/devtools/build/lib/packages",
+        "//src/main/java/com/google/devtools/build/lib/skylarkbuildapi",
         "//third_party:guava",
         "//third_party:jsr305",
     ],
@@ -2035,6 +2037,7 @@
     name = "skylark/starlark_config",
     srcs = ["skylark/StarlarkConfig.java"],
     deps = [
+        ":config/execution_transition_factory",
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/packages:type",
         "//src/main/java/com/google/devtools/build/lib/skylarkbuildapi",
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/DependencyResolver.java b/src/main/java/com/google/devtools/build/lib/analysis/DependencyResolver.java
index 6dbb7f0..b84e2f9 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/DependencyResolver.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/DependencyResolver.java
@@ -22,6 +22,7 @@
 import com.google.devtools.build.lib.analysis.AspectCollection.AspectCycleOnPathException;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.analysis.config.ConfigMatchingProvider;
+import com.google.devtools.build.lib.analysis.config.ExecutionTransitionFactory;
 import com.google.devtools.build.lib.analysis.config.Fragment;
 import com.google.devtools.build.lib.analysis.config.HostTransition;
 import com.google.devtools.build.lib.analysis.config.TransitionResolver;
@@ -291,7 +292,8 @@
           @Nullable Rule fromRule,
           ConfiguredAttributeMapper attributeMap,
           @Nullable ToolchainCollection<ToolchainContext> toolchainContexts,
-          Iterable<Aspect> aspects) {
+          Iterable<Aspect> aspects)
+          throws EvalException {
     OrderedSetMultimap<DependencyKind, PartiallyResolvedDependency> partiallyResolvedDeps =
         OrderedSetMultimap.create();
 
@@ -335,13 +337,28 @@
           aspects, attribute.getName(), entry.getKey().getOwningAspect(), propagatingAspects);
 
       Label executionPlatformLabel = null;
-      // TODO(b/151742236): support transitions to other ({@link ExecGroup defined}) execution
-      // platforms
-      if (toolchainContexts != null
-          && toolchainContexts.getDefaultToolchainContext().executionPlatform() != null) {
-        executionPlatformLabel =
-            toolchainContexts.getDefaultToolchainContext().executionPlatform().label();
+      if (toolchainContexts != null) {
+        if (attribute.getTransitionFactory() instanceof ExecutionTransitionFactory) {
+          String execGroup =
+              ((ExecutionTransitionFactory) attribute.getTransitionFactory()).getExecGroup();
+          if (!toolchainContexts.hasToolchainContext(execGroup)) {
+            String error =
+                String.format(
+                    "Attr '%s' declares a transition for non-existent exec group '%s'",
+                    attribute.getName(), execGroup);
+            if (fromRule != null) {
+              throw new EvalException(fromRule.getLocation(), error);
+            } else {
+              throw new EvalException(error);
+            }
+          }
+          if (toolchainContexts.getToolchainContext(execGroup).executionPlatform() != null) {
+            executionPlatformLabel =
+                toolchainContexts.getToolchainContext(execGroup).executionPlatform().label();
+          }
+        }
       }
+
       AttributeTransitionData attributeTransitionData =
           AttributeTransitionData.builder()
               .attributes(attributeMap)
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/ToolchainCollection.java b/src/main/java/com/google/devtools/build/lib/analysis/ToolchainCollection.java
index 2cbe240..206a8d6 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/ToolchainCollection.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/ToolchainCollection.java
@@ -24,7 +24,7 @@
 import java.util.Map;
 
 /**
- * A wrapper class for a map of exec_group names to their relevent ToolchainContext.
+ * A wrapper class for a map of exec_group names to their relevant ToolchainContext.
  *
  * @param <T> any class that extends ToolchainContext. This generic allows ToolchainCollection to be
  *     used, e.g., both before and after toolchain resolution.
@@ -73,6 +73,10 @@
     return toolchainContexts.get(DEFAULT_EXEC_GROUP_NAME);
   }
 
+  boolean hasToolchainContext(String execGroup) {
+    return toolchainContexts.containsKey(execGroup);
+  }
+
   public T getToolchainContext(String execGroup) {
     return toolchainContexts.get(execGroup);
   }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/config/ExecutionTransitionFactory.java b/src/main/java/com/google/devtools/build/lib/analysis/config/ExecutionTransitionFactory.java
index 643f6fb..28a5bf5 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/config/ExecutionTransitionFactory.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/config/ExecutionTransitionFactory.java
@@ -14,6 +14,8 @@
 
 package com.google.devtools.build.lib.analysis.config;
 
+import static com.google.devtools.build.lib.analysis.ToolchainCollection.DEFAULT_EXEC_GROUP_NAME;
+
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.analysis.PlatformOptions;
@@ -22,6 +24,7 @@
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.packages.AttributeTransitionData;
+import com.google.devtools.build.lib.skylarkbuildapi.StarlarkConfigApi.ExecTransitionFactoryApi;
 import javax.annotation.Nullable;
 
 /**
@@ -29,26 +32,29 @@
  * transition to a configuration suitable for building dependencies for the execution platform of
  * the depending target.
  */
-public class ExecutionTransitionFactory implements TransitionFactory<AttributeTransitionData> {
+public class ExecutionTransitionFactory
+    implements TransitionFactory<AttributeTransitionData>, ExecTransitionFactoryApi {
 
-  private ExecutionTransitionFactory() {}
+  private final String execGroup;
 
-  /** Returns a new {@link ExecutionTransitionFactory}. */
-  public static ExecutionTransitionFactory create() {
-    return new ExecutionTransitionFactory();
+  private ExecutionTransitionFactory(String execGroup) {
+    this.execGroup = execGroup;
   }
 
   /**
-   * Returns either a new {@link ExecutionTransitionFactory} or a factory for the {@link
-   * HostTransition}, depending on the value of {@code enableExecutionTransition}.
+   * Returns a new {@link ExecutionTransitionFactory} for the default {@link
+   * com.google.devtools.build.lib.packages.ExecGroup}.
    */
-  public static TransitionFactory<AttributeTransitionData> create(
-      boolean enableExecutionTransition) {
-    if (enableExecutionTransition) {
-      return create();
-    } else {
-      return HostTransition.createFactory();
-    }
+  public static ExecutionTransitionFactory create() {
+    return new ExecutionTransitionFactory(DEFAULT_EXEC_GROUP_NAME);
+  }
+
+  /**
+   * Returns a new {@link ExecutionTransitionFactory} for the given {@link
+   * com.google.devtools.build.lib.packages.ExecGroup}.
+   */
+  public static ExecutionTransitionFactory create(String execGroup) {
+    return new ExecutionTransitionFactory(execGroup);
   }
 
   @Override
@@ -56,6 +62,10 @@
     return new ExecutionTransition(data.executionPlatform());
   }
 
+  public String getExecGroup() {
+    return execGroup;
+  }
+
   @Override
   public boolean isHost() {
     return false;
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkAttr.java b/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkAttr.java
index 5fa7eb8..bfb7f3f 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkAttr.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkAttr.java
@@ -257,6 +257,8 @@
         builder.cfg(HostTransition.createFactory());
       } else if (trans.equals("exec")) {
         builder.cfg(ExecutionTransitionFactory.create());
+      } else if (trans instanceof ExecutionTransitionFactory) {
+        builder.cfg((ExecutionTransitionFactory) trans);
       } else if (trans instanceof SplitTransition) {
         builder.cfg(TransitionFactories.of((SplitTransition) trans));
       } else if (trans instanceof TransitionFactory) {
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/skylark/StarlarkConfig.java b/src/main/java/com/google/devtools/build/lib/analysis/skylark/StarlarkConfig.java
index d345023..5b197c6 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/skylark/StarlarkConfig.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/skylark/StarlarkConfig.java
@@ -19,9 +19,11 @@
 import static com.google.devtools.build.lib.packages.Type.STRING;
 import static com.google.devtools.build.lib.packages.Type.STRING_LIST;
 
+import com.google.devtools.build.lib.analysis.config.ExecutionTransitionFactory;
 import com.google.devtools.build.lib.packages.BuildSetting;
 import com.google.devtools.build.lib.skylarkbuildapi.StarlarkConfigApi;
 import com.google.devtools.build.lib.syntax.Printer;
+import com.google.devtools.build.lib.syntax.Starlark;
 
 /** Starlark namespace for creating build settings. */
 // TODO(juliexxia): Consider adding more types of build settings, specifically other label types.
@@ -48,6 +50,13 @@
   }
 
   @Override
+  public ExecutionTransitionFactory exec(Object execGroupUnchecked) {
+    return execGroupUnchecked == Starlark.NONE
+        ? ExecutionTransitionFactory.create()
+        : ExecutionTransitionFactory.create((String) execGroupUnchecked);
+  }
+
+  @Override
   public void repr(Printer printer) {
     printer.append("<config>");
   }
diff --git a/src/main/java/com/google/devtools/build/lib/skylarkbuildapi/SkylarkAttrApi.java b/src/main/java/com/google/devtools/build/lib/skylarkbuildapi/SkylarkAttrApi.java
index f0618e5..badbe68 100644
--- a/src/main/java/com/google/devtools/build/lib/skylarkbuildapi/SkylarkAttrApi.java
+++ b/src/main/java/com/google/devtools/build/lib/skylarkbuildapi/SkylarkAttrApi.java
@@ -89,7 +89,7 @@
           + "attribute.";
 
   String CONFIGURATION_ARG = "cfg";
-  // TODO(bazel-team): Update when new Starlark-based configuration framework is implemented.
+  // TODO(b/151742236): Update when new Starlark-based configuration framework is implemented.
   String CONFIGURATION_DOC =
       "<a href=\"../rules.$DOC_EXT#configurations\">Configuration</a> of the attribute. It can be "
           + "either <code>\"host\"</code> or <code>\"target\"</code>.";
diff --git a/src/main/java/com/google/devtools/build/lib/skylarkbuildapi/StarlarkConfigApi.java b/src/main/java/com/google/devtools/build/lib/skylarkbuildapi/StarlarkConfigApi.java
index e4c32e1..53a91b5 100644
--- a/src/main/java/com/google/devtools/build/lib/skylarkbuildapi/StarlarkConfigApi.java
+++ b/src/main/java/com/google/devtools/build/lib/skylarkbuildapi/StarlarkConfigApi.java
@@ -18,6 +18,7 @@
 import com.google.devtools.build.lib.skylarkinterface.SkylarkCallable;
 import com.google.devtools.build.lib.skylarkinterface.SkylarkModule;
 import com.google.devtools.build.lib.skylarkinterface.SkylarkModuleCategory;
+import com.google.devtools.build.lib.syntax.StarlarkSemantics.FlagIdentifier;
 import com.google.devtools.build.lib.syntax.StarlarkValue;
 
 /**
@@ -116,4 +117,29 @@
               + "key-value map of settings like {'cpu': 'ppc', 'copt': '-DFoo'}, this describes a "
               + "single entry in that map.")
   interface BuildSettingApi extends StarlarkValue {}
+
+  @SkylarkCallable(
+      name = "exec",
+      doc = "<i>experimental</i> Creates an execution transition.",
+      enableOnlyWithFlag = FlagIdentifier.EXPERIMENTAL_EXEC_GROUPS,
+      parameters = {
+        @Param(
+            name = "exec_group",
+            type = String.class,
+            named = true,
+            noneable = true,
+            defaultValue = "None",
+            doc =
+                "The name of the exec group whose execution platform this transition will use. If"
+                    + " not provided, this exec transition will use the target's default execution"
+                    + " platform.")
+      })
+  ExecTransitionFactoryApi exec(Object execGroupUnchecked);
+
+  /** The api for exec transitions. */
+  @SkylarkModule(
+      name = "ExecTransitionFactory",
+      category = SkylarkModuleCategory.BUILTIN,
+      doc = "<i>experimental</i> an execution transition.")
+  interface ExecTransitionFactoryApi extends StarlarkValue {}
 }