Add support to Polymorphic strategy for grandchild classes that have no codec, but whose parent classes have codecs. This is ok because the polymorphic strategy doesn't need an instance of the grandchild class: the parent class is fine, so long as it has a codec.

PiperOrigin-RevId: 185200943

2 files changed, 52 insertions, 10 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/PolymorphicHelper.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/PolymorphicHelper.java
index a325e7ea96..11689bade6 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/PolymorphicHelper.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/PolymorphicHelper.java
@@ -36,12 +36,16 @@ public class PolymorphicHelper {
    */
   @SuppressWarnings("unchecked")
   public static void serialize(
-      Object input, CodedOutputStream codedOut, @Nullable Optional<?> dependency)
+      Object input,
+      Class<?> baseClass,
+      CodedOutputStream codedOut,
+      @Nullable Optional<?> dependency)
       throws IOException, SerializationException {
     if (input != null) {
-      Class<?> clazz = input.getClass();
       try {
-        Object codec = getCodec(clazz);
+        ClassAndCodec classAndCodec = getClassAndCodecInClassHierarchy(input, baseClass);
+        Class<?> clazz = classAndCodec.clazz;
+        Object codec = classAndCodec.codec;
         codedOut.writeBoolNoTag(true);
         StringCodecs.asciiOptimized().serialize(clazz.getName(), codedOut);
         if (codec instanceof ObjectCodec) {
@@ -100,11 +104,46 @@ public class PolymorphicHelper {
     return deserialized;
   }
 
+  private static ClassAndCodec getClassAndCodecInClassHierarchy(Object input, Class<?> baseClass)
+      throws SerializationException, IllegalAccessException {
+    Class<?> clazz = input.getClass();
+    Field codecField = null;
+    while (!clazz.equals(baseClass)) {
+      try {
+        codecField = getCodecField(clazz);
+        break;
+      } catch (NoSuchFieldException e) {
+        clazz = clazz.getSuperclass();
+      }
+    }
+    if (clazz.equals(baseClass)) {
+      throw new SerializationException("Could not find codec for " + input.getClass());
+    }
+    return new ClassAndCodec(clazz, getValueOfField(codecField));
+  }
+
+  private static Field getCodecField(Class<?> clazz) throws NoSuchFieldException {
+    return clazz.getDeclaredField("CODEC");
+  }
+
+  private static Object getValueOfField(Field field) throws IllegalAccessException {
+    field.setAccessible(true);
+    return field.get(null);
+  }
+
   /** Returns the static CODEC instance for {@code clazz}. */
   private static Object getCodec(Class<?> clazz)
       throws NoSuchFieldException, IllegalAccessException {
-    Field codecField = clazz.getDeclaredField("CODEC");
-    codecField.setAccessible(true);
-    return codecField.get(null);
+    return getValueOfField(getCodecField(clazz));
+  }
+
+  private static class ClassAndCodec {
+    final Class<?> clazz;
+    final Object codec;
+
+    ClassAndCodec(Class<?> clazz, Object codec) {
+      this.clazz = clazz;
+      this.codec = codec;
+    }
   }
 }
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 c576f0c52c..ff99931c8e 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
@@ -596,7 +596,7 @@ public class AutoCodecProcessor extends AbstractProcessor {
     return Optional.empty();
   }
 
-  private static TypeSpec.Builder buildClassWithPolymorphicStrategy(
+  private TypeSpec.Builder buildClassWithPolymorphicStrategy(
       TypeElement encodedType, @Nullable TypeElement dependency) {
     if (!encodedType.getModifiers().contains(Modifier.ABSTRACT)) {
       throw new IllegalArgumentException(
@@ -609,16 +609,19 @@ public class AutoCodecProcessor extends AbstractProcessor {
     return codecClassBuilder;
   }
 
-  private static MethodSpec buildPolymorphicSerializeMethod(
+  private MethodSpec buildPolymorphicSerializeMethod(
       TypeElement encodedType, @Nullable TypeElement dependency) {
     MethodSpec.Builder builder =
         AutoCodecUtil.initializeSerializeMethodBuilder(encodedType, dependency);
+    TypeName polyClass = TypeName.get(env.getTypeUtils().erasure(encodedType.asType()));
     if (dependency == null) {
-      builder.addStatement("$T.serialize(input, codedOut, null)", PolymorphicHelper.class);
+      builder.addStatement(
+          "$T.serialize(input, $T.class, codedOut, null)", PolymorphicHelper.class, polyClass);
     } else {
       builder.addStatement(
-          "$T.serialize(input, codedOut, $T.ofNullable(dependency))",
+          "$T.serialize(input, $T.class, codedOut, $T.ofNullable(dependency))",
           PolymorphicHelper.class,
+          polyClass,
           Optional.class);
     }
     return builder.build();