Makes ImmutableMap codecs DeferredObjectCodec.

* Fixes a soundness error with ImmutableMapCodec which would cause a crash if
  the declared type is ImmutableSortedMap by splitting into 2 codecs.
* This necessarily includes ImmutableBiMapCodec and
  ImmutableClassToInstanceMapCodec because those two previously had
  dependencies on ImmutableMapCodec.

PiperOrigin-RevId: 596078698
Change-Id: I23b928307a3f1b1722d31b2907366fe46352ee6f

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

@@ -14,10 +14,14 @@
 
 package com.google.devtools.build.lib.skyframe.serialization;
 
+import static com.google.devtools.build.lib.skyframe.serialization.MapHelpers.deserializeMapEntries;
+import static com.google.devtools.build.lib.skyframe.serialization.MapHelpers.serializeMapEntries;
+
 import com.google.common.collect.ImmutableBiMap;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
+import java.util.function.Supplier;
 
 /**
  * Encodes an {@link ImmutableBiMap}. The iteration order of the deserialized map is the same as the
@@ -37,41 +41,61 @@
  * include the class name of the entry's value. Errors that occur while serializing an entry key are
  * not affected.
  */
-class ImmutableBiMapCodec<K, V> implements ObjectCodec<ImmutableBiMap<K, V>> {
+@SuppressWarnings({"unchecked", "rawtypes"})
+class ImmutableBiMapCodec extends DeferredObjectCodec<ImmutableBiMap> {
 
-  @SuppressWarnings("unchecked")
   @Override
-  public Class<ImmutableBiMap<K, V>> getEncodedClass() {
-    return (Class<ImmutableBiMap<K, V>>) ((Class<?>) ImmutableBiMap.class);
+  public Class<ImmutableBiMap> getEncodedClass() {
+    return ImmutableBiMap.class;
   }
 
   @Override
   public void serialize(
-      SerializationContext context, ImmutableBiMap<K, V> map, CodedOutputStream codedOut)
+      SerializationContext context, ImmutableBiMap map, CodedOutputStream codedOut)
       throws SerializationException, IOException {
-
     codedOut.writeInt32NoTag(map.size());
-    ImmutableMapCodec.serializeEntries(context, map.entrySet(), codedOut);
+    serializeMapEntries(context, map, codedOut);
   }
 
   @Override
-  public ImmutableBiMap<K, V> deserialize(DeserializationContext context, CodedInputStream codedIn)
+  public Supplier<ImmutableBiMap> deserializeDeferred(
+      AsyncDeserializationContext context, CodedInputStream codedIn)
       throws SerializationException, IOException {
+    int size = codedIn.readInt32();
+    if (size < 0) {
+      throw new SerializationException("Expected non-negative size: " + size);
+    }
 
-    int length = codedIn.readInt32();
-    if (length < 0) {
-      throw new SerializationException("Expected non-negative length: " + length);
+    if (size == 0) {
+      return ImmutableBiMap::of;
     }
 
-    ImmutableBiMap.Builder<K, V> builder =
-        ImmutableMapCodec.deserializeEntries(
-            ImmutableBiMap.builderWithExpectedSize(length), length, context, codedIn);
+    EntryBuffer buffer = new EntryBuffer(size);
+    deserializeMapEntries(
+        context, codedIn, /* requiresFullValueDeserialization= */ true, buffer.keys, buffer.values);
+    return buffer;
+  }
+
+  private static class EntryBuffer implements Supplier<ImmutableBiMap> {
+    final Object[] keys;
+    final Object[] values;
+
+    private EntryBuffer(int size) {
+      this.keys = new Object[size];
+      this.values = new Object[size];
+    }
+
+    int size() {
+      return keys.length;
+    }
 
-    try {
+    @Override
+    public ImmutableBiMap get() {
+      ImmutableBiMap.Builder builder = ImmutableBiMap.builderWithExpectedSize(size());
+      for (int i = 0; i < size(); i++) {
+        builder.put(keys[i], values[i]);
+      }
       return builder.buildOrThrow();
-    } catch (IllegalArgumentException e) {
-      throw new SerializationException(
-          "Duplicate keys during ImmutableBiMapCodec deserialization", e);
     }
   }
 }

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

@@ -11,14 +11,17 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-
 package com.google.devtools.build.lib.skyframe.serialization;
 
+import static com.google.devtools.build.lib.skyframe.serialization.MapHelpers.deserializeMapEntries;
+import static com.google.devtools.build.lib.skyframe.serialization.MapHelpers.serializeMapEntries;
+
 import com.google.common.collect.ImmutableClassToInstanceMap;
 import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
+import java.util.function.Supplier;
 
 /**
  * Encodes an {@link ImmutableClassToInstanceMap}. The iteration order of the deserialized map is
@@ -36,48 +39,64 @@
  * include the class name of the entry's value. Errors that occur while serializing an entry key are
  * not affected.
  */
-class ImmutableClassToInstanceMapCodec<B> implements ObjectCodec<ImmutableClassToInstanceMap<B>> {
+@SuppressWarnings({"unchecked", "rawtypes"})
+class ImmutableClassToInstanceMapCodec extends DeferredObjectCodec<ImmutableClassToInstanceMap> {
 
-  @SuppressWarnings("unchecked")
   @Override
-  public Class<ImmutableClassToInstanceMap<B>> getEncodedClass() {
-    return (Class<ImmutableClassToInstanceMap<B>>) ((Class<?>) ImmutableClassToInstanceMap.class);
+  public Class<ImmutableClassToInstanceMap> getEncodedClass() {
+    return ImmutableClassToInstanceMap.class;
   }
 
   @Override
   public void serialize(
-      SerializationContext context, ImmutableClassToInstanceMap<B> map, CodedOutputStream codedOut)
+      SerializationContext context, ImmutableClassToInstanceMap map, CodedOutputStream codedOut)
       throws SerializationException, IOException {
-
     codedOut.writeInt32NoTag(map.size());
-    ImmutableMapCodec.serializeEntries(context, map.entrySet(), codedOut);
+    serializeMapEntries(context, map, codedOut);
   }
 
   @Override
-  public ImmutableClassToInstanceMap<B> deserialize(
-      DeserializationContext context, CodedInputStream codedIn)
+  public Supplier<ImmutableClassToInstanceMap> deserializeDeferred(
+      AsyncDeserializationContext context, CodedInputStream codedIn)
       throws SerializationException, IOException {
-
-    int length = codedIn.readInt32();
-    if (length < 0) {
-      throw new SerializationException("Expected non-negative length: " + length);
+    int size = codedIn.readInt32();
+    if (size < 0) {
+      throw new SerializationException("Expected non-negative length: " + size);
+    }
+    if (size == 0) {
+      return ImmutableClassToInstanceMap::of;
     }
 
-    ImmutableMap.Builder<Class<? extends B>, B> classKeyedBuilder =
-        ImmutableMapCodec.deserializeEntries(ImmutableMap.builder(), length, context, codedIn);
+    EntryBuffer buffer = new EntryBuffer(size);
+    deserializeMapEntries(
+        context,
+        codedIn,
+        /* requiresFullValueDeserialization= */ false,
+        buffer.keys,
+        buffer.values);
+    return buffer;
+  }
+
+  private static class EntryBuffer implements Supplier<ImmutableClassToInstanceMap> {
+    final Object[] keys;
+    final Object[] values;
 
-    ImmutableMap<Class<? extends B>, B> classKeyedMap;
-    try {
-      classKeyedMap = classKeyedBuilder.buildOrThrow();
-    } catch (IllegalArgumentException e) {
-      throw new SerializationException(
-          "Duplicate keys during ImmutableClassToInstanceMapCodec deserialization", e);
+    private EntryBuffer(int size) {
+      this.keys = new Object[size];
+      this.values = new Object[size];
     }
-    try {
-      return ImmutableClassToInstanceMap.copyOf(classKeyedMap);
-    } catch (ClassCastException e) {
-      throw new SerializationException(
-          "Key-value type mismatch during ImmutableClassToInstanceMapCodec deserialization", e);
+
+    @Override
+    public ImmutableClassToInstanceMap get() {
+      ImmutableClassToInstanceMap.Builder builder = ImmutableClassToInstanceMap.builder();
+      for (int i = 0; i < size(); i++) {
+        builder.put((Class) keys[i], values[i]);
+      }
+      return builder.build();
+    }
+
+    private int size() {
+      return keys.length;
     }
   }
 }