fix: expand_location binary target inputs

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

@@ -1133,19 +1133,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 runfiles 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

@@ -50,6 +50,7 @@
 import com.google.devtools.build.lib.analysis.configuredtargets.RuleConfiguredTarget;
 import com.google.devtools.build.lib.analysis.starlark.Args;
 import com.google.devtools.build.lib.analysis.starlark.StarlarkRuleContext;
+import com.google.devtools.build.lib.analysis.TransitiveInfoCollection;
 import com.google.devtools.build.lib.analysis.util.BuildViewTestCase;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
@@ -658,6 +659,78 @@ 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 {