Integrate memoization into standard serialization. This involves a number of large changes:

1. SerializationContext and DeserializationContext become the owners of the Memoizer if requested. They produce new versions of themselves on demand that are memoization-aware. Because of intricacies of Skylark that I do not fully understand, we inject a Mutability object when starting memoization, and so to be conservative, that injection starts up a new memoization frame, nested if we were already memoizing, just like before. It would be nice to decouple this injection from memoization in the future.

2. MemoizingCodec is deleted, but really ObjectCodec becomes MemoizingCodec, so it lives on. BaseCodec is deleted since it now has only one implementation.

3. The simplified model of registering MemoizingCodecs is adopted for ObjectCodecs: all codecs are registered based on their #getEncodedClass and #additionalEncodedSubclasses. This also allows us to register codecs that are defined using tricky parameter types, since we're no longer trying to reflectively examine them. This required a clean-up of such codecs, and the addition of ArrayListCodec to stop NullableListCodec from making lists unmodifiable when they shouldn't be.

4. @AutoCodec is extended to allow users to specify that memoization should start with this codec.

To ensure bit-equivalence, SkyKeySerializer disables memoization.

PiperOrigin-RevId: 188918251

1 files changed, 43 insertions, 13 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMapCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMapCodec.java
index 9be4046daf..2c1e11c92b 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMapCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMapCodec.java
@@ -24,22 +24,38 @@ import java.util.Comparator;
 import java.util.Map;
 
 /**
- * Encodes an {@link ImmutableMap}, which may be an ImmutableSortedMap. We handle both here because
- * we cannot handle ImmutableSortedMap with any ordering other than the default, and so we degrade
- * to ImmutableMap in that case, hoping that the caller does not notice.
+ * Encodes an {@link ImmutableMap}, which may be an {@link ImmutableSortedMap}. The iteration order
+ * of the deserialized map is the same as the original map's.
+ *
+ * <p>We handle both {@link ImmutableMap} and {@link ImmutableSortedMap} here because we cannot
+ * handle {@link ImmutableSortedMap} with any ordering other than the default, and so we degrade to
+ * ImmutableMap in that case, hoping that the caller does not notice. Since the ordering is
+ * preserved, the caller should not be sensitive to this change unless the caller's field is
+ * declared to be an {@link ImmutableSortedMap}, in which case we will crash at runtime.
+ *
+ * <p>Any {@link SerializationException} or {@link IOException} that arises while serializing or
+ * deserializing a map entry's value (not its key) will be wrapped in a new {@link
+ * SerializationException} using {@link SerializationException#propagate}. (Note that this preserves
+ * the type of {@link SerializationException.NoCodecException} exceptions.) The message will include
+ * the {@code toString()} of the entry's key. For errors that occur while serializing, it will also
+ * include the class name of the entry's value. Errors that occur while serializing an entry key are
+ * not affected.
+ *
+ * <p>Because of the ambiguity around the key type (Comparable in the case of {@link
+ * ImmutableSortedMap}, arbitrary otherwise, we avoid specifying the key type as a parameter.
  */
-class ImmutableMapCodec implements ObjectCodec<ImmutableMap<?, ?>> {
+class ImmutableMapCodec<V> implements ObjectCodec<ImmutableMap<?, V>> {
   @SuppressWarnings("unchecked")
   @Override
-  public Class<ImmutableMap<?, ?>> getEncodedClass() {
+  public Class<ImmutableMap<?, V>> getEncodedClass() {
     // Because Java disallows converting from Class<ImmutableMap> to Class<ImmutableMap<?, ?>>
     // directly.
-    return (Class<ImmutableMap<?, ?>>) ((Class<?>) ImmutableMap.class);
+    return (Class<ImmutableMap<?, V>>) ((Class<?>) ImmutableMap.class);
   }
 
   @Override
   public void serialize(
-      SerializationContext context, ImmutableMap<?, ?> map, CodedOutputStream codedOut)
+      SerializationContext context, ImmutableMap<?, V> map, CodedOutputStream codedOut)
       throws SerializationException, IOException {
     codedOut.writeInt32NoTag(map.size());
     boolean serializeAsSortedMap = false;
@@ -52,12 +68,20 @@ class ImmutableMapCodec implements ObjectCodec<ImmutableMap<?, ?>> {
     codedOut.writeBoolNoTag(serializeAsSortedMap);
     for (Map.Entry<?, ?> entry : map.entrySet()) {
       context.serialize(entry.getKey(), codedOut);
-      context.serialize(entry.getValue(), codedOut);
+      try {
+        context.serialize(entry.getValue(), codedOut);
+      } catch (SerializationException | IOException e) {
+        throw SerializationException.propagate(
+            String.format(
+                "Exception while serializing value of type %s for key '%s'",
+                entry.getValue().getClass().getName(), entry.getKey()),
+            e);
+      }
     }
   }
 
   @Override
-  public ImmutableMap<?, ?> deserialize(DeserializationContext context, CodedInputStream codedIn)
+  public ImmutableMap<?, V> deserialize(DeserializationContext context, CodedInputStream codedIn)
       throws SerializationException, IOException {
     int length = codedIn.readInt32();
     if (length < 0) {
@@ -70,15 +94,21 @@ class ImmutableMapCodec implements ObjectCodec<ImmutableMap<?, ?>> {
     }
   }
 
-  private static <T> ImmutableMap<T, Object> buildMap(
-      ImmutableMap.Builder<T, Object> builder,
+  private static <K, V> ImmutableMap<K, V> buildMap(
+      ImmutableMap.Builder<K, V> builder,
       int length,
       DeserializationContext context,
       CodedInputStream codedIn)
       throws IOException, SerializationException {
     for (int i = 0; i < length; i++) {
-      T key = context.deserialize(codedIn);
-      Object value = context.deserialize(codedIn);
+      K key = context.deserialize(codedIn);
+      V value;
+      try {
+        value = context.deserialize(codedIn);
+      } catch (SerializationException | IOException e) {
+        throw SerializationException.propagate(
+            String.format("Exception while deserializing value for key '%s'", key), e);
+      }
       builder.put(key, value);
     }
     try {