1 files changed, 63 insertions, 24 deletions

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 e3c7898c58..7a01498151 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
@@ -22,6 +22,8 @@ import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.Maps;
+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.strings.StringCodecs;
 import com.google.protobuf.GeneratedMessage;
@@ -32,7 +34,10 @@ import java.util.List;
 import java.util.Map;
 import java.util.regex.Pattern;
 import javax.annotation.processing.ProcessingEnvironment;
+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.DeclaredType;
 import javax.lang.model.type.TypeMirror;
 
@@ -61,31 +66,16 @@ class Marshallers {
   }
 
   private Marshaller getMatchingMarshaller(DeclaredType type) {
-    return marshallers.stream().filter(marshaller -> marshaller.matches(type)).findFirst().get();
+    return marshallers
+        .stream()
+        .filter(marshaller -> marshaller.matches(type))
+        .findFirst()
+        .orElseThrow(
+            () ->
+                new IllegalArgumentException(
+                    "No marshaller for: " + ((TypeElement) type.asElement()).getQualifiedName()));
   }
 
-  private static final Marshaller CODEC_MARSHALLER =
-      new Marshaller() {
-        @Override
-        public boolean matches(DeclaredType type) {
-          // TODO(shahan): check for getCodec or CODEC.
-          // CODEC is the final fallback for all Marshallers so this returns true.
-          return true;
-        }
-
-        @Override
-        public void addSerializationCode(Context context) {
-          context.builder.addStatement(
-              "$T.CODEC.serialize($L, codedOut)", context.getTypeName(), context.name);
-        }
-
-        @Override
-        public void addDeserializationCode(Context context) {
-          context.builder.addStatement(
-              "$L = $T.CODEC.deserialize(codedIn)", context.name, context.getTypeName());
-        }
-      };
-
   private final Marshaller enumMarshaller =
       new Marshaller() {
         @Override
@@ -450,6 +440,39 @@ class Marshallers {
         }
       };
 
+  private final Marshaller codecMarshaller =
+      new Marshaller() {
+        @Override
+        public boolean matches(DeclaredType type) {
+          return getCodec(type).isPresent();
+        }
+
+        @Override
+        public void addSerializationCode(Context context) {
+          context.builder.addStatement(
+              "$T.CODEC.serialize($L, codedOut)", context.getTypeName(), context.name);
+        }
+
+        @Override
+        public void addDeserializationCode(Context context) {
+          TypeMirror codecType = getCodec(context.type).get().asType();
+          if (isSubtypeErased(codecType, ObjectCodec.class)) {
+            context.builder.addStatement(
+                "$L = $T.CODEC.deserialize(codedIn)", context.name, context.getTypeName());
+          } else if (isSubtypeErased(codecType, InjectingObjectCodec.class)) {
+            context.builder.addStatement(
+                "$L = $T.CODEC.deserialize(dependency, codedIn)",
+                context.name,
+                context.getTypeName());
+          } else {
+            throw new IllegalArgumentException(
+                "CODEC field of "
+                    + ((TypeElement) context.type.asElement()).getQualifiedName()
+                    + " is neither ObjectCodec nor InjectingCodec");
+          }
+        }
+      };
+
   private final ImmutableList<Marshaller> marshallers =
       ImmutableList.of(
           enumMarshaller,
@@ -462,7 +485,7 @@ class Marshallers {
           multimapMarshaller,
           patternMarshaller,
           protoMarshaller,
-          CODEC_MARSHALLER);
+          codecMarshaller);
 
   /** True when {@code type} has the same type as {@code clazz}. */
   private boolean matchesType(TypeMirror type, Class<?> clazz) {
@@ -480,8 +503,24 @@ class Marshallers {
         .isSameType(env.getTypeUtils().erasure(type), env.getTypeUtils().erasure(getType(clazz)));
   }
 
+  /** True when erasure of {@code type} is a subtype of the erasure of {@code clazz}. */
+  private boolean isSubtypeErased(TypeMirror type, Class<?> clazz) {
+    return env.getTypeUtils()
+        .isSubtype(env.getTypeUtils().erasure(type), env.getTypeUtils().erasure(getType(clazz)));
+  }
+
   /** Returns the TypeMirror corresponding to {@code clazz}. */
   private TypeMirror getType(Class<?> clazz) {
     return env.getElementUtils().getTypeElement((clazz.getCanonicalName())).asType();
   }
+
+  private static java.util.Optional<? extends Element> getCodec(DeclaredType type) {
+    return type.asElement()
+        .getEnclosedElements()
+        .stream()
+        .filter(t -> t.getModifiers().contains(Modifier.STATIC))
+        .filter(t -> t.getSimpleName().contentEquals("CODEC"))
+        .filter(t -> t.getKind() == ElementKind.FIELD)
+        .findAny();
+  }
 }