Use specialized fastutil maps in serialization to avoid boxing of primitive ints.

PiperOrigin-RevId: 587006554
Change-Id: I006548544656f40a976a038cddd259d57ec15611

Author: justinhorvitz

src/main/java/com/google/devtools/build/lib/skyframe/serialization/BUILD

@@ -39,6 +39,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/unsafe:string",
         "//src/main/java/com/google/devtools/build/lib/unsafe:unsafe-provider",
         "//third_party:error_prone_annotations",
+        "//third_party:fastutil",
         "//third_party:flogger",
         "//third_party:guava",
         "//third_party:jsr305",

src/main/java/com/google/devtools/build/lib/skyframe/serialization/Memoizer.java

@@ -19,11 +19,11 @@
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec.MemoizationStrategy;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
+import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
+import it.unimi.dsi.fastutil.objects.Reference2IntOpenHashMap;
 import java.io.IOException;
 import java.util.ArrayDeque;
 import java.util.Deque;
-import java.util.HashMap;
-import java.util.IdentityHashMap;
 import javax.annotation.Nullable;
 
 /**
@@ -157,9 +157,8 @@ <T> void serialize(
       }
     }
 
-    @Nullable
-    Integer getMemoizedIndex(Object obj) {
-      return memo.lookupNullable(obj);
+    int getMemoizedIndex(Object obj) {
+      return memo.lookup(obj);
     }
 
     // Corresponds to MemoContent in the abstract grammar.
@@ -170,52 +169,55 @@ private <T> void serializeMemoContent(
         CodedOutputStream codedOut,
         MemoizationStrategy strategy)
         throws SerializationException, IOException {
-      switch(strategy) {
-        case MEMOIZE_BEFORE: {
-          int id = memo.memoize(obj);
-          codedOut.writeInt32NoTag(id);
+      switch (strategy) {
+        case MEMOIZE_BEFORE:
+          {
+            int id = memo.memoize(obj);
+            codedOut.writeInt32NoTag(id);
             codec.serialize(context, obj, codedOut);
-          break;
-        }
-        case MEMOIZE_AFTER: {
+            break;
+          }
+        case MEMOIZE_AFTER:
+          {
             codec.serialize(context, obj, codedOut);
-          // If serializing the children caused the parent object itself to be serialized due to a
-          // cycle, then there's now a memo entry for the parent. Don't overwrite it with a new id.
-          Integer cylicallyCreatedId = memo.lookupNullable(obj);
-          int id = (cylicallyCreatedId != null) ? cylicallyCreatedId : memo.memoize(obj);
-          codedOut.writeInt32NoTag(id);
-          break;
-        }
+            // If serializing the children caused the parent object itself to be serialized due to a
+            // cycle, then there's now a memo entry for the parent. Don't overwrite it with a new
+            // id.
+            int cylicallyCreatedId = memo.lookup(obj);
+            int id = (cylicallyCreatedId != -1) ? cylicallyCreatedId : memo.memoize(obj);
+            codedOut.writeInt32NoTag(id);
+            break;
+          }
         default:
           throw new AssertionError("Unreachable (strategy=" + strategy + ")");
       }
     }
 
     private static class SerializingMemoTable {
-      private final IdentityHashMap<Object, Integer> table = new IdentityHashMap<>();
+      private final Reference2IntOpenHashMap<Object> table = new Reference2IntOpenHashMap<>();
+
+      SerializingMemoTable() {
+        table.defaultReturnValue(-1);
+      }
 
       /**
        * Adds a new value to the memo table and returns its id. The value must not already be
        * present.
        */
       private int memoize(Object value) {
         Preconditions.checkArgument(
-            !table.containsKey(value),
-            "Tried to memoize object '%s' multiple times", value);
+            !table.containsKey(value), "Tried to memoize object '%s' multiple times", value);
         // Ids count sequentially from 0.
         int newId = table.size();
         table.put(value, newId);
         return newId;
       }
 
       /**
-       * If the value is already memoized, return its on-the-wire id; otherwise return null. Opaque.
-       *
-       * <p>Beware accidental unboxing of a null result.
+       * If the value is already memoized, return its on-the-wire id; otherwise returns {@code -1}.
        */
-      @Nullable
-      private Integer lookupNullable(Object value) {
-        return table.get(value);
+      private int lookup(Object value) {
+        return table.getInt(value);
       }
     }
   }
@@ -225,7 +227,7 @@ private Integer lookupNullable(Object value) {
    */
   static class Deserializer {
     private final DeserializingMemoTable memo = new DeserializingMemoTable();
-    @Nullable private Integer tagForMemoizedBefore = null;
+    private int tagForMemoizedBefore = -1;
     private final Deque<Object> memoizedBeforeStackForSanityChecking = new ArrayDeque<>();
 
     /**
@@ -241,7 +243,7 @@ <T> T deserialize(
         CodedInputStream codedIn)
         throws SerializationException, IOException {
       Preconditions.checkState(
-          tagForMemoizedBefore == null,
+          tagForMemoizedBefore == -1,
           "non-null memoized-before tag %s (%s)",
           tagForMemoizedBefore,
           codec);
@@ -302,10 +304,10 @@ private static <T> T castedDeserialize(
     }
 
     <T> void registerInitialValue(T initialValue) {
-      int tag =
-          Preconditions.checkNotNull(
-              tagForMemoizedBefore, " Not called with memoize before: %s", initialValue);
-      tagForMemoizedBefore = null;
+      Preconditions.checkState(
+          tagForMemoizedBefore != -1, "Not called with memoize before: %s", initialValue);
+      int tag = tagForMemoizedBefore;
+      tagForMemoizedBefore = -1;
       memo.memoize(tag, initialValue);
       memoizedBeforeStackForSanityChecking.addLast(initialValue);
     }
@@ -348,18 +350,21 @@ private <T> T deserializeMemoAfterContent(
 
     private static class DeserializingMemoTable {
 
-      private HashMap<Integer, Object> table = new HashMap<>();
+      private final Int2ObjectOpenHashMap<Object> table = new Int2ObjectOpenHashMap<>();
 
       /**
        * Adds a new id -> object maplet to the memo table. The value must not already be present.
        */
       private void memoize(int id, Object value) {
         Preconditions.checkNotNull(value);
         Object prev = table.put(id, value);
-        Preconditions.checkArgument(
-            prev == null,
-            "Tried to memoize id %s to object '%s', when it is already memoized to object '%s'",
-            id, value, prev);
+        if (prev != null) { // Avoid boxing int with checkArgument.
+          throw new IllegalArgumentException(
+              String.format(
+                  "Tried to memoize id %s to object '%s', when it is already memoized to object"
+                      + " '%s'",
+                  id, value, prev));
+        }
       }
 
       /** If the id has been memoized, returns its corresponding object. Otherwise returns null. */

src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java

@@ -278,8 +278,8 @@ private ObjectCodecRegistry.CodecDescriptor recordAndGetDescriptorIfNotConstantM
       return null;
     }
     if (serializer != null) {
-      Integer memoizedIndex = serializer.getMemoizedIndex(object);
-      if (memoizedIndex != null) {
+      int memoizedIndex = serializer.getMemoizedIndex(object);
+      if (memoizedIndex != -1) {
         // Subtract 1 so it will be negative and not collide with null.
         codedOut.writeSInt32NoTag(-memoizedIndex - 1);
         return null;