Add value constants to ObjectCodecRegistry. Value constants are to be used when there may not be a canonical instance of the object we want (or the canonical instance isn't available at registry construction time), but we can reasonably cheaply do value equality comparisons. Strings are a good example.

PiperOrigin-RevId: 192354865

5 files changed, 96 insertions, 28 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/AutoRegistry.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/AutoRegistry.java
index 7785493aa5..6c81ccfe92 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/AutoRegistry.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/AutoRegistry.java
@@ -21,7 +21,10 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Ordering;
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashMap;
 
 /**
  * A lazy, automatically populated registry.
@@ -46,7 +49,7 @@ public class AutoRegistry {
   private static final ImmutableList<String> EXTERNAL_CLASS_NAMES_TO_REGISTER =
       ImmutableList.of("java.io.FileNotFoundException", "java.io.IOException");
 
-  private static final ImmutableList<Object> CONSTANTS_TO_REGISTER =
+  private static final ImmutableList<Object> REFERENCE_CONSTANTS_TO_REGISTER =
       ImmutableList.of(
           Predicates.alwaysTrue(),
           Predicates.alwaysFalse(),
@@ -57,6 +60,14 @@ public class AutoRegistry {
           Comparator.naturalOrder(),
           Ordering.natural());
 
+  private static final ImmutableList<Object> VALUE_CONSTANTS_TO_REGISTER =
+      ImmutableList.of(
+          "",
+          Boolean.FALSE,
+          Boolean.TRUE,
+          Collections.unmodifiableMap(new HashMap<>()),
+          Collections.unmodifiableList(new ArrayList<>()));
+
   public static ObjectCodecRegistry get() {
     return SUPPLIER.get();
   }
@@ -68,8 +79,11 @@ public class AutoRegistry {
       for (String className : EXTERNAL_CLASS_NAMES_TO_REGISTER) {
         registry.addClassName(className);
       }
-      for (Object constant : CONSTANTS_TO_REGISTER) {
-        registry.addConstant(constant);
+      for (Object constant : REFERENCE_CONSTANTS_TO_REGISTER) {
+        registry.addReferenceConstant(constant);
+      }
+      for (Object constant : VALUE_CONSTANTS_TO_REGISTER) {
+        registry.addValueConstant(constant);
       }
       return registry.build();
     } catch (IOException | ReflectiveOperationException e) {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecScanner.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecScanner.java
index a8f50a8ad1..dc93c5817c 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecScanner.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecScanner.java
@@ -129,7 +129,8 @@ public class CodecScanner {
           e);
     }
     try {
-      builder.addConstant(Preconditions.checkNotNull(field.get(null), "%s %s", field, type));
+      builder.addReferenceConstant(
+          Preconditions.checkNotNull(field.get(null), "%s %s", field, type));
     } catch (IllegalAccessException e) {
       throw new IllegalStateException("Could not access field " + field + " for " + type, e);
     }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java
index be324bcc05..5c144112af 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java
@@ -29,7 +29,6 @@ import java.util.Comparator;
 import java.util.HashMap;
 import java.util.IdentityHashMap;
 import java.util.Map;
-import java.util.Set;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
@@ -49,9 +48,13 @@ public class ObjectCodecRegistry {
   private final ImmutableMap<Class<?>, CodecDescriptor> classMappedCodecs;
   private final ImmutableList<CodecDescriptor> tagMappedCodecs;
 
-  private final int constantsStartTag;
-  private final IdentityHashMap<Object, Integer> constantsMap;
-  private final ImmutableList<Object> constants;
+  private final int referenceConstantsStartTag;
+  private final IdentityHashMap<Object, Integer> referenceConstantsMap;
+  private final ImmutableList<Object> referenceConstants;
+
+  private final int valueConstantsStartTag;
+  private final ImmutableMap<Class<?>, ImmutableMap<Object, Integer>> valueConstantsMap;
+  private final ImmutableList<Object> valueConstants;
 
   /** This is sorted, but we need index-based access. */
   private final ImmutableList<String> classNames;
@@ -59,8 +62,9 @@ public class ObjectCodecRegistry {
   private final IdentityHashMap<String, Supplier<CodecDescriptor>> dynamicCodecs;
 
   private ObjectCodecRegistry(
-      Set<ObjectCodec<?>> memoizingCodecs,
-      ImmutableList<Object> constants,
+      ImmutableSet<ObjectCodec<?>> memoizingCodecs,
+      ImmutableList<Object> referenceConstants,
+      ImmutableList<Object> valueConstants,
       ImmutableSortedSet<String> classNames,
       boolean allowDefaultCodec) {
     this.allowDefaultCodec = allowDefaultCodec;
@@ -77,13 +81,29 @@ public class ObjectCodecRegistry {
     this.classMappedCodecs = memoizingCodecsBuilder.build();
     this.tagMappedCodecs = tagMappedMemoizingCodecsBuilder.build();
 
-    constantsStartTag = nextTag;
-    constantsMap = new IdentityHashMap<>();
-    for (Object constant : constants) {
-      constantsMap.put(constant, nextTag++);
+    referenceConstantsStartTag = nextTag;
+    referenceConstantsMap = new IdentityHashMap<>();
+    for (Object constant : referenceConstants) {
+      referenceConstantsMap.put(constant, nextTag++);
     }
-    this.constants = constants;
+    this.referenceConstants = referenceConstants;
+
+    valueConstantsStartTag = nextTag;
 
+    HashMap<Class<?>, HashMap<Object, Integer>> valuesBuilder = new HashMap<>();
+    for (Object constant : valueConstants) {
+      valuesBuilder
+          .computeIfAbsent(constant.getClass(), k -> new HashMap<>())
+          .put(constant, nextTag++);
+    }
+    this.valueConstantsMap =
+        valuesBuilder
+            .entrySet()
+            .stream()
+            .collect(
+                ImmutableMap.toImmutableMap(
+                    Map.Entry::getKey, e -> ImmutableMap.copyOf(e.getValue())));
+    this.valueConstants = valueConstants;
     this.classNames = classNames.asList();
     this.dynamicCodecs = createDynamicCodecs(classNames, nextTag);
   }
@@ -124,14 +144,27 @@ public class ObjectCodecRegistry {
 
   @Nullable
   Object maybeGetConstantByTag(int tag) {
-    return tag < constantsStartTag || tag - constantsStartTag >= constants.size()
-        ? null
-        : constants.get(tag - constantsStartTag);
+    if (referenceConstantsStartTag <= tag
+        && tag < referenceConstantsStartTag + referenceConstants.size()) {
+      return referenceConstants.get(tag - referenceConstantsStartTag);
+    }
+    if (valueConstantsStartTag <= tag && tag < valueConstantsStartTag + valueConstants.size()) {
+      return valueConstants.get(tag - valueConstantsStartTag);
+    }
+    return null;
   }
 
   @Nullable
   Integer maybeGetTagForConstant(Object object) {
-    return constantsMap.get(object);
+    Integer result = referenceConstantsMap.get(object);
+    if (result != null) {
+      return result;
+    }
+    ImmutableMap<Object, Integer> valueConstantsForClass = valueConstantsMap.get(object.getClass());
+    if (valueConstantsForClass == null) {
+      return null;
+    }
+    return valueConstantsForClass.get(object);
   }
 
   /** Returns the {@link CodecDescriptor} associated with the supplied tag. */
@@ -146,7 +179,8 @@ public class ObjectCodecRegistry {
     }
 
     tagOffset -= tagMappedCodecs.size();
-    tagOffset -= constants.size();
+    tagOffset -= referenceConstants.size();
+    tagOffset -= valueConstants.size();
     if (!allowDefaultCodec || tagOffset < 0 || tagOffset >= classNames.size()) {
       throw new SerializationException.NoCodecException("No codec available for tag " + tag);
     }
@@ -166,8 +200,12 @@ public class ObjectCodecRegistry {
       builder.add(entry.getValue().getCodec());
     }
 
-    for (Object constant : constants) {
-      builder.addConstant(constant);
+    for (Object constant : referenceConstants) {
+      builder.addReferenceConstant(constant);
+    }
+
+    for (Object constant : valueConstants) {
+      builder.addValueConstant(constant);
     }
 
     for (String className : classNames) {
@@ -245,7 +283,8 @@ public class ObjectCodecRegistry {
   /** Builder for {@link ObjectCodecRegistry}. */
   public static class Builder {
     private final Map<Class<?>, ObjectCodec<?>> codecs = new HashMap<>();
-    private final ImmutableList.Builder<Object> constantsBuilder = ImmutableList.builder();
+    private final ImmutableList.Builder<Object> referenceConstantsBuilder = ImmutableList.builder();
+    private final ImmutableList.Builder<Object> valueConstantsBuilder = ImmutableList.builder();
     private final ImmutableSortedSet.Builder<String> classNames = ImmutableSortedSet.naturalOrder();
     private boolean allowDefaultCodec = true;
 
@@ -288,8 +327,21 @@ public class ObjectCodecRegistry {
      * not serialize bit-for-bit identical disregarding this list of constants, since the list
      * object's codec will be different for the two objects.
      */
-    public Builder addConstant(Object object) {
-      constantsBuilder.add(object);
+    public Builder addReferenceConstant(Object object) {
+      referenceConstantsBuilder.add(object);
+      return this;
+    }
+
+    /**
+     * Adds a constant value. Any value encountered during serialization which has the same class as
+     * {@code object} and {@link Object#equals} {@code object} will be replaced by {@code object}
+     * upon deserialization. These objects should therefore be indistinguishable, and unequal
+     * objects should quickly compare unequal (it is ok for equal objects to be relatively expensive
+     * to compare equal, if that is still less expensive than the cost of serializing the object).
+     * Short {@link String} objects are ideal for value constants.
+     */
+    public Builder addValueConstant(Object object) {
+      valueConstantsBuilder.add(object);
       return this;
     }
 
@@ -301,7 +353,8 @@ public class ObjectCodecRegistry {
     public ObjectCodecRegistry build() {
       return new ObjectCodecRegistry(
           ImmutableSet.copyOf(codecs.values()),
-          constantsBuilder.build(),
+          referenceConstantsBuilder.build(),
+          valueConstantsBuilder.build(),
           classNames.build(),
           allowDefaultCodec);
     }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializationTester.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializationTester.java
index b45d443f6c..1a6841fff7 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializationTester.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializationTester.java
@@ -111,7 +111,7 @@ public class SerializationTester {
     ImmutableMap<Class<?>, Object> dependencies = dependenciesBuilder.build();
     ObjectCodecRegistry.Builder registryBuilder = registry.getBuilder();
     for (Object val : dependencies.values()) {
-      registryBuilder.addConstant(val);
+      registryBuilder.addReferenceConstant(val);
     }
     for (ObjectCodec<?> codec : additionalCodecs) {
       registryBuilder.add(codec);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/TestUtils.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/TestUtils.java
index 657765f87f..2ed9db4290 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/TestUtils.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/TestUtils.java
@@ -84,7 +84,7 @@ public class TestUtils {
       throws IOException, SerializationException {
     ObjectCodecRegistry.Builder builder = AutoRegistry.get().getBuilder();
     for (Object constant : dependencies.values()) {
-      builder.addConstant(constant);
+      builder.addReferenceConstant(constant);
     }
     ObjectCodecRegistry registry = builder.build();
     return new DeserializationContext(registry, dependencies)