Extend StarlarkRuleContextApi.runfiles to accept depset for symlinks and root symlinks

This is useful for efficiently creating a new instance of Runfiles from an existing one since the RunfilesApi returns a depset of symlinks

PiperOrigin-RevId: 408853831

Author: hvadehra

src/main/java/com/google/devtools/build/lib/analysis/Runfiles.java

@@ -106,10 +106,10 @@ public Iterable<PathFragment> getExtraPaths(Set<PathFragment> manifestPaths) {
   // first one prevents the last one from appearing on the iterator.
   //
   // The lack of a .equals() method prevents this by making the first entry in the above case not
-  // equals to the third one if they are not the same instance (which they almost never are)
+  // equal to the third one if they are not the same instance (which they almost never are).
   //
   // Goodnight, prince(ss)?, and sweet dreams.
-  private static final class SymlinkEntry implements SymlinkEntryApi {
+  public static final class SymlinkEntry implements SymlinkEntryApi {
 
     static final Depset.ElementType TYPE = Depset.ElementType.of(SymlinkEntry.class);
 

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

@@ -40,6 +40,7 @@
 import com.google.devtools.build.lib.analysis.LocationExpander;
 import com.google.devtools.build.lib.analysis.RuleContext;
 import com.google.devtools.build.lib.analysis.Runfiles;
+import com.google.devtools.build.lib.analysis.Runfiles.SymlinkEntry;
 import com.google.devtools.build.lib.analysis.RunfilesProvider;
 import com.google.devtools.build.lib.analysis.ShToolchain;
 import com.google.devtools.build.lib.analysis.TransitiveInfoCollection;
@@ -54,6 +55,7 @@
 import com.google.devtools.build.lib.analysis.test.InstrumentedFilesInfo;
 import com.google.devtools.build.lib.cmdline.Label;
 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;
 import com.google.devtools.build.lib.collect.nestedset.Order;
 import com.google.devtools.build.lib.packages.Aspect;
@@ -912,9 +914,9 @@ public Runfiles runfiles(
       Object transitiveFiles,
       Boolean collectData,
       Boolean collectDefault,
-      Dict<?, ?> symlinks,
-      Dict<?, ?> rootSymlinks)
-      throws EvalException {
+      Object symlinks,
+      Object rootSymlinks)
+      throws EvalException, TypeException {
     checkMutable("runfiles");
     Runfiles.Builder builder =
         new Runfiles.Builder(
@@ -941,15 +943,19 @@ public Runfiles runfiles(
       }
       builder.addTransitiveArtifacts(transitiveArtifacts);
     }
-    if (!symlinks.isEmpty()) {
+    if (isDepset(symlinks)) {
+      builder.addSymlinks(((Depset) symlinks).getSet(SymlinkEntry.class));
+    } else if (isNonEmptyDict(symlinks)) {
       // If Starlark code directly manipulates symlinks, activate more stringent validity checking.
       checkConflicts = true;
       for (Map.Entry<String, Artifact> entry :
           Dict.cast(symlinks, String.class, Artifact.class, "symlinks").entrySet()) {
         builder.addSymlink(PathFragment.create(entry.getKey()), entry.getValue());
       }
     }
-    if (!rootSymlinks.isEmpty()) {
+    if (isDepset(rootSymlinks)) {
+      builder.addRootSymlinks(((Depset) rootSymlinks).getSet(SymlinkEntry.class));
+    } else if (isNonEmptyDict(rootSymlinks)) {
       checkConflicts = true;
       for (Map.Entry<String, Artifact> entry :
           Dict.cast(rootSymlinks, String.class, Artifact.class, "root_symlinks").entrySet()) {
@@ -963,6 +969,14 @@ public Runfiles runfiles(
     return runfiles;
   }
 
+  private static boolean isNonEmptyDict(Object o) {
+    return o instanceof Dict && !((Dict<?, ?>) o).isEmpty();
+  }
+
+  private static boolean isDepset(Object o) {
+    return o instanceof Depset;
+  }
+
   @Override
   public Tuple resolveCommand(
       String command,

src/main/java/com/google/devtools/build/lib/starlarkbuildapi/StarlarkRuleContextApi.java

@@ -18,6 +18,7 @@
 import com.google.devtools.build.docgen.annot.DocCategory;
 import com.google.devtools.build.lib.cmdline.Label;
 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.starlarkbuildapi.core.ProviderApi;
 import com.google.devtools.build.lib.starlarkbuildapi.core.StructApi;
 import com.google.devtools.build.lib.starlarkbuildapi.core.TransitiveInfoCollectionApi;
@@ -564,27 +565,36 @@ String expandLocation(String input, Sequence<?> targets, StarlarkThread thread)
             name = "symlinks",
             defaultValue = "{}",
             named = true,
+            allowedTypes = {
+              @ParamType(type = Dict.class),
+              @ParamType(type = Depset.class, generic1 = SymlinkEntryApi.class)
+            },
             doc =
-                "The map of symlinks to be added to the runfiles, prefixed by workspace name. See "
+                "Either a SymlinkEntry depset or the map of symlinks, prefixed by workspace name, "
+                    + "to be added to the runfiles. See "
                     + "<a href=\"../rules.$DOC_EXT#runfiles-symlinks\">Runfiles symlinks</a> in "
                     + "the rules guide."),
         @Param(
             name = "root_symlinks",
             defaultValue = "{}",
             named = true,
+            allowedTypes = {
+              @ParamType(type = Dict.class),
+              @ParamType(type = Depset.class, generic1 = SymlinkEntryApi.class)
+            },
             doc =
-                "The map of symlinks to be added to the runfiles. See "
-                    + "<a href=\"../rules.$DOC_EXT#runfiles-symlinks\">Runfiles symlinks</a> in "
-                    + "the rules guide.")
+                "Either a SymlinkEntry depset or a map of symlinks to be added to the runfiles. "
+                    + "See <a href=\"../rules.$DOC_EXT#runfiles-symlinks\">Runfiles symlinks</a> "
+                    + "in the rules guide.")
       })
   RunfilesApi runfiles(
       Sequence<?> files,
       Object transitiveFiles,
       Boolean collectData,
       Boolean collectDefault,
-      Dict<?, ?> symlinks,
-      Dict<?, ?> rootSymlinks)
-      throws EvalException;
+      Object symlinks,
+      Object rootSymlinks)
+      throws EvalException, TypeException;
 
   @StarlarkMethod(
       name = "resolve_command",

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

@@ -1688,6 +1688,78 @@ public void testForwardingDefaultInfoRetainsDataRunfiles() throws Exception {
     assertThat(forwardedRunfiles.get(0).getPath().getBaseName()).isEqualTo("i_am_a_runfile");
   }
 
+  @Test
+  public void testAccessingRunfilesSymlinksAsDepsets() throws Exception {
+    // Arrange
+    scratch.file("test/a.py");
+    scratch.file("test/b.py");
+    scratch.file(
+        "test/rule.bzl",
+        "def symlink_impl(ctx):",
+        "  symlinks = {",
+        "    'symlink_' + f.short_path: f",
+        "    for f in ctx.files.symlink",
+        "  }",
+        "  root_symlinks = {",
+        "    'root_symlink_' + f.short_path: f",
+        "    for f in ctx.files.symlink",
+        "  }",
+        "  runfiles_from_dict = ctx.runfiles(",
+        "    symlinks=symlinks,",
+        "    root_symlinks=root_symlinks,",
+        "  )",
+        "  runfiles_from_depset = ctx.runfiles(",
+        "    symlinks = runfiles_from_dict.symlinks,",
+        "    root_symlinks = runfiles_from_dict.root_symlinks,",
+        "  )",
+        "   ",
+        "  return DefaultInfo(runfiles = runfiles_from_depset,)",
+        "symlink_rule = rule(",
+        "  implementation = symlink_impl,",
+        "  attrs = {",
+        "    'symlink': attr.label(allow_files=True),",
+        "  },",
+        ")");
+    scratch.file(
+        "test/BUILD",
+        "load('//test:rule.bzl', 'symlink_rule')",
+        "symlink_rule(name = 'lib_with_symlink', symlink = ':a.py')",
+        "sh_binary(",
+        "  name = 'test_with_symlink',",
+        "  srcs = ['test/b.py'],",
+        "  data = [':lib_with_symlink'],",
+        ")");
+    setRuleContext(createRuleContext("//test:test_with_symlink"));
+
+    // Act
+    Object symlinkPaths =
+        ev.eval("[s.path for s in ruleContext.attr.data[0].data_runfiles.symlinks.to_list()]");
+    Object rootSymlinkPaths =
+        ev.eval("[s.path for s in ruleContext.attr.data[0].data_runfiles.root_symlinks.to_list()]");
+
+    // Assert
+    assertThat(symlinkPaths).isInstanceOf(Sequence.class);
+    Sequence<?> symlinkPathsList = (Sequence) symlinkPaths;
+    assertThat(symlinkPathsList).containsExactly("symlink_test/a.py").inOrder();
+    Object symlinkFilenames =
+        ev.eval(
+            "[s.target_file.short_path for s in"
+                + " ruleContext.attr.data[0].data_runfiles.symlinks.to_list()]");
+    assertThat(symlinkFilenames).isInstanceOf(Sequence.class);
+    Sequence<?> symlinkFilenamesList = (Sequence) symlinkFilenames;
+    assertThat(symlinkFilenamesList).containsExactly("test/a.py").inOrder();
+    assertThat(rootSymlinkPaths).isInstanceOf(Sequence.class);
+    Sequence<?> rootSymlinkPathsList = (Sequence) rootSymlinkPaths;
+    assertThat(rootSymlinkPathsList).containsExactly("root_symlink_test/a.py").inOrder();
+    Object rootSymlinkFilenames =
+        ev.eval(
+            "[s.target_file.short_path for s in"
+                + " ruleContext.attr.data[0].data_runfiles.root_symlinks.to_list()]");
+    assertThat(rootSymlinkFilenames).isInstanceOf(Sequence.class);
+    Sequence<?> rootSymlinkFilenamesList = (Sequence) rootSymlinkFilenames;
+    assertThat(rootSymlinkFilenamesList).containsExactly("test/a.py").inOrder();
+  }
+
   @Test
   public void runfiles_merge() throws Exception {
     scratch.file("test/a.py");