Add context argument to ObjectCodec.{serialize,deserialize}

Context implementations are currently empty, just doing the plumbing in this
change. Once this is in we can start passing along the ObjectCodecRegistry, which
will allow runtime codec resolution for classes not known at compile time.
We'll also inevitably add some memoization helpers, allowing us to optimize the
serialization process further.

PiperOrigin-RevId: 185305674

1 files changed, 32 insertions, 18 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodec.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodec.java
index 072b936fcc..475e93b7ab 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodec.java
@@ -16,10 +16,12 @@ package com.google.devtools.build.lib.collect.nestedset;
 import com.google.common.base.Preconditions;
 import com.google.common.hash.Hashing;
 import com.google.common.hash.HashingOutputStream;
+import com.google.devtools.build.lib.skyframe.serialization.DeserializationContext;
 import com.google.devtools.build.lib.skyframe.serialization.EnumCodec;
 import com.google.devtools.build.lib.skyframe.serialization.InjectingObjectCodec;
 import com.google.devtools.build.lib.skyframe.serialization.InjectingObjectCodecAdapter;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
+import com.google.devtools.build.lib.skyframe.serialization.SerializationContext;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.CodedInputStream;
@@ -61,20 +63,20 @@ public class NestedSetCodec<T> implements ObjectCodec<NestedSet<T>> {
   }
 
   @Override
-  public void serialize(NestedSet<T> obj, CodedOutputStream codedOut)
+  public void serialize(SerializationContext context, NestedSet<T> obj, CodedOutputStream codedOut)
       throws SerializationException, IOException {
     // Topo sort the nested set to ensure digests are available for children at time of writing
     Collection<Object> topoSortedChildren = getTopologicallySortedChildren(obj);
     Map<Object, byte[]> childToDigest = new IdentityHashMap<>();
     codedOut.writeInt32NoTag(topoSortedChildren.size());
-    orderCodec.serialize(obj.getOrder(), codedOut);
+    orderCodec.serialize(context, obj.getOrder(), codedOut);
     for (Object children : topoSortedChildren) {
-      serializeOneNestedSet(children, codedOut, childToDigest);
+      serializeOneNestedSet(context, children, codedOut, childToDigest);
     }
   }
 
   @Override
-  public NestedSet<T> deserialize(CodedInputStream codedIn)
+  public NestedSet<T> deserialize(DeserializationContext context, CodedInputStream codedIn)
       throws SerializationException, IOException {
     Map<ByteString, Object> digestToChild = new HashMap<>();
     int nestedSetCount = codedIn.readInt32();
@@ -82,17 +84,20 @@ public class NestedSetCodec<T> implements ObjectCodec<NestedSet<T>> {
         nestedSetCount >= 1,
         "Should have at least serialized one nested set, got: %d",
         nestedSetCount);
-    Order order = orderCodec.deserialize(codedIn);
+    Order order = orderCodec.deserialize(context, codedIn);
     Object children = null;
     for (int i = 0; i < nestedSetCount; ++i) {
       // Update var, the last one in the list is the top level nested set
-      children = deserializeOneNestedSet(codedIn, digestToChild);
+      children = deserializeOneNestedSet(context, codedIn, digestToChild);
     }
     return createNestedSet(order, children);
   }
 
   private void serializeOneNestedSet(
-      Object children, CodedOutputStream codedOut, Map<Object, byte[]> childToDigest)
+      SerializationContext context,
+      Object children,
+      CodedOutputStream codedOut,
+      Map<Object, byte[]> childToDigest)
       throws IOException, SerializationException {
     // Serialize nested set into an inner byte array so we can take its digest
     ByteArrayOutputStream childOutputStream = new ByteArrayOutputStream();
@@ -100,9 +105,9 @@ public class NestedSetCodec<T> implements ObjectCodec<NestedSet<T>> {
         new HashingOutputStream(Hashing.md5(), childOutputStream);
     CodedOutputStream childCodedOut = CodedOutputStream.newInstance(hashingOutputStream);
     if (children instanceof Object[]) {
-      serializeMultiItemChildArray((Object[]) children, childToDigest, childCodedOut);
+      serializeMultiItemChildArray(context, (Object[]) children, childToDigest, childCodedOut);
     } else if (children != NestedSet.EMPTY_CHILDREN) {
-      serializeSingleItemChildArray(children, childCodedOut);
+      serializeSingleItemChildArray(context, children, childCodedOut);
     } else {
       // Empty set
       childCodedOut.writeInt32NoTag(0);
@@ -116,7 +121,10 @@ public class NestedSetCodec<T> implements ObjectCodec<NestedSet<T>> {
   }
 
   private void serializeMultiItemChildArray(
-      Object[] children, Map<Object, byte[]> childToDigest, CodedOutputStream childCodedOut)
+      SerializationContext context,
+      Object[] children,
+      Map<Object, byte[]> childToDigest,
+      CodedOutputStream childCodedOut)
       throws IOException, SerializationException {
     childCodedOut.writeInt32NoTag(children.length);
     for (Object child : children) {
@@ -129,20 +137,23 @@ public class NestedSetCodec<T> implements ObjectCodec<NestedSet<T>> {
         childCodedOut.writeByteArrayNoTag(digest);
       } else {
         childCodedOut.writeBoolNoTag(false);
-        objectCodec.serialize(cast(child), childCodedOut);
+        objectCodec.serialize(context, cast(child), childCodedOut);
       }
     }
   }
 
-  private void serializeSingleItemChildArray(Object children, CodedOutputStream childCodedOut)
+  private void serializeSingleItemChildArray(
+      SerializationContext context, Object children, CodedOutputStream childCodedOut)
       throws IOException, SerializationException {
     childCodedOut.writeInt32NoTag(1);
     T singleChild = cast(children);
-    objectCodec.serialize(singleChild, childCodedOut);
+    objectCodec.serialize(context, singleChild, childCodedOut);
   }
 
   private Object deserializeOneNestedSet(
-      CodedInputStream codedIn, Map<ByteString, Object> digestToChild)
+      DeserializationContext context,
+      CodedInputStream codedIn,
+      Map<ByteString, Object> digestToChild)
       throws SerializationException, IOException {
     ByteString digest = codedIn.readBytes();
     CodedInputStream childCodedIn = codedIn.readBytes().newCodedInput();
@@ -150,9 +161,9 @@ public class NestedSetCodec<T> implements ObjectCodec<NestedSet<T>> {
     int childCount = childCodedIn.readInt32();
     final Object result;
     if (childCount > 1) {
-      result = deserializeMultipleItemChildArray(digestToChild, childCodedIn, childCount);
+      result = deserializeMultipleItemChildArray(context, digestToChild, childCodedIn, childCount);
     } else if (childCount == 1) {
-      result = objectCodec.deserialize(childCodedIn);
+      result = objectCodec.deserialize(context, childCodedIn);
     } else {
       result = NestedSet.EMPTY_CHILDREN;
     }
@@ -161,7 +172,10 @@ public class NestedSetCodec<T> implements ObjectCodec<NestedSet<T>> {
   }
 
   private Object deserializeMultipleItemChildArray(
-      Map<ByteString, Object> digestToChild, CodedInputStream childCodedIn, int childCount)
+      DeserializationContext context,
+      Map<ByteString, Object> digestToChild,
+      CodedInputStream childCodedIn,
+      int childCount)
       throws IOException, SerializationException {
     Object[] children = new Object[childCount];
     for (int i = 0; i < childCount; ++i) {
@@ -171,7 +185,7 @@ public class NestedSetCodec<T> implements ObjectCodec<NestedSet<T>> {
         children[i] =
             Preconditions.checkNotNull(digestToChild.get(digest), "Transitive nested set missing");
       } else {
-        children[i] = objectCodec.deserialize(childCodedIn);
+        children[i] = objectCodec.deserialize(context, childCodedIn);
       }
     }
     return children;