Add functionality for AEGs when toolchain parameter is a Label

ctx.actions.{run,run_shell} were missing functionality for toolchain parameter when toolchain is an instance of Label. Moreover, I've added appropriate parsing of toolchain when it's a String (relative to the package). I've also added tests for these changes. PiperOrigin-RevId: 516802154 Change-Id: I50b219a88ad6c78b8f18f42eee13ab66c6bb4cab
bazelbuild · Mar 15, 2023 · f556833 · f556833
1 parent e93a8a2
commit f556833
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 8 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkActionFactory.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkActionFactory.java
@@ -53,6 +53,7 @@
 import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.analysis.platform.PlatformInfo;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
 import com.google.devtools.build.lib.collect.nestedset.Depset;
 import com.google.devtools.build.lib.collect.nestedset.Depset.TypeException;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
@@ -475,9 +476,9 @@ private void verifyExecGroupExists(String execGroup, RuleContext ctx) throws Eva
     }
   }
 
-  private void verifyAutomaticExecGroupExists(String execGroup, RuleContext ctx)
+  private void verifyAutomaticExecGroupExists(String execGroup, RuleContext ruleContext)
       throws EvalException {
-    if (!ctx.hasToolchainContext(execGroup)) {
+    if (!ruleContext.hasToolchainContext(execGroup)) {
       throw Starlark.errorf("Action declared for non-existent toolchain '%s'.", execGroup);
     }
   }
@@ -736,14 +737,26 @@ private void registerStarlarkAction(
       }
     }
 
+    Label toolchainLabel = null;
+    if (toolchainUnchecked instanceof Label) {
+      toolchainLabel = (Label) toolchainUnchecked;
+    } else if (toolchainUnchecked instanceof String) {
+      try {
+        toolchainLabel =
+            Label.parseWithPackageContext(
+                (String) toolchainUnchecked, ruleContext.getPackageContext());
+      } catch (LabelSyntaxException e) {
+        throw Starlark.errorf("%s", e.getMessage());
+      }
+    }
+
     if (execGroupUnchecked != Starlark.NONE) {
       String execGroup = (String) execGroupUnchecked;
       verifyExecGroupExists(execGroup, ruleContext);
       checkValidGroupName(execGroup);
 
       // If toolchain and exec_groups are both defined, verify they are compatible.
-      if (useAutoExecGroups && toolchainUnchecked instanceof String) {
-        Label toolchainLabel = Label.parseCanonicalUnchecked((String) toolchainUnchecked);
+      if (useAutoExecGroups && toolchainLabel != null) {
         if (ruleContext.getExecGroups().getExecGroup(execGroup).toolchainTypes().stream()
             .map(ToolchainTypeRequirement::toolchainType)
             .noneMatch(toolchainLabel::equals)) {
@@ -755,10 +768,9 @@ private void registerStarlarkAction(
       }
 
       builder.setExecGroup(execGroup);
-    } else if (useAutoExecGroups && toolchainUnchecked instanceof String) {
-      String toolchain = (String) toolchainUnchecked;
-      verifyAutomaticExecGroupExists(toolchain, ruleContext);
-      builder.setExecGroup(toolchain);
+    } else if (useAutoExecGroups && toolchainLabel != null) {
+      verifyAutomaticExecGroupExists(toolchainLabel.toString(), ruleContext);
+      builder.setExecGroup(toolchainLabel.toString());
     } else {
       builder.setExecGroup(ExecGroup.DEFAULT_EXEC_GROUP_NAME);
     }

diff --git a/src/test/java/com/google/devtools/build/lib/analysis/AutoExecGroupsTest.java b/src/test/java/com/google/devtools/build/lib/analysis/AutoExecGroupsTest.java
@@ -659,6 +659,73 @@ public void defaultExecGroupHasBasicExecutionPlatform(String action) throws Exce
         .isEqualTo(Label.parseCanonical("//platforms:platform_1"));
   }
 
+  @Test
+  @TestParameters({
+    "{action: ctx.actions.run}",
+    "{action: ctx.actions.run_shell}",
+  })
+  public void toolchainParameterAsLabel_correctParsingOfToolchain(String action) throws Exception {
+    createCustomRule(
+        /* action= */ action,
+        /* actionParameters= */ "toolchain = Label('@//rule:toolchain_type_1'),",
+        /* extraAttributes= */ "",
+        /* toolchains= */ "['//rule:toolchain_type_1', '//rule:toolchain_type_2']",
+        /* execGroups= */ "",
+        /* execCompatibleWith= */ "");
+    useConfiguration("--incompatible_auto_exec_groups");
+
+    ConfiguredTarget target = getConfiguredTarget("//test:custom_rule_name");
+    Action generatedAction = getGeneratingAction(target, "test/custom_rule_name_dummy_output.jar");
+
+    assertThat(generatedAction.getOwner().getExecutionPlatform().label())
+        .isEqualTo(Label.parseCanonical("//platforms:platform_1"));
+  }
+
+  @Test
+  @TestParameters({
+    "{action: ctx.actions.run}",
+    "{action: ctx.actions.run_shell}",
+  })
+  public void toolchainParameterAsString_correctParsingOfToolchain(String action) throws Exception {
+    createCustomRule(
+        /* action= */ action,
+        /* actionParameters= */ "toolchain = '@//rule:toolchain_type_1',",
+        /* extraAttributes= */ "",
+        /* toolchains= */ "['//rule:toolchain_type_1', '//rule:toolchain_type_2']",
+        /* execGroups= */ "",
+        /* execCompatibleWith= */ "");
+    useConfiguration("--incompatible_auto_exec_groups");
+
+    ConfiguredTarget target = getConfiguredTarget("//test:custom_rule_name");
+    Action generatedAction = getGeneratingAction(target, "test/custom_rule_name_dummy_output.jar");
+
+    assertThat(generatedAction.getOwner().getExecutionPlatform().label())
+        .isEqualTo(Label.parseCanonical("//platforms:platform_1"));
+  }
+
+  @Test
+  @TestParameters({
+    "{action: ctx.actions.run}",
+    "{action: ctx.actions.run_shell}",
+  })
+  public void toolchainParameterAsString_syntaxErrorInParsingOfToolchain(String action)
+      throws Exception {
+    createCustomRule(
+        /* action= */ action,
+        /* actionParameters= */ "toolchain = 'rule:toolchain_type_1',",
+        /* extraAttributes= */ "",
+        /* toolchains= */ "['//rule:toolchain_type_1', '//rule:toolchain_type_2']",
+        /* execGroups= */ "",
+        /* execCompatibleWith= */ "");
+    useConfiguration("--incompatible_auto_exec_groups");
+
+    reporter.removeHandler(failFastHandler);
+    getConfiguredTarget("//test:custom_rule_name");
+
+    assertContainsEvent(
+        "invalid label 'rule:toolchain_type_1': absolute label must begin with '@'" + " or '//'");
+  }
+
   @Test
   @TestParameters({
     "{action: ctx.actions.run}",