@AutoCodec features and refactoring.

* Features - ImmutableList, Map, ImmutableSortedMap, and ProtocolMessageEnum types.
* Refactoring - creates an abstraction, "Marshaller", to replace a big switch over supported types.
* Cleanup - implements equals for most of the test classes.

PiperOrigin-RevId: 180024241

4 files changed, 432 insertions, 136 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 c285cb1fe1..95b9706187 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
@@ -16,8 +16,6 @@ package com.google.devtools.build.lib.skyframe.serialization.autocodec;
 
 import com.google.auto.service.AutoService;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.ImmutableSortedSet;
-import com.google.common.collect.Maps;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
 import com.google.devtools.build.lib.skyframe.serialization.strings.StringCodecs;
 import com.google.protobuf.CodedInputStream;
@@ -30,10 +28,7 @@ import java.io.IOException;
 import java.lang.reflect.Field;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
-import java.util.ArrayList;
-import java.util.Comparator;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
@@ -43,7 +38,6 @@ import javax.annotation.processing.Processor;
 import javax.annotation.processing.RoundEnvironment;
 import javax.lang.model.SourceVersion;
 import javax.lang.model.element.Element;
-import javax.lang.model.element.ElementKind;
 import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
@@ -68,6 +62,7 @@ public class AutoCodecProcessor extends AbstractProcessor {
   private static final String PRINT_GENERATED_OPTION = "autocodec_print_generated";
 
   private ProcessingEnvironment env; // Captured from `init` method.
+  private Marshallers marshallers;
 
   @Override
   public Set<String> getSupportedOptions() {
@@ -88,6 +83,7 @@ public class AutoCodecProcessor extends AbstractProcessor {
   public synchronized void init(ProcessingEnvironment processingEnv) {
     super.init(processingEnv);
     this.env = processingEnv;
+    this.marshallers = new Marshallers(processingEnv);
   }
 
   @Override
@@ -196,7 +192,9 @@ public class AutoCodecProcessor extends AbstractProcessor {
           serializeBuilder.addStatement("codedOut.writeInt32NoTag($L)", paramAccessor);
           break;
         case DECLARED:
-          buildSerializeBody(serializeBuilder, (DeclaredType) parameter.asType(), paramAccessor);
+          marshallers.writeSerializationCode(
+              new Marshaller.Context(
+                  serializeBuilder, (DeclaredType) parameter.asType(), paramAccessor));
           break;
         default:
           throw new IllegalArgumentException("Unimplemented or invalid kind: " + typeKind);
@@ -206,51 +204,6 @@ public class AutoCodecProcessor extends AbstractProcessor {
   }
 
   /**
-   * Appends code statements to {@code builder} to serialize a pre-declared variable named {@code
-   * accessor}.
-   *
-   * @param type the type of {@code accessor}
-   * @param depth recursion depth of buildSerializeBody
-   */
-  private void buildSerializeBody(
-      MethodSpec.Builder builder, DeclaredType type, String accessor, int depth) {
-    builder.beginControlFlow("if ($L != null)", accessor); // Begin if not null block.
-    builder.addStatement("codedOut.writeBoolNoTag(true)");
-    // TODO(shahan): Add support for more types.
-    if (isEnum(type)) {
-      builder.addStatement("codedOut.writeInt32NoTag($L.ordinal())", accessor);
-    } else if (matchesType(type, String.class)) {
-      builder.addStatement(
-          "$T.asciiOptimized().serialize($L, codedOut)", StringCodecs.class, accessor);
-    } else if (matchesErased(type, Map.Entry.class)) {
-      DeclaredType keyType = (DeclaredType) type.getTypeArguments().get(0);
-      buildSerializeBody(builder, keyType, accessor + ".getKey()", depth + 1);
-      DeclaredType valueType = (DeclaredType) type.getTypeArguments().get(1);
-      buildSerializeBody(builder, valueType, accessor + ".getValue()", depth + 1);
-    } else if (matchesErased(type, List.class) || matchesErased(type, ImmutableSortedSet.class)) {
-      // Writes the target count to the stream so deserialization knows when to stop.
-      builder.addStatement("codedOut.writeInt32NoTag($L.size())", accessor);
-      DeclaredType repeatedType = (DeclaredType) type.getTypeArguments().get(0);
-      String repeatedName = "repeated" + depth;
-      builder.beginControlFlow(
-          "for ($T $L : $L)", TypeName.get(repeatedType), repeatedName, accessor);
-      buildSerializeBody(builder, repeatedType, repeatedName, depth + 1);
-      builder.endControlFlow();
-    } else {
-      // Otherwise use the type's codec.
-      builder.addStatement("$T.CODEC.serialize($L, codedOut)", TypeName.get(type), accessor);
-    }
-    builder.nextControlFlow("else");
-    builder.addStatement("codedOut.writeBoolNoTag(false)");
-    builder.endControlFlow(); // End if not null.
-  }
-
-  /** Convenience overload for depth = 0. */
-  private void buildSerializeBody(MethodSpec.Builder builder, DeclaredType type, String accessor) {
-    buildSerializeBody(builder, type, accessor, /*depth=*/ 0);
-  }
-
-  /**
    * Adds a body to the deserialize method that extracts serialized parameters.
    *
    * <p>Parameter values are extracted into local variables with the same name as the parameter
@@ -270,7 +223,8 @@ public class AutoCodecProcessor extends AbstractProcessor {
           builder.addStatement("int $L = codedIn.readInt32()", paramName);
           break;
         case DECLARED:
-          buildDeserializeBody(builder, (DeclaredType) parameter.asType(), paramName);
+          marshallers.writeDeserializationCode(
+              new Marshaller.Context(builder, (DeclaredType) parameter.asType(), paramName));
           break;
         default:
           throw new IllegalArgumentException("Unimplemented or invalid kind: " + typeKind);
@@ -279,74 +233,6 @@ public class AutoCodecProcessor extends AbstractProcessor {
   }
 
   /**
-   * Appends code statements to {@code builder}, declaring a variable called {@code name} and
-   * initializing it with deserialization.
-   *
-   * @param type the type of {@code name}.
-   * @param depth recursion depth of buildDeserializeBody.
-   */
-  private void buildDeserializeBody(
-      MethodSpec.Builder builder, DeclaredType type, String name, int depth) {
-    builder.addStatement("$T $L = null", TypeName.get(type), name);
-    builder.beginControlFlow("if (codedIn.readBool())"); // Begin null-handling block.
-    // TODO(shahan): Add support for more types.
-    if (isEnum(type)) {
-      // TODO(shahan): memoize this expensive call to values().
-      builder.addStatement("$L = $T.values()[codedIn.readInt32()]", name, TypeName.get(type));
-    } else if (matchesType(type, String.class)) {
-      builder.addStatement(
-          "$L = $T.asciiOptimized().deserialize(codedIn)", name, StringCodecs.class);
-    } else if (matchesErased(type, Map.Entry.class)) {
-      DeclaredType keyType = (DeclaredType) type.getTypeArguments().get(0);
-      String keyName = "key" + depth;
-      buildDeserializeBody(builder, keyType, keyName, depth + 1);
-      DeclaredType valueType = (DeclaredType) type.getTypeArguments().get(1);
-      String valueName = "value" + depth;
-      buildDeserializeBody(builder, valueType, valueName, depth + 1);
-      builder.addStatement("$L = $T.immutableEntry($L, $L)", name, Maps.class, keyName, valueName);
-    } else if (matchesErased(type, List.class)) {
-      builder.addStatement("$L = new $T<>()", name, ArrayList.class);
-      String lengthName = "length" + depth;
-      builder.addStatement("int $L = codedIn.readInt32()", lengthName);
-      String indexName = "i" + depth;
-      builder.beginControlFlow(
-          "for (int $L = 0; $L < $L; ++$L)", indexName, indexName, lengthName, indexName);
-      DeclaredType repeatedType = (DeclaredType) type.getTypeArguments().get(0);
-      String repeatedName = "repeated" + depth;
-      buildDeserializeBody(builder, repeatedType, repeatedName, depth + 1);
-      builder.addStatement("$L.add($L)", name, repeatedName);
-      builder.endControlFlow();
-    } else if (matchesErased(type, ImmutableSortedSet.class)) {
-      DeclaredType repeatedType = (DeclaredType) type.getTypeArguments().get(0);
-      builder.addStatement(
-          "$T<$T> builder = new $T<>($T.naturalOrder())",
-          ImmutableSortedSet.Builder.class,
-          TypeName.get(repeatedType),
-          ImmutableSortedSet.Builder.class,
-          Comparator.class);
-      String lengthName = "length" + depth;
-      builder.addStatement("int $L = codedIn.readInt32()", lengthName);
-      String indexName = "i" + depth;
-      builder.beginControlFlow(
-          "for (int $L = 0; $L < $L; ++$L)", indexName, indexName, lengthName, indexName);
-      String repeatedName = "repeated" + depth;
-      buildDeserializeBody(builder, repeatedType, repeatedName, depth + 1);
-      builder.addStatement("builder.add($L)", repeatedName);
-      builder.endControlFlow();
-      builder.addStatement("$L = builder.build()", name);
-    } else {
-      // Otherwise, use the type's codec.
-      builder.addStatement("$L = $T.CODEC.deserialize(codedIn)", name, TypeName.get(type));
-    }
-    builder.endControlFlow(); // End null-handling block.
-  }
-
-  /** Overload of above, for common case of 0 depth. */
-  private void buildDeserializeBody(MethodSpec.Builder builder, DeclaredType type, String name) {
-    buildDeserializeBody(builder, type, name, /*depth=*/ 0);
-  }
-
-  /**
    * Invokes the constructor and returns the value.
    *
    * <p>Used by the {@link AutoCodec.Strategy.CONSTRUCTOR} strategy.
@@ -453,20 +339,6 @@ public class AutoCodecProcessor extends AbstractProcessor {
             type, env.getElementUtils().getTypeElement((clazz.getCanonicalName())).asType());
   }
 
-  /** True when erasure of {@code type} matches erasure of {@code clazz}. */
-  private boolean matchesErased(TypeMirror type, Class<?> clazz) {
-    return env.getTypeUtils()
-        .isSameType(
-            env.getTypeUtils().erasure(type),
-            env.getTypeUtils()
-                .erasure(
-                    env.getElementUtils().getTypeElement((clazz.getCanonicalName())).asType()));
-  }
-
-  private boolean isEnum(TypeMirror type) {
-    return env.getTypeUtils().asElement(type).getKind() == ElementKind.ENUM;
-  }
-
   /** Emits a note to BUILD log during annotation processing for debugging. */
   private void note(String note) {
     env.getMessager().printMessage(Diagnostic.Kind.NOTE, note);
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 4f98384702..1116067e28 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
@@ -37,7 +37,11 @@ java_plugin(
 # @AutoCodec annotation processor implementation.
 java_library(
     name = "autocodec-processor",
-    srcs = ["AutoCodecProcessor.java"],
+    srcs = [
+        "AutoCodecProcessor.java",
+        "Marshaller.java",
+        "Marshallers.java",
+    ],
     deps = [
         ":autocodec-annotation",
         ":autocodec-util",
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/Marshaller.java
new file mode 100644
index 0000000000..ad4b9d7878
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/Marshaller.java
@@ -0,0 +1,80 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.skyframe.serialization.autocodec;
+
+import com.squareup.javapoet.MethodSpec;
+import com.squareup.javapoet.TypeName;
+import javax.annotation.processing.ProcessingEnvironment;
+import javax.lang.model.type.DeclaredType;
+
+/**
+ * Generates serialize and deserialize code fragments for matching types.
+ *
+ * <p>All methods are logically static and take the {@link ProcessingEnvironment} as a parameter.
+ */
+interface Marshaller {
+  static class Context {
+    /** Builder for the method. */
+    public final MethodSpec.Builder builder;
+    /** Type of {@code name}. */
+    public final DeclaredType type;
+    /** Name of variable. */
+    public final String name;
+    /**
+     * Recursion depth.
+     *
+     * <p>Recursion is used to traverse generic types.
+     */
+    public final int depth;
+
+    Context(MethodSpec.Builder builder, DeclaredType type, String name) {
+      this(builder, type, name, 0);
+    }
+
+    private Context(MethodSpec.Builder builder, DeclaredType type, String name, int depth) {
+      this.builder = builder;
+      this.type = type;
+      this.name = name;
+      this.depth = depth;
+    }
+
+    /** Returns a new context with a new type and name at the next recursion depth. */
+    Context with(DeclaredType newType, String newName) {
+      return new Context(builder, newType, newName, depth + 1);
+    }
+
+    TypeName getTypeName() {
+      return TypeName.get(type);
+    }
+
+    /**
+     * Returns a depth-qualified name.
+     *
+     * <p>This helps to avoid name collisions when recursion reuses symbol names.
+     */
+    String makeName(String name) {
+      return name + depth;
+    }
+  };
+
+  /** 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);
+}
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
new file mode 100644
index 0000000000..4a34cd6e81
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/Marshallers.java
@@ -0,0 +1,340 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.skyframe.serialization.autocodec;
+
+import com.google.common.collect.ImmutableList;
+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.autocodec.Marshaller.Context;
+import com.google.devtools.build.lib.skyframe.serialization.strings.StringCodecs;
+import com.google.protobuf.ProtocolMessageEnum;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.processing.ProcessingEnvironment;
+import javax.lang.model.element.ElementKind;
+import javax.lang.model.type.DeclaredType;
+import javax.lang.model.type.TypeMirror;
+
+/** Class containing all {@link Marshaller} instances. */
+class Marshallers {
+  private final ProcessingEnvironment env;
+
+  Marshallers(ProcessingEnvironment env) {
+    this.env = env;
+  }
+
+  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();
+  }
+
+  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();
+  }
+
+  private Marshaller getMatchingMarshaller(DeclaredType type) {
+    return marshallers.stream().filter(marshaller -> marshaller.matches(type)).findFirst().get();
+  }
+
+  private static final Marshaller CODEC_MARSHALLER =
+      new Marshaller() {
+        @Override
+        public boolean matches(DeclaredType type) {
+          // 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
+        public boolean matches(DeclaredType type) {
+          return env.getTypeUtils().asElement(type).getKind() == ElementKind.ENUM;
+        }
+
+        @Override
+        public void addSerializationCode(Context context) {
+          if (isProtoEnum(context.type)) {
+            context.builder.addStatement("codedOut.writeInt32NoTag($L.getNumber())", context.name);
+          } else {
+            context.builder.addStatement("codedOut.writeInt32NoTag($L.ordinal())", context.name);
+          }
+        }
+
+        @Override
+        public void addDeserializationCode(Context context) {
+          if (isProtoEnum(context.type)) {
+            context.builder.addStatement(
+                "$L = $T.forNumber(codedIn.readInt32())", context.name, context.getTypeName());
+          } else {
+            // TODO(shahan): memoize this expensive call to values().
+            context.builder.addStatement(
+                "$L = $T.values()[codedIn.readInt32()]", context.name, context.getTypeName());
+          }
+        }
+
+        private boolean isProtoEnum(DeclaredType type) {
+          return env.getTypeUtils()
+              .isSubtype(
+                  type,
+                  env.getElementUtils()
+                      .getTypeElement(ProtocolMessageEnum.class.getCanonicalName())
+                      .asType());
+        }
+      };
+
+  private final Marshaller stringMarshaller =
+      new Marshaller() {
+        @Override
+        public boolean matches(DeclaredType type) {
+          return matchesType(type, String.class);
+        }
+
+        @Override
+        public void addSerializationCode(Context context) {
+          context.builder.addStatement(
+              "$T.asciiOptimized().serialize($L, codedOut)", StringCodecs.class, context.name);
+        }
+
+        @Override
+        public void addDeserializationCode(Context context) {
+          context.builder.addStatement(
+              "$L = $T.asciiOptimized().deserialize(codedIn)", context.name, StringCodecs.class);
+        }
+      };
+
+  private final Marshaller mapEntryMarshaller =
+      new Marshaller() {
+        @Override
+        public boolean matches(DeclaredType type) {
+          return matchesErased(type, Map.Entry.class);
+        }
+
+        @Override
+        public void addSerializationCode(Context context) {
+          DeclaredType keyType = (DeclaredType) context.type.getTypeArguments().get(0);
+          writeSerializationCode(context.with(keyType, context.name + ".getKey()"));
+          DeclaredType valueType = (DeclaredType) context.type.getTypeArguments().get(1);
+          writeSerializationCode(context.with(valueType, context.name + ".getValue()"));
+        }
+
+        @Override
+        public void addDeserializationCode(Context context) {
+          DeclaredType keyType = (DeclaredType) context.type.getTypeArguments().get(0);
+          String keyName = context.makeName("key");
+          writeDeserializationCode(context.with(keyType, keyName));
+          DeclaredType valueType = (DeclaredType) context.type.getTypeArguments().get(1);
+          String valueName = context.makeName("value");
+          writeDeserializationCode(context.with(valueType, valueName));
+          context.builder.addStatement(
+              "$L = $T.immutableEntry($L, $L)", context.name, Maps.class, keyName, valueName);
+        }
+      };
+
+  private final Marshaller listMarshaller =
+      new Marshaller() {
+        @Override
+        public boolean matches(DeclaredType type) {
+          // TODO(shahan): List is more general than ImmutableList. Consider whether these
+          // two should have distinct marshallers.
+          return matchesErased(type, List.class) || matchesErased(type, ImmutableList.class);
+        }
+
+        @Override
+        public void addSerializationCode(Context context) {
+          // Writes the target count to the stream so deserialization knows when to stop.
+          context.builder.addStatement("codedOut.writeInt32NoTag($L.size())", context.name);
+          Context repeated =
+              context.with(
+                  (DeclaredType) context.type.getTypeArguments().get(0),
+                  context.makeName("repeated"));
+          context.builder.beginControlFlow(
+              "for ($T $L : $L)", repeated.getTypeName(), repeated.name, context.name);
+          writeSerializationCode(repeated);
+          context.builder.endControlFlow();
+        }
+
+        @Override
+        public void addDeserializationCode(Context context) {
+          Context repeated =
+              context.with(
+                  (DeclaredType) context.type.getTypeArguments().get(0),
+                  context.makeName("repeated"));
+          String builderName = context.makeName("builder");
+          context.builder.addStatement(
+              "$T<$T> $L = new $T<>()",
+              ImmutableList.Builder.class,
+              repeated.getTypeName(),
+              builderName,
+              ImmutableList.Builder.class);
+          String lengthName = context.makeName("length");
+          context.builder.addStatement("int $L = codedIn.readInt32()", 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.add($L)", builderName, repeated.name);
+          context.builder.endControlFlow();
+          context.builder.addStatement("$L = $L.build()", context.name, builderName);
+        }
+      };
+
+  private final Marshaller immutableSortedSetMarshaller =
+      new Marshaller() {
+        @Override
+        public boolean matches(DeclaredType type) {
+          return matchesErased(type, ImmutableSortedSet.class);
+        }
+
+        @Override
+        public void addSerializationCode(Context context) {
+          listMarshaller.addSerializationCode(context);
+        }
+
+        @Override
+        public void addDeserializationCode(Context context) {
+          Context repeated =
+              context.with(
+                  (DeclaredType) context.type.getTypeArguments().get(0),
+                  context.makeName("repeated"));
+          String builderName = context.makeName("builder");
+          context.builder.addStatement(
+              "$T<$T> $L = new $T<>($T.naturalOrder())",
+              ImmutableSortedSet.Builder.class,
+              repeated.getTypeName(),
+              builderName,
+              ImmutableSortedSet.Builder.class,
+              Comparator.class);
+          String lengthName = context.makeName("length");
+          context.builder.addStatement("int $L = codedIn.readInt32()", 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.add($L)", builderName, repeated.name);
+          context.builder.endControlFlow();
+          context.builder.addStatement("$L = $L.build()", context.name, builderName);
+        }
+      };
+
+  private final Marshaller mapMarshaller =
+      new Marshaller() {
+        @Override
+        public boolean matches(DeclaredType type) {
+          // TODO(shahan): since Map is a bit more general than ImmutableSortedMap, consider
+          // splitting these.
+          return matchesErased(type, Map.class) || matchesErased(type, ImmutableSortedMap.class);
+        }
+
+        @Override
+        public void addSerializationCode(Context context) {
+          context.builder.addStatement("codedOut.writeInt32NoTag($L.size())", context.name);
+          String entryName = context.makeName("entry");
+          Context key =
+              context.with(
+                  (DeclaredType) context.type.getTypeArguments().get(0), entryName + ".getKey()");
+          Context value =
+              context.with(
+                  (DeclaredType) context.type.getTypeArguments().get(1), entryName + ".getValue()");
+          context.builder.beginControlFlow(
+              "for ($T<$T, $T> $L : $L.entrySet())",
+              Map.Entry.class,
+              key.getTypeName(),
+              value.getTypeName(),
+              entryName,
+              context.name);
+          writeSerializationCode(key);
+          writeSerializationCode(value);
+          context.builder.endControlFlow();
+        }
+
+        @Override
+        public void addDeserializationCode(Context context) {
+          Context key =
+              context.with(
+                  (DeclaredType) context.type.getTypeArguments().get(0), context.makeName("key"));
+          Context value =
+              context.with(
+                  (DeclaredType) context.type.getTypeArguments().get(1), context.makeName("value"));
+          String builderName = context.makeName("builder");
+          context.builder.addStatement(
+              "$T<$T, $T> $L = new $T<>($T.naturalOrder())",
+              ImmutableSortedMap.Builder.class,
+              key.getTypeName(),
+              value.getTypeName(),
+              builderName,
+              ImmutableSortedMap.Builder.class,
+              Comparator.class);
+          String lengthName = context.makeName("length");
+          context.builder.addStatement("int $L = codedIn.readInt32()", lengthName);
+          String indexName = context.makeName("i");
+          context.builder.beginControlFlow(
+              "for (int $L = 0; $L < $L; ++$L)", indexName, indexName, lengthName, indexName);
+          writeDeserializationCode(key);
+          writeDeserializationCode(value);
+          context.builder.addStatement("$L.put($L, $L)", builderName, key.name, value.name);
+          context.builder.endControlFlow();
+          context.builder.addStatement("$L = $L.build()", context.name, builderName);
+        }
+      };
+
+  private final ImmutableList<Marshaller> marshallers =
+      ImmutableList.of(
+          enumMarshaller,
+          stringMarshaller,
+          mapEntryMarshaller,
+          listMarshaller,
+          immutableSortedSetMarshaller,
+          mapMarshaller,
+          CODEC_MARSHALLER);
+
+  /** True when {@code type} has the same type as {@code clazz}. */
+  private boolean matchesType(TypeMirror type, Class<?> clazz) {
+    return env.getTypeUtils()
+        .isSameType(
+            type, env.getElementUtils().getTypeElement((clazz.getCanonicalName())).asType());
+  }
+
+  /** True when erasure of {@code type} matches erasure of {@code clazz}. */
+  private boolean matchesErased(TypeMirror type, Class<?> clazz) {
+    return env.getTypeUtils()
+        .isSameType(
+            env.getTypeUtils().erasure(type),
+            env.getTypeUtils()
+                .erasure(
+                    env.getElementUtils().getTypeElement((clazz.getCanonicalName())).asType()));
+  }
+}