1 files changed, 33 insertions, 24 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[]) {