Integrate memoization into standard serialization. This involves a number of large changes:

1. SerializationContext and DeserializationContext become the owners of the Memoizer if requested. They produce new versions of themselves on demand that are memoization-aware. Because of intricacies of Skylark that I do not fully understand, we inject a Mutability object when starting memoization, and so to be conservative, that injection starts up a new memoization frame, nested if we were already memoizing, just like before. It would be nice to decouple this injection from memoization in the future.

2. MemoizingCodec is deleted, but really ObjectCodec becomes MemoizingCodec, so it lives on. BaseCodec is deleted since it now has only one implementation.

3. The simplified model of registering MemoizingCodecs is adopted for ObjectCodecs: all codecs are registered based on their #getEncodedClass and #additionalEncodedSubclasses. This also allows us to register codecs that are defined using tricky parameter types, since we're no longer trying to reflectively examine them. This required a clean-up of such codecs, and the addition of ArrayListCodec to stop NullableListCodec from making lists unmodifiable when they shouldn't be.

4. @AutoCodec is extended to allow users to specify that memoization should start with this codec.

To ensure bit-equivalence, SkyKeySerializer disables memoization.

PiperOrigin-RevId: 188918251

24 files changed, 659 insertions, 873 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/packages/Attribute.java b/src/main/java/com/google/devtools/build/lib/packages/Attribute.java
index 045520aa6d..61878dd4ef 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/Attribute.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/Attribute.java
@@ -57,6 +57,7 @@ import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.Nullable;
@@ -2251,6 +2252,51 @@ public final class Attribute implements Comparable<Attribute> {
     return name.compareTo(other.name);
   }
 
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    Attribute attribute = (Attribute) o;
+    return Objects.equals(name, attribute.name)
+        && Objects.equals(type, attribute.type)
+        && Objects.equals(propertyFlags, attribute.propertyFlags)
+        && Objects.equals(defaultValue, attribute.defaultValue)
+        && Objects.equals(configTransition, attribute.configTransition)
+        && Objects.equals(splitTransitionProvider, attribute.splitTransitionProvider)
+        && Objects.equals(allowedRuleClassesForLabels, attribute.allowedRuleClassesForLabels)
+        && Objects.equals(
+            allowedRuleClassesForLabelsWarning, attribute.allowedRuleClassesForLabelsWarning)
+        && Objects.equals(allowedFileTypesForLabels, attribute.allowedFileTypesForLabels)
+        && Objects.equals(validityPredicate, attribute.validityPredicate)
+        && Objects.equals(condition, attribute.condition)
+        && Objects.equals(allowedValues, attribute.allowedValues)
+        && Objects.equals(requiredProviders, attribute.requiredProviders)
+        && Objects.equals(aspects, attribute.aspects);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(
+        name,
+        type,
+        propertyFlags,
+        defaultValue,
+        configTransition,
+        splitTransitionProvider,
+        allowedRuleClassesForLabels,
+        allowedRuleClassesForLabelsWarning,
+        allowedFileTypesForLabels,
+        validityPredicate,
+        condition,
+        allowedValues,
+        requiredProviders,
+        aspects);
+  }
+
   /**
    * Returns a replica builder of this Attribute.
    */
diff --git a/src/main/java/com/google/devtools/build/lib/packages/RequiredProviders.java b/src/main/java/com/google/devtools/build/lib/packages/RequiredProviders.java
index 1551d70f35..2f026c84a2 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/RequiredProviders.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/RequiredProviders.java
@@ -20,6 +20,7 @@ import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec.VisibleForSerialization;
+import java.util.Objects;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import javax.annotation.Nullable;
@@ -303,6 +304,25 @@ public final class RequiredProviders {
     }
   }
 
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    RequiredProviders that = (RequiredProviders) o;
+    return constraint == that.constraint
+        && Objects.equals(nativeProviders, that.nativeProviders)
+        && Objects.equals(skylarkProviders, that.skylarkProviders);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(constraint, nativeProviders, skylarkProviders);
+  }
+
   /**
    * A builder for {@link RequiredProviders} that accepts any dependency
    * unless restriction provider sets are added.
diff --git a/src/main/java/com/google/devtools/build/lib/packages/SkylarkDefinedAspect.java b/src/main/java/com/google/devtools/build/lib/packages/SkylarkDefinedAspect.java
index 14ec9d646e..e8cd11ae55 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/SkylarkDefinedAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/SkylarkDefinedAspect.java
@@ -29,6 +29,7 @@ import com.google.devtools.build.lib.syntax.EvalException;
 import com.google.devtools.build.lib.syntax.Type;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Objects;
 
 /** A Skylark value that is a result of an 'aspect(..)' function call. */
 @AutoCodec
@@ -224,4 +225,42 @@ public class SkylarkDefinedAspect implements SkylarkExportable, SkylarkAspect {
     }
     attrBuilder.aspect(this, loc);
   }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    SkylarkDefinedAspect that = (SkylarkDefinedAspect) o;
+    return Objects.equals(implementation, that.implementation)
+        && Objects.equals(attributeAspects, that.attributeAspects)
+        && Objects.equals(attributes, that.attributes)
+        && Objects.equals(requiredAspectProviders, that.requiredAspectProviders)
+        && Objects.equals(provides, that.provides)
+        && Objects.equals(paramAttributes, that.paramAttributes)
+        && Objects.equals(fragments, that.fragments)
+        && Objects.equals(hostTransition, that.hostTransition)
+        && Objects.equals(hostFragments, that.hostFragments)
+        && Objects.equals(requiredToolchains, that.requiredToolchains)
+        && Objects.equals(aspectClass, that.aspectClass);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(
+        implementation,
+        attributeAspects,
+        attributes,
+        requiredAspectProviders,
+        provides,
+        paramAttributes,
+        fragments,
+        hostTransition,
+        hostFragments,
+        requiredToolchains,
+        aspectClass);
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/RuleConfiguredTargetValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/RuleConfiguredTargetValue.java
index abe1c70e66..b48f247feb 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/RuleConfiguredTargetValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/RuleConfiguredTargetValue.java
@@ -26,15 +26,14 @@ import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec.Memoization;
 import java.util.ArrayList;
 import javax.annotation.Nullable;
 
-/**
- * A configured target in the context of a Skyframe graph.
- */
+/** A configured target in the context of a Skyframe graph. */
 @Immutable
 @ThreadSafe
-@AutoCodec
+@AutoCodec(memoization = Memoization.START_MEMOIZING)
 @VisibleForTesting
 public final class RuleConfiguredTargetValue extends ActionLookupValue
     implements ConfiguredTargetValue {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/BaseCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/BaseCodec.java
deleted file mode 100644
index abec3f1007..0000000000
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/BaseCodec.java
+++ /dev/null
@@ -1,51 +0,0 @@
-// 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;
-
-import com.google.common.collect.ImmutableList;
-import com.google.devtools.build.skyframe.SkyKey;
-import com.google.devtools.build.skyframe.SkyValue;
-import java.util.List;
-
-/** An opaque interface for codecs that just reveals the {@link Class} of its objects. */
-interface BaseCodec<T> {
-  /**
-   * Returns the class of the objects serialized/deserialized by this codec.
-   *
-   * <p>This is useful for automatically dispatching to the correct codec, e.g. in {@link
-   * ObjectCodecs} and {@link BaseCodecMap}. It may also be useful for automatically registering
-   * codecs for {@link SkyKey}s and {@link SkyValue}s instead of using the current manual mapping
-   * (b/26186886).
-   */
-  Class<T> getEncodedClass();
-
-  /**
-   * Returns additional subtypes of {@code T} that may be serialized/deserialized using this codec
-   * without loss of information.
-   *
-   * <p>This method is intended for when {@code T} has multiple concrete implementations whose
-   * details are known to the codec but not to the codec dispatching mechanism. It signals that the
-   * dispatcher may choose to use this codec for the subtype, rather than raise {@link
-   * SerializationException.NoCodecException}.
-   *
-   * <p>This method should not be used if the codec's serialization and deserialization methods
-   * perform their own dispatching to other codecs for subtypes of {@code T}.
-   *
-   * <p>{@code T} itself should not be included in the returned list.
-   */
-  default List<Class<? extends T>> additionalEncodedSubclasses() {
-    return ImmutableList.of();
-  }
-}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecScanner.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecScanner.java
index 0a5f6b9db1..1adb15047f 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecScanner.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecScanner.java
@@ -17,8 +17,6 @@ package com.google.devtools.build.lib.skyframe.serialization;
 import com.google.common.base.Preconditions;
 import com.google.common.reflect.ClassPath;
 import com.google.common.reflect.ClassPath.ClassInfo;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer.MemoizingCodec;
-import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecRegistry.Builder;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.RegisteredSingletonDoNotUse;
 import java.io.IOException;
 import java.lang.reflect.Constructor;
@@ -35,7 +33,6 @@ import java.util.logging.Level;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-import javax.annotation.Nullable;
 
 /**
  * Scans the classpath to find {@link ObjectCodec} and {@link CodecRegisterer} instances.
@@ -61,18 +58,15 @@ public class CodecScanner {
       throws IOException, ReflectiveOperationException {
     log.info("Building ObjectCodecRegistry");
     ArrayList<Class<? extends ObjectCodec<?>>> codecs = new ArrayList<>();
-    ArrayList<Class<? extends MemoizingCodec<?>>> memoizingCodecs = new ArrayList<>();
     ArrayList<Class<? extends CodecRegisterer<?>>> registerers = new ArrayList<>();
     List<ClassInfo> classInfos = getClassInfos(packagePrefix).collect(Collectors.toList());
     getCodecs(classInfos)
         .forEach(
             type -> {
-              if (!ObjectCodec.class.equals(type) && ObjectCodec.class.isAssignableFrom(type)) {
-                codecs.add((Class<? extends ObjectCodec<?>>) type);
-              } else if (!MemoizingCodec.class.equals(type)
-                  && MemoizingCodec.class.isAssignableFrom(type)
+              if (!ObjectCodec.class.equals(type)
+                  && ObjectCodec.class.isAssignableFrom(type)
                   && !Modifier.isAbstract(type.getModifiers())) {
-                memoizingCodecs.add((Class<? extends MemoizingCodec<?>>) type);
+                codecs.add((Class<? extends ObjectCodec<?>>) type);
               } else if (!CodecRegisterer.class.equals(type)
                   && CodecRegisterer.class.isAssignableFrom(type)) {
                 registerers.add((Class<? extends CodecRegisterer<?>>) type);
@@ -117,7 +111,6 @@ public class CodecScanner {
         runRegisterers(builder, registerers);
 
     applyDefaultRegistration(builder, alreadyRegistered, codecs);
-    applyMemoizingCodecRegistration(builder, memoizingCodecs);
     return builder;
   }
 
@@ -142,24 +135,6 @@ public class CodecScanner {
     return registered;
   }
 
-  private static void applyMemoizingCodecRegistration(
-      Builder builder, List<Class<? extends MemoizingCodec<?>>> memoizingCodecs)
-      throws ReflectiveOperationException {
-    for (Class<? extends MemoizingCodec<?>> codecType : memoizingCodecs) {
-      try {
-        Constructor<? extends MemoizingCodec<?>> constructor = codecType.getDeclaredConstructor();
-        constructor.setAccessible(true);
-        builder.addMemoizing(constructor.newInstance());
-      } catch (NoSuchMethodException e) {
-        log.log(
-            Level.FINE,
-            "Skipping registration of " + codecType + " because it had no default constructor.");
-      } catch (InstantiationException e) {
-        throw new IllegalStateException("Couldn't instantiate " + codecType, e);
-      }
-    }
-  }
-
   @SuppressWarnings({"rawtypes", "unchecked"})
   private static void applyDefaultRegistration(
       ObjectCodecRegistry.Builder builder,
@@ -170,19 +145,10 @@ public class CodecScanner {
       if (alreadyRegistered.contains(codecType)) {
         continue;
       }
-      Class<?> serializedType = getSerializedType(codecType);
-      if (serializedType == null) {
-        log.log(
-            Level.FINE,
-            "Skipping registration of "
-                + codecType
-                + " because its generic parameter is not reified.");
-        continue;
-      }
       try {
         Constructor constructor = codecType.getDeclaredConstructor();
         constructor.setAccessible(true);
-        builder.add((Class) serializedType, (ObjectCodec) constructor.newInstance());
+        builder.add((ObjectCodec<?>) constructor.newInstance());
       } catch (NoSuchMethodException e) {
         log.log(
             Level.FINE,
@@ -193,7 +159,7 @@ public class CodecScanner {
 
   @SuppressWarnings("unchecked")
   private static Class<? extends ObjectCodec<?>> getObjectCodecType(
-      Class<? extends CodecRegisterer<?>> registererType) throws ReflectiveOperationException {
+      Class<? extends CodecRegisterer<?>> registererType) {
     Type typeArg =
         ((ParameterizedType)
                 registererType.getGenericInterfaces()[getCodecRegistererIndex(registererType)])
@@ -220,37 +186,6 @@ public class CodecScanner {
     throw new IllegalStateException(registererType + " doesn't directly implement CodecRegisterer");
   }
 
-  /** Returns the serialized class if available and null otherwise. */
-  @Nullable
-  private static Class<?> getSerializedType(Class<? extends ObjectCodec<?>> codecType) {
-    int objectCodecIndex = getObjectCodecIndex(codecType);
-    if (objectCodecIndex < 0) {
-      // This occurs if ObjectCodec is implemented by a superclass of codecType. There's no clear
-      // way to get the generic type parameter in this case.
-      return null;
-    }
-    Type typeArg =
-        ((ParameterizedType) codecType.getGenericInterfaces()[objectCodecIndex])
-            .getActualTypeArguments()[0];
-    if (typeArg instanceof ParameterizedType) {
-      typeArg = ((ParameterizedType) typeArg).getRawType();
-    }
-    if (!(typeArg instanceof Class)) {
-      return null;
-    }
-    return (Class<?>) typeArg;
-  }
-
-  private static int getObjectCodecIndex(Class<? extends ObjectCodec<?>> codecType) {
-    Class<?>[] interfaces = codecType.getInterfaces();
-    for (int i = 0; i < interfaces.length; ++i) {
-      if (ObjectCodec.class.equals(interfaces[i])) {
-        return i;
-      }
-    }
-    return -1;
-  }
-
   /** Return the {@link ClassInfo} objects matching {@code packagePrefix}, sorted by name. */
   private static Stream<ClassInfo> getClassInfos(String packagePrefix) throws IOException {
     return ClassPath.from(ClassLoader.getSystemClassLoader())
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/DeserializationContext.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/DeserializationContext.java
index 3948b64118..54b5820d31 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/DeserializationContext.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/DeserializationContext.java
@@ -14,35 +14,92 @@
 
 package com.google.devtools.build.lib.skyframe.serialization;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.skyframe.serialization.Memoizer.Deserializer;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer.MemoizingAfterObjectCodecAdapter;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer.MemoizingCodec;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer.Serializer;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer.Strategy;
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecRegistry.CodecDescriptor;
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecs.MemoizationPermission;
 import com.google.protobuf.CodedInputStream;
-import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
-import javax.annotation.Nullable;
+import java.util.function.Function;
+import javax.annotation.CheckReturnValue;
 
-/** Stateful class for providing additional context to a single deserialization "session". */
-// TODO(bazel-team): This class is just a shell, fill in.
+/**
+ * Stateful class for providing additional context to a single deserialization "session". This class
+ * is thread-safe so long as {@link #deserializer} is null. If it is not null, this class is not
+ * thread-safe and should only be accessed on a single thread for deserializing one serialized
+ * object (that may contain other serialized objects inside it).
+ */
 public class DeserializationContext {
-
   private final ObjectCodecRegistry registry;
   private final ImmutableMap<Class<?>, Object> dependencies;
+  private final MemoizationPermission memoizationPermission;
+  private final Memoizer.Deserializer deserializer;
 
-  public DeserializationContext(
-      ObjectCodecRegistry registry, ImmutableMap<Class<?>, Object> dependencies) {
+  private DeserializationContext(
+      ObjectCodecRegistry registry,
+      ImmutableMap<Class<?>, Object> dependencies,
+      MemoizationPermission memoizationPermission,
+      Deserializer deserializer) {
     this.registry = registry;
     this.dependencies = dependencies;
+    this.memoizationPermission = memoizationPermission;
+    this.deserializer = deserializer;
   }
 
+  @VisibleForTesting
+  public DeserializationContext(
+      ObjectCodecRegistry registry, ImmutableMap<Class<?>, Object> dependencies) {
+    this(registry, dependencies, MemoizationPermission.ALLOWED, /*deserializer=*/ null);
+  }
+
+  @VisibleForTesting
   public DeserializationContext(ImmutableMap<Class<?>, Object> dependencies) {
     this(AutoRegistry.get(), dependencies);
   }
 
+  DeserializationContext disableMemoization() {
+    Preconditions.checkState(
+        memoizationPermission == MemoizationPermission.ALLOWED, "memoization already disabled");
+    Preconditions.checkState(deserializer == null, "deserializer already present");
+    return new DeserializationContext(
+        registry, dependencies, MemoizationPermission.DISABLED, deserializer);
+  }
+
+  /**
+   * Returns a {@link DeserializationContext} that will memoize values it encounters (using
+   * reference equality), the inverse of the memoization performed by a {@link SerializationContext}
+   * returned by {@link SerializationContext#newMemoizingContext}. The context returned here should
+   * be used instead of the original: memoization may only occur when using the returned context.
+   *
+   * <p>A new {@link DeserializationContext} will be created for each call of this method, since it
+   * is assumed that each codec that calls this method has data that it needs to inject into its
+   * children that should mask its parents' injected data. Thus, over-eagerly calling this method
+   * will reduce the effectiveness of memoization.
+   */
+  @CheckReturnValue
+  public DeserializationContext newMemoizingContext(Object additionalData) {
+    Preconditions.checkNotNull(additionalData);
+    Preconditions.checkState(
+        memoizationPermission == MemoizationPermission.ALLOWED, "memoization disabled");
+    // TODO(janakr,brandjon): De we want to support a single context with data that gets pushed and
+    // popped.
+    return new DeserializationContext(
+        this.registry, this.dependencies, memoizationPermission, new Deserializer(additionalData));
+  }
+
+  /**
+   * Construct an initial value for the currently deserializing value from the additional data
+   * stored in a memoizing deserializer. The additional data must have type {@code klass}. The given
+   * {@code initialValueFunction} should be a hermetic function that does not interact with this
+   * context or {@code codedIn in any way}: it should be a pure function of an element of type
+   * {@link A}.
+   */
+  public <A, T> T makeInitialValue(Function<A, T> initialValueFunction, Class<A> klass) {
+    return deserializer.makeInitialValue(initialValueFunction, klass);
+  }
+
   // TODO(shahan): consider making codedIn a member of this class.
   @SuppressWarnings({"TypeParameterUnusedInFormals", "unchecked"})
   public <T> T deserialize(CodedInputStream codedIn) throws IOException, SerializationException {
@@ -51,69 +108,17 @@ public class DeserializationContext {
       return null;
     }
     T constant = (T) registry.maybeGetConstantByTag(tag);
-    return constant == null
-        ? (T) registry.getCodecDescriptorByTag(tag).deserialize(this, codedIn)
-        : constant;
-  }
-
-  /**
-   * Returns a {@link MemoizingCodec} appropriate for deserializing the next object on the wire.
-   * Only for use by {@link Memoizer.Deserializer}.
-   */
-  @SuppressWarnings("unchecked")
-  public MemoizingCodec<?> getMemoizingCodecRecordedInInput(CodedInputStream codedIn)
-      throws IOException, SerializationException {
-    int tag = codedIn.readSInt32();
-    if (tag == 0) {
-      return new InitializedMemoizingCodec<>(null);
-    }
-    Object constant = registry.maybeGetConstantByTag(tag);
     if (constant != null) {
-      return new InitializedMemoizingCodec<>(constant);
+      return constant;
     }
-    MemoizingCodec<?> codec = registry.maybeGetMemoizingCodecByTag(tag);
-    if (codec != null) {
-      return codec;
+    CodecDescriptor codecDescriptor = registry.getCodecDescriptorByTag(tag);
+    if (deserializer == null) {
+      return (T) codecDescriptor.deserialize(this, codedIn);
+    } else {
+      return deserializer.deserialize(this, (ObjectCodec<T>) codecDescriptor.getCodec(), codedIn);
     }
-    return new MemoizingAfterObjectCodecAdapter<>(registry.getCodecDescriptorByTag(tag).getCodec());
   }
 
-  private static class InitializedMemoizingCodec<T> implements MemoizingCodec<T> {
-    private final T obj;
-    private boolean deserialized = false;
-
-    private InitializedMemoizingCodec(T obj) {
-      this.obj = obj;
-    }
-
-    @Override
-    public Strategy getStrategy() {
-      return Strategy.DO_NOT_MEMOIZE;
-    }
-
-    @Override
-    public void serializePayload(
-        SerializationContext context, T obj, CodedOutputStream codedOut, Serializer serializer) {
-      throw new UnsupportedOperationException("Unexpected serialize: " + obj);
-    }
-
-    @Override
-    public T deserializePayload(
-        DeserializationContext context,
-        @Nullable T initial,
-        CodedInputStream codedIn,
-        Deserializer deserializer) {
-      Preconditions.checkState(!deserialized, "Already deserialized: %s", obj);
-      deserialized = true;
-      return obj;
-    }
-
-    @SuppressWarnings("unchecked")
-    @Override
-    public Class<T> getEncodedClass() {
-      return (Class<T>) obj.getClass();
-    }
-  }
   @SuppressWarnings("unchecked")
   public <T> T getDependency(Class<T> type) {
     Preconditions.checkNotNull(type);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMapCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMapCodec.java
index 9be4046daf..2c1e11c92b 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMapCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMapCodec.java
@@ -24,22 +24,38 @@ import java.util.Comparator;
 import java.util.Map;
 
 /**
- * Encodes an {@link ImmutableMap}, which may be an ImmutableSortedMap. We handle both here because
- * we cannot handle ImmutableSortedMap with any ordering other than the default, and so we degrade
- * to ImmutableMap in that case, hoping that the caller does not notice.
+ * Encodes an {@link ImmutableMap}, which may be an {@link ImmutableSortedMap}. The iteration order
+ * of the deserialized map is the same as the original map's.
+ *
+ * <p>We handle both {@link ImmutableMap} and {@link ImmutableSortedMap} here because we cannot
+ * handle {@link ImmutableSortedMap} with any ordering other than the default, and so we degrade to
+ * ImmutableMap in that case, hoping that the caller does not notice. Since the ordering is
+ * preserved, the caller should not be sensitive to this change unless the caller's field is
+ * declared to be an {@link ImmutableSortedMap}, in which case we will crash at runtime.
+ *
+ * <p>Any {@link SerializationException} or {@link IOException} that arises while serializing or
+ * deserializing a map entry's value (not its key) will be wrapped in a new {@link
+ * SerializationException} using {@link SerializationException#propagate}. (Note that this preserves
+ * the type of {@link SerializationException.NoCodecException} exceptions.) The message will include
+ * the {@code toString()} of the entry's key. For errors that occur while serializing, it will also
+ * include the class name of the entry's value. Errors that occur while serializing an entry key are
+ * not affected.
+ *
+ * <p>Because of the ambiguity around the key type (Comparable in the case of {@link
+ * ImmutableSortedMap}, arbitrary otherwise, we avoid specifying the key type as a parameter.
  */
-class ImmutableMapCodec implements ObjectCodec<ImmutableMap<?, ?>> {
+class ImmutableMapCodec<V> implements ObjectCodec<ImmutableMap<?, V>> {
   @SuppressWarnings("unchecked")
   @Override
-  public Class<ImmutableMap<?, ?>> getEncodedClass() {
+  public Class<ImmutableMap<?, V>> getEncodedClass() {
     // Because Java disallows converting from Class<ImmutableMap> to Class<ImmutableMap<?, ?>>
     // directly.
-    return (Class<ImmutableMap<?, ?>>) ((Class<?>) ImmutableMap.class);
+    return (Class<ImmutableMap<?, V>>) ((Class<?>) ImmutableMap.class);
   }
 
   @Override
   public void serialize(
-      SerializationContext context, ImmutableMap<?, ?> map, CodedOutputStream codedOut)
+      SerializationContext context, ImmutableMap<?, V> map, CodedOutputStream codedOut)
       throws SerializationException, IOException {
     codedOut.writeInt32NoTag(map.size());
     boolean serializeAsSortedMap = false;
@@ -52,12 +68,20 @@ class ImmutableMapCodec implements ObjectCodec<ImmutableMap<?, ?>> {
     codedOut.writeBoolNoTag(serializeAsSortedMap);
     for (Map.Entry<?, ?> entry : map.entrySet()) {
       context.serialize(entry.getKey(), codedOut);
-      context.serialize(entry.getValue(), codedOut);
+      try {
+        context.serialize(entry.getValue(), codedOut);
+      } catch (SerializationException | IOException e) {
+        throw SerializationException.propagate(
+            String.format(
+                "Exception while serializing value of type %s for key '%s'",
+                entry.getValue().getClass().getName(), entry.getKey()),
+            e);
+      }
     }
   }
 
   @Override
-  public ImmutableMap<?, ?> deserialize(DeserializationContext context, CodedInputStream codedIn)
+  public ImmutableMap<?, V> deserialize(DeserializationContext context, CodedInputStream codedIn)
       throws SerializationException, IOException {
     int length = codedIn.readInt32();
     if (length < 0) {
@@ -70,15 +94,21 @@ class ImmutableMapCodec implements ObjectCodec<ImmutableMap<?, ?>> {
     }
   }
 
-  private static <T> ImmutableMap<T, Object> buildMap(
-      ImmutableMap.Builder<T, Object> builder,
+  private static <K, V> ImmutableMap<K, V> buildMap(
+      ImmutableMap.Builder<K, V> builder,
       int length,
       DeserializationContext context,
       CodedInputStream codedIn)
       throws IOException, SerializationException {
     for (int i = 0; i < length; i++) {
-      T key = context.deserialize(codedIn);
-      Object value = context.deserialize(codedIn);
+      K key = context.deserialize(codedIn);
+      V value;
+      try {
+        value = context.deserialize(codedIn);
+      } catch (SerializationException | IOException e) {
+        throw SerializationException.propagate(
+            String.format("Exception while deserializing value for key '%s'", key), e);
+      }
       builder.put(key, value);
     }
     try {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/Memoizer.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/Memoizer.java
index c31269ef5c..0276a847ce 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/Memoizer.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/Memoizer.java
@@ -14,28 +14,23 @@
 
 package com.google.devtools.build.lib.skyframe.serialization;
 
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableSet;
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec.MemoizationStrategy;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.Deque;
 import java.util.HashMap;
 import java.util.IdentityHashMap;
+import java.util.function.Function;
 import javax.annotation.Nullable;
 
 /**
  * A framework for serializing and deserializing with memo tables. Memoization is useful both for
  * performance and, in the case of cyclic data structures, to help avoid infinite recursion.
  *
- * <p><b>Usage:</b> To add support for a new type of value, create a class that implements {@link
- * MemoizingCodec}. To serialize, instantiate a new {@link Serializer} and call its {@link
- * Serializer#serialize serialize} method with the value and a codec for that value; and similarly
- * to deserialize, instantiate a {@link Deserializer} and call {@link Deserializer#deserialize
- * deserialize}. The memo table lives for the lifetime of the {@code Serializer}/{@code
- * Deserializer} instance. Do not call {@link MemoizingCodec#serializePayload} and {@link
- * MemoizingCodec#deserializePayload} directly, since that will bypass the memoization logic for the
- * top-level value to be encoded or decoded.
- *
  * <p>The memo table associates each value with an integer identifier.
  *
  * <ul>
@@ -135,7 +130,7 @@ import javax.annotation.Nullable;
 // that instead of failing with a stack overflow, we detect the cycle and throw
 // SerializationException. This requires just a little extra memo tracking for the MEMOIZE_AFTER
 // case.
-public class Memoizer {
+class Memoizer {
 
   /**
    * Constants used in the wire format to signal whether the next bytes are content or a
@@ -150,139 +145,9 @@ public class Memoizer {
 
   private Memoizer() {}
 
-  /** Indicates how a {@link MemoizingCodec} uses the memo table. */
-  public enum Strategy {
-    /**
-     * Indicates that the memo table is not directly used by this codec.
-     *
-     * <p>Codecs with this strategy will always serialize payloads, never backreferences, even if
-     * the same value has been serialized before. This does not apply to other codecs that are
-     * delegated to via {@link Serializer#serialize}. Deserialization behaves
-     * analogously.
-     *
-     * <p>This strategy is useful for codecs that write very little data themselves, but that still
-     * delegate to other codecs.
-     */
-    DO_NOT_MEMOIZE,
-
-    /**
-     * Indicates that the value is memoized before recursing to its children, so that it is
-     * available to form cyclic references from its children. If this strategy is used, {@link
-     * MemoizingCodec#makeInitialValue} must be overridden.
-     *
-     * <p>This should be used for all types where it is feasible to provide an initial value. Any
-     * cycle that does not go through at least one {@code MEMOIZE_BEFORE} type of value (e.g., a
-     * pathological self-referential tuple) is unserializable.
-     */
-    MEMOIZE_BEFORE,
-
-    /**
-     * Indicates that the value is memoized after recursing to its children, so that it cannot be
-     * referred to until after it has been constructed (regardless of whether its children are still
-     * under construction).
-     *
-     * <p>This is typically used for immutable types, since they cannot be created by mutating an
-     * initial value.
-     */
-    MEMOIZE_AFTER
-  }
-
-  /**
-   * A codec that knows how to serialize/deserialize {@code T} using a memo strategy.
-   *
-   * <p>Codecs do not interact with memo tables directly; rather, they just define a memo strategy
-   * and let {@link Serializer} and {@link Deserializer} take care of the memo logic and metadata. A
-   * codec can dispatch to a subcodec to handle components of {@code T} compositionally; this
-   * dispatching occurs via the {@code Serializer} or {@code Deserializer} so that the memo table is
-   * consulted.
-   *
-   * <p>Codecs may access additional data from the deserialization context. For instance, Skylark
-   * codecs may retrieve a {@link com.google.devtools.build.lib.syntax.Mutability} object to
-   * associate with the value they construct. The type of this additional data is checked
-   * dynamically, and a mistyping will result in a {@link SerializationException}.
-   */
-  // It is technically possible to work the type of the additional data into the generic type of
-  // MemoizingCodec and Deserializer. But it makes the types ugly, and probably isn't worth the
-  // damage to readability.
-  public interface MemoizingCodec<T> extends BaseCodec<T> {
-
-    /**
-     * Returns the memoization strategy for this codec.
-     *
-     * <p>If set to {@link Strategy#MEMOIZE_BEFORE}, then {@link #makeInitialValue} must be
-     * implemented.
-     *
-     * <p>Implementations of this method should just return a constant, since the choice of strategy
-     * is usually intrinsic to {@code T}.
-     */
-    Strategy getStrategy();
-
-    /**
-     * If {@link #getStrategy} returns {@link Strategy#MEMOIZE_BEFORE}, then this method is called
-     * to obtain a new initial / empty instance of {@code T} for deserialization. The returned
-     * instance will be populated by {@link #deserializePayload}.
-     *
-     * <p>If any other strategy is used, this method is ignored. The default implementation throws
-     * {@link UnsupportedOperationException}.
-     *
-     * <p>The passed in deserialization context can provide additional data that is used to help
-     * initialize the value. For instance, it may contain a Skylark {@link
-     * com.google.devtools.build.lib.syntax.Mutability}.
-     */
-    default T makeInitialValue(Deserializer deserializer) throws SerializationException {
-      throw new UnsupportedOperationException(
-          "No initial value was provided for codec class " + getClass().getName());
-    }
-
-    /**
-     * Serializes an object's payload.
-     *
-     * <p>To delegate to a subcodec, an implementation should call {@code serializer}'s {@link
-     * Serializer#serialize serialize} method with the appropriate piece of {@code obj} and the
-     * subcodec. Do not call the subcodec's {@code serializePayload} method directly, as that would
-     * bypass the memo table.
-     */
-    void serializePayload(
-        SerializationContext context,
-        T obj,
-        CodedOutputStream codedOut,
-        Serializer serializer)
-        throws SerializationException, IOException;
-
-    /**
-     * Deserializes an object's payload.
-     *
-     * <p>To delegate to a subcodec, an implementation should call {@code deserializer}'s {@link
-     * Deserializer#deserialize deserialize} method with the appropriate subcodec. Do not call the
-     * subcodec's {@code deserializePayload} method directly, as that would bypass the memo table.
-     *
-     * If this codec uses the {@link Strategy#MEMOIZE_BEFORE} strategy (as determined by {@link
-     * #getStrategy}), then the {@code initial} parameter will be a new value of {@code T} that can
-     * be mutated into the result. In that case, this function must return {@code initial}. If any
-     * other strategy is used, then {@code initial} will be null and this function must instantiate
-     * the value to return.
-     */
-    T deserializePayload(
-        DeserializationContext context,
-        @Nullable T initial,
-        CodedInputStream codedIn,
-        Deserializer deserializer)
-        throws SerializationException, IOException;
-  }
-
-  /** A context for serializing; wraps a memo table. */
-  public static class Serializer {
+  /** A context for serializing; wraps a memo table. Not thread-safe. */
+  static class Serializer {
     private final SerializingMemoTable memo = new SerializingMemoTable();
-
-    public <T> void serialize(SerializationContext context, T obj, CodedOutputStream codedOut)
-        throws IOException, SerializationException {
-      MemoizingCodec<? super T> memoizingCodec =
-          context.recordAndMaybeGetMemoizingCodec(obj, codedOut);
-      if (memoizingCodec != null) {
-        serialize(context, obj, memoizingCodec, codedOut);
-      }
-    }
-
     /**
      * Serializes an object using the given codec and current memo table state.
      *
@@ -292,12 +157,12 @@ public class Memoizer {
     <T> void serialize(
         SerializationContext context,
         T obj,
-        MemoizingCodec<? super T> codec,
+        ObjectCodec<? super T> codec,
         CodedOutputStream codedOut)
         throws SerializationException, IOException {
-      Strategy strategy = codec.getStrategy();
-      if (strategy == Strategy.DO_NOT_MEMOIZE) {
-        codec.serializePayload(context, obj, codedOut, this);
+      MemoizationStrategy strategy = codec.getStrategy();
+      if (strategy == MemoizationStrategy.DO_NOT_MEMOIZE) {
+        codec.serialize(context, obj, codedOut);
       } else {
         Integer id = memo.lookupNullable(obj);
         if (id != null) {
@@ -314,19 +179,19 @@ public class Memoizer {
     private <T> void serializeMemoContent(
         SerializationContext context,
         T obj,
-        MemoizingCodec<T> codec,
+        ObjectCodec<T> codec,
         CodedOutputStream codedOut,
-        Strategy strategy)
+        MemoizationStrategy strategy)
         throws SerializationException, IOException {
       switch(strategy) {
         case MEMOIZE_BEFORE: {
           int id = memo.memoize(obj);
           codedOut.writeInt32NoTag(id);
-          codec.serializePayload(context, obj, codedOut, this);
+            codec.serialize(context, obj, codedOut);
           break;
         }
         case MEMOIZE_AFTER: {
-          codec.serializePayload(context, obj, codedOut, this);
+            codec.serialize(context, obj, codedOut);
           // If serializing the children caused the parent object itself to be serialized due to a
           // cycle, then there's now a memo entry for the parent. Don't overwrite it with a new id.
           Integer cylicallyCreatedId = memo.lookupNullable(obj);
@@ -368,9 +233,13 @@ public class Memoizer {
     }
   }
 
-  /** A context for deserializing; wraps a memo table and possibly additional data. */
-  public static class Deserializer {
+  /**
+   * A context for deserializing; wraps a memo table and possibly additional data. Not thread-safe.
+   */
+  static class Deserializer {
     private final DeserializingMemoTable memo = new DeserializingMemoTable();
+    @Nullable private Integer tagForMemoizedBefore = null;
+    private final Deque<Object> memoizedBeforeStackForSanityChecking = new ArrayDeque<>();
 
     /**
      * Additional data is dynamically typed, and retrieved using a type token.
@@ -378,15 +247,9 @@ public class Memoizer {
      * <p>If we need to support multiple kinds of additional data in the future, this could become
      * a mapping.
      */
-    @Nullable
     private final Object additionalData;
 
-    public Deserializer() {
-      this.additionalData = null;
-    }
-
-    @VisibleForTesting
-    public Deserializer(Object additionalData) {
+    Deserializer(Object additionalData) {
       Preconditions.checkNotNull(additionalData);
       this.additionalData = additionalData;
     }
@@ -399,9 +262,6 @@ public class Memoizer {
      * @throws IllegalArgumentException if the additional data is not an instance of {@code klass}
      */
     <T> T getAdditionalData(Class<T> klass) {
-      Preconditions.checkNotNull(additionalData,
-          "Codec requires additional data of type %s, but no additional data is defined",
-          klass.getName());
       try {
         return klass.cast(additionalData);
       } catch (ClassCastException e) {
@@ -414,41 +274,27 @@ public class Memoizer {
     }
 
     /**
-     * For use with generic types where a {@link Class} object is not available at runtime, or where
-     * {@link Object} is the desired return type anyway.
-     */
-    public Object deserialize(DeserializationContext context, CodedInputStream codedIn)
-        throws IOException, SerializationException {
-      return deserialize(context, codedIn, Object.class);
-    }
-
-    public <T> T deserialize(
-        DeserializationContext context, CodedInputStream codedIn, Class<T> deserializedClass)
-        throws SerializationException, IOException {
-      @SuppressWarnings("unchecked")
-      MemoizingCodec<? extends T> codec =
-          (MemoizingCodec<? extends T>) context.getMemoizingCodecRecordedInInput(codedIn);
-      return deserializedClass.cast(deserialize(context, codec, codedIn));
-    }
-
-    /**
      * Deserializes an object using the given codec and current memo table state.
      *
      * @throws SerializationException on a logical error during deserialization
      * @throws IOException on {@link IOException} during deserialization
      */
     <T> T deserialize(
-        DeserializationContext context, MemoizingCodec<? extends T> codec, CodedInputStream codedIn)
+        DeserializationContext context, ObjectCodec<? extends T> codec, CodedInputStream codedIn)
         throws SerializationException, IOException {
-      Strategy strategy = codec.getStrategy();
-      if (strategy == Strategy.DO_NOT_MEMOIZE) {
-        return codec.deserializePayload(context, null, codedIn, this);
+      Preconditions.checkState(
+          tagForMemoizedBefore == null,
+          "non-null memoized-before tag %s (%s)",
+          tagForMemoizedBefore,
+          codec);
+      MemoizationStrategy strategy = codec.getStrategy();
+      if (strategy == MemoizationStrategy.DO_NOT_MEMOIZE) {
+        return codec.deserialize(context, codedIn);
       } else {
-        MemoEntry memoEntry =
-            memoEntryCodec.deserialize(context, codedIn);
+        MemoEntry memoEntry = memoEntryCodec.deserialize(context, codedIn);
         if (memoEntry == MemoEntry.BACKREF) {
           int id = codedIn.readInt32();
-          return lookupBackreference(id, codec.getEncodedClass());
+          return lookupBackreference(id, codec);
         } else if (memoEntry == MemoEntry.NEW_VALUE) {
           switch (strategy) {
             case MEMOIZE_BEFORE:
@@ -464,33 +310,71 @@ public class Memoizer {
       }
     }
 
+    private <T> T safeCast(Object obj, ObjectCodec<T> codec) throws SerializationException {
+      if (obj == null) {
+        return null;
+      }
+      ImmutableSet<Class<? extends T>> expectedTypes =
+          codec.additionalEncodedClasses().isEmpty()
+              ? ImmutableSet.of(codec.getEncodedClass())
+              : ImmutableSet.<Class<? extends T>>builderWithExpectedSize(
+                      codec.additionalEncodedClasses().size() + 1)
+                  .add(codec.getEncodedClass())
+                  .addAll(codec.additionalEncodedClasses())
+                  .build();
+      Class<?> objectClass = obj.getClass();
+      if (expectedTypes.contains(objectClass)) {
+        @SuppressWarnings("unchecked")
+        T checkedResult = (T) obj;
+        return checkedResult;
+      }
+      for (Class<? extends T> expectedType : expectedTypes) {
+        if (expectedType.isAssignableFrom(objectClass)) {
+          return expectedType.cast(obj);
+        }
+      }
+      throw new SerializationException(
+          "Object "
+              + obj
+              + ") has type "
+              + objectClass.getName()
+              + " but expected type one of "
+              + expectedTypes);
+    }
+
+    private <T> T castedDeserialize(
+        DeserializationContext context, ObjectCodec<T> codec, CodedInputStream codedIn)
+        throws IOException, SerializationException {
+      return safeCast(codec.deserialize(context, codedIn), codec);
+    }
     /** Retrieves a memo entry, validating that it exists and has the expected type. */
-    private <T> T lookupBackreference(int id, Class<T> expectedType) throws SerializationException {
+    private <T> T lookupBackreference(int id, ObjectCodec<T> codec) throws SerializationException {
       Object savedUnchecked = memo.lookup(id);
       if (savedUnchecked == null) {
         throw new SerializationException(
             "Found backreference to non-existent memo id (" + id + ")");
       }
-      try {
-        return expectedType.cast(savedUnchecked);
-      } catch (ClassCastException e) {
-        throw new SerializationException(
-            "Backreference (to memo id " + id + ") has type "
-                + savedUnchecked.getClass().getName()
-                + " but expected type " + expectedType.getName());
-      }
+      return safeCast(savedUnchecked, codec);
+    }
+
+    <A, T> T makeInitialValue(Function<A, T> initialValueFunction, Class<A> klass) {
+      T initial = initialValueFunction.apply(getAdditionalData(klass));
+      int tag =
+          Preconditions.checkNotNull(
+              tagForMemoizedBefore, " Not called with memoize before: %s", initial);
+      tagForMemoizedBefore = null;
+      memo.memoize(tag, initial);
+      memoizedBeforeStackForSanityChecking.addLast(initial);
+      return initial;
     }
 
     // Corresponds to MemoBeforeContent in the abstract grammar.
     private <T> T deserializeMemoBeforeContent(
-        DeserializationContext context,
-        MemoizingCodec<T> codec,
-        CodedInputStream codedIn)
+        DeserializationContext context, ObjectCodec<T> codec, CodedInputStream codedIn)
         throws SerializationException, IOException {
-      int id = codedIn.readInt32();
-      T initial = codec.makeInitialValue(this);
-      memo.memoize(id, initial);
-      T value = codec.deserializePayload(context, initial, codedIn, this);
+      this.tagForMemoizedBefore = codedIn.readInt32();
+      T value = castedDeserialize(context, codec, codedIn);
+      Object initial = memoizedBeforeStackForSanityChecking.removeLast();
       if (value != initial) {
         // This indicates a bug in the particular codec subclass.
         throw new SerializationException("doDeserialize did not return the initial instance");
@@ -500,11 +384,9 @@ public class Memoizer {
 
     // Corresponds to MemoAfterContent in the abstract grammar.
     private <T> T deserializeMemoAfterContent(
-        DeserializationContext context,
-        MemoizingCodec<T> codec,
-        CodedInputStream codedIn)
+        DeserializationContext context, ObjectCodec<T> codec, CodedInputStream codedIn)
         throws SerializationException, IOException {
-      T value = codec.deserializePayload(context, null, codedIn, this);
+      T value = castedDeserialize(context, codec, codedIn);
       int id = codedIn.readInt32();
       // If deserializing the children caused the parent object itself to be deserialized due to
       // a cycle, then there's now a memo entry for the parent. Reuse that object, discarding
@@ -512,18 +394,7 @@ public class Memoizer {
       // the object graph.
       Object cyclicallyCreatedObject = memo.lookup(id);
       if (cyclicallyCreatedObject != null) {
-        Class<T> expectedType = codec.getEncodedClass();
-        try {
-          return expectedType.cast(cyclicallyCreatedObject);
-        } catch (ClassCastException e) {
-          // This indicates either some kind of type reification mismatch, or a memo id
-          // collision.
-          throw new SerializationException(
-              "While trying to deserialize a value of type " + expectedType.getName()
-                  + " marked with memo id " + id + ", another value having that id was "
-                  + "recursively created; but that value had unexpected type "
-                  + cyclicallyCreatedObject.getClass().getName());
-        }
+        return safeCast(cyclicallyCreatedObject, codec);
       } else {
         memo.memoize(id, value);
         return value;
@@ -553,99 +424,4 @@ public class Memoizer {
       }
     }
   }
-
-  /**
-   * An adaptor for treating an {@link ObjectCodec} as a {@link MemoizingCodec} that does not use
-   * memoization.
-   */
-  static class ObjectCodecAdaptor<T> implements MemoizingCodec<T> {
-
-    private final ObjectCodec<T> codec;
-
-    ObjectCodecAdaptor(ObjectCodec<T> codec) {
-      this.codec = codec;
-    }
-
-    @Override
-    public Class<T> getEncodedClass() {
-      return codec.getEncodedClass();
-    }
-
-    @Override
-    public Strategy getStrategy() {
-      return Strategy.DO_NOT_MEMOIZE;
-    }
-
-    @Override
-    public void serializePayload(
-        SerializationContext context,
-        T obj,
-        CodedOutputStream codedOut,
-        Serializer serializer)
-        throws SerializationException, IOException {
-      codec.serialize(context, obj, codedOut);
-    }
-
-    @Override
-    public T deserializePayload(DeserializationContext context,
-        T initial,
-        CodedInputStream codedIn,
-        Deserializer deserializer)
-        throws SerializationException, IOException {
-      return codec.deserialize(context, codedIn);
-    }
-  }
-
-  /** An adapter for {@link ObjectCodec} that allows the values to be memoized after creation. */
-  static class MemoizingAfterObjectCodecAdapter<T> extends ObjectCodecAdaptor<T> {
-    MemoizingAfterObjectCodecAdapter(ObjectCodec<T> codec) {
-      super(codec);
-    }
-
-    @Override
-    public Strategy getStrategy() {
-      return Strategy.MEMOIZE_AFTER;
-    }
-  }
-
-  /**
-   * An adaptor for treating a {@link MemoizingCodec} as an {@link ObjectCodec}. By default, each
-   * call to {@link ObjectCodec#serialize} or {@link ObjectCodec#deserialize} uses a fresh memo
-   * table.
-   *
-   * <p>A subclass can override the {@link #getDeserializer} hook to add additional context
-   * information, e.g. a Skylark {@link com.google.devtools.build.lib.syntax.Mutability} to share
-   * among all deserialized Skylark values. Note that if a {@code Deserializer} is reused across
-   * multiple calls to this hook, this adaptor may become stateful.
-   */
-  static class MemoizingCodecAdaptor<T> implements ObjectCodec<T> {
-
-    private final MemoizingCodec<T> codec;
-
-    MemoizingCodecAdaptor(MemoizingCodec<T> codec) {
-      this.codec = codec;
-    }
-
-    /** Provides the context used in {@link #deserialize}. */
-    Deserializer getDeserializer() {
-      return new Deserializer();
-    }
-
-    @Override
-    public Class<T> getEncodedClass() {
-      return codec.getEncodedClass();
-    }
-
-    @Override
-    public void serialize(SerializationContext context, T obj, CodedOutputStream codedOut)
-        throws SerializationException, IOException {
-      new Serializer().serialize(context, obj, codec, codedOut);
-    }
-
-    @Override
-    public T deserialize(DeserializationContext context, CodedInputStream codedIn)
-        throws SerializationException, IOException {
-      return getDeserializer().deserialize(context, codec, codedIn);
-    }
-  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodec.java
index a17cf3d83d..4b47b4bad0 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodec.java
@@ -14,15 +14,42 @@
 
 package com.google.devtools.build.lib.skyframe.serialization;
 
+import com.google.common.collect.ImmutableList;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
+import java.util.List;
 
 /**
  * Generic object serialization/deserialization. Implementations should serialize values
  * deterministically.
  */
-public interface ObjectCodec<T> extends BaseCodec<T> {
+public interface ObjectCodec<T> {
+  /**
+   * Returns the class of the objects serialized/deserialized by this codec.
+   *
+   * <p>This is useful for automatically dispatching to the correct codec, e.g. in {@link
+   * ObjectCodecs}.
+   */
+  Class<T> getEncodedClass();
+
+  /**
+   * Returns additional subtypes of {@code T} that may be serialized/deserialized using this codec
+   * without loss of information.
+   *
+   * <p>This method is intended for when {@code T} has multiple concrete implementations whose
+   * details are known to the codec but not to the codec dispatching mechanism. It signals that the
+   * dispatcher may choose to use this codec for the subtype, rather than raise {@link
+   * SerializationException.NoCodecException}.
+   *
+   * <p>This method should not be used if the codec's serialization and deserialization methods
+   * perform their own dispatching to other codecs for subtypes of {@code T}.
+   *
+   * <p>{@code T} itself should not be included in the returned list.
+   */
+  default List<Class<? extends T>> additionalEncodedClasses() {
+    return ImmutableList.of();
+  }
 
   /**
    * Serializes {@code obj}, inverse of {@link #deserialize}.
@@ -50,4 +77,55 @@ public interface ObjectCodec<T> extends BaseCodec<T> {
    */
   T deserialize(DeserializationContext context, CodedInputStream codedIn)
       throws SerializationException, IOException;
+
+  /**
+   * Returns the memoization strategy for this codec.
+   *
+   * <p>If set to {@link MemoizationStrategy#MEMOIZE_BEFORE}, then {@link
+   * DeserializationContext#makeInitialValue} must be called first in the {@link #deserialize}
+   * method, before delegating to any other codecs.
+   *
+   * <p>Implementations of this method should just return a constant, since the choice of strategy
+   * is usually intrinsic to {@link T}.
+   */
+  default MemoizationStrategy getStrategy() {
+    return MemoizationStrategy.MEMOIZE_AFTER;
+  }
+
+  /** Indicates how an {@link ObjectCodec} is memoized. */
+  enum MemoizationStrategy {
+    /**
+     * Indicates that memoization is not directly used by this codec.
+     *
+     * <p>Codecs with this strategy will always serialize payloads, never backreferences, even if
+     * the same value has been serialized before. This does not apply to other codecs that are
+     * delegated to within this codec. Deserialization behaves analogously.
+     *
+     * <p>This strategy is useful for codecs that write very little data themselves, but that still
+     * delegate to other codecs.
+     */
+    DO_NOT_MEMOIZE,
+
+    /**
+     * Indicates that the value is memoized before recursing to its children, so that it is
+     * available to form cyclic references from its children. If this strategy is used, {@link
+     * DeserializationContext#makeInitialValue} must be called during the {@link #deserialize}
+     * method.
+     *
+     * <p>This should be used for all types where it is feasible to provide an initial value. Any
+     * cycle that does not go through at least one {@code MEMOIZE_BEFORE} type of value (e.g., a
+     * pathological self-referential tuple) is unserializable.
+     */
+    MEMOIZE_BEFORE,
+
+    /**
+     * Indicates that the value is memoized after recursing to its children, so that it cannot be
+     * referred to until after it has been constructed (regardless of whether its children are still
+     * under construction).
+     *
+     * <p>This is typically used for immutable types, since they cannot be created by mutating an
+     * initial value.
+     */
+    MEMOIZE_AFTER
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java
index 80ec3281ae..95590143db 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java
@@ -19,11 +19,9 @@ import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer.MemoizingCodec;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
-import java.util.Collection;
 import java.util.Comparator;
 import java.util.IdentityHashMap;
 import java.util.Map;
@@ -45,44 +43,33 @@ public class ObjectCodecRegistry {
   @Nullable
   private final CodecDescriptor defaultCodecDescriptor;
 
-  private final int memoizingCodecsStartTag;
-  private final ImmutableMap<Class<?>, MemoizingCodecDescriptor<?>> classMappedMemoizingCodecs;
-  private final ImmutableList<MemoizingCodecDescriptor<?>> tagMappedMemoizingCodecs;
-
   private final IdentityHashMap<Object, Integer> constantsMap;
   private final ImmutableList<Object> constants;
   private final int constantsStartTag;
 
   private ObjectCodecRegistry(
-      Map<Class<?>, CodecHolder> codecs,
-      ImmutableSet<MemoizingCodec<?>> memoizingCodecs,
+      ImmutableSet<ObjectCodec<?>> memoizingCodecs,
       ImmutableList<Object> constants,
       boolean allowDefaultCodec) {
-    ImmutableMap.Builder<Class<?>, CodecDescriptor> codecMappingsBuilder = ImmutableMap.builder();
     int nextTag = 1; // 0 is reserved for null.
-    for (Class<?> clazz :
-        ImmutableList.sortedCopyOf(Comparator.comparing(Class::getName), codecs.keySet())) {
-      codecMappingsBuilder.put(clazz, codecs.get(clazz).createDescriptor(nextTag));
-      nextTag++;
-    }
-    this.classMappedCodecs = codecMappingsBuilder.build();
+    ImmutableMap.Builder<Class<?>, CodecDescriptor> memoizingCodecsBuilder =
+        ImmutableMap.builderWithExpectedSize(memoizingCodecs.size());
+    ImmutableList.Builder<CodecDescriptor> tagMappedMemoizingCodecsBuilder =
+        ImmutableList.builderWithExpectedSize(memoizingCodecs.size());
+    nextTag =
+        processCodecs(
+            memoizingCodecs, nextTag, tagMappedMemoizingCodecsBuilder, memoizingCodecsBuilder);
 
     this.defaultCodecDescriptor =
         allowDefaultCodec
             ? new TypedCodecDescriptor<>(nextTag++, new JavaSerializableCodec())
             : null;
-    this.tagMappedCodecs = makeTagMappedCodecs(classMappedCodecs.values(), defaultCodecDescriptor);
+    if (allowDefaultCodec) {
+      tagMappedMemoizingCodecsBuilder.add(defaultCodecDescriptor);
+    }
 
-    this.memoizingCodecsStartTag = nextTag;
-    ImmutableMap.Builder<Class<?>, MemoizingCodecDescriptor<?>> memoizingCodecsBuilder =
-        ImmutableMap.builderWithExpectedSize(memoizingCodecs.size());
-    ImmutableList.Builder<MemoizingCodecDescriptor<?>> tagMappedMemoizingCodecsBuilder =
-        ImmutableList.builderWithExpectedSize(memoizingCodecs.size());
-    nextTag =
-        processMemoizingCodecs(
-            memoizingCodecs, nextTag, tagMappedMemoizingCodecsBuilder, memoizingCodecsBuilder);
-    this.classMappedMemoizingCodecs = memoizingCodecsBuilder.build();
-    this.tagMappedMemoizingCodecs = tagMappedMemoizingCodecsBuilder.build();
+    this.classMappedCodecs = memoizingCodecsBuilder.build();
+    this.tagMappedCodecs = tagMappedMemoizingCodecsBuilder.build();
 
     constantsStartTag = nextTag;
     constantsMap = new IdentityHashMap<>();
@@ -115,20 +102,6 @@ public class ObjectCodecRegistry {
   }
 
   @Nullable
-  public <T> MemoizingCodecDescriptor<? super T> getMemoizingCodecDescriptor(Class<T> type) {
-    for (Class<?> nextType = type; nextType != null; nextType = nextType.getSuperclass()) {
-      MemoizingCodecDescriptor<?> result = classMappedMemoizingCodecs.get(nextType);
-      if (result != null) {
-        @SuppressWarnings("unchecked")
-        MemoizingCodecDescriptor<? super T> castResult =
-            (MemoizingCodecDescriptor<? super T>) result;
-        return castResult;
-      }
-    }
-    return null;
-  }
-
-  @Nullable
   Object maybeGetConstantByTag(int tag) {
     return tag < constantsStartTag || tag - constantsStartTag >= constants.size()
         ? null
@@ -156,15 +129,6 @@ public class ObjectCodecRegistry {
     }
   }
 
-  @Nullable
-  public MemoizingCodec<?> maybeGetMemoizingCodecByTag(int tag) {
-    int tagOffset = tag - memoizingCodecsStartTag;
-    if (tagOffset < 0 || tagOffset > tagMappedMemoizingCodecs.size()) {
-      return null;
-    }
-    return tagMappedMemoizingCodecs.get(tagOffset).getMemoizingCodec();
-  }
-
   /**
    * Creates a builder using the current contents of this registry.
    *
@@ -175,12 +139,7 @@ public class ObjectCodecRegistry {
     Builder builder = newBuilder();
     builder.setAllowDefaultCodec(defaultCodecDescriptor != null);
     for (Map.Entry<Class<?>, CodecDescriptor> entry : classMappedCodecs.entrySet()) {
-      // Cast here is safe because the original #add in the Builder was type-checked.
-      builder.add(entry.getKey(), getObjectCodec(entry.getValue()));
-    }
-
-    for (MemoizingCodecDescriptor<?> descriptor : tagMappedMemoizingCodecs) {
-      builder.addMemoizing(descriptor.getMemoizingCodec());
+      builder.add(entry.getValue().getCodec());
     }
 
     for (Object constant : constants) {
@@ -189,11 +148,6 @@ public class ObjectCodecRegistry {
     return builder;
   }
 
-  @SuppressWarnings("unchecked")
-  private static ObjectCodec<Object> getObjectCodec(CodecDescriptor descriptor) {
-    return (ObjectCodec<Object>) descriptor.getCodec();
-  }
-
   /** Describes encoding logic. */
   interface CodecDescriptor {
     void serialize(SerializationContext context, Object obj, CodedOutputStream codedOut)
@@ -218,32 +172,6 @@ public class ObjectCodecRegistry {
     ObjectCodec<?> getCodec();
   }
 
-  static class MemoizingCodecDescriptor<T> {
-    private final int tag;
-    private final MemoizingCodec<T> memoizingCodec;
-
-    MemoizingCodecDescriptor(int tag, MemoizingCodec<T> memoizingCodec) {
-      this.tag = tag;
-      this.memoizingCodec = memoizingCodec;
-    }
-
-    int getTag() {
-      return tag;
-    }
-
-    MemoizingCodec<T> getMemoizingCodec() {
-      return memoizingCodec;
-    }
-
-    @Override
-    public String toString() {
-      return MoreObjects.toStringHelper(this)
-          .add("tag", tag)
-          .add("codec", memoizingCodec)
-          .toString();
-    }
-  }
-
   private static class TypedCodecDescriptor<T> implements CodecDescriptor {
     private final int tag;
     private final ObjectCodec<T> codec;
@@ -275,46 +203,21 @@ public class ObjectCodecRegistry {
     public ObjectCodec<T> getCodec() {
       return codec;
     }
-  }
-
-  private interface CodecHolder {
-    CodecDescriptor createDescriptor(int tag);
-  }
-
-  private static class TypedCodecHolder<T> implements CodecHolder {
-    private final ObjectCodec<T> codec;
-
-    private TypedCodecHolder(ObjectCodec<T> codec) {
-      this.codec = codec;
-    }
-
-    @Override
-    public CodecDescriptor createDescriptor(int tag) {
-      return new TypedCodecDescriptor<T>(tag, codec);
-    }
 
     @Override
     public String toString() {
-      return MoreObjects.toStringHelper(this).add("codec", codec).toString();
+      return MoreObjects.toStringHelper(this).add("codec", codec).add("tag", tag).toString();
     }
   }
 
   /** Builder for {@link ObjectCodecRegistry}. */
   public static class Builder {
-    private final ImmutableMap.Builder<Class<?>, CodecHolder> codecsBuilder =
-        ImmutableMap.builder();
-    private final ImmutableSet.Builder<MemoizingCodec<?>> memoizingCodecBuilder =
-        ImmutableSet.builder();
+    private final ImmutableSet.Builder<ObjectCodec<?>> codecBuilder = ImmutableSet.builder();
     private final ImmutableList.Builder<Object> constantsBuilder = ImmutableList.builder();
     private boolean allowDefaultCodec = true;
 
-    public <T> Builder add(Class<? extends T> type, ObjectCodec<T> codec) {
-      codecsBuilder.put(type, new TypedCodecHolder<>(codec));
-      return this;
-    }
-
-    public Builder addMemoizing(MemoizingCodec<?> memoizingCodec) {
-      memoizingCodecBuilder.add(memoizingCodec);
+    public Builder add(ObjectCodec<?> codec) {
+      codecBuilder.add(codec);
       return this;
     }
 
@@ -333,41 +236,23 @@ public class ObjectCodecRegistry {
 
     public ObjectCodecRegistry build() {
       return new ObjectCodecRegistry(
-          codecsBuilder.build(),
-          memoizingCodecBuilder.build(),
-          constantsBuilder.build(),
-          allowDefaultCodec);
-    }
-  }
-
-  private static ImmutableList<CodecDescriptor> makeTagMappedCodecs(
-      Collection<CodecDescriptor> codecs, @Nullable CodecDescriptor defaultCodecDescriptor) {
-    CodecDescriptor[] codecTable =
-        new CodecDescriptor[codecs.size() + (defaultCodecDescriptor != null ? 1 : 0)];
-    for (CodecDescriptor codecDescriptor : codecs) {
-      codecTable[codecDescriptor.getTag() - 1] = codecDescriptor;
-    }
-
-    if (defaultCodecDescriptor != null) {
-      codecTable[defaultCodecDescriptor.getTag() - 1] = defaultCodecDescriptor;
+          codecBuilder.build(), constantsBuilder.build(), allowDefaultCodec);
     }
-    return ImmutableList.copyOf(codecTable);
   }
 
-  private static int processMemoizingCodecs(
-      Iterable<? extends MemoizingCodec<?>> memoizingCodecs,
+  private static int processCodecs(
+      Iterable<? extends ObjectCodec<?>> memoizingCodecs,
       int nextTag,
-      ImmutableList.Builder<MemoizingCodecDescriptor<?>> tagMappedMemoizingCodecsBuilder,
-      ImmutableMap.Builder<Class<?>, MemoizingCodecDescriptor<?>> memoizingCodecsBuilder) {
-    for (MemoizingCodec<?> memoizingCodec :
+      ImmutableList.Builder<CodecDescriptor> tagMappedCodecsBuilder,
+      ImmutableMap.Builder<Class<?>, CodecDescriptor> codecsBuilder) {
+    for (ObjectCodec<?> codec :
         ImmutableList.sortedCopyOf(
             Comparator.comparing(o -> o.getEncodedClass().getName()), memoizingCodecs)) {
-      MemoizingCodecDescriptor<?> codecDescriptor =
-          new MemoizingCodecDescriptor<>(nextTag++, memoizingCodec);
-      tagMappedMemoizingCodecsBuilder.add(codecDescriptor);
-      memoizingCodecsBuilder.put(memoizingCodec.getEncodedClass(), codecDescriptor);
-      for (Class<?> otherClass : memoizingCodec.additionalEncodedSubclasses()) {
-        memoizingCodecsBuilder.put(otherClass, codecDescriptor);
+      CodecDescriptor codecDescriptor = new TypedCodecDescriptor<>(nextTag++, codec);
+      tagMappedCodecsBuilder.add(codecDescriptor);
+      codecsBuilder.put(codec.getEncodedClass(), codecDescriptor);
+      for (Class<?> otherClass : codec.additionalEncodedClasses()) {
+        codecsBuilder.put(otherClass, codecDescriptor);
       }
     }
     return nextTag;
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java
index 1d0fbe66b6..e34ecc1a14 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java
@@ -53,21 +53,41 @@ public class ObjectCodecs {
     }
   }
 
-  public void serialize(Object subject, CodedOutputStream codedOut) throws SerializationException {
+  public void serialize(
+      Object subject, CodedOutputStream codedOut, MemoizationPermission memoizationPermission)
+      throws SerializationException {
+    SerializationContext context = serializationContext;
+    if (memoizationPermission == MemoizationPermission.DISABLED) {
+      context = context.disableMemoization();
+    }
     try {
-      serializationContext.serialize(subject, codedOut);
+      context.serialize(subject, codedOut);
     } catch (IOException e) {
       throw new SerializationException("Failed to serialize " + subject, e);
     }
   }
 
+  /**
+   * Controls whether memoization can occur for serialization/deserialization. Should be allowed
+   * unless bit-equivalence is needed.
+   */
+  public enum MemoizationPermission {
+    ALLOWED,
+    DISABLED
+  }
+
   public Object deserialize(ByteString data) throws SerializationException {
-    return deserialize(data.newCodedInput());
+    return deserialize(data.newCodedInput(), MemoizationPermission.ALLOWED);
   }
 
-  public Object deserialize(CodedInputStream codedIn) throws SerializationException {
+  public Object deserialize(CodedInputStream codedIn, MemoizationPermission memoizationPermission)
+      throws SerializationException {
+    DeserializationContext context = deserializationContext;
+    if (memoizationPermission == MemoizationPermission.DISABLED) {
+      context = context.disableMemoization();
+    }
     try {
-      return deserializationContext.deserialize(codedIn);
+      return context.deserialize(codedIn);
     } catch (IOException e) {
       throw new SerializationException("Failed to deserialize data", e);
     }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java
index b106a63c44..2eb22ba2a3 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java
@@ -14,32 +14,77 @@
 
 package com.google.devtools.build.lib.skyframe.serialization;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer.MemoizingAfterObjectCodecAdapter;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer.MemoizingCodec;
-import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecRegistry.MemoizingCodecDescriptor;
+import com.google.devtools.build.lib.skyframe.serialization.Memoizer.Serializer;
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecs.MemoizationPermission;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationException.NoCodecException;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
+import javax.annotation.CheckReturnValue;
 import javax.annotation.Nullable;
 
-/** Stateful class for providing additional context to a single serialization "session". */
+/**
+ * Stateful class for providing additional context to a single serialization "session". This class
+ * is thread-safe so long as {@link #serializer} is null. If it is not null, this class is not
+ * thread-safe and should only be accessed on a single thread for serializing one object (that may
+ * involve serializing other objects contained in it).
+ */
 public class SerializationContext {
-
   private final ObjectCodecRegistry registry;
   private final ImmutableMap<Class<?>, Object> dependencies;
+  private final MemoizationPermission memoizationPermission;
+  @Nullable private final Memoizer.Serializer serializer;
 
-  public SerializationContext(
-      ObjectCodecRegistry registry, ImmutableMap<Class<?>, Object> dependencies) {
+  private SerializationContext(
+      ObjectCodecRegistry registry,
+      ImmutableMap<Class<?>, Object> dependencies,
+      MemoizationPermission memoizationPermission,
+      @Nullable Serializer serializer) {
     this.registry = registry;
     this.dependencies = dependencies;
+    this.serializer = serializer;
+    this.memoizationPermission = memoizationPermission;
+    Preconditions.checkState(
+        serializer == null || memoizationPermission == MemoizationPermission.ALLOWED);
   }
 
+  @VisibleForTesting
+  public SerializationContext(
+      ObjectCodecRegistry registry, ImmutableMap<Class<?>, Object> dependencies) {
+    this(registry, dependencies, MemoizationPermission.ALLOWED, /*serializer=*/ null);
+  }
+
+  @VisibleForTesting
   public SerializationContext(ImmutableMap<Class<?>, Object> dependencies) {
     this(AutoRegistry.get(), dependencies);
   }
 
+  SerializationContext disableMemoization() {
+    Preconditions.checkState(
+        memoizationPermission == MemoizationPermission.ALLOWED, "memoization already disabled");
+    Preconditions.checkState(serializer == null, "serializer already present");
+    return new SerializationContext(
+        registry, dependencies, MemoizationPermission.DISABLED, serializer);
+  }
+
+  /**
+   * Returns a {@link SerializationContext} that will memoize values it encounters (using reference
+   * equality) in a new memoization table. The returned context should be used instead of the
+   * original: memoization may only occur when using the returned context. Calls must be in pairs
+   * with {@link DeserializationContext#newMemoizingContext} in the corresponding deserialization
+   * code. Over-eagerly calling this method will reduce the effectiveness of memoization, since any
+   * previous memo state is inaccessible to the new context.
+   */
+  @CheckReturnValue
+  public SerializationContext newMemoizingContext() {
+    Preconditions.checkState(
+        memoizationPermission == MemoizationPermission.ALLOWED, "memoization disabled");
+    return new SerializationContext(
+        this.registry, this.dependencies, memoizationPermission, new Memoizer.Serializer());
+  }
+
   private boolean writeNullOrConstant(@Nullable Object object, CodedOutputStream codedOut)
       throws IOException {
     if (object == null) {
@@ -71,30 +116,15 @@ public class SerializationContext {
     ObjectCodecRegistry.CodecDescriptor descriptor =
         recordAndGetDescriptorIfNotConstantOrNull(object, codedOut);
     if (descriptor != null) {
-      descriptor.serialize(this, object, codedOut);
+      if (serializer == null) {
+        descriptor.serialize(this, object, codedOut);
+      } else {
+        @SuppressWarnings("unchecked")
+        ObjectCodec<Object> castCodec = (ObjectCodec<Object>) descriptor.getCodec();
+        serializer.serialize(this, object, castCodec, codedOut);
+      }
     }
-  }
-
-  @Nullable
-  public <T> MemoizingCodec<? super T> recordAndMaybeGetMemoizingCodec(
-      T object, CodedOutputStream codedOut) throws IOException, NoCodecException {
-    if (writeNullOrConstant(object, codedOut)) {
-      return null;
     }
-    @SuppressWarnings("unchecked")
-    Class<T> clazz = (Class<T>) object.getClass();
-    MemoizingCodecDescriptor<? super T> memoizingCodecDescriptor =
-        registry.getMemoizingCodecDescriptor(clazz);
-    if (memoizingCodecDescriptor != null) {
-      codedOut.writeSInt32NoTag(memoizingCodecDescriptor.getTag());
-      return memoizingCodecDescriptor.getMemoizingCodec();
-    }
-    ObjectCodecRegistry.CodecDescriptor codecDescriptor = registry.getCodecDescriptor(clazz);
-    codedOut.writeSInt32NoTag(codecDescriptor.getTag());
-    @SuppressWarnings("unchecked")
-    ObjectCodec<? super T> objectCodec = (ObjectCodec<? super T>) codecDescriptor.getCodec();
-    return new MemoizingAfterObjectCodecAdapter<>(objectCodec);
-  }
 
   @SuppressWarnings("unchecked")
   public <T> T getDependency(Class<T> type) {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/UnmodifiableListCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/UnmodifiableListCodec.java
deleted file mode 100644
index 648b1a5f42..0000000000
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/UnmodifiableListCodec.java
+++ /dev/null
@@ -1,60 +0,0 @@
-// Copyright 2018 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;
-
-import com.google.protobuf.CodedInputStream;
-import com.google.protobuf.CodedOutputStream;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.LinkedList;
-import java.util.List;
-
-/** A {@link ObjectCodec} for objects produced by {@link Collections#unmodifiableList}. */
-class UnmodifiableListCodec implements ObjectCodec<List<Object>> {
-
-  private static final Class<?> RANDOM_ACCESS_TYPE =
-      Collections.unmodifiableList(new ArrayList<Object>()).getClass();
-  private static final Class<?> SEQUENTIAL_ACCESS_TYPE =
-      Collections.unmodifiableList(new LinkedList<Object>()).getClass();
-
-  @Override
-  public Class<List<Object>> getEncodedClass() {
-    return null; // No reasonable value here.
-  }
-
-  @Override
-  public void serialize(
-      SerializationContext context, List<Object> object, CodedOutputStream output) {
-    // TODO(shahan): Stub. Replace with actual implementation, which requires the registry to be
-    // added to the context.
-  }
-
-  @Override
-  public List<Object> deserialize(DeserializationContext context, CodedInputStream input) {
-    // TODO(shahan): Stub. Replace with actual implementation, which requires the registry to be
-    // added to the context.
-    return null;
-  }
-
-  static class UnmodifiableListCodecRegisterer implements CodecRegisterer<UnmodifiableListCodec> {
-    @SuppressWarnings({"rawtypes", "unchecked"})
-    @Override
-    public void register(ObjectCodecRegistry.Builder builder) {
-      UnmodifiableListCodec codec = new UnmodifiableListCodec();
-      builder.add((Class) RANDOM_ACCESS_TYPE, (ObjectCodec) codec);
-      builder.add((Class) SEQUENTIAL_ACCESS_TYPE, (ObjectCodec) codec);
-    }
-  }
-}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodec.java
index e6a687d773..3d9d68bdb9 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodec.java
@@ -83,6 +83,20 @@ public @interface AutoCodec {
 
   Strategy strategy() default Strategy.INSTANTIATOR;
 
+  /** Whether to start memoizing values below this codec. */
+  enum Memoization {
+    /** Do not start memoization, but also do not disable memoization if it is already happening. */
+    UNCHANGED,
+    /**
+     * Start memoizing. Memoization is assumed to always need a Skylark "Mutability" object. If this
+     * package does not have access to the {@link com.google.devtools.build.lib.syntax.Mutability}
+     * class, memoization cannot be started here.
+     */
+    START_MEMOIZING
+  }
+
+  Memoization memoization() default Memoization.UNCHANGED;
+
   /**
    * Signals that the annotated element is only visible for use by serialization. It should not be
    * used by other callers.
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 dd7ad40fc0..268a87cf5e 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
@@ -24,6 +24,7 @@ import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.skyframe.serialization.CodecScanningConstants;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec.Memoization;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationCodeGenerator.Marshaller;
 import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.FieldSpec;
@@ -101,12 +102,13 @@ public class AutoCodecProcessor extends AbstractProcessor {
       if (element instanceof TypeElement) {
         TypeElement encodedType = (TypeElement) element;
         TypeSpec.Builder codecClassBuilder;
+        boolean startMemoizing = annotation.memoization() == Memoization.START_MEMOIZING;
         switch (annotation.strategy()) {
           case INSTANTIATOR:
-            codecClassBuilder = buildClassWithInstantiatorStrategy(encodedType);
+            codecClassBuilder = buildClassWithInstantiatorStrategy(encodedType, startMemoizing);
             break;
           case PUBLIC_FIELDS:
-            codecClassBuilder = buildClassWithPublicFieldsStrategy(encodedType);
+            codecClassBuilder = buildClassWithPublicFieldsStrategy(encodedType, startMemoizing);
             break;
           default:
             throw new IllegalArgumentException("Unknown strategy: " + annotation.strategy());
@@ -166,7 +168,8 @@ public class AutoCodecProcessor extends AbstractProcessor {
         .build();
   }
 
-  private TypeSpec.Builder buildClassWithInstantiatorStrategy(TypeElement encodedType) {
+  private TypeSpec.Builder buildClassWithInstantiatorStrategy(
+      TypeElement encodedType, boolean startMemoizing) {
     ExecutableElement constructor = selectInstantiator(encodedType);
     List<? extends VariableElement> fields = constructor.getParameters();
 
@@ -175,16 +178,17 @@ public class AutoCodecProcessor extends AbstractProcessor {
 
     if (encodedType.getAnnotation(AutoValue.class) == null) {
       initializeUnsafeOffsets(codecClassBuilder, encodedType, fields);
-      codecClassBuilder.addMethod(buildSerializeMethodWithInstantiator(encodedType, fields));
+      codecClassBuilder.addMethod(
+          buildSerializeMethodWithInstantiator(encodedType, fields, startMemoizing));
     } else {
       codecClassBuilder.addMethod(
-          buildSerializeMethodWithInstantiatorForAutoValue(encodedType, fields));
+          buildSerializeMethodWithInstantiatorForAutoValue(encodedType, fields, startMemoizing));
     }
 
     MethodSpec.Builder deserializeBuilder =
-        AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType, env);
+        AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType, startMemoizing, env);
     buildDeserializeBody(deserializeBuilder, fields);
-    addReturnNew(deserializeBuilder, encodedType, constructor, env);
+    addReturnNew(deserializeBuilder, encodedType, constructor, startMemoizing, env);
     codecClassBuilder.addMethod(deserializeBuilder.build());
 
     return codecClassBuilder;
@@ -287,9 +291,9 @@ public class AutoCodecProcessor extends AbstractProcessor {
   }
 
   private MethodSpec buildSerializeMethodWithInstantiator(
-      TypeElement encodedType, List<? extends VariableElement> fields) {
+      TypeElement encodedType, List<? extends VariableElement> fields, boolean startMemoizing) {
     MethodSpec.Builder serializeBuilder =
-        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType, env);
+        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType, startMemoizing, env);
     for (VariableElement parameter : fields) {
       Optional<FieldValueAndClass> hasField =
           getFieldByNameRecursive(encodedType, parameter.getSimpleName().toString());
@@ -388,16 +392,17 @@ public class AutoCodecProcessor extends AbstractProcessor {
   }
 
   private MethodSpec buildSerializeMethodWithInstantiatorForAutoValue(
-      TypeElement encodedType, List<? extends VariableElement> fields) {
+      TypeElement encodedType, List<? extends VariableElement> fields, boolean startMemoizing) {
     MethodSpec.Builder serializeBuilder =
-        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType, env);
+        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType, startMemoizing, env);
     for (VariableElement parameter : fields) {
       addSerializeParameterWithGetter(encodedType, parameter, serializeBuilder);
     }
     return serializeBuilder.build();
   }
 
-  private TypeSpec.Builder buildClassWithPublicFieldsStrategy(TypeElement encodedType) {
+  private TypeSpec.Builder buildClassWithPublicFieldsStrategy(
+      TypeElement encodedType, boolean startMemoizing) {
     TypeSpec.Builder codecClassBuilder =
         AutoCodecUtil.initializeCodecClassBuilder(encodedType, env);
     ImmutableList<? extends VariableElement> publicFields =
@@ -405,9 +410,10 @@ public class AutoCodecProcessor extends AbstractProcessor {
             .stream()
             .filter(this::isPublicField)
             .collect(toImmutableList());
-    codecClassBuilder.addMethod(buildSerializeMethodWithPublicFields(encodedType, publicFields));
+    codecClassBuilder.addMethod(
+        buildSerializeMethodWithPublicFields(encodedType, publicFields, startMemoizing));
     MethodSpec.Builder deserializeBuilder =
-        AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType, env);
+        AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType, startMemoizing, env);
     buildDeserializeBody(deserializeBuilder, publicFields);
     addInstantiatePopulateFieldsAndReturn(deserializeBuilder, encodedType, publicFields);
     codecClassBuilder.addMethod(deserializeBuilder.build());
@@ -423,9 +429,9 @@ public class AutoCodecProcessor extends AbstractProcessor {
   }
 
   private MethodSpec buildSerializeMethodWithPublicFields(
-      TypeElement encodedType, List<? extends VariableElement> fields) {
+      TypeElement encodedType, List<? extends VariableElement> fields, boolean startMemoizing) {
     MethodSpec.Builder serializeBuilder =
-        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType, env);
+        AutoCodecUtil.initializeSerializeMethodBuilder(encodedType, startMemoizing, env);
     for (VariableElement parameter : fields) {
       String paramAccessor = "input." + parameter.getSimpleName();
       marshallers.writeSerializationCode(
@@ -459,7 +465,13 @@ public class AutoCodecProcessor extends AbstractProcessor {
       MethodSpec.Builder builder,
       TypeElement type,
       ExecutableElement instantiator,
+      boolean startMemoizing,
       ProcessingEnvironment env) {
+    // TODO(janakr): When injection of additional data and memoizing are separated, rework this so
+    // that mutability isn't deeply embedded in AutoCodec.
+    if (startMemoizing) {
+      builder.addStatement("$L.close()", AutoCodecUtil.MUTABILITY_VARIABLE_NAME);
+    }
     List<? extends TypeMirror> allThrown = instantiator.getThrownTypes();
     if (!allThrown.isEmpty()) {
       builder.beginControlFlow("try");
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 102c1dfa4e..f0bbaf2270 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
@@ -38,7 +38,7 @@ class AutoCodecUtil {
   // Synthesized classes will have `_AutoCodec` suffix added.
   public static final String GENERATED_CLASS_NAME_SUFFIX = "AutoCodec";
   static final Class<AutoCodec> ANNOTATION = AutoCodec.class;
-
+  static final String MUTABILITY_VARIABLE_NAME = "mutabilityForMemoizeAdditionalData";
   /**
    * Initializes a builder for a class of the appropriate type.
    *
@@ -68,39 +68,54 @@ class AutoCodecUtil {
    * Initializes the deserialize method.
    *
    * @param encodedType type being serialized
+   * @param startMemoizing whether memoization should start in this method.
    */
   static MethodSpec.Builder initializeSerializeMethodBuilder(
-      TypeElement encodedType, ProcessingEnvironment env) {
+      TypeElement encodedType, boolean startMemoizing, ProcessingEnvironment env) {
     MethodSpec.Builder builder =
         MethodSpec.methodBuilder("serialize")
             .addModifiers(Modifier.PUBLIC)
             .returns(void.class)
             .addAnnotation(Override.class)
             .addException(SerializationException.class)
-            .addException(IOException.class);
-    return builder
-        .addParameter(SerializationContext.class, "context")
-        .addParameter(TypeName.get(env.getTypeUtils().erasure(encodedType.asType())), "input")
-        .addParameter(CodedOutputStream.class, "codedOut");
+            .addException(IOException.class)
+            .addParameter(SerializationContext.class, "context")
+            .addParameter(TypeName.get(env.getTypeUtils().erasure(encodedType.asType())), "input")
+            .addParameter(CodedOutputStream.class, "codedOut");
+    if (startMemoizing) {
+      builder.addStatement("context = context.newMemoizingContext()");
+    }
+    return builder;
   }
 
   /**
    * Initializes the deserialize method.
    *
    * @param encodedType type being serialized
+   * @param startMemoizing whether memoization should start in this method.
    */
   static MethodSpec.Builder initializeDeserializeMethodBuilder(
-      TypeElement encodedType, ProcessingEnvironment env) {
+      TypeElement encodedType, boolean startMemoizing, ProcessingEnvironment env) {
     MethodSpec.Builder builder =
         MethodSpec.methodBuilder("deserialize")
             .addModifiers(Modifier.PUBLIC)
             .returns(TypeName.get(env.getTypeUtils().erasure(encodedType.asType())))
             .addAnnotation(Override.class)
             .addException(SerializationException.class)
-            .addException(IOException.class);
-    return builder
-        .addParameter(DeserializationContext.class, "context")
-        .addParameter(CodedInputStream.class, "codedIn");
+            .addException(IOException.class)
+            .addParameter(DeserializationContext.class, "context")
+            .addParameter(CodedInputStream.class, "codedIn");
+    if (startMemoizing) {
+      // We can't directly use the Mutability class here because there are @AutoCodec'ed classes
+      // that Mutability depends on. Those classes won't start memoization, of course, but the code
+      // generator doesn't know that.
+      builder.addStatement(
+          "com.google.devtools.build.lib.syntax.Mutability $L = "
+              + "com.google.devtools.build.lib.syntax.Mutability.create(\"deserialize skylark\")",
+          MUTABILITY_VARIABLE_NAME);
+      builder.addStatement("context = context.newMemoizingContext($L)", MUTABILITY_VARIABLE_NAME);
+    }
+    return builder;
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/strings/FastStringCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/strings/FastStringCodec.java
index 8e1a8fc641..0cd658a8c6 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/strings/FastStringCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/strings/FastStringCodec.java
@@ -65,6 +65,15 @@ class FastStringCodec implements ObjectCodec<String> {
   }
 
   @Override
+  public MemoizationStrategy getStrategy() {
+    // Don't memoize strings inside memoizing serialization, to preserve current behavior.
+    // TODO(janakr,brandjon,michajlo): Is it actually a problem to memoize strings? Doubt there
+    // would be much performance impact from increasing the size of the identity map, and we
+    // could potentially drop our string tables in the future.
+    return MemoizationStrategy.DO_NOT_MEMOIZE;
+  }
+
+  @Override
   public void serialize(SerializationContext context, String string, CodedOutputStream codedOut)
       throws IOException {
     codedOut.writeStringNoTag(string);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/strings/StringCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/strings/StringCodec.java
index 30901ed95a..aa0e169849 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/strings/StringCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/strings/StringCodec.java
@@ -30,6 +30,15 @@ class StringCodec implements ObjectCodec<String> {
   }
 
   @Override
+  public MemoizationStrategy getStrategy() {
+    // Don't memoize strings inside memoizing serialization, to preserve current behavior.
+    // TODO(janakr,brandjon,michajlo): Is it actually a problem to memoize strings? Doubt there
+    // would be much performance impact from increasing the size of the identity map, and we
+    // could potentially drop our string tables in the future.
+    return MemoizationStrategy.DO_NOT_MEMOIZE;
+  }
+
+  @Override
   public void serialize(SerializationContext context, String str, CodedOutputStream codedOut)
       throws IOException {
     codedOut.writeStringNoTag(str);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/strings/StringCodecs.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/strings/StringCodecs.java
index aea5a4c67d..a5516a75a6 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/strings/StringCodecs.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/strings/StringCodecs.java
@@ -76,7 +76,7 @@ public final class StringCodecs {
     @Override
     public void register(ObjectCodecRegistry.Builder builder) {
       if (!supportsOptimizedAscii()) {
-        builder.add(String.class, simple());
+        builder.add(simple());
       }
     }
   }
@@ -90,7 +90,7 @@ public final class StringCodecs {
     @Override
     public void register(ObjectCodecRegistry.Builder builder) {
       if (supportsOptimizedAscii()) {
-        builder.add(String.class, asciiOptimized());
+        builder.add(asciiOptimized());
       }
     }
   }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializationTester.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializationTester.java
index b0af31a1a8..32a42b2a27 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializationTester.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializationTester.java
@@ -22,9 +22,8 @@ import com.google.common.base.Stopwatch;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.skyframe.serialization.AutoRegistry;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer.Deserializer;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer.MemoizingCodec;
+import com.google.devtools.build.lib.skyframe.serialization.DeserializationContext;
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecRegistry;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecs;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
@@ -62,11 +61,7 @@ public class SerializationTester {
 
   private final ImmutableList<Object> subjects;
   private final ImmutableMap.Builder<Class<?>, Object> dependenciesBuilder;
-  private final ArrayList<MemoizingCodec<?>> memoizingCodecs = new ArrayList<>();
-  // TODO(janakr): in future, there may be tests that want to start with default serialization but
-  // need memoizing codecs, so they should be able to set this to false even after adding a
-  // memoizing codec.
-  private boolean useMemoization = false;
+  private final ArrayList<ObjectCodec<?>> additionalCodecs = new ArrayList<>();
   private Object additionalDeserializationData;
 
   @SuppressWarnings("rawtypes")
@@ -90,9 +85,8 @@ public class SerializationTester {
     return this;
   }
 
-  public SerializationTester addMemoizingCodec(MemoizingCodec<?> memoizingCodec) {
-    memoizingCodecs.add(memoizingCodec);
-    useMemoization = true;
+  public SerializationTester addCodec(ObjectCodec<?> codec) {
+    additionalCodecs.add(codec);
     return this;
   }
 
@@ -122,8 +116,8 @@ public class SerializationTester {
     for (Object val : dependencies.values()) {
       registryBuilder.addConstant(val);
     }
-    for (MemoizingCodec<?> memoizingCodec : memoizingCodecs) {
-      registryBuilder.addMemoizing(memoizingCodec);
+    for (ObjectCodec<?> codec : additionalCodecs) {
+      registryBuilder.add(codec);
     }
     ObjectCodecs codecs = new ObjectCodecs(registryBuilder.build(), dependencies);
     testSerializeDeserialize(codecs);
@@ -133,26 +127,26 @@ public class SerializationTester {
 
   private ByteString serialize(Object subject, ObjectCodecs codecs)
       throws SerializationException, IOException {
-    if (!useMemoization) {
+    if (additionalDeserializationData == null) {
       return codecs.serialize(subject);
     }
     ByteString.Output output = ByteString.newOutput();
     CodedOutputStream codedOut = CodedOutputStream.newInstance(output);
-    new Memoizer.Serializer()
-        .serialize(codecs.getSerializationContextForTesting(), subject, codedOut);
+    codecs.getSerializationContextForTesting().newMemoizingContext().serialize(subject, codedOut);
     codedOut.flush();
     return output.toByteString();
   }
 
   private Object deserialize(ByteString serialized, ObjectCodecs codecs)
       throws SerializationException, IOException {
-    if (!useMemoization) {
+    if (additionalDeserializationData == null) {
       return codecs.deserialize(serialized);
     }
-    return (additionalDeserializationData == null
-            ? new Deserializer()
-            : new Memoizer.Deserializer(additionalDeserializationData))
-        .deserialize(codecs.getDeserializationContextForTesting(), serialized.newCodedInput());
+    DeserializationContext context =
+        codecs
+            .getDeserializationContextForTesting()
+            .newMemoizingContext(additionalDeserializationData);
+    return context.deserialize(serialized.newCodedInput());
   }
 
   /** Runs serialization/deserialization tests. */
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/TestUtils.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/TestUtils.java
index 2d5d7c9a96..a7478d4fa7 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/TestUtils.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/TestUtils.java
@@ -18,16 +18,11 @@ import static com.google.common.truth.Truth.assertThat;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.skyframe.serialization.AutoRegistry;
-import com.google.devtools.build.lib.skyframe.serialization.CodecRegisterer;
 import com.google.devtools.build.lib.skyframe.serialization.DeserializationContext;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer.Deserializer;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer.MemoizingCodec;
-import com.google.devtools.build.lib.skyframe.serialization.Memoizer.Serializer;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecRegistry;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationContext;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
-import com.google.devtools.build.lib.skyframe.serialization.strings.StringCodecs;
 import com.google.devtools.build.lib.syntax.Environment.GlobalFrame;
 import com.google.devtools.build.lib.syntax.Mutability;
 import com.google.protobuf.ByteString;
@@ -54,10 +49,10 @@ public class TestUtils {
 
   public static <T> ByteString toBytes(T value, ImmutableMap<Class<?>, Object> dependencies)
       throws IOException, SerializationException {
-    return toBytes(value, new SerializationContext(dependencies));
+    return toBytes(new SerializationContext(dependencies), value);
   }
 
-  private static <T> ByteString toBytes(T value, SerializationContext serializationContext)
+  public static <T> ByteString toBytes(SerializationContext serializationContext, T value)
       throws IOException, SerializationException {
     ByteString.Output output = ByteString.newOutput();
     CodedOutputStream codedOut = CodedOutputStream.newInstance(output);
@@ -66,6 +61,11 @@ public class TestUtils {
     return output.toByteString();
   }
 
+  public static Object fromBytes(DeserializationContext deserializationContext, ByteString bytes)
+      throws IOException, SerializationException {
+    return deserializationContext.deserialize(bytes.newCodedInput());
+  }
+
   /** Deserialize a value from a byte array. */
   public static <T> T fromBytes(DeserializationContext context, ObjectCodec<T> codec, byte[] bytes)
       throws SerializationException, IOException {
@@ -81,7 +81,11 @@ public class TestUtils {
     ObjectCodecRegistry registry = builder.build();
     return new DeserializationContext(registry, dependencies)
         .deserialize(
-            toBytes(value, new SerializationContext(registry, dependencies)).newCodedInput());
+            toBytes(new SerializationContext(registry, dependencies), value).newCodedInput());
+  }
+
+  public static <T> T roundTrip(T value) throws IOException, SerializationException {
+    return TestUtils.roundTrip(value, ImmutableMap.of());
   }
 
   /**
@@ -106,81 +110,50 @@ public class TestUtils {
       throws IOException, SerializationException {
     ByteString.Output output = ByteString.newOutput();
     CodedOutputStream codedOut = CodedOutputStream.newInstance(output);
-    new Serializer()
-        .serialize(new SerializationContext(registry, ImmutableMap.of()), original, codedOut);
+    new SerializationContext(registry, ImmutableMap.of())
+        .newMemoizingContext()
+        .serialize(original, codedOut);
     codedOut.flush();
     return output.toByteString();
   }
 
   public static Object fromBytesMemoized(ByteString bytes, ObjectCodecRegistry registry)
       throws IOException, SerializationException {
-    return new Deserializer()
-        .deserialize(
-            new DeserializationContext(registry, ImmutableMap.of()), bytes.newCodedInput());
+    return new DeserializationContext(registry, ImmutableMap.of())
+        .newMemoizingContext(Mutability.create("deserialize skylark"))
+        .deserialize(bytes.newCodedInput());
   }
 
   public static <T> T roundTripMemoized(T original, ObjectCodecRegistry registry)
       throws IOException, SerializationException {
-    return roundTripMemoized(original, /*mutability=*/ null, registry);
+    return roundTripMemoized(original, Mutability.create("deserialize skylark"), registry);
   }
 
-  @SuppressWarnings("unchecked") // Cast to T in return.
   public static <T> T roundTripMemoized(
       T original, @Nullable Mutability mutability, ObjectCodecRegistry registry)
       throws IOException, SerializationException {
-    return (T)
-        (mutability == null ? new Deserializer() : new Deserializer(mutability))
-            .deserialize(
-                new DeserializationContext(registry, ImmutableMap.of()),
-                toBytesMemoized(original, registry).newCodedInput());
+    return new DeserializationContext(registry, ImmutableMap.of())
+        .newMemoizingContext(mutability)
+        .deserialize(toBytesMemoized(original, registry).newCodedInput());
   }
 
-  public static <T> T roundTripMemoized(T original, MemoizingCodec<?>... codecs)
+  public static <T> T roundTripMemoized(T original, ObjectCodec<?>... codecs)
       throws IOException, SerializationException {
-    return roundTripMemoized(original, getBuilderWithMemoized(codecs).build());
+    return roundTripMemoized(original, getBuilderWithAdditionalCodecs(codecs).build());
   }
 
   public static <T> T roundTripMemoized(
-      T original, @Nullable Mutability mutability, MemoizingCodec<?>... codecs)
+      T original, @Nullable Mutability mutability, ObjectCodec<?>... codecs)
       throws IOException, SerializationException {
-    return roundTripMemoized(original, mutability, getBuilderWithMemoized(codecs).build());
+    return roundTripMemoized(original, mutability, getBuilderWithAdditionalCodecs(codecs).build());
   }
 
-  public static ObjectCodecRegistry.Builder getBuilderWithMemoized(MemoizingCodec<?>... codecs) {
+  public static ObjectCodecRegistry.Builder getBuilderWithAdditionalCodecs(
+      ObjectCodec<?>... codecs) {
     ObjectCodecRegistry.Builder builder = AutoRegistry.get().getBuilder();
-    for (MemoizingCodec<?> codec : codecs) {
-      builder.addMemoizing(codec);
+    for (ObjectCodec<?> codec : codecs) {
+      builder.add(codec);
     }
     return builder;
   }
-
-  /**
-   * Fake string codec that replaces all input and output string values with the constant "dummy".
-   */
-  public static class ConstantStringCodec implements ObjectCodec<String> {
-
-    private static final ObjectCodec<String> stringCodec = StringCodecs.simple();
-
-    @Override
-    public Class<String> getEncodedClass() {
-      return String.class;
-    }
-
-    @Override
-    public void serialize(SerializationContext context, String value, CodedOutputStream codedOut)
-        throws SerializationException, IOException {
-      stringCodec.serialize(context, "dummy", codedOut);
-    }
-
-    @Override
-    public String deserialize(DeserializationContext context, CodedInputStream codedIn)
-        throws SerializationException, IOException {
-      stringCodec.deserialize(context, codedIn);
-      return "dummy";
-    }
-
-    /** Disables auto-registration of ConstantStringCodec. */
-    private static class ConstantStringCodecRegisterer
-        implements CodecRegisterer<ConstantStringCodec> {}
-  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Environment.java b/src/main/java/com/google/devtools/build/lib/syntax/Environment.java
index eb06a10862..86d9c369ee 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/Environment.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/Environment.java
@@ -24,6 +24,8 @@ import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.events.EventKind;
 import com.google.devtools.build.lib.events.Location;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec.Memoization;
 import com.google.devtools.build.lib.syntax.Mutability.Freezable;
 import com.google.devtools.build.lib.syntax.Mutability.MutabilityException;
 import com.google.devtools.build.lib.util.Fingerprint;
@@ -473,6 +475,9 @@ public final class Environment implements Freezable {
 
   /** An Extension to be imported with load() into a BUILD or .bzl file. */
   @Immutable
+  // TODO(janakr,brandjon): Do Extensions actually have to start their own memoization? Or can we
+  // have a node higher up in the hierarchy inject the mutability?
+  @AutoCodec(memoization = Memoization.START_MEMOIZING)
   public static final class Extension {
 
     private final ImmutableMap<String, Object> bindings;
@@ -485,6 +490,7 @@ public final class Environment implements Freezable {
     private final String transitiveContentHashCode;
 
     /** Constructs with the given hash code and bindings. */
+    @AutoCodec.Instantiator
     public Extension(ImmutableMap<String, Object> bindings, String transitiveContentHashCode) {
       this.bindings = bindings;
       this.transitiveContentHashCode = transitiveContentHashCode;
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/UserDefinedFunction.java b/src/main/java/com/google/devtools/build/lib/syntax/UserDefinedFunction.java
index 485bac26e8..fb60ca6df0 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/UserDefinedFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/UserDefinedFunction.java
@@ -18,13 +18,15 @@ import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.Location;
 import com.google.devtools.build.lib.profiler.Profiler;
 import com.google.devtools.build.lib.profiler.ProfilerTask;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.skylarkinterface.SkylarkPrinter;
 import com.google.devtools.build.lib.syntax.Environment.LexicalFrame;
 
 /**
- * The actual function registered in the environment. This function is defined in the
- * parsed code using {@link FunctionDefStatement}.
+ * The actual function registered in the environment. This function is defined in the parsed code
+ * using {@link FunctionDefStatement}.
  */
+@AutoCodec
 public class UserDefinedFunction extends BaseFunction {
 
   private final ImmutableList<Statement> statements;
@@ -34,11 +36,11 @@ public class UserDefinedFunction extends BaseFunction {
 
   public UserDefinedFunction(
       String name,
-      Location loc,
+      Location location,
       FunctionSignature.WithValues<Object, SkylarkType> signature,
       ImmutableList<Statement> statements,
       Environment.GlobalFrame definitionGlobals) {
-    super(name, signature, loc);
+    super(name, signature, location);
     this.statements = statements;
     this.definitionGlobals = definitionGlobals;
   }