Support array fields in AutoCodec. To do this, introduce

SerializationCodeGenerator, which is a generalization of Marshaller that
supports primitive and array values.

PiperOrigin-RevId: 182053617

4 files changed, 208 insertions, 48 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecProcessor.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecProcessor.java
index bbb3d3bca3..052aa242b2 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecProcessor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecProcessor.java
@@ -22,6 +22,7 @@ import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
 import com.google.devtools.build.lib.skyframe.serialization.PolymorphicHelper;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationCodeGenerator.Marshaller;
 import com.squareup.javapoet.JavaFile;
 import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.TypeName;
@@ -252,6 +253,18 @@ public class AutoCodecProcessor extends AbstractProcessor {
               UnsafeProvider.class,
               parameter.getSimpleName());
           break;
+        case ARRAY:
+          serializeBuilder.addStatement(
+              "$T unsafe_$L = ($T)$T.getInstance().getObject(input, $L_offset)",
+              field.asType(),
+              parameter.getSimpleName(),
+              field.asType(),
+              UnsafeProvider.class,
+              parameter.getSimpleName());
+          marshallers.writeSerializationCode(
+              new Marshaller.Context(
+                  serializeBuilder, parameter.asType(), "unsafe_" + parameter.getSimpleName()));
+          break;
         case DECLARED:
           serializeBuilder.addStatement(
               "$T unsafe_$L = ($T)$T.getInstance().getObject(input, $L_offset)",
@@ -316,6 +329,10 @@ public class AutoCodecProcessor extends AbstractProcessor {
         case INT:
           serializeBuilder.addStatement("codedOut.writeInt32NoTag($L)", paramAccessor);
           break;
+        case ARRAY:
+          marshallers.writeSerializationCode(
+              new Marshaller.Context(serializeBuilder, parameter.asType(), paramAccessor));
+          break;
         case DECLARED:
           marshallers.writeSerializationCode(
               new Marshaller.Context(
@@ -347,6 +364,10 @@ public class AutoCodecProcessor extends AbstractProcessor {
         case INT:
           builder.addStatement("int $L = codedIn.readInt32()", paramName);
           break;
+        case ARRAY:
+          marshallers.writeDeserializationCode(
+              new Marshaller.Context(builder, parameter.asType(), paramName));
+          break;
         case DECLARED:
           marshallers.writeDeserializationCode(
               new Marshaller.Context(builder, (DeclaredType) parameter.asType(), paramName));
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/BUILD b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/BUILD
index da038b425e..d18a8f0738 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/BUILD
@@ -46,8 +46,8 @@ java_library(
     srcs = [
         "AutoCodecProcessor.java",
         "AutoCodecUtil.java",
-        "Marshaller.java",
         "Marshallers.java",
+        "SerializationCodeGenerator.java",
     ],
     deps = [
         ":autocodec-annotation",
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/Marshallers.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/Marshallers.java
index 81f836472e..be28995d29 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/Marshallers.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/Marshallers.java
@@ -25,7 +25,9 @@ import com.google.common.collect.Maps;
 import com.google.common.hash.HashCode;
 import com.google.devtools.build.lib.skyframe.serialization.InjectingObjectCodec;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
-import com.google.devtools.build.lib.skyframe.serialization.autocodec.Marshaller.Context;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationCodeGenerator.Context;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationCodeGenerator.Marshaller;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationCodeGenerator.PrimitiveValueSerializationCodeGenerator;
 import com.google.devtools.build.lib.skyframe.serialization.strings.StringCodecs;
 import com.google.protobuf.AbstractMessage;
 import com.google.protobuf.ExtensionRegistryLite;
@@ -42,7 +44,10 @@ import javax.lang.model.element.Element;
 import javax.lang.model.element.ElementKind;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
+import javax.lang.model.type.ArrayType;
 import javax.lang.model.type.DeclaredType;
+import javax.lang.model.type.PrimitiveType;
+import javax.lang.model.type.TypeKind;
 import javax.lang.model.type.TypeMirror;
 
 /** Class containing all {@link Marshaller} instances. */
@@ -54,32 +59,129 @@ class Marshallers {
   }
 
   void writeSerializationCode(Context context) {
-    context.builder.beginControlFlow("if ($L != null)", context.name);
-    context.builder.addStatement("codedOut.writeBoolNoTag(true)");
-    getMatchingMarshaller(context.type).addSerializationCode(context);
-    context.builder.nextControlFlow("else");
-    context.builder.addStatement("codedOut.writeBoolNoTag(false)");
-    context.builder.endControlFlow();
+    if (context.canBeNull()) {
+      context.builder.beginControlFlow("if ($L != null)", context.name);
+      context.builder.addStatement("codedOut.writeBoolNoTag(true)");
+    }
+    getMatchingCodeGenerator(context.type).addSerializationCode(context);
+    if (context.canBeNull()) {
+      context.builder.nextControlFlow("else");
+      context.builder.addStatement("codedOut.writeBoolNoTag(false)");
+      context.builder.endControlFlow();
+    }
   }
 
   void writeDeserializationCode(Context context) {
-    context.builder.addStatement("$T $L = null", context.getTypeName(), context.name);
-    context.builder.beginControlFlow("if (codedIn.readBool())");
-    getMatchingMarshaller(context.type).addDeserializationCode(context);
-    context.builder.endControlFlow();
+    if (context.canBeNull()) {
+      context.builder.addStatement("$T $L = null", context.getTypeName(), context.name);
+      context.builder.beginControlFlow("if (codedIn.readBool())");
+    } else {
+      context.builder.addStatement("$T $L", context.getTypeName(), context.name);
+    }
+    getMatchingCodeGenerator(context.type).addDeserializationCode(context);
+    if (requiresNullityCheck(context)) {
+      context.builder.endControlFlow();
+    }
   }
 
-  private Marshaller getMatchingMarshaller(DeclaredType type) {
+  private static boolean requiresNullityCheck(Context context) {
+    return !(context.type instanceof PrimitiveType);
+  }
+
+  private SerializationCodeGenerator getMatchingCodeGenerator(TypeMirror type) {
+    if (type.getKind() == TypeKind.ARRAY) {
+      return arrayCodeGenerator;
+    }
+
+    if (type instanceof PrimitiveType) {
+      PrimitiveType primitiveType = (PrimitiveType) type;
+      return primitiveGenerators
+          .stream()
+          .filter(generator -> generator.matches((PrimitiveType) type))
+          .findFirst()
+          .orElseThrow(() -> new IllegalArgumentException("No generator for: " + primitiveType));
+    }
+
+    // TODO(cpeyser): Refactor primitive handling from AutoCodecProcessor.java
+    if (!(type instanceof DeclaredType)) {
+      throw new IllegalArgumentException(
+          "Can only serialize primitive, array or declared fields, found " + type);
+    }
+    DeclaredType declaredType = (DeclaredType) type;
+
     return marshallers
         .stream()
-        .filter(marshaller -> marshaller.matches(type))
+        .filter(marshaller -> marshaller.matches(declaredType))
         .findFirst()
         .orElseThrow(
             () ->
                 new IllegalArgumentException(
-                    "No marshaller for: " + ((TypeElement) type.asElement()).getQualifiedName()));
+                    "No marshaller for: "
+                        + ((TypeElement) declaredType.asElement()).getQualifiedName()));
   }
 
+  private final SerializationCodeGenerator arrayCodeGenerator =
+      new SerializationCodeGenerator() {
+        @Override
+        public void addSerializationCode(Context context) {
+          String length = context.makeName("length");
+          context.builder.addStatement("int $L = $L.length", length, context.name);
+          context.builder.addStatement("codedOut.writeInt32NoTag($L)", length);
+          Context repeated =
+              context.with(
+                  ((ArrayType) context.type).getComponentType(), context.makeName("repeated"));
+          String indexName = context.makeName("i");
+          context.builder.beginControlFlow(
+              "for(int $L = 0; $L < $L; ++$L)", indexName, indexName, length, indexName);
+          context.builder.addStatement(
+              "$T $L = $L[$L]", repeated.getTypeName(), repeated.name, context.name, indexName);
+          writeSerializationCode(repeated);
+          context.builder.endControlFlow();
+        }
+
+        @Override
+        public void addDeserializationCode(Context context) {
+          Context repeated =
+              context.with(
+                  ((ArrayType) context.type).getComponentType(), context.makeName("repeated"));
+          String lengthName = context.makeName("length");
+          context.builder.addStatement("int $L = codedIn.readInt32()", lengthName);
+
+          String resultName = context.makeName("result");
+          context.builder.addStatement(
+              "$T[] $L = new $T[$L]",
+              repeated.getTypeName(),
+              resultName,
+              repeated.getTypeName(),
+              lengthName);
+          String indexName = context.makeName("i");
+          context.builder.beginControlFlow(
+              "for (int $L = 0; $L < $L; ++$L)", indexName, indexName, lengthName, indexName);
+          writeDeserializationCode(repeated);
+          context.builder.addStatement("$L[$L] = $L", resultName, indexName, repeated.name);
+          context.builder.endControlFlow();
+          context.builder.addStatement("$L = $L", context.name, resultName);
+        }
+      };
+
+  private final PrimitiveValueSerializationCodeGenerator intCodeGenerator =
+      new PrimitiveValueSerializationCodeGenerator() {
+        @Override
+        public boolean matches(PrimitiveType type) {
+          return type.getKind() == TypeKind.INT;
+        }
+
+        @Override
+        public void addSerializationCode(Context context) {
+          context.builder.addStatement("codedOut.writeInt32NoTag($L)", context.name);
+        }
+
+        @Override
+        public void addDeserializationCode(Context context) {
+          context.builder.addStatement("$L = codedIn.readInt32()", context.name);
+        }
+      };
+
   private final Marshaller enumMarshaller =
       new Marshaller() {
         @Override
@@ -89,7 +191,7 @@ class Marshallers {
 
         @Override
         public void addSerializationCode(Context context) {
-          if (isProtoEnum(context.type)) {
+          if (isProtoEnum(context.getDeclaredType())) {
             context.builder.addStatement("codedOut.writeInt32NoTag($L.getNumber())", context.name);
           } else {
             context.builder.addStatement("codedOut.writeInt32NoTag($L.ordinal())", context.name);
@@ -98,7 +200,7 @@ class Marshallers {
 
         @Override
         public void addDeserializationCode(Context context) {
-          if (isProtoEnum(context.type)) {
+          if (isProtoEnum(context.getDeclaredType())) {
             context.builder.addStatement(
                 "$L = $T.forNumber(codedIn.readInt32())", context.name, context.getTypeName());
           } else {
@@ -147,13 +249,15 @@ class Marshallers {
 
         @Override
         public void addSerializationCode(Context context) {
-          DeclaredType optionalType = (DeclaredType) context.type.getTypeArguments().get(0);
+          DeclaredType optionalType =
+              (DeclaredType) context.getDeclaredType().getTypeArguments().get(0);
           writeSerializationCode(context.with(optionalType, context.name + ".orNull()"));
         }
 
         @Override
         public void addDeserializationCode(Context context) {
-          DeclaredType optionalType = (DeclaredType) context.type.getTypeArguments().get(0);
+          DeclaredType optionalType =
+              (DeclaredType) context.getDeclaredType().getTypeArguments().get(0);
           String optionalName = context.makeName("optional");
           writeDeserializationCode(context.with(optionalType, optionalName));
           context.builder.addStatement(
@@ -170,18 +274,20 @@ class Marshallers {
 
         @Override
         public void addSerializationCode(Context context) {
-          DeclaredType keyType = (DeclaredType) context.type.getTypeArguments().get(0);
+          DeclaredType keyType = (DeclaredType) context.getDeclaredType().getTypeArguments().get(0);
           writeSerializationCode(context.with(keyType, context.name + ".getKey()"));
-          DeclaredType valueType = (DeclaredType) context.type.getTypeArguments().get(1);
+          DeclaredType valueType =
+              (DeclaredType) context.getDeclaredType().getTypeArguments().get(1);
           writeSerializationCode(context.with(valueType, context.name + ".getValue()"));
         }
 
         @Override
         public void addDeserializationCode(Context context) {
-          DeclaredType keyType = (DeclaredType) context.type.getTypeArguments().get(0);
+          DeclaredType keyType = (DeclaredType) context.getDeclaredType().getTypeArguments().get(0);
           String keyName = context.makeName("key");
           writeDeserializationCode(context.with(keyType, keyName));
-          DeclaredType valueType = (DeclaredType) context.type.getTypeArguments().get(1);
+          DeclaredType valueType =
+              (DeclaredType) context.getDeclaredType().getTypeArguments().get(1);
           String valueName = context.makeName("value");
           writeDeserializationCode(context.with(valueType, valueName));
           context.builder.addStatement(
@@ -206,7 +312,7 @@ class Marshallers {
           context.builder.addStatement("codedOut.writeInt32NoTag($L.size())", context.name);
           Context repeated =
               context.with(
-                  (DeclaredType) context.type.getTypeArguments().get(0),
+                  (DeclaredType) context.getDeclaredType().getTypeArguments().get(0),
                   context.makeName("repeated"));
           context.builder.beginControlFlow(
               "for ($T $L : $L)", repeated.getTypeName(), repeated.name, context.name);
@@ -218,7 +324,7 @@ class Marshallers {
         public void addDeserializationCode(Context context) {
           Context repeated =
               context.with(
-                  (DeclaredType) context.type.getTypeArguments().get(0),
+                  (DeclaredType) context.getDeclaredType().getTypeArguments().get(0),
                   context.makeName("repeated"));
           String builderName = context.makeName("builder");
           context.builder.addStatement(
@@ -255,7 +361,7 @@ class Marshallers {
         public void addDeserializationCode(Context context) {
           Context repeated =
               context.with(
-                  (DeclaredType) context.type.getTypeArguments().get(0),
+                  (DeclaredType) context.getDeclaredType().getTypeArguments().get(0),
                   context.makeName("repeated"));
           String builderName = context.makeName("builder");
           context.builder.addStatement(
@@ -291,10 +397,12 @@ class Marshallers {
           String entryName = context.makeName("entry");
           Context key =
               context.with(
-                  (DeclaredType) context.type.getTypeArguments().get(0), entryName + ".getKey()");
+                  (DeclaredType) context.getDeclaredType().getTypeArguments().get(0),
+                  entryName + ".getKey()");
           Context value =
               context.with(
-                  (DeclaredType) context.type.getTypeArguments().get(1), entryName + ".getValue()");
+                  (DeclaredType) context.getDeclaredType().getTypeArguments().get(1),
+                  entryName + ".getValue()");
           context.builder.addStatement(
               "$T<$T, $T> $L = null", Map.class, key.getTypeName(), value.getTypeName(), mapName);
           context.builder.beginControlFlow("if ($L instanceof $T)", context.name, SortedMap.class);
@@ -347,10 +455,12 @@ class Marshallers {
       Context context, String builderName, boolean isImmutableMap) {
           Context key =
               context.with(
-                  (DeclaredType) context.type.getTypeArguments().get(0), context.makeName("key"));
+                  (DeclaredType) context.getDeclaredType().getTypeArguments().get(0),
+                  context.makeName("key"));
           Context value =
               context.with(
-                  (DeclaredType) context.type.getTypeArguments().get(1), context.makeName("value"));
+                  (DeclaredType) context.getDeclaredType().getTypeArguments().get(1),
+                  context.makeName("value"));
     if (isImmutableMap) {
       context.builder.addStatement(
           "$T<$T, $T> $L = new $T<>($T.naturalOrder())",
@@ -388,10 +498,12 @@ class Marshallers {
           String entryName = context.makeName("entry");
           Context key =
               context.with(
-                  (DeclaredType) context.type.getTypeArguments().get(0), entryName + ".getKey()");
+                  (DeclaredType) context.getDeclaredType().getTypeArguments().get(0),
+                  entryName + ".getKey()");
           Context value =
               context.with(
-                  (DeclaredType) context.type.getTypeArguments().get(1), entryName + ".getValue()");
+                  (DeclaredType) context.getDeclaredType().getTypeArguments().get(1),
+                  entryName + ".getValue()");
           context.builder.beginControlFlow(
               "for ($T<$T, $T> $L : $L.entries())",
               Map.Entry.class,
@@ -408,10 +520,12 @@ class Marshallers {
         public void addDeserializationCode(Context context) {
           Context key =
               context.with(
-                  (DeclaredType) context.type.getTypeArguments().get(0), context.makeName("key"));
+                  (DeclaredType) context.getDeclaredType().getTypeArguments().get(0),
+                  context.makeName("key"));
           Context value =
               context.with(
-                  (DeclaredType) context.type.getTypeArguments().get(1), context.makeName("value"));
+                  (DeclaredType) context.getDeclaredType().getTypeArguments().get(1),
+                  context.makeName("value"));
           String builderName = context.makeName("builder");
           context.builder.addStatement(
               "$T<$T, $T> $L = new $T<>()",
@@ -517,7 +631,7 @@ class Marshallers {
 
         @Override
         public void addSerializationCode(Context context) {
-          TypeMirror codecType = getCodec(context.type).get().asType();
+          TypeMirror codecType = getCodec(context.getDeclaredType()).get().asType();
           if (isSubtypeErased(codecType, ObjectCodec.class)) {
             context.builder.addStatement(
                 "$T.CODEC.serialize($L, codedOut)", context.getTypeName(), context.name);
@@ -529,14 +643,14 @@ class Marshallers {
           } else {
             throw new IllegalArgumentException(
                 "CODEC field of "
-                    + ((TypeElement) context.type.asElement()).getQualifiedName()
+                    + ((TypeElement) context.getDeclaredType().asElement()).getQualifiedName()
                     + " is neither ObjectCodec nor InjectingCodec");
           }
         }
 
         @Override
         public void addDeserializationCode(Context context) {
-          TypeMirror codecType = getCodec(context.type).get().asType();
+          TypeMirror codecType = getCodec(context.getDeclaredType()).get().asType();
           if (isSubtypeErased(codecType, ObjectCodec.class)) {
             context.builder.addStatement(
                 "$L = $T.CODEC.deserialize(codedIn)", context.name, context.getTypeName());
@@ -548,12 +662,15 @@ class Marshallers {
           } else {
             throw new IllegalArgumentException(
                 "CODEC field of "
-                    + ((TypeElement) context.type.asElement()).getQualifiedName()
+                    + ((TypeElement) context.getDeclaredType().asElement()).getQualifiedName()
                     + " is neither ObjectCodec nor InjectingCodec");
           }
         }
       };
 
+  private final ImmutableList<PrimitiveValueSerializationCodeGenerator> primitiveGenerators =
+      ImmutableList.of(intCodeGenerator);
+
   private final ImmutableList<Marshaller> marshallers =
       ImmutableList.of(
           enumMarshaller,
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/Marshaller.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/SerializationCodeGenerator.java
index ad4b9d7878..94690393ab 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/Marshaller.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/SerializationCodeGenerator.java
@@ -14,22 +14,25 @@
 
 package com.google.devtools.build.lib.skyframe.serialization.autocodec;
 
+import com.google.common.base.Preconditions;
 import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.TypeName;
 import javax.annotation.processing.ProcessingEnvironment;
 import javax.lang.model.type.DeclaredType;
+import javax.lang.model.type.PrimitiveType;
+import javax.lang.model.type.TypeMirror;
 
 /**
- * Generates serialize and deserialize code fragments for matching types.
+ * Generates serialize and deserialize code fragments.
  *
  * <p>All methods are logically static and take the {@link ProcessingEnvironment} as a parameter.
  */
-interface Marshaller {
-  static class Context {
+interface SerializationCodeGenerator {
+  class Context {
     /** Builder for the method. */
     public final MethodSpec.Builder builder;
     /** Type of {@code name}. */
-    public final DeclaredType type;
+    public final TypeMirror type;
     /** Name of variable. */
     public final String name;
     /**
@@ -39,11 +42,11 @@ interface Marshaller {
      */
     public final int depth;
 
-    Context(MethodSpec.Builder builder, DeclaredType type, String name) {
+    Context(MethodSpec.Builder builder, TypeMirror type, String name) {
       this(builder, type, name, 0);
     }
 
-    private Context(MethodSpec.Builder builder, DeclaredType type, String name, int depth) {
+    private Context(MethodSpec.Builder builder, TypeMirror type, String name, int depth) {
       this.builder = builder;
       this.type = type;
       this.name = name;
@@ -51,7 +54,7 @@ interface Marshaller {
     }
 
     /** Returns a new context with a new type and name at the next recursion depth. */
-    Context with(DeclaredType newType, String newName) {
+    Context with(TypeMirror newType, String newName) {
       return new Context(builder, newType, newName, depth + 1);
     }
 
@@ -59,6 +62,16 @@ interface Marshaller {
       return TypeName.get(type);
     }
 
+    DeclaredType getDeclaredType() {
+      Preconditions.checkState(type instanceof DeclaredType, "Expected DeclaredType, was " + type);
+      return (DeclaredType) type;
+    }
+
+    /** Returns true if this Context represents a type that can be null */
+    boolean canBeNull() {
+      return !(type instanceof PrimitiveType);
+    }
+
     /**
      * Returns a depth-qualified name.
      *
@@ -69,12 +82,21 @@ interface Marshaller {
     }
   };
 
-  /** Returns true if `type` is handled by this. */
-  boolean matches(DeclaredType type);
-
   /** Appends code statements to serialize a pre-declared variable. */
   void addSerializationCode(Context context);
 
   /** Appends code statements to initialize the pre-declared variable with deserialization. */
   void addDeserializationCode(Context context);
+
+  /** A {@link SerializationCodeGenerator} for a particular declared type. */
+  interface Marshaller extends SerializationCodeGenerator {
+    /** Returns true if {@code type} is handled by this. */
+    boolean matches(DeclaredType type);
+  }
+
+  /** A {@link SerializationCodeGenerator} for primitive values. */
+  interface PrimitiveValueSerializationCodeGenerator extends SerializationCodeGenerator {
+    /** Returns true if {@code type} is handled by this. */
+    boolean matches(PrimitiveType type);
+  }
 }