Add support for parameterized types to AutoCodec.

AutoCodec still cannnot handle types with generic elements, e.g.

class Foo<T> {
   private T member;
}

PiperOrigin-RevId: 187052487

2 files changed, 41 insertions, 25 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 0287117121..41f28835c7 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
@@ -103,14 +103,15 @@ public class AutoCodecProcessor extends AbstractProcessor {
           codecClassBuilder = buildClassWithPublicFieldsStrategy(encodedType);
           break;
         case SINGLETON:
-          codecClassBuilder = buildClassWithSingletonStrategy(encodedType);
+          codecClassBuilder = buildClassWithSingletonStrategy(encodedType, env);
           break;
         default:
           throw new IllegalArgumentException("Unknown strategy: " + annotation.strategy());
       }
       codecClassBuilder.addMethod(
-          AutoCodecUtil.initializeGetEncodedClassMethod(encodedType)
-              .addStatement("return $T.class", TypeName.get(encodedType.asType()))
+          AutoCodecUtil.initializeGetEncodedClassMethod(encodedType, env)
+              .addStatement(
+                  "return $T.class", TypeName.get(env.getTypeUtils().erasure(encodedType.asType())))
               .build());
       String packageName =
           env.getElementUtils().getPackageOf(encodedType).getQualifiedName().toString();
@@ -133,7 +134,8 @@ public class AutoCodecProcessor extends AbstractProcessor {
     ExecutableElement constructor = selectInstantiator(encodedType);
     List<? extends VariableElement> fields = constructor.getParameters();
 
-    TypeSpec.Builder codecClassBuilder = AutoCodecUtil.initializeCodecClassBuilder(encodedType);
+    TypeSpec.Builder codecClassBuilder =
+        AutoCodecUtil.initializeCodecClassBuilder(encodedType, env);
 
     if (encodedType.getAnnotation(AutoValue.class) == null) {
       initializeUnsafeOffsets(codecClassBuilder, encodedType, fields);
@@ -144,9 +146,9 @@ public class AutoCodecProcessor extends AbstractProcessor {
     }
 
     MethodSpec.Builder deserializeBuilder =
-        AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType);
+        AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType, env);
     buildDeserializeBody(deserializeBuilder, fields);
-    addReturnNew(deserializeBuilder, encodedType, constructor);
+    addReturnNew(deserializeBuilder, encodedType, constructor, env);
     codecClassBuilder.addMethod(deserializeBuilder.build());
 
     return codecClassBuilder;
@@ -200,7 +202,7 @@ public class AutoCodecProcessor extends AbstractProcessor {
   private MethodSpec buildSerializeMethodWithInstantiator(
       TypeElement encodedType, List<? extends VariableElement> fields) {
     MethodSpec.Builder serializeBuilder =
-        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType);
+        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType, env);
     for (VariableElement parameter : fields) {
       Optional<FieldValueAndClass> hasField =
           getFieldByNameRecursive(encodedType, parameter.getSimpleName().toString());
@@ -285,7 +287,7 @@ public class AutoCodecProcessor extends AbstractProcessor {
   private MethodSpec buildSerializeMethodWithInstantiatorForAutoValue(
       TypeElement encodedType, List<? extends VariableElement> fields) {
     MethodSpec.Builder serializeBuilder =
-        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType);
+        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType, env);
     for (VariableElement parameter : fields) {
       addSerializeParameterWithGetter(encodedType, parameter, serializeBuilder);
     }
@@ -293,7 +295,8 @@ public class AutoCodecProcessor extends AbstractProcessor {
   }
 
   private TypeSpec.Builder buildClassWithPublicFieldsStrategy(TypeElement encodedType) {
-    TypeSpec.Builder codecClassBuilder = AutoCodecUtil.initializeCodecClassBuilder(encodedType);
+    TypeSpec.Builder codecClassBuilder =
+        AutoCodecUtil.initializeCodecClassBuilder(encodedType, env);
     ImmutableList<? extends VariableElement> publicFields =
         ElementFilter.fieldsIn(env.getElementUtils().getAllMembers(encodedType))
             .stream()
@@ -301,7 +304,7 @@ public class AutoCodecProcessor extends AbstractProcessor {
             .collect(toImmutableList());
     codecClassBuilder.addMethod(buildSerializeMethodWithPublicFields(encodedType, publicFields));
     MethodSpec.Builder deserializeBuilder =
-        AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType);
+        AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType, env);
     buildDeserializeBody(deserializeBuilder, publicFields);
     addInstantiatePopulateFieldsAndReturn(deserializeBuilder, encodedType, publicFields);
     codecClassBuilder.addMethod(deserializeBuilder.build());
@@ -319,7 +322,7 @@ public class AutoCodecProcessor extends AbstractProcessor {
   private MethodSpec buildSerializeMethodWithPublicFields(
       TypeElement encodedType, List<? extends VariableElement> fields) {
     MethodSpec.Builder serializeBuilder =
-        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType);
+        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType, env);
     for (VariableElement parameter : fields) {
       String paramAccessor = "input." + parameter.getSimpleName();
       marshallers.writeSerializationCode(
@@ -350,7 +353,10 @@ public class AutoCodecProcessor extends AbstractProcessor {
    * <p>Used by the {@link AutoCodec.Strategy#INSTANTIATOR} strategy.
    */
   private static void addReturnNew(
-      MethodSpec.Builder builder, TypeElement type, ExecutableElement instantiator) {
+      MethodSpec.Builder builder,
+      TypeElement type,
+      ExecutableElement instantiator,
+      ProcessingEnvironment env) {
     List<? extends TypeMirror> allThrown = instantiator.getThrownTypes();
     if (!allThrown.isEmpty()) {
       builder.beginControlFlow("try");
@@ -362,7 +368,8 @@ public class AutoCodecProcessor extends AbstractProcessor {
             .map(AutoCodecProcessor::handleFromParameter)
             .collect(Collectors.joining(", "));
     if (instantiator.getKind().equals(ElementKind.CONSTRUCTOR)) {
-      builder.addStatement("return new $T($L)", TypeName.get(type.asType()), parameters);
+      builder.addStatement(
+          "return new $T($L)", TypeName.get(env.getTypeUtils().erasure(type.asType())), parameters);
     } else { // Otherwise, it's a factory method.
       builder.addStatement(
           "return $T.$L($L)",
@@ -473,13 +480,15 @@ public class AutoCodecProcessor extends AbstractProcessor {
     return Optional.empty();
   }
 
-  private static TypeSpec.Builder buildClassWithSingletonStrategy(TypeElement encodedType) {
-    TypeSpec.Builder codecClassBuilder = AutoCodecUtil.initializeCodecClassBuilder(encodedType);
+  private static TypeSpec.Builder buildClassWithSingletonStrategy(
+      TypeElement encodedType, ProcessingEnvironment env) {
+    TypeSpec.Builder codecClassBuilder =
+        AutoCodecUtil.initializeCodecClassBuilder(encodedType, env);
     // Serialization is a no-op.
     codecClassBuilder.addMethod(
-        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType).build());
+        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType, env).build());
     MethodSpec.Builder deserializeMethodBuilder =
-        AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType);
+        AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType, env);
     deserializeMethodBuilder.addStatement("return $T.INSTANCE", TypeName.get(encodedType.asType()));
     codecClassBuilder.addMethod(deserializeMethodBuilder.build());
     return codecClassBuilder;
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecUtil.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecUtil.java
index 610badc465..8adcddf7d7 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecUtil.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecUtil.java
@@ -28,6 +28,7 @@ import com.squareup.javapoet.TypeName;
 import com.squareup.javapoet.TypeSpec;
 import java.io.IOException;
 import java.util.stream.Collectors;
+import javax.annotation.processing.ProcessingEnvironment;
 import javax.lang.model.element.Element;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
@@ -43,20 +44,24 @@ class AutoCodecUtil {
    *
    * @param encodedType type being serialized
    */
-  static TypeSpec.Builder initializeCodecClassBuilder(TypeElement encodedType) {
+  static TypeSpec.Builder initializeCodecClassBuilder(
+      TypeElement encodedType, ProcessingEnvironment env) {
     TypeSpec.Builder builder = TypeSpec.classBuilder(getCodecName(encodedType));
     return builder.addSuperinterface(
         ParameterizedTypeName.get(
-            ClassName.get(ObjectCodec.class), TypeName.get(encodedType.asType())));
+            ClassName.get(ObjectCodec.class),
+            TypeName.get(env.getTypeUtils().erasure(encodedType.asType()))));
   }
 
-  static MethodSpec.Builder initializeGetEncodedClassMethod(TypeElement encodedType) {
+  static MethodSpec.Builder initializeGetEncodedClassMethod(
+      TypeElement encodedType, ProcessingEnvironment env) {
     return MethodSpec.methodBuilder("getEncodedClass")
         .addModifiers(Modifier.PUBLIC)
         .addAnnotation(Override.class)
         .returns(
             ParameterizedTypeName.get(
-                ClassName.get(Class.class), TypeName.get(encodedType.asType())));
+                ClassName.get(Class.class),
+                TypeName.get(env.getTypeUtils().erasure(encodedType.asType()))));
   }
 
   /**
@@ -64,7 +69,8 @@ class AutoCodecUtil {
    *
    * @param encodedType type being serialized
    */
-  static MethodSpec.Builder initializeSerializeMethodBuilder(TypeElement encodedType) {
+  static MethodSpec.Builder initializeSerializeMethodBuilder(
+      TypeElement encodedType, ProcessingEnvironment env) {
     MethodSpec.Builder builder =
         MethodSpec.methodBuilder("serialize")
             .addModifiers(Modifier.PUBLIC)
@@ -74,7 +80,7 @@ class AutoCodecUtil {
             .addException(IOException.class);
     return builder
         .addParameter(SerializationContext.class, "context")
-        .addParameter(TypeName.get(encodedType.asType()), "input")
+        .addParameter(TypeName.get(env.getTypeUtils().erasure(encodedType.asType())), "input")
         .addParameter(CodedOutputStream.class, "codedOut");
   }
 
@@ -83,11 +89,12 @@ class AutoCodecUtil {
    *
    * @param encodedType type being serialized
    */
-  static MethodSpec.Builder initializeDeserializeMethodBuilder(TypeElement encodedType) {
+  static MethodSpec.Builder initializeDeserializeMethodBuilder(
+      TypeElement encodedType, ProcessingEnvironment env) {
     MethodSpec.Builder builder =
         MethodSpec.methodBuilder("deserialize")
             .addModifiers(Modifier.PUBLIC)
-            .returns(TypeName.get(encodedType.asType()))
+            .returns(TypeName.get(env.getTypeUtils().erasure(encodedType.asType())))
             .addAnnotation(Override.class)
             .addException(SerializationException.class)
             .addException(IOException.class);