Save nested set size after initial expand to avoid GC from resizing list when replaying.

We stash the size in the same field as the order using bit packing. The small additional cost of masking out the int should be less than the resizing when replaying large nested sets.

The upper size bound on nested sets will be 512M objects as a result. If Java had unsigned ints we could push it to 1G.

RELNOTES: None
PiperOrigin-RevId: 200417105

2 files changed, 41 insertions, 28 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSet.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSet.java
index 7ff6da4a3f..52df58e36a 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSet.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSet.java
@@ -16,6 +16,7 @@ package com.google.devtools.build.lib.collect.nestedset;
 import static java.util.stream.Collectors.joining;
 
 import com.google.common.base.Function;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.collect.compacthashset.CompactHashSet;
@@ -39,7 +40,15 @@ import javax.annotation.Nullable;
 @AutoCodec
 public final class NestedSet<E> implements Iterable<E> {
 
-  private final Order order;
+  /**
+   * Order and size of set packed into one int.
+   *
+   * <p>Bits 31-2: size, bits 1-0: order enum ordinal. The order is assigned on construction time,
+   * the size is computed on the first expansion and set afterwards so it's available for {@link
+   * #replay}.
+   */
+  private int orderAndSize;
+
   private final Object children;
   private byte[] memo;
 
@@ -50,13 +59,13 @@ public final class NestedSet<E> implements Iterable<E> {
    * Construct an empty NestedSet.  Should only be called by Order's class initializer.
    */
   NestedSet(Order order) {
-    this.order = order;
+    this.orderAndSize = order.ordinal();
     this.children = EMPTY_CHILDREN;
     this.memo = LEAF_MEMO;
   }
 
   NestedSet(Order order, Set<E> direct, Set<NestedSet<E>> transitive) {
-    this.order = order;
+    this.orderAndSize = order.ordinal();
 
     // The iteration order of these collections is the order in which we add the items.
     Collection<E> directOrder = direct;
@@ -143,7 +152,7 @@ public final class NestedSet<E> implements Iterable<E> {
   // Only used by deserialization
   @AutoCodec.Instantiator
   NestedSet(Order order, Object children) {
-    this.order = order;
+    this.orderAndSize = order.ordinal();
     this.children = children;
     boolean hasChildren =
         children instanceof Object[]
@@ -155,7 +164,7 @@ public final class NestedSet<E> implements Iterable<E> {
    * Returns the ordering of this nested set.
    */
   public Order getOrder() {
-    return order;
+    return Order.getOrder(orderAndSize & 3);
   }
 
   /**
@@ -202,7 +211,7 @@ public final class NestedSet<E> implements Iterable<E> {
     if (isEmpty()) {
       return ImmutableList.of();
     }
-    return order == Order.LINK_ORDER ? expand().reverse() : expand();
+    return getOrder() == Order.LINK_ORDER ? expand().reverse() : expand();
   }
 
   /**
@@ -229,9 +238,10 @@ public final class NestedSet<E> implements Iterable<E> {
       return true;
     }
     return other != null
-        && order == other.order
+        && getOrder() == other.getOrder()
         && (children.equals(other.children)
-            || (!isSingleton() && !other.isSingleton()
+            || (!isSingleton()
+                && !other.isSingleton()
                 && Arrays.equals((Object[]) children, (Object[]) other.children)));
   }
 
@@ -246,8 +256,8 @@ public final class NestedSet<E> implements Iterable<E> {
    */
   public int shallowHashCode() {
     return isSingleton()
-        ? Objects.hash(order, children)
-        : Objects.hash(order, Arrays.hashCode((Object[]) children));
+        ? Objects.hash(getOrder(), children)
+        : Objects.hash(getOrder(), Arrays.hashCode((Object[]) children));
   }
 
   @Override
@@ -294,10 +304,7 @@ public final class NestedSet<E> implements Iterable<E> {
       return ImmutableList.copyOf(members);
     }
     Object[] children = (Object[]) this.children;
-    // TODO:  We could record the exact size (inside memo, or by making order an int with two bits
-    // for Order.ordinal()) and avoid an array copy here.  It's not directly visible in profiles but
-    // it would reduce garbage generated.
-    ImmutableList.Builder<E> output = ImmutableList.builder();
+    ImmutableList.Builder<E> output = ImmutableList.builderWithExpectedSize(orderAndSize >> 2);
     replay(output, children, memo, 0);
     return output.build();
   }
@@ -338,18 +345,20 @@ public final class NestedSet<E> implements Iterable<E> {
     if (bytes <= memo.length - 16) {
       memo = Arrays.copyOf(memo, bytes);
     }
+    Preconditions.checkState(members.size() < (Integer.MAX_VALUE >> 2));
+    orderAndSize |= (members.size()) << 2;
     return members;
   }
 
   /**
-   * Perform a depth-first traversal of {@code children}, tracking visited
-   * arrays in {@code sets} and visited leaves in {@code members}.  We also
-   * record which edges were taken in {@code this.memo} starting at {@code pos}.
+   * Perform a depth-first traversal of {@code children}, tracking visited arrays in {@code sets}
+   * and visited leaves in {@code members}. We also record which edges were taken in {@code
+   * this.memo} starting at {@code pos}.
    *
-   * Returns the final value of {@code pos}.
+   * <p>Returns the final value of {@code pos}.
    */
-  private int walk(CompactHashSet<Object> sets, CompactHashSet<E> members,
-                   Object[] children, int pos) {
+  private int walk(
+      CompactHashSet<Object> sets, CompactHashSet<E> members, Object[] children, int pos) {
     for (Object child : children) {
       if ((pos >> 3) >= memo.length) {
         memo = Arrays.copyOf(memo, memo.length * 2);
@@ -381,11 +390,11 @@ public final class NestedSet<E> implements Iterable<E> {
   }
 
   /**
-   * Repeat a previous traversal of {@code children} performed by {@link #walk}
-   * and recorded in {@code memo}, appending leaves to {@code output}.
+   * Repeat a previous traversal of {@code children} performed by {@link #walk} and recorded in
+   * {@code memo}, appending leaves to {@code output}.
    */
-  private static <E> int replay(ImmutableList.Builder<E> output, Object[] children,
-                                byte[] memo, int pos) {
+  private static <E> int replay(
+      ImmutableList.Builder<E> output, Object[] children, byte[] memo, int pos) {
     for (Object child : children) {
       if ((memo[pos >> 3] & (1 << (pos & 7))) != 0) {
         if (child instanceof Object[]) {
diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/Order.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/Order.java
index 0fb2ae4d8c..967dd5e299 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/Order.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/Order.java
@@ -107,6 +107,7 @@ public enum Order {
   NAIVE_LINK_ORDER("preorder");
 
   private static final ImmutableMap<String, Order> VALUES;
+  private static final Order[] ORDINALS;
 
   private final String skylarkName;
   private final NestedSet<?> emptySet;
@@ -180,14 +181,17 @@ public enum Order {
    * Indexes all possible values by name and stores the results in a {@code ImmutableMap}
    */
   static {
-    Order[] tmpValues = Order.values();
+    ORDINALS = values();
+    HashMap<String, Order> entries = Maps.newHashMapWithExpectedSize(ORDINALS.length);
 
-    HashMap<String, Order> entries = Maps.newHashMapWithExpectedSize(tmpValues.length);
-
-    for (Order current : tmpValues) {
+    for (Order current : ORDINALS) {
       entries.put(current.getSkylarkName(), current);
     }
 
     VALUES = ImmutableMap.copyOf(entries);
   }
+
+  static Order getOrder(int ordinal) {
+    return ORDINALS[ordinal];
+  }
 }