Attribute parameterize all starlark attribute transition implementation functions.

The impl functions of these transitions currently take one argument but the final API calls for a second. For now, support both 1-param and 2-param transition impl functions to not break current users.

Though this param looks like the attr param of the starlark rule `ctx` object, it's not exactly the same. For dependency-typed attributes, `ctx` can actually access the providers of those dependencies. For this attr object, we just see them as the representative labels. Current usage of transitions has no need to actually see providers so leaving as a TODO for now and we can reconsider when we find a use case.

PiperOrigin-RevId: 227543951
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/config/StarlarkDefinedConfigTransition.java b/src/main/java/com/google/devtools/build/lib/analysis/config/StarlarkDefinedConfigTransition.java
index 7604ca2..89a33e6 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/config/StarlarkDefinedConfigTransition.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/config/StarlarkDefinedConfigTransition.java
@@ -18,6 +18,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.events.Location;
+import com.google.devtools.build.lib.packages.StructImpl;
 import com.google.devtools.build.lib.skylarkbuildapi.config.ConfigurationTransitionApi;
 import com.google.devtools.build.lib.skylarkinterface.SkylarkPrinter;
 import com.google.devtools.build.lib.syntax.BaseFunction;
@@ -91,7 +92,8 @@
    * @throws InterruptedException if evaluating the transition is interrupted
    */
   public abstract ImmutableList<Map<String, Object>> getChangedSettings(
-      Map<String, Object> previousSettings) throws EvalException, InterruptedException;
+      Map<String, Object> previousSettings, StructImpl attributeMap)
+      throws EvalException, InterruptedException;
 
   public static StarlarkDefinedConfigTransition newRegularTransition(
       BaseFunction impl,
@@ -122,7 +124,7 @@
 
     @Override
     public ImmutableList<Map<String, Object>> getChangedSettings(
-        Map<String, Object> previousSettings) {
+        Map<String, Object> previousSettings, StructImpl attributeMapper) {
       return ImmutableList.of(changedSettings);
     }
 
@@ -156,12 +158,23 @@
 
     @Override
     public ImmutableList<Map<String, Object>> getChangedSettings(
-        Map<String, Object> previousSettings) throws EvalException, InterruptedException {
+        Map<String, Object> previousSettings, StructImpl attributeMapper)
+        throws EvalException, InterruptedException {
       Object result;
       try {
-        result = evalFunction(impl, previousSettings);
+        result = evalFunction(impl, ImmutableList.of(previousSettings, attributeMapper));
       } catch (EvalException e) {
-        throw new EvalException(impl.getLocation(), e.getMessage());
+        // TODO(b/121134880): Still support the one-param syntax since we have users using that.
+        // Deprecate when this API will stop changing.
+        if (e.getMessage().contains("too many (2) positional arguments in call to")) {
+          try {
+            result = evalFunction(impl, ImmutableList.of(previousSettings));
+          } catch (EvalException e2) {
+            throw new EvalException(impl.getLocation(), e2.getMessage());
+          }
+        } else {
+          throw new EvalException(impl.getLocation(), e.getMessage());
+        }
       }
 
       if (!(result instanceof SkylarkDict<?, ?>)) {
@@ -209,7 +222,7 @@
     }
 
     /** Evaluate the input function with the given argument, and return the return value. */
-    private Object evalFunction(BaseFunction function, Object arg)
+    private Object evalFunction(BaseFunction function, ImmutableList<Object> args)
         throws InterruptedException, EvalException {
       try (Mutability mutability = Mutability.create("eval_transition_function")) {
         Environment env =
@@ -218,7 +231,7 @@
                 .setEventHandler(eventHandler)
                 .build();
 
-        return function.call(ImmutableList.of(arg), ImmutableMap.of(), null, env);
+        return function.call(args, ImmutableMap.of(), null, env);
       }
     }
   }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/skylark/FunctionSplitTransitionProvider.java b/src/main/java/com/google/devtools/build/lib/analysis/skylark/FunctionSplitTransitionProvider.java
index 11ece1c..d374a75 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/skylark/FunctionSplitTransitionProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/skylark/FunctionSplitTransitionProvider.java
@@ -14,9 +14,11 @@
 
 package com.google.devtools.build.lib.analysis.skylark;
 
+import static com.google.devtools.build.lib.analysis.skylark.SkylarkAttributesCollection.ERROR_MESSAGE_FOR_NO_ATTR;
 import static java.nio.charset.StandardCharsets.US_ASCII;
 
 import com.google.common.base.Joiner;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -27,29 +29,38 @@
 import com.google.devtools.build.lib.analysis.config.FragmentOptions;
 import com.google.devtools.build.lib.analysis.config.StarlarkDefinedConfigTransition;
 import com.google.devtools.build.lib.analysis.config.transitions.SplitTransition;
+import com.google.devtools.build.lib.packages.Attribute;
 import com.google.devtools.build.lib.packages.Attribute.SplitTransitionProvider;
 import com.google.devtools.build.lib.packages.AttributeMap;
+import com.google.devtools.build.lib.packages.ConfiguredAttributeMapper;
+import com.google.devtools.build.lib.packages.StructImpl;
+import com.google.devtools.build.lib.packages.StructProvider;
+import com.google.devtools.build.lib.syntax.Environment;
 import com.google.devtools.build.lib.syntax.EvalException;
 import com.google.devtools.build.lib.syntax.Mutability;
+import com.google.devtools.build.lib.syntax.Runtime;
 import com.google.devtools.build.lib.syntax.Runtime.NoneType;
 import com.google.devtools.build.lib.syntax.SkylarkDict;
+import com.google.devtools.build.lib.syntax.SkylarkType;
 import com.google.devtools.common.options.OptionDefinition;
 import com.google.devtools.common.options.OptionsParser;
 import com.google.devtools.common.options.OptionsParsingException;
 import java.lang.reflect.Field;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
+import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 
 /**
  * This class implements a split transition provider that takes a Skylark transition function as
- * input.  The transition function takes a settings argument, which is a dictionary containing the
- * current option values.  It either returns a dictionary mapping option name to new option value
+ * input. The transition function takes a settings argument, which is a dictionary containing the
+ * current option values. It either returns a dictionary mapping option name to new option value
  * (for a patch transition), or a dictionary of such dictionaries (for a split transition).
  *
- * Currently the implementation ignores the attributes provided by the containing function.
+ * <p>TODO(bazel-team): Consider allowing dependency-typed attributes to actually return providers
+ * instead of just labels (see {@link SkylarkAttributesCollection#addAttribute}).
  */
 public class FunctionSplitTransitionProvider implements SplitTransitionProvider {
 
@@ -64,15 +75,30 @@
 
   @Override
   public SplitTransition apply(AttributeMap attributeMap) {
-    return new FunctionSplitTransition(starlarkDefinedConfigTransition);
+    return new FunctionSplitTransition(starlarkDefinedConfigTransition, attributeMap);
   }
 
   private static class FunctionSplitTransition implements SplitTransition {
     private final StarlarkDefinedConfigTransition starlarkDefinedConfigTransition;
+    private final StructImpl attrObject;
 
-    public FunctionSplitTransition(
-        StarlarkDefinedConfigTransition starlarkDefinedConfigTransition) {
+    FunctionSplitTransition(
+        StarlarkDefinedConfigTransition starlarkDefinedConfigTransition,
+        AttributeMap attributeMap) {
+      Preconditions.checkArgument(attributeMap instanceof ConfiguredAttributeMapper);
       this.starlarkDefinedConfigTransition = starlarkDefinedConfigTransition;
+
+      ConfiguredAttributeMapper configuredAttributeMapper =
+          (ConfiguredAttributeMapper) attributeMap;
+      LinkedHashMap<String, Object> attributes = new LinkedHashMap<>();
+      for (String attribute : attributeMap.getAttributeNames()) {
+        Object val =
+            configuredAttributeMapper.get(attribute, attributeMap.getAttributeType(attribute));
+        attributes.put(
+            Attribute.getSkylarkName(attribute),
+            val == null ? Runtime.NONE : SkylarkType.convertToSkylark(val, (Environment) null));
+      }
+      attrObject = StructProvider.STRUCT.create(attributes, ERROR_MESSAGE_FOR_NO_ATTR);
     }
 
     @Override
@@ -87,7 +113,7 @@
         ImmutableList.Builder<BuildOptions> splitBuildOptions = ImmutableList.builder();
 
         ImmutableList<Map<String, Object>> transitions =
-            starlarkDefinedConfigTransition.getChangedSettings(settings);
+            starlarkDefinedConfigTransition.getChangedSettings(settings, attrObject);
         // TODO(juliexxia): Validate that the output values correctly match the output types.
         validateFunctionOutputs(transitions, starlarkDefinedConfigTransition.getOutputs());
 
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkAttributesCollection.java b/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkAttributesCollection.java
index fbf027c..c000612 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkAttributesCollection.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkAttributesCollection.java
@@ -48,6 +48,9 @@
   private final ImmutableMap<Artifact, FilesToRunProvider> executableRunfilesMap;
   private final String ruleClassName;
 
+  static final String ERROR_MESSAGE_FOR_NO_ATTR =
+      "No attribute '%s' in attr. Make sure you declared a rule attribute with this name.";
+
   private SkylarkAttributesCollection(
       SkylarkRuleContext skylarkRuleContext,
       String ruleClassName,
@@ -58,10 +61,7 @@
       ImmutableMap<Artifact, FilesToRunProvider> executableRunfilesMap) {
     this.skylarkRuleContext = skylarkRuleContext;
     this.ruleClassName = ruleClassName;
-    attrObject =
-        StructProvider.STRUCT.create(
-            attrs,
-            "No attribute '%s' in attr. Make sure you declared a rule attribute with this name.");
+    attrObject = StructProvider.STRUCT.create(attrs, ERROR_MESSAGE_FOR_NO_ATTR);
     executableObject =
         StructProvider.STRUCT.create(
             executables,
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/FunctionSplitTransitionProviderTest.java b/src/test/java/com/google/devtools/build/lib/analysis/FunctionSplitTransitionProviderTest.java
index 8e5d867..6aab61a 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/FunctionSplitTransitionProviderTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/FunctionSplitTransitionProviderTest.java
@@ -18,7 +18,9 @@
 import static org.junit.Assert.fail;
 
 import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.ListMultimap;
+import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.analysis.util.BuildViewTestCase;
 import com.google.devtools.build.lib.packages.util.BazelMockAndroidSupport;
 import java.util.List;
@@ -83,6 +85,63 @@
   }
 
   @Test
+  @SuppressWarnings("unchecked")
+  public void testParameterizedTransition() throws Exception {
+    setSkylarkSemanticsOptions("--experimental_starlark_config_transitions=true");
+
+    writeWhitelistFile();
+    scratch.file(
+        "test/skylark/my_rule.bzl",
+        "def _transition_impl(settings, attr):",
+        "  if (attr.big_foot_is_real):",
+        "    return {'t0': {'//command_line_option:cpu': 'x86_64'}}",
+        "  else:",
+        "    return {'t0': {'//command_line_option:cpu': 'armeabi-v7a'}}",
+        "",
+        "my_transition = transition(implementation = _transition_impl, inputs = [],",
+        "  outputs = ['//command_line_option:cpu'])",
+        "def _my_rule_impl(ctx): ",
+        "  return struct(",
+        "    split_attr_dep = ctx.split_attr.dep,",
+        "    attr_dep = ctx.attr.dep)",
+        "my_rule = rule(",
+        "  implementation = _my_rule_impl,",
+        "  attrs = {",
+        "    'big_foot_is_real': attr.bool(default = False),",
+        "    'dep':  attr.label(cfg = my_transition),",
+        "    '_whitelist_function_transition': attr.label(",
+        "        default = '//tools/whitelists/function_transition_whitelist',",
+        "    ),",
+        "  })",
+        "def _simple_impl(ctx):",
+        "  return []",
+        "simple_rule = rule(implementation = _simple_impl)");
+
+    scratch.file(
+        "test/skylark/BUILD",
+        "load('//test/skylark:my_rule.bzl', 'my_rule', 'simple_rule')",
+        "my_rule(name = 'lies', dep = ':dep')",
+        "my_rule(name = 'the-truth', dep = ':dep', big_foot_is_real = True)",
+        "simple_rule(name = 'dep')");
+
+    useConfiguration("--cpu=k8");
+
+    BuildConfiguration liesDepConfiguration =
+        getConfiguration(
+            Iterables.getOnlyElement(
+                (List<ConfiguredTarget>)
+                    getConfiguredTarget("//test/skylark:lies").get("attr_dep")));
+    assertThat(liesDepConfiguration.getCpu()).isEqualTo("armeabi-v7a");
+
+    BuildConfiguration theTruthDepConfiguration =
+        getConfiguration(
+            Iterables.getOnlyElement(
+                (List<ConfiguredTarget>)
+                    getConfiguredTarget("//test/skylark:the-truth").get("attr_dep")));
+    assertThat(theTruthDepConfiguration.getCpu()).isEqualTo("x86_64");
+  }
+
+  @Test
   public void testFunctionSplitTransitionCheckSplitAttrDeps() throws Exception {
     writeBasicTestFiles();
     testSplitTransitionCheckSplitAttrDeps(getConfiguredTarget("//test/skylark:test"));