From da23370dcdf6ea19545002fb86bd5d3e6519cdf6 Mon Sep 17 00:00:00 2001
From: Googler <shahan@google.com>
Date: Thu, 13 Jul 2023 12:03:03 -0700
Subject: [PATCH] Fix crash for top level aspects on targets with
 non-idempotent rule transitions.

PiperOrigin-RevId: 547876237
Change-Id: I236e26e8403cc9d02c86695284abe4103297a876
---
 .../ToplevelStarlarkAspectFunction.java       |  2 +-
 .../lib/starlark/StarlarkIntegrationTest.java | 68 +++++++++++++++++++
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ToplevelStarlarkAspectFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ToplevelStarlarkAspectFunction.java
index ec613bf8cc6a6d..f66b8e239144cc 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ToplevelStarlarkAspectFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ToplevelStarlarkAspectFunction.java
@@ -75,7 +75,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
         baseConfiguredTargetValue.getConfiguredTarget().getConfigurationKey();
     if (!Objects.equals(realConfiguration, baseConfiguredTargetKey.getConfigurationKey())) {
       baseConfiguredTargetKey =
-          baseConfiguredTargetKey.toBuilder().setConfigurationKey(realConfiguration).build();
+          ConfiguredTargetKey.fromConfiguredTarget(baseConfiguredTargetValue.getConfiguredTarget());
     }
 
     Collection<AspectKey> aspectsKeys =
diff --git a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkIntegrationTest.java b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkIntegrationTest.java
index d912579ab53711..17f291f0cca955 100644
--- a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkIntegrationTest.java
+++ b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkIntegrationTest.java
@@ -13,6 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.lib.starlark;
 
+import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static com.google.devtools.build.lib.analysis.OutputGroupInfo.INTERNAL_SUFFIX;
@@ -22,8 +23,10 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.eventbus.EventBus;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.util.ActionsTestUtil;
+import com.google.devtools.build.lib.analysis.AnalysisResult;
 import com.google.devtools.build.lib.analysis.ConfiguredRuleClassProvider;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.FileProvider;
@@ -53,6 +56,7 @@
 import com.google.devtools.build.lib.packages.StructImpl;
 import com.google.devtools.build.lib.packages.Type;
 import com.google.devtools.build.lib.rules.objc.ObjcProvider;
+import com.google.devtools.build.lib.skyframe.AspectKeyCreator.AspectKey;
 import com.google.devtools.build.lib.skyframe.ConfiguredTargetAndData;
 import com.google.devtools.build.lib.testutil.TestRuleClassProvider;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
@@ -3924,4 +3928,68 @@ public void identicalPrintStatementsOnSameLineNotDeduplicated_ruleImplementation
     update("//foo:all", /*loadingPhaseThreads=*/ 1, /*doAnalysis=*/ true);
     assertContainsEventWithFrequency("this is a print statement", 2);
   }
+
+  @Test
+  public void topLevelAspectOnTargetWithNonIdempotentRuleTransition() throws Exception {
+    scratch.file(
+        "test/aspect.bzl",
+        "def _impl(target, ctx):",
+        "   print('This aspect does nothing')",
+        "   return struct()",
+        "MyAspect = aspect(implementation=_impl)");
+
+    scratch.overwriteFile(
+        "tools/allowlists/function_transition_allowlist/BUILD",
+        "package_group(",
+        "    name = 'function_transition_allowlist',",
+        "    packages = [",
+        "        '//test/...',",
+        "    ],",
+        ")");
+
+    scratch.file(
+        "test/rules.bzl",
+        "def _transition_impl(settings, attr):",
+        "  print('printing from transition impl', settings['//command_line_option:foo'])",
+        "  return {'//command_line_option:foo': " + "settings['//command_line_option:foo']+'meow'}",
+        "my_transition = transition(",
+        "  implementation = _transition_impl,",
+        "  inputs = ['//command_line_option:foo'],",
+        "  outputs = ['//command_line_option:foo'],",
+        ")",
+        "def _rule_impl(ctx):",
+        "  return []",
+        "my_rule = rule(",
+        "  implementation = _rule_impl,",
+        "  cfg = my_transition,",
+        "  attrs = {",
+        "    'dep': attr.label(cfg = my_transition),",
+        "    '_allowlist_function_transition': attr.label(",
+        "        default = '//tools/allowlists/function_transition_allowlist',",
+        "    ),",
+        "  }",
+        ")");
+
+    scratch.file(
+        "test/BUILD",
+        "load('//test:rules.bzl', 'my_rule')",
+        "my_rule(name = 'test', dep = ':dep')",
+        "my_rule(name = 'dep')");
+
+    useConfiguration("--foo=meow");
+
+    AnalysisResult analysisResult =
+        update(
+            ImmutableList.of("//test:test"),
+            ImmutableList.of("test/aspect.bzl%MyAspect"),
+            /* keepGoing= */ true,
+            /* loadingPhaseThreads= */ 1,
+            /* doAnalysis= */ true,
+            new EventBus());
+    assertThat(getOnlyElement(analysisResult.getTargetsToBuild()).getLabel().toString())
+        .isEqualTo("//test:test");
+    AspectKey aspectKey = getOnlyElement(analysisResult.getAspectsMap().keySet());
+    assertThat(aspectKey.getAspectClass().getName()).isEqualTo("//test:aspect.bzl%MyAspect");
+    assertThat(aspectKey.getLabel().toString()).isEqualTo("//test:test");
+  }
 }