fix: expand_location binary target inputs

The previous behaviour was that location expansion on native rules would accept binary targets as inputs and resolve to a single location. This was not an option on custom rules. This patch aligns the behaviour of Starlark's ctx.expand_location to the native rules. E.g. a py_binary can be expanded using location/execpath/rootpath instead of previously only locations/execpaths/rootpaths.

For more context, see bazelbuild/rules_python#846. To be able to add the multi-toolchain feature to rules_python, I had to trick `ctx.expand_location` using the `tools` attribute while this is not fixed.

Closes #16381.

PiperOrigin-RevId: 502466606
Change-Id: I3977f2dd479a55308bdda982588697fb04abcedf

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

@@ -1123,19 +1123,34 @@ private static void checkDeprecated(String newApi, String oldApi, StarlarkSemant
   }
 
   /**
-   * Builds a map: Label -> List of files from the given labels
+   * Builds a map: Label -> List of files from the given labels. It first looks into the
+   * files-to-run and then into files. If the label being iterated is an executable, then it will
+   * have the same behaviour as native rules special attributes (data, tools) in which only the
+   * executable file is mapped to the label. This allows executable targets to have the location
+   * expansion work with the singular form of location/execpath/rootpath.
    *
    * @param knownLabels List of known labels
    * @return Immutable map with immutable collections as values
    */
   public static ImmutableMap<Label, ImmutableCollection<Artifact>> makeLabelMap(
       Iterable<TransitiveInfoCollection> knownLabels) {
+
     ImmutableMap.Builder<Label, ImmutableCollection<Artifact>> builder = ImmutableMap.builder();
 
-    for (TransitiveInfoCollection current : knownLabels) {
-      builder.put(
-          AliasProvider.getDependencyLabel(current),
-          current.getProvider(FileProvider.class).getFilesToBuild().toList());
+    for (TransitiveInfoCollection label : knownLabels) {
+      FilesToRunProvider filesToRun = label.getProvider(FilesToRunProvider.class);
+      if (filesToRun != null) {
+        Artifact executableArtifact = filesToRun.getExecutable();
+        builder.put(
+            AliasProvider.getDependencyLabel(label),
+            executableArtifact != null
+                ? ImmutableList.of(executableArtifact)
+                : filesToRun.getFilesToRun().toList());
+      } else {
+        builder.put(
+            AliasProvider.getDependencyLabel(label),
+            label.getProvider(FileProvider.class).getFilesToBuild().toList());
+      }
     }
 
     return builder.buildOrThrow();

src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleImplementationFunctionsTest.java

@@ -42,6 +42,7 @@
 import com.google.devtools.build.lib.analysis.FileProvider;
 import com.google.devtools.build.lib.analysis.FilesToRunProvider;
 import com.google.devtools.build.lib.analysis.Runfiles;
+import com.google.devtools.build.lib.analysis.TransitiveInfoCollection;
 import com.google.devtools.build.lib.analysis.actions.FileWriteAction;
 import com.google.devtools.build.lib.analysis.actions.ParameterFileWriteAction;
 import com.google.devtools.build.lib.analysis.actions.SpawnAction;
@@ -658,6 +659,76 @@ public void testExpandLocationWithShortPaths() throws Exception {
     assertThat(loc).isEqualTo("foo/libjl.jar");
   }
 
+  /**
+   * Asserts that a custom rule has the same behaviour as native rules when expanding executable
+   * target locations.
+   */
+  @Test
+  public void testExpandLocationExecutableTargets() throws Exception {
+    scratch.file(
+        "test/defs.bzl",
+        "def _my_binary_impl(ctx):",
+        "  executable = ctx.actions.declare_file(ctx.attr.name)",
+        "  ctx.actions.write(executable, '', is_executable = True)",
+        "  return [DefaultInfo(executable = executable, files = depset(ctx.files.srcs), runfiles ="
+            + " ctx.runfiles(ctx.files.srcs))]",
+        "my_binary = rule(",
+        "    implementation = _my_binary_impl,",
+        "    attrs = {'srcs': attr.label_list(allow_files=True)},",
+        "    executable = True,",
+        ")",
+        "def _expand_location_env_rule_impl(ctx):",
+        "  env = {}",
+        "  for k, v in ctx.attr.env.items():",
+        "    env[k] = ctx.expand_location(v, targets = ctx.attr.data)",
+        "  env_file = ctx.actions.declare_file('env_file')",
+        "  ctx.actions.write(env_file, str(env))",
+        "  return [DefaultInfo(files = depset([env_file]))]",
+        "expand_location_env_rule = rule(",
+        "    implementation = _expand_location_env_rule_impl,",
+        "    attrs = {",
+        "       'data': attr.label_list(allow_files=True),",
+        "       'env': attr.string_dict(),",
+        "    },",
+        ")");
+
+    scratch.file("test/file1", "");
+    scratch.file("test/file2", "");
+
+    scratch.file(
+        "test/BUILD",
+        "load('//test:defs.bzl', 'expand_location_env_rule', 'my_binary')",
+        "my_binary(name = 'main')",
+        "filegroup(name = 'files', srcs = ['file1', 'file2'])",
+        "expand_location_env_rule(",
+        "  name = 'expand_execpath_env',",
+        "  data = [':main'],",
+        "  env = {'MAIN_EXECPATH': '$(execpath :main)'},",
+        ")",
+        "expand_location_env_rule(",
+        "  name = 'expand_execpaths_env',",
+        "  data = [':files'],",
+        "  env = {'FILES': '$(execpaths :files)'},",
+        ")");
+
+    TransitiveInfoCollection expandExecpathEnv = getConfiguredTarget("//test:expand_execpath_env");
+    Artifact artifact =
+        Iterables.getOnlyElement(
+            expandExecpathEnv.getProvider(FileProvider.class).getFilesToBuild().toList());
+    FileWriteAction action = (FileWriteAction) getGeneratingAction(artifact);
+    assertMatches(
+        "env file contains expanded location of runfile",
+        "\\{\"MAIN_EXECPATH\": \"[\\w-_/]+/test/main\"\\}",
+        action.getFileContents());
+
+    expandExecpathEnv = getConfiguredTarget("//test:expand_execpaths_env");
+    artifact =
+        Iterables.getOnlyElement(
+            expandExecpathEnv.getProvider(FileProvider.class).getFilesToBuild().toList());
+    action = (FileWriteAction) getGeneratingAction(artifact);
+    assertThat(action.getFileContents()).isEqualTo("{\"FILES\": \"test/file1 test/file2\"}");
+  }
+
   /** Regression test to check that expand_location allows ${var} and $$. */
   @Test
   public void testExpandLocationWithDollarSignsAndCurlys() throws Exception {