From 64297d6b861098da36205860d0fe75c46be7e421 Mon Sep 17 00:00:00 2001
From: Fabian Meumertzheim <fabian@meumertzhe.im>
Date: Thu, 31 Aug 2023 16:40:10 +0200
Subject: [PATCH] Also apply `NestedSet` optimizations to `Depset`

If the construction of a `NestedSet` involves only a single non-empty
transitive set, this set may be returned directly as an optimization.
This commit checks for this case when constructing a `Depset` and reuses
the corresponding `Depset` if possible.
---
 .../build/lib/collect/nestedset/Depset.java      | 16 +++++++++++++++-
 .../build/lib/collect/nestedset/DepsetTest.java  |  8 ++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/Depset.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/Depset.java
index 13dd5cfa85ceb7..33fe5dba93e382 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/Depset.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/Depset.java
@@ -21,6 +21,7 @@
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
 import java.util.List;
+import java.util.Optional;
 import javax.annotation.Nullable;
 import net.starlark.java.annot.Param;
 import net.starlark.java.annot.ParamType;
@@ -365,6 +366,7 @@ public static Depset fromDirectAndTransitive(
     builder.addAll(direct);
 
     // Add transitive sets, checking that type is equal to elements already added.
+    long nonEmptyTransitives = 0;
     for (Depset x : transitive) {
       if (!x.isEmpty()) {
         type = checkType(type, x.elemClass);
@@ -374,10 +376,22 @@ public static Depset fromDirectAndTransitive(
               order.getStarlarkName(), x.getOrder().getStarlarkName());
         }
         builder.addTransitive(x.getSet());
+        nonEmptyTransitives++;
       }
     }
 
-    return new Depset(type, builder.build());
+    NestedSet<Object> set = builder.build();
+    // If the nested set was optimized to the only non-empty transitive element, reuse the
+    // corresponding depset.
+    if (nonEmptyTransitives == 1) {
+      Optional<Depset> unwrappedTransitive =
+          transitive.stream().filter(x -> x.getSet() == set).findFirst();
+      if (unwrappedTransitive.isPresent()) {
+        return unwrappedTransitive.get();
+      }
+    }
+
+    return new Depset(type, set);
   }
 
   /** An exception thrown when validation fails on the type of elements of a nested set. */
diff --git a/src/test/java/com/google/devtools/build/lib/collect/nestedset/DepsetTest.java b/src/test/java/com/google/devtools/build/lib/collect/nestedset/DepsetTest.java
index 44b6a51b0fd140..cd24f87d7f9f7f 100644
--- a/src/test/java/com/google/devtools/build/lib/collect/nestedset/DepsetTest.java
+++ b/src/test/java/com/google/devtools/build/lib/collect/nestedset/DepsetTest.java
@@ -406,4 +406,12 @@ ev.new Scenario()
                 + " NoneType'",
             "depset(direct='hello')");
   }
+
+  @Test
+  public void testSingleNonEmptyTransitiveUnwrapped() throws Exception {
+    ev.exec(
+        "inner = depset([1, 2, 3])",
+        "outer = depset(transitive = [depset(), inner, depset()])");
+    assertThat(ev.lookup("outer")).isSameInstanceAs(ev.lookup("inner"));
+  }
 }