Add runfiles.merge_all() for merging a sequence of runfiles objects.

We've seen some pathological cases where runfiles.merge() run in a loop results
in an extraordinarily deep and inefficient NestedSet representing the
runfiles's artifacts - and attempting to walk such a NestedSet could lead to a
stack overflow.

Merging a list of runfiles at once produces a NestedSet which is much shallower
(by two orders of magnitude in some cases).

In a followup to this change, we intend to introduce checks to make sure
NestedSet depth does not become excessive.

RELNOTES: Add runfiles.merge_all() for merging a sequence of runfiles objects.
PiperOrigin-RevId: 364399853

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

@@ -48,7 +48,9 @@
 import java.util.Objects;
 import java.util.Set;
 import javax.annotation.Nullable;
+import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Printer;
+import net.starlark.java.eval.Sequence;
 import net.starlark.java.syntax.Location;
 
 /**
@@ -1076,6 +1078,43 @@ public Runfiles merge(RunfilesApi other) {
     return new Runfiles.Builder(suffix, false).merge(this).merge(o).build();
   }
 
+  @Override
+  public Runfiles mergeAll(Sequence<?> sequence) throws EvalException {
+    // The delayed initialization of the Builder is not just a memory / performance optimization.
+    // The Builder requires a valid suffix, but the {@code Runfiles.EMPTY} singleton has an invalid
+    // one, which must not be used to construct a Runfiles.Builder.
+    Builder builder = null;
+    // When merging exactly one non-empty Runfiles object, we want to return that object and avoid a
+    // Builder. This is a memory optimization and provides identical behavior for `x.merge_all([y])`
+    // and `x.merge(y)` in Starlark.
+    Runfiles uniqueNonEmptyMergee = null;
+    if (!this.isEmpty()) {
+      builder = new Builder(suffix, false).merge(this);
+      uniqueNonEmptyMergee = this;
+    }
+
+    Sequence<Runfiles> runfilesSequence = Sequence.cast(sequence, Runfiles.class, "param");
+    for (Runfiles runfiles : runfilesSequence) {
+      if (!runfiles.isEmpty()) {
+        if (builder == null) {
+          builder = new Builder(runfiles.suffix, /* legacyExternalRunfiles = */ false);
+          uniqueNonEmptyMergee = runfiles;
+        } else {
+          uniqueNonEmptyMergee = null;
+        }
+        builder.merge(runfiles);
+      }
+    }
+
+    if (uniqueNonEmptyMergee != null) {
+      return uniqueNonEmptyMergee;
+    } else if (builder != null) {
+      return builder.build();
+    } else {
+      return EMPTY;
+    }
+  }
+
   /**
    * Fingerprint this {@link Runfiles} tree.
    */

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

@@ -17,8 +17,11 @@
 import com.google.devtools.build.docgen.annot.DocCategory;
 import com.google.devtools.build.lib.collect.nestedset.Depset;
 import net.starlark.java.annot.Param;
+import net.starlark.java.annot.ParamType;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.annot.StarlarkMethod;
+import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.Sequence;
 import net.starlark.java.eval.StarlarkValue;
 
 /** An interface for a set of runfiles. */
@@ -55,7 +58,11 @@ public interface RunfilesApi extends StarlarkValue {
       name = "merge",
       doc =
           "Returns a new runfiles object that includes all the contents of this one and the "
-              + "argument.",
+              + "argument."
+              + "<p><i>Note:</i> When you have many runfiles objects to merge, use <a "
+              + "href='#merge_all'><code>merge_all()</code></a> rather than calling <code>merge"
+              + "</code> in a loop. This avoids constructing deep depset structures which can "
+              + "cause build failures.",
       parameters = {
         @Param(
             name = "other",
@@ -64,4 +71,21 @@ public interface RunfilesApi extends StarlarkValue {
             doc = "The runfiles object to merge into this."),
       })
   RunfilesApi merge(RunfilesApi other);
+
+  @StarlarkMethod(
+      name = "merge_all",
+      doc =
+          "Returns a new runfiles object that includes all the contents of this one and of the "
+              + "runfiles objects in the argument.",
+      parameters = {
+        @Param(
+            name = "other",
+            allowedTypes = {
+              @ParamType(type = Sequence.class, generic1 = RunfilesApi.class),
+            },
+            positional = true,
+            named = false,
+            doc = "The sequence of runfiles objects to merge into this."),
+      })
+  RunfilesApi mergeAll(Sequence<?> sequence) throws EvalException;
 }

src/test/java/com/google/devtools/build/lib/analysis/RunfilesTest.java

@@ -36,6 +36,7 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 import javax.annotation.Nullable;
+import net.starlark.java.eval.StarlarkList;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -480,6 +481,24 @@ public void testMergeWithSymlinks() {
     assertThat(runfilesC.getSymlinksAsMap(null).get(sympathB)).isEqualTo(artifactB);
   }
 
+  @Test
+  public void mergeAll_symlinks() throws Exception {
+    ArtifactRoot root = ArtifactRoot.asSourceRoot(Root.fromPath(scratch.resolve("/workspace")));
+    Artifact artifactA = ActionsTestUtil.createArtifact(root, "a/target");
+    Artifact artifactB = ActionsTestUtil.createArtifact(root, "b/target");
+    Artifact artifactC = ActionsTestUtil.createArtifact(root, "c/target");
+    PathFragment sympathA = PathFragment.create("a/symlink");
+    PathFragment sympathB = PathFragment.create("b/symlink");
+    PathFragment sympathC = PathFragment.create("c/symlink");
+    Runfiles runfilesA = new Runfiles.Builder("TESTING").addSymlink(sympathA, artifactA).build();
+    Runfiles runfilesB = new Runfiles.Builder("TESTING").addSymlink(sympathB, artifactB).build();
+    Runfiles runfilesC = new Runfiles.Builder("TESTING").addSymlink(sympathC, artifactC).build();
+
+    Runfiles runfilesMerged = runfilesA.mergeAll(StarlarkList.immutableOf(runfilesB, runfilesC));
+    assertThat(runfilesMerged.getSymlinksAsMap(null))
+        .containsExactly(sympathA, artifactA, sympathB, artifactB, sympathC, artifactC);
+  }
+
   @Test
   public void testMergeEmptyWithNonEmpty() {
     ArtifactRoot root = ArtifactRoot.asSourceRoot(Root.fromPath(scratch.resolve("/workspace")));
@@ -489,6 +508,31 @@ public void testMergeEmptyWithNonEmpty() {
     assertThat(runfilesB.merge(Runfiles.EMPTY)).isSameInstanceAs(runfilesB);
   }
 
+  @Test
+  public void mergeAll_emptyWithNonEmpty() throws Exception {
+    ArtifactRoot root = ArtifactRoot.asSourceRoot(Root.fromPath(scratch.resolve("/workspace")));
+    Artifact artifact = ActionsTestUtil.createArtifact(root, "target");
+    Runfiles nonEmpty = new Runfiles.Builder("TESTING").addArtifact(artifact).build();
+
+    assertThat(Runfiles.EMPTY.mergeAll(StarlarkList.immutableOf(nonEmpty)))
+        .isSameInstanceAs(nonEmpty);
+    assertThat(
+            Runfiles.EMPTY.mergeAll(
+                StarlarkList.immutableOf(Runfiles.EMPTY, nonEmpty, Runfiles.EMPTY)))
+        .isSameInstanceAs(nonEmpty);
+    assertThat(nonEmpty.mergeAll(StarlarkList.immutableOf(Runfiles.EMPTY, Runfiles.EMPTY)))
+        .isSameInstanceAs(nonEmpty);
+    assertThat(nonEmpty.mergeAll(StarlarkList.immutableOf())).isSameInstanceAs(nonEmpty);
+  }
+
+  @Test
+  public void mergeAll_emptyWithEmpty() throws Exception {
+    assertThat(Runfiles.EMPTY.mergeAll(StarlarkList.immutableOf()))
+        .isSameInstanceAs(Runfiles.EMPTY);
+    assertThat(Runfiles.EMPTY.mergeAll(StarlarkList.immutableOf(Runfiles.EMPTY, Runfiles.EMPTY)))
+        .isSameInstanceAs(Runfiles.EMPTY);
+  }
+
   @Test
   public void testOnlyExtraMiddlemenNotConsideredEmpty() {
     ArtifactRoot root =

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

@@ -1825,6 +1825,95 @@ public void testAccessingRunfilesRootSymlinks() throws Exception {
     assertThat(rootSymlinkFilenamesList).containsExactly("test/a.py").inOrder();
   }
 
+  @Test
+  public void runfiles_merge() throws Exception {
+    scratch.file("test/a.py");
+    scratch.file("test/b.py");
+    scratch.file("test/other.py");
+    scratch.file(
+        "test/rule.bzl",
+        "def symlink_merge_impl(ctx):",
+        "  runfiles = ctx.runfiles(symlinks = {",
+        "    'symlink_' + ctx.file.symlink.short_path: ctx.file.symlink",
+        "  })",
+        "  if ctx.attr.dep:",
+        "    runfiles = runfiles.merge(ctx.attr.dep[DefaultInfo].default_runfiles)",
+        "  return DefaultInfo(",
+        "    runfiles = runfiles",
+        "  )",
+        "symlink_merge_rule = rule(",
+        "  implementation = symlink_merge_impl,",
+        "  attrs = {",
+        "    'symlink': attr.label(allow_single_file=True),",
+        "    'dep': attr.label(),",
+        "  },",
+        ")");
+    scratch.file(
+        "test/BUILD",
+        "load('//test:rule.bzl', 'symlink_merge_rule')",
+        "symlink_merge_rule(name = 'lib_a', symlink = ':a.py', dep = 'lib_b')",
+        "symlink_merge_rule(name = 'lib_b', symlink = ':b.py')",
+        "sh_binary(",
+        "  name = 'test',",
+        "  srcs = ['test/other.py'],",
+        "  data = [':lib_a'],",
+        ")");
+    setRuleContext(createRuleContext("//test:test"));
+    Object symlinkPaths =
+        ev.eval("[s.path for s in ruleContext.attr.data[0].data_runfiles.symlinks.to_list()]");
+    assertThat(symlinkPaths).isInstanceOf(Sequence.class);
+    Sequence<?> symlinkPathsList = (Sequence) symlinkPaths;
+    assertThat(symlinkPathsList)
+        .containsExactly("symlink_test/a.py", "symlink_test/b.py")
+        .inOrder();
+  }
+
+  @Test
+  public void runfiles_mergeAll() throws Exception {
+    scratch.file("test/a.py");
+    scratch.file("test/b.py");
+    scratch.file("test/c.py");
+    scratch.file("test/other.py");
+    scratch.file(
+        "test/rule.bzl",
+        "def symlink_merge_all_impl(ctx):",
+        "  runfiles = ctx.runfiles(symlinks = {",
+        "    'symlink_' + ctx.file.symlink.short_path: ctx.file.symlink",
+        "  })",
+        "  if ctx.attr.deps:",
+        "    runfiles = runfiles.merge_all([dep[DefaultInfo].default_runfiles",
+        "                                   for dep in ctx.attr.deps])",
+        "  return DefaultInfo(",
+        "    runfiles = runfiles",
+        "  )",
+        "symlink_merge_all_rule = rule(",
+        "  implementation = symlink_merge_all_impl,",
+        "  attrs = {",
+        "    'symlink': attr.label(allow_single_file=True),",
+        "    'deps': attr.label_list(),",
+        "  },",
+        ")");
+    scratch.file(
+        "test/BUILD",
+        "load('//test:rule.bzl', 'symlink_merge_all_rule')",
+        "symlink_merge_all_rule(name = 'lib_a', symlink = ':a.py', deps = [':lib_b', ':lib_c'])",
+        "symlink_merge_all_rule(name = 'lib_b', symlink = ':b.py')",
+        "symlink_merge_all_rule(name = 'lib_c', symlink = ':c.py')",
+        "sh_binary(",
+        "  name = 'test',",
+        "  srcs = ['test/other.py'],",
+        "  data = [':lib_a'],",
+        ")");
+    setRuleContext(createRuleContext("//test:test"));
+    Object symlinkPaths =
+        ev.eval("[s.path for s in ruleContext.attr.data[0].data_runfiles.symlinks.to_list()]");
+    assertThat(symlinkPaths).isInstanceOf(Sequence.class);
+    Sequence<?> symlinkPathsList = Sequence.cast(symlinkPaths, String.class, "symlinkPaths");
+    assertThat(symlinkPathsList)
+        .containsExactly("symlink_test/a.py", "symlink_test/b.py", "symlink_test/c.py")
+        .inOrder();
+  }
+
   @Test
   public void testExternalShortPath() throws Exception {
     scratch.file("/bar/WORKSPACE");