Make ImmutableTableCodec a DeferredObjectCodec.

PiperOrigin-RevId: 596051669
Change-Id: Ibfa7b49e4a7dc3e2282a554a00d37586ab6ebb67

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

@@ -11,54 +11,87 @@
 // 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 sun.misc.Unsafe.ARRAY_OBJECT_BASE_OFFSET;
+import static sun.misc.Unsafe.ARRAY_OBJECT_INDEX_SCALE;
+
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableTable;
 import com.google.common.collect.Table.Cell;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
-import java.util.Set;
+import java.util.function.Supplier;
 
 /** Codec for {@link ImmutableTable}. */
-public class ImmutableTableCodec<R, C, V> implements ObjectCodec<ImmutableTable<R, C, V>> {
+@SuppressWarnings({"rawtypes", "unchecked"})
+public class ImmutableTableCodec extends DeferredObjectCodec<ImmutableTable> {
 
-  @SuppressWarnings("unchecked")
   @Override
-  public Class<ImmutableTable<R, C, V>> getEncodedClass() {
-    // Compiler doesn't like to do a direct cast.
-    return (Class<ImmutableTable<R, C, V>>) ((Class<?>) ImmutableTable.class);
+  public Class<ImmutableTable> getEncodedClass() {
+    return ImmutableTable.class;
   }
 
   @Override
   public void serialize(
-      SerializationContext context, ImmutableTable<R, C, V> object, CodedOutputStream codedOut)
+      SerializationContext context, ImmutableTable object, CodedOutputStream codedOut)
       throws SerializationException, IOException {
-    Set<Cell<R, C, V>> cellSet = object.cellSet();
+    ImmutableSet<Cell> cellSet = object.cellSet();
     codedOut.writeInt32NoTag(cellSet.size());
-    for (Cell<R, C, V> cell : cellSet) {
+    for (Cell cell : cellSet) {
       context.serialize(cell.getRowKey(), codedOut);
       context.serialize(cell.getColumnKey(), codedOut);
       context.serialize(cell.getValue(), codedOut);
     }
   }
 
   @Override
-  public ImmutableTable<R, C, V> deserialize(
-      DeserializationContext context, CodedInputStream codedIn)
+  public Supplier<ImmutableTable> 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 size: " + size);
+    }
+    if (size == 0) {
+      return ImmutableTable::of;
+    }
+
+    EntryBuffer buffer = new EntryBuffer(size);
+    long offset = ARRAY_OBJECT_BASE_OFFSET;
+    for (int i = 0; i < size; i++) {
+      context.deserializeFully(codedIn, buffer.rowKeys, offset);
+      context.deserializeFully(codedIn, buffer.columnKeys, offset);
+      context.deserialize(codedIn, buffer.values, offset);
+      offset += ARRAY_OBJECT_INDEX_SCALE;
+    }
+    return buffer;
+  }
+
+  private static final class EntryBuffer implements Supplier<ImmutableTable> {
+    private final Object[] rowKeys;
+    private final Object[] columnKeys;
+    private final Object[] values;
+
+    private EntryBuffer(int size) {
+      this.rowKeys = new Object[size];
+      this.columnKeys = new Object[size];
+      this.values = new Object[size];
     }
-    ImmutableTable.Builder<R, C, V> builder = ImmutableTable.builder();
-    for (int i = 0; i < length; i++) {
-      builder.put(
-          /*rowKey=*/ context.deserialize(codedIn),
-          /*columnKey=*/ context.deserialize(codedIn),
-          /*value=*/ context.deserialize(codedIn));
+
+    @Override
+    public ImmutableTable get() {
+      ImmutableTable.Builder builder = ImmutableTable.builder();
+      for (int i = 0; i < size(); i++) {
+        builder.put(
+            /* rowKey= */ rowKeys[i], /* columnKey= */ columnKeys[i], /* value= */ values[i]);
+      }
+      return builder.buildOrThrow();
+    }
+
+    private int size() {
+      return rowKeys.length;
     }
-    return builder.buildOrThrow();
   }
 }