Makes the decision to memoize at the SkyValue level.

PiperOrigin-RevId: 192158825

6 files changed, 17 insertions, 26 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ASTFileLookupValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/ASTFileLookupValue.java
index 2bb6d1ab9f..e42ba4efac 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ASTFileLookupValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ASTFileLookupValue.java
@@ -45,7 +45,8 @@ public abstract class ASTFileLookupValue implements NotComparableSkyValue {
   public abstract String getErrorMsg();
 
   /** If the file is found, this class encapsulates the parsed AST. */
-  private static class ASTLookupWithFile extends ASTFileLookupValue {
+  @AutoCodec.VisibleForSerialization
+  public static class ASTLookupWithFile extends ASTFileLookupValue {
     private final BuildFileAST ast;
 
     private ASTLookupWithFile(BuildFileAST ast) {
@@ -71,7 +72,8 @@ public abstract class ASTFileLookupValue implements NotComparableSkyValue {
   }
 
   /** If the file isn't found, this class encapsulates a message with the reason. */
-  private static class ASTLookupNoFile extends ASTFileLookupValue {
+  @AutoCodec.VisibleForSerialization
+  public static class ASTLookupNoFile extends ASTFileLookupValue {
     private final String errorMsg;
 
     private ASTLookupNoFile(String errorMsg) {
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 1e025af063..2080517e1f 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
@@ -104,12 +104,9 @@ public class DeserializationContext {
    *
    * <p>This method is idempotent: calling it on an already memoizing context will return the same
    * context.
-   *
-   * <p><em>This is public for testing and {@link
-   * com.google.devtools.build.lib.packages.PackageSerializer} only.</em>
    */
   @CheckReturnValue
-  public DeserializationContext getMemoizingContext() {
+  DeserializationContext getMemoizingContext() {
     if (deserializer != null) {
       return this;
     }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/DynamicCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/DynamicCodec.java
index d8218da3ef..c4301a8fe4 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/DynamicCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/DynamicCodec.java
@@ -58,8 +58,6 @@ public class DynamicCodec implements ObjectCodec<Object> {
   @Override
   public void serialize(SerializationContext context, Object obj, CodedOutputStream codedOut)
       throws SerializationException, IOException {
-    // TODO(janakr,shahan): Remove when memoization is on by default.
-    context = context.getMemoizingContext();
     for (Map.Entry<Field, Long> entry : offsets.entrySet()) {
       serializeField(context, codedOut, obj, entry.getKey().getType(), entry.getValue());
     }
@@ -142,11 +140,6 @@ public class DynamicCodec implements ObjectCodec<Object> {
       throw new SerializationException("Could not instantiate object of type: " + type, e);
     }
     context.registerInitialValue(instance);
-    // We start memoizing if we weren't already doing so. We can't start before registering the
-    // initial value because the memoizer would get confused, since it doesn't have a tag for this
-    // object.
-    // TODO(janakr,shahan): Remove when memoization is on by default.
-    context = context.getMemoizingContext();
     for (Map.Entry<Field, Long> entry : offsets.entrySet()) {
       deserializeField(context, codedIn, instance, entry.getKey().getType(), entry.getValue());
     }
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 3674901826..557ed17255 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
@@ -38,6 +38,10 @@ public class ObjectCodecs {
     deserializationContext = new DeserializationContext(codecRegistry, dependencies);
   }
 
+  public ObjectCodecs(ObjectCodecRegistry codecRegistry) {
+    this(codecRegistry, ImmutableMap.of());
+  }
+
   public ByteString serialize(Object subject) throws SerializationException {
     return serializeToByteString(subject, this::serialize);
   }
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 a8110eaf8b..30b02517ee 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
@@ -86,12 +86,9 @@ public class SerializationContext {
    *
    * <p>This method is idempotent: calling it on an already memoizing context will return the same
    * context.
-   *
-   * <p><em>This is public for testing and {@link
-   * com.google.devtools.build.lib.packages.PackageSerializer} only.</em>
    */
   @CheckReturnValue
-  public SerializationContext getMemoizingContext() {
+  SerializationContext getMemoizingContext() {
     if (serializer != null) {
       return this;
     }
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 63f9d35c15..657765f87f 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
@@ -21,6 +21,7 @@ import com.google.devtools.build.lib.skyframe.serialization.AutoRegistry;
 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.SerializationContext;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
 import com.google.devtools.build.lib.syntax.Environment.GlobalFrame;
@@ -117,25 +118,22 @@ public class TestUtils {
       throws IOException, SerializationException {
     ByteString.Output output = ByteString.newOutput();
     CodedOutputStream codedOut = CodedOutputStream.newInstance(output);
-    new SerializationContext(registry, ImmutableMap.of())
-        .getMemoizingContext()
-        .serialize(original, codedOut);
+    new ObjectCodecs(registry).serializeMemoized(original, codedOut);
     codedOut.flush();
     return output.toByteString();
   }
 
   public static Object fromBytesMemoized(ByteString bytes, ObjectCodecRegistry registry)
       throws IOException, SerializationException {
-    return new DeserializationContext(registry, ImmutableMap.of())
-        .getMemoizingContext()
-        .deserialize(bytes.newCodedInput());
+    return new ObjectCodecs(registry).deserializeMemoized(bytes.newCodedInput());
   }
 
+  @SuppressWarnings("unchecked")
   public static <T> T roundTripMemoized(T original, ObjectCodecRegistry registry)
       throws IOException, SerializationException {
-    return new DeserializationContext(registry, ImmutableMap.of())
-        .getMemoizingContext()
-        .deserialize(toBytesMemoized(original, registry).newCodedInput());
+    return (T)
+        new ObjectCodecs(registry)
+            .deserializeMemoized(toBytesMemoized(original, registry).newCodedInput());
   }
 
   public static <T> T roundTripMemoized(T original, ObjectCodec<?>... codecs)