Serialize singleton NestedSets directly in NestedSetCodecWithStore.

PiperOrigin-RevId: 194232982

4 files changed, 61 insertions, 38 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecWithStore.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecWithStore.java
index 93df1e378d..018d360990 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecWithStore.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecWithStore.java
@@ -62,7 +62,16 @@ public class NestedSetCodecWithStore<T> implements ObjectCodec<NestedSet<T>> {
     if (obj.isEmpty()) {
       return;
     }
-    ByteString fingerprint = nestedSetStore.computeFingerprintAndStore(obj.rawChildren(), context);
+
+    // If the NestedSet is a singleton, we serialize directly as an optimization.
+    codedOut.writeBoolNoTag(obj.isSingleton());
+    if (obj.isSingleton()) {
+      context.serialize(obj.rawChildren(), codedOut);
+      return;
+    }
+
+    ByteString fingerprint =
+        nestedSetStore.computeFingerprintAndStore((Object[]) obj.rawChildren(), context);
     codedOut.writeByteArrayNoTag(fingerprint.toByteArray());
   }
 
@@ -80,6 +89,12 @@ public class NestedSetCodecWithStore<T> implements ObjectCodec<NestedSet<T>> {
       return NestedSetBuilder.emptySet(order);
     }
 
+    boolean isSingleton = codedIn.readBool();
+    if (isSingleton) {
+      T contents = context.deserialize(codedIn);
+      return new NestedSet<T>(order, contents);
+    }
+
     ByteString fingerprint = ByteString.copyFrom(codedIn.readByteArray());
     Object members = nestedSetStore.getContentsAndDeserialize(fingerprint, context);
     return new NestedSet<>(order, members);
diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetStore.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetStore.java
index 9ad55efee5..7377daa336 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetStore.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetStore.java
@@ -84,14 +84,14 @@ public class NestedSetStore {
 
   /** An in-memory cache for fingerprint <-> NestedSet associations. */
   private static class NestedSetCache {
-    private final Map<ByteString, Object> fingerprintToContents =
+    private final Map<ByteString, Object[]> fingerprintToContents =
         new MapMaker()
             .concurrencyLevel(SerializationConstants.DESERIALIZATION_POOL_SIZE)
             .weakValues()
             .makeMap();
 
     /** Object/Object[] contents to fingerprint. Maintained for fast fingerprinting. */
-    private final Map<Object, ByteString> contentsToFingerprint =
+    private final Map<Object[], ByteString> contentsToFingerprint =
         new MapMaker()
             .concurrencyLevel(SerializationConstants.DESERIALIZATION_POOL_SIZE)
             .weakKeys()
@@ -102,7 +102,7 @@ public class NestedSetStore {
      * fingerprint is not known.
      */
     @Nullable
-    public Object contentsForFingerprint(ByteString fingerprint) {
+    public Object[] contentsForFingerprint(ByteString fingerprint) {
       return fingerprintToContents.get(fingerprint);
     }
 
@@ -111,12 +111,12 @@ public class NestedSetStore {
      * contents are not known.
      */
     @Nullable
-    public ByteString fingerprintForContents(Object contents) {
+    public ByteString fingerprintForContents(Object[] contents) {
       return contentsToFingerprint.get(contents);
     }
 
     /** Associates the provided fingerprint and NestedSet contents. */
-    public void put(ByteString fingerprint, Object contents) {
+    public void put(ByteString fingerprint, Object[] contents) {
       contentsToFingerprint.put(contents, fingerprint);
       fingerprintToContents.put(fingerprint, contents);
     }
@@ -132,8 +132,8 @@ public class NestedSetStore {
    * store. Recursively does the same for all transitive members (i.e. Object[] members) of the
    * provided contents.
    */
-  ByteString computeFingerprintAndStore(Object contents, SerializationContext serializationContext)
-      throws SerializationException {
+  ByteString computeFingerprintAndStore(
+      Object[] contents, SerializationContext serializationContext) throws SerializationException {
     ByteString priorFingerprint = nestedSetCache.fingerprintForContents(contents);
     if (priorFingerprint != null) {
       return priorFingerprint;
@@ -150,24 +150,19 @@ public class NestedSetStore {
     CodedOutputStream codedOutputStream = CodedOutputStream.newInstance(byteArrayOutputStream);
 
     try {
-      if (contents instanceof Object[]) {
-        Object[] contentsArray = (Object[]) contents;
-        codedOutputStream.writeInt32NoTag(contentsArray.length);
-        for (Object child : contentsArray) {
-          if (child instanceof Object[]) {
-            ByteString fingerprint = computeFingerprintAndStore(child, serializationContext);
-            newSerializationContext.serialize(fingerprint, codedOutputStream);
-          } else {
-            newSerializationContext.serialize(child, codedOutputStream);
-          }
+      codedOutputStream.writeInt32NoTag(contents.length);
+      for (Object child : contents) {
+        if (child instanceof Object[]) {
+          ByteString fingerprint =
+              computeFingerprintAndStore((Object[]) child, serializationContext);
+          newSerializationContext.serialize(fingerprint, codedOutputStream);
+        } else {
+          newSerializationContext.serialize(child, codedOutputStream);
         }
-      } else {
-        codedOutputStream.writeInt32NoTag(1);
-        newSerializationContext.serialize(contents, codedOutputStream);
       }
       codedOutputStream.flush();
     } catch (IOException e) {
-      throw new SerializationException("Could not serialize " + contents, e);
+      throw new SerializationException("Could not serialize NestedSet contents", e);
     }
 
     byte[] serializedBytes = byteArrayOutputStream.toByteArray();
@@ -181,10 +176,10 @@ public class NestedSetStore {
   }
 
   /** Retrieves and deserializes the NestedSet contents associated with the given fingerprint. */
-  public Object getContentsAndDeserialize(
+  public Object[] getContentsAndDeserialize(
       ByteString fingerprint, DeserializationContext deserializationContext)
       throws IOException, SerializationException {
-    Object contents = nestedSetCache.contentsForFingerprint(fingerprint);
+    Object[] contents = nestedSetCache.contentsForFingerprint(fingerprint);
     if (contents != null) {
       return contents;
     }
@@ -199,20 +194,14 @@ public class NestedSetStore {
         deserializationContext.getNewMemoizingContext();
 
     int numberOfElements = codedIn.readInt32();
-    Object dereferencedContents;
-    if (numberOfElements > 1) {
-      Object[] dereferencedContentsArray = new Object[numberOfElements];
-      for (int i = 0; i < numberOfElements; i++) {
-        Object deserializedElement = newDeserializationContext.deserialize(codedIn);
-        dereferencedContentsArray[i] =
-            deserializedElement instanceof ByteString
-                ? getContentsAndDeserialize(
-                    (ByteString) deserializedElement, deserializationContext)
-                : deserializedElement;
-      }
-      dereferencedContents = dereferencedContentsArray;
-    } else {
-      dereferencedContents = newDeserializationContext.deserialize(codedIn);
+    Object[] dereferencedContents = new Object[numberOfElements];
+    for (int i = 0; i < numberOfElements; i++) {
+      Object deserializedElement = newDeserializationContext.deserialize(codedIn);
+      dereferencedContents[i] =
+          deserializedElement instanceof ByteString
+              ? getContentsAndDeserialize(
+                  (ByteString) deserializedElement, deserializationContext)
+              : deserializedElement;
     }
 
     nestedSetCache.put(fingerprint, dereferencedContents);
diff --git a/src/test/java/com/google/devtools/build/lib/BUILD b/src/test/java/com/google/devtools/build/lib/BUILD
index 91c8a1d236..3ac58e14db 100644
--- a/src/test/java/com/google/devtools/build/lib/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/BUILD
@@ -225,6 +225,7 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization:constants",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils",
+        "//third_party:mockito",
         "//third_party/protobuf:protobuf_java",
     ],
 )
diff --git a/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTest.java b/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTest.java
index 920eec7c55..4ad1f0494b 100644
--- a/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTest.java
+++ b/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTest.java
@@ -25,6 +25,7 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
 
 /** Tests for {@link NestedSet} serialization. */
 @RunWith(JUnit4.class)
@@ -64,6 +65,23 @@ public class NestedSetCodecTest {
     checkCodec(objectCodecs);
   }
 
+  @Test
+  public void testSingletonNestedSetSerializedWithoutStore() throws Exception {
+    NestedSetStore mockNestedSetStore = Mockito.mock(NestedSetStore.class);
+    Mockito.when(mockNestedSetStore.computeFingerprintAndStore(Mockito.any(), Mockito.any()))
+        .thenThrow(new AssertionError("NestedSetStore should not have been used"));
+
+    ObjectCodecs objectCodecs =
+        new ObjectCodecs(
+            AutoRegistry.get()
+                .getBuilder()
+                .setAllowDefaultCodec(true)
+                .add(new NestedSetCodecWithStore<>(mockNestedSetStore))
+                .build());
+    NestedSet<String> singletonNestedSet = new NestedSet<>(Order.STABLE_ORDER, "a");
+    objectCodecs.serialize(singletonNestedSet);
+  }
+
   private void checkCodec(ObjectCodecs objectCodecs) throws Exception {
     new SerializationTester(
             NestedSetBuilder.emptySet(Order.STABLE_ORDER),