Clean up unnecessary "additional data" from memoizing deserialization. Since memoization is now a simple on-off switch, change semantics to have at most one memoizing frame: starting memoization is now an idempotent operation.

PiperOrigin-RevId: 189993914

8 files changed, 51 insertions, 111 deletions

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 55908b9e64..86aea0408f 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
@@ -22,7 +22,6 @@ import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecRegistry.
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecs.MemoizationPermission;
 import com.google.protobuf.CodedInputStream;
 import java.io.IOException;
-import java.util.function.Function;
 import javax.annotation.CheckReturnValue;
 
 /**
@@ -70,42 +69,29 @@ public class DeserializationContext {
   /**
    * 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
+   * returned by {@link SerializationContext#getMemoizingContext}. 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.
+   * <p>This method is idempotent: calling it on an already memoizing context will return the same
+   * context.
    */
   @CheckReturnValue
-  public DeserializationContext newMemoizingContext(Object additionalData) {
-    Preconditions.checkNotNull(additionalData);
+  public DeserializationContext getMemoizingContext() {
     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.
+    if (deserializer != null) {
+      return this;
+    }
     return new DeserializationContext(
-        this.registry, this.dependencies, memoizationPermission, new Deserializer(additionalData));
+        this.registry, this.dependencies, memoizationPermission, new Deserializer());
   }
 
   /**
-   * 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}.
+   * Register an initial value for the currently deserializing value, for use by child objects that
+   * may have references to it. Only for use during memoizing deserialization.
    */
-  public <A, T> T makeInitialValue(Function<A, T> initialValueFunction, Class<A> klass) {
-    // TODO(janakr): Possibly remove the "additional data" mechanism entirely now that it's no
-    // longer needed to thread Skylark Mutability's around. Instead of initial value functions, we
-    // can just have codecs construct their initial values directly. The only reasons to keep
-    // additional data around are 1) if we think something else might need it in the future (but
-    // YAGNI), or 2) if we want to share the same Mutability across multiple deserializations
-    // (within a single thread) to reduce garbage. For 2), we could even access the Mutability
-    // through dedicated methods instead of a generic "additional data" mechanism, though it might
-    // require subclassing DeserializationContext to avoid adding a dependency on Skylark.
-    return deserializer.makeInitialValue(initialValueFunction, klass);
+  public <T> void registerInitialValue(T initialValue) {
+    deserializer.registerInitialValue(initialValue);
   }
 
   // TODO(shahan): consider making codedIn a member of this class.
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 0276a847ce..2ffc339db1 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
@@ -24,7 +24,6 @@ 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;
 
 /**
@@ -242,38 +241,6 @@ class Memoizer {
     private final Deque<Object> memoizedBeforeStackForSanityChecking = new ArrayDeque<>();
 
     /**
-     * Additional data is dynamically typed, and retrieved using a type token.
-     *
-     * <p>If we need to support multiple kinds of additional data in the future, this could become
-     * a mapping.
-     */
-    private final Object additionalData;
-
-    Deserializer(Object additionalData) {
-      Preconditions.checkNotNull(additionalData);
-      this.additionalData = additionalData;
-    }
-
-    /**
-     * If this {@code Deserializer} was constructed with (non-null) additional data, and if its type
-     * satisfies the given type token {@code klass}, returns that additional data.
-     *
-     * @throws NullPointerException if no additional data is present
-     * @throws IllegalArgumentException if the additional data is not an instance of {@code klass}
-     */
-    <T> T getAdditionalData(Class<T> klass) {
-      try {
-        return klass.cast(additionalData);
-      } catch (ClassCastException e) {
-        throw new IllegalArgumentException(String.format(
-            "Codec requires additional data of type %s, but the available additional data has type "
-            + "%s",
-            klass.getName(),
-            additionalData.getClass().getName()));
-      }
-    }
-
-    /**
      * Deserializes an object using the given codec and current memo table state.
      *
      * @throws SerializationException on a logical error during deserialization
@@ -357,15 +324,13 @@ class Memoizer {
       return safeCast(savedUnchecked, codec);
     }
 
-    <A, T> T makeInitialValue(Function<A, T> initialValueFunction, Class<A> klass) {
-      T initial = initialValueFunction.apply(getAdditionalData(klass));
+    <T> void registerInitialValue(T initialValue) {
       int tag =
           Preconditions.checkNotNull(
-              tagForMemoizedBefore, " Not called with memoize before: %s", initial);
+              tagForMemoizedBefore, " Not called with memoize before: %s", initialValue);
       tagForMemoizedBefore = null;
-      memo.memoize(tag, initial);
-      memoizedBeforeStackForSanityChecking.addLast(initial);
-      return initial;
+      memo.memoize(tag, initialValue);
+      memoizedBeforeStackForSanityChecking.addLast(initialValue);
     }
 
     // Corresponds to MemoBeforeContent in the abstract grammar.
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 a754246db6..ab004def86 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
@@ -82,7 +82,7 @@ public interface ObjectCodec<T> {
    * 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}
+   * DeserializationContext#registerInitialValue} 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
@@ -109,7 +109,7 @@ public interface ObjectCodec<T> {
     /**
      * 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}
+     * DeserializationContext#registerInitialValue} 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
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 2eb22ba2a3..6e0d6811d9 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
@@ -73,14 +73,19 @@ public class SerializationContext {
    * 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.
+   * with {@link DeserializationContext#getMemoizingContext} in the corresponding deserialization
+   * code.
+   *
+   * <p>This method is idempotent: calling it on an already memoizing context will return the same
+   * context.
    */
   @CheckReturnValue
-  public SerializationContext newMemoizingContext() {
+  public SerializationContext getMemoizingContext() {
     Preconditions.checkState(
         memoizationPermission == MemoizationPermission.ALLOWED, "memoization disabled");
+    if (serializer != null) {
+      return this;
+    }
     return new SerializationContext(
         this.registry, this.dependencies, memoizationPermission, new Memoizer.Serializer());
   }
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 631bce8729..eb8e6c268c 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
@@ -193,8 +193,7 @@ public class AutoCodecProcessor extends AbstractProcessor {
     MethodSpec.Builder deserializeBuilder =
         AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType, startMemoizing, env);
     buildDeserializeBody(deserializeBuilder, fields);
-    addReturnNew(
-        deserializeBuilder, encodedType, constructor, /*builderVar=*/ null, startMemoizing, env);
+    addReturnNew(deserializeBuilder, encodedType, constructor, /*builderVar=*/ null, env);
     codecClassBuilder.addMethod(deserializeBuilder.build());
 
     return codecClassBuilder;
@@ -223,7 +222,7 @@ public class AutoCodecProcessor extends AbstractProcessor {
     String builderVarName =
         buildDeserializeBodyWithBuilder(
             encodedType, builderType, deserializeBuilder, getters, builderCreationMethod);
-    addReturnNew(deserializeBuilder, encodedType, buildMethod, builderVarName, startMemoizing, env);
+    addReturnNew(deserializeBuilder, encodedType, buildMethod, builderVarName, env);
     codecClassBuilder.addMethod(deserializeBuilder.build());
 
     return codecClassBuilder;
@@ -695,13 +694,7 @@ public class AutoCodecProcessor extends AbstractProcessor {
       TypeElement type,
       ExecutableElement instantiator,
       Object builderVar,
-      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 dd4706ab15..2ba6d9b2c2 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
@@ -37,9 +37,9 @@ import javax.lang.model.type.TypeMirror;
 /** Static utilities for AutoCodec processors. */
 class AutoCodecUtil {
   // Synthesized classes will have `_AutoCodec` suffix added.
-  public static final String GENERATED_CLASS_NAME_SUFFIX = "AutoCodec";
+  private 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.
    *
@@ -84,7 +84,7 @@ class AutoCodecUtil {
             .addParameter(TypeName.get(env.getTypeUtils().erasure(encodedType.asType())), "input")
             .addParameter(CodedOutputStream.class, "codedOut");
     if (startMemoizing) {
-      builder.addStatement("context = context.newMemoizingContext()");
+      builder.addStatement("context = context.getMemoizingContext()");
     }
     return builder;
   }
@@ -107,14 +107,7 @@ class AutoCodecUtil {
             .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);
+      builder.addStatement("context = context.getMemoizingContext()");
     }
     return builder;
   }
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 32a42b2a27..91926dbc8e 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
@@ -62,7 +62,7 @@ public class SerializationTester {
   private final ImmutableList<Object> subjects;
   private final ImmutableMap.Builder<Class<?>, Object> dependenciesBuilder;
   private final ArrayList<ObjectCodec<?>> additionalCodecs = new ArrayList<>();
-  private Object additionalDeserializationData;
+  private boolean memoize;
 
   @SuppressWarnings("rawtypes")
   private VerificationFunction verificationFunction =
@@ -90,9 +90,8 @@ public class SerializationTester {
     return this;
   }
 
-  public SerializationTester setAdditionalDeserializationData(
-      Object additionalDeserializationData) {
-    this.additionalDeserializationData = additionalDeserializationData;
+  public SerializationTester makeMemoizing() {
+    this.memoize = true;
     return this;
   }
 
@@ -127,25 +126,23 @@ public class SerializationTester {
 
   private ByteString serialize(Object subject, ObjectCodecs codecs)
       throws SerializationException, IOException {
-    if (additionalDeserializationData == null) {
+    if (!memoize) {
       return codecs.serialize(subject);
     }
     ByteString.Output output = ByteString.newOutput();
     CodedOutputStream codedOut = CodedOutputStream.newInstance(output);
-    codecs.getSerializationContextForTesting().newMemoizingContext().serialize(subject, codedOut);
+    codecs.getSerializationContextForTesting().getMemoizingContext().serialize(subject, codedOut);
     codedOut.flush();
     return output.toByteString();
   }
 
   private Object deserialize(ByteString serialized, ObjectCodecs codecs)
       throws SerializationException, IOException {
-    if (additionalDeserializationData == null) {
+    if (!memoize) {
       return codecs.deserialize(serialized);
     }
     DeserializationContext context =
-        codecs
-            .getDeserializationContextForTesting()
-            .newMemoizingContext(additionalDeserializationData);
+        codecs.getDeserializationContextForTesting().getMemoizingContext();
     return context.deserialize(serialized.newCodedInput());
   }
 
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 a7478d4fa7..63f9d35c15 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
@@ -72,6 +72,13 @@ public class TestUtils {
     return codec.deserialize(context, CodedInputStream.newInstance(bytes));
   }
 
+  public static <T> T roundTrip(T value, ObjectCodecRegistry registry)
+      throws IOException, SerializationException {
+    return new DeserializationContext(registry, ImmutableMap.of())
+        .deserialize(
+            toBytes(new SerializationContext(registry, ImmutableMap.of()), value).newCodedInput());
+  }
+
   public static <T> T roundTrip(T value, ImmutableMap<Class<?>, Object> dependencies)
       throws IOException, SerializationException {
     ObjectCodecRegistry.Builder builder = AutoRegistry.get().getBuilder();
@@ -111,7 +118,7 @@ public class TestUtils {
     ByteString.Output output = ByteString.newOutput();
     CodedOutputStream codedOut = CodedOutputStream.newInstance(output);
     new SerializationContext(registry, ImmutableMap.of())
-        .newMemoizingContext()
+        .getMemoizingContext()
         .serialize(original, codedOut);
     codedOut.flush();
     return output.toByteString();
@@ -120,20 +127,14 @@ public class TestUtils {
   public static Object fromBytesMemoized(ByteString bytes, ObjectCodecRegistry registry)
       throws IOException, SerializationException {
     return new DeserializationContext(registry, ImmutableMap.of())
-        .newMemoizingContext(Mutability.create("deserialize skylark"))
+        .getMemoizingContext()
         .deserialize(bytes.newCodedInput());
   }
 
   public static <T> T roundTripMemoized(T original, ObjectCodecRegistry registry)
       throws IOException, SerializationException {
-    return roundTripMemoized(original, Mutability.create("deserialize skylark"), registry);
-  }
-
-  public static <T> T roundTripMemoized(
-      T original, @Nullable Mutability mutability, ObjectCodecRegistry registry)
-      throws IOException, SerializationException {
     return new DeserializationContext(registry, ImmutableMap.of())
-        .newMemoizingContext(mutability)
+        .getMemoizingContext()
         .deserialize(toBytesMemoized(original, registry).newCodedInput());
   }
 
@@ -145,7 +146,7 @@ public class TestUtils {
   public static <T> T roundTripMemoized(
       T original, @Nullable Mutability mutability, ObjectCodec<?>... codecs)
       throws IOException, SerializationException {
-    return roundTripMemoized(original, mutability, getBuilderWithAdditionalCodecs(codecs).build());
+    return roundTripMemoized(original, getBuilderWithAdditionalCodecs(codecs).build());
   }
 
   public static ObjectCodecRegistry.Builder getBuilderWithAdditionalCodecs(