Use ClassToInstanceMap for serialization dependencies.

    Serialization deps are mappings from a class to an instance of that class. Using `ClassToInstanceMap` provides compile-time enforcement that the value is in fact an instance of the key.

    Also use `checkNotNull` in `getDependency` to throw when a dependency is not present, so that callers don't have to check.

    PiperOrigin-RevId: 369249742

dataset/GitHub_Java/bazelbuild.bazel/src/main/java/com/google/devtools/build/lib/packages/NativeAspectClass.java

@@ -14,58 +14,55 @@
 
 package com.google.devtools.build.lib.packages;
 
+import com.google.common.base.Preconditions;
+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.SerializationContext;
+import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.CodedOutputStream;
+import java.io.IOException;
+
 /**
  * A class of aspects that are implemented natively in Bazel.
  *
  * <p>This class just wraps a {@link java.lang.Class} implementing the
  * aspect factory. All wrappers of the same class are equal.
  */
-public final class NativeAspectClass<T extends NativeAspectClass.NativeAspectFactory>
-    implements AspectClass {
-  private final Class<? extends T> nativeClass;
-
-  public NativeAspectClass(Class<? extends T> nativeClass) {
-    this.nativeClass = nativeClass;
-  }
+public abstract class NativeAspectClass implements AspectClass {
+  public static final ObjectCodec<NativeAspectClass> CODEC = new Codec();
 
   @Override
   public String getName() {
-    return nativeClass.getSimpleName();
+    return getClass().getSimpleName();
   }
 
-  @Override
-  public AspectDefinition getDefinition() {
-    return newInstance().getDefinition();
-  }
+  public abstract AspectDefinition getDefinition(AspectParameters aspectParameters);
 
-  public T newInstance() {
-    try {
-      return nativeClass.newInstance();
-    } catch (Exception e) {
-      throw new IllegalStateException(e);
+  private static class Codec implements ObjectCodec<NativeAspectClass> {
+    @Override
+    public Class<NativeAspectClass> getEncodedClass() {
+      return NativeAspectClass.class;
     }
-  }
 
-  @Override
-  public int hashCode() {
-    return nativeClass.hashCode();
-  }
-
-  @Override
-  public boolean equals(Object obj) {
-    if (!(obj instanceof NativeAspectClass)) {
-      return false;
+    @Override
+    public void serialize(
+        SerializationContext context, NativeAspectClass obj, CodedOutputStream codedOut)
+        throws SerializationException, IOException {
+      RuleClassProvider ruleClassProvider = context.getDependency(RuleClassProvider.class);
+      NativeAspectClass storedAspect = ruleClassProvider.getNativeAspectClass(obj.getKey());
+      Preconditions.checkState(
+          obj == storedAspect, "Not stored right: %s %s %s", obj, storedAspect, ruleClassProvider);
+      context.serialize(obj.getKey(), codedOut);
     }
-    return nativeClass.equals(((NativeAspectClass) obj).nativeClass);
-  }
 
-  /**
-   * Every native aspect should implement this interface.
-   */
-  public interface NativeAspectFactory {
-    /**
-     * Returns the definition of the aspect.
-     */
-    AspectDefinition getDefinition();
+    @Override
+    public NativeAspectClass deserialize(DeserializationContext context, CodedInputStream codedIn)
+        throws SerializationException, IOException {
+      String aspectKey = context.deserialize(codedIn);
+      return Preconditions.checkNotNull(
+          context.getDependency(RuleClassProvider.class).getNativeAspectClass(aspectKey),
+          aspectKey);
+    }
   }
 }

dataset/GitHub_Java/bazelbuild.bazel/src/main/java/com/google/devtools/build/lib/skyframe/serialization/DeserializationContext.java

@@ -14,33 +14,134 @@
 
 package com.google.devtools.build.lib.skyframe.serialization;
 
-import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableMap;
+import static com.google.common.base.Preconditions.checkNotNull;
 
-/** Stateful class for providing additional context to a single deserialization "session". */
-// TODO(bazel-team): This class is just a shell, fill in.
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableClassToInstanceMap;
+import com.google.devtools.build.lib.skyframe.serialization.Memoizer.Deserializer;
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec.MemoizationStrategy;
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecRegistry.CodecDescriptor;
+import com.google.protobuf.CodedInputStream;
+import java.io.IOException;
+import javax.annotation.CheckReturnValue;
+import javax.annotation.Nullable;
+
+/**
+ * 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 ImmutableClassToInstanceMap<Object> dependencies;
+  private final Memoizer.Deserializer deserializer;
+
+  private DeserializationContext(
+      ObjectCodecRegistry registry,
+      ImmutableClassToInstanceMap<Object> dependencies,
+      Deserializer deserializer) {
+    this.registry = registry;
+    this.dependencies = dependencies;
+    this.deserializer = deserializer;
+  }
+
+  @VisibleForTesting
+  public DeserializationContext(
+      ObjectCodecRegistry registry, ImmutableClassToInstanceMap<Object> dependencies) {
+    this(registry, dependencies, /*deserializer=*/ null);
+  }
+
+  @VisibleForTesting
+  public DeserializationContext(ImmutableClassToInstanceMap<Object> dependencies) {
+    this(AutoRegistry.get(), dependencies);
+  }
+
+  // TODO(shahan): consider making codedIn a member of this class.
+  @SuppressWarnings({"TypeParameterUnusedInFormals"})
+  public <T> T deserialize(CodedInputStream codedIn) throws IOException, SerializationException {
+    return deserializeInternal(codedIn, /*customMemoizationStrategy=*/ null);
+  }
+
+  @SuppressWarnings({"TypeParameterUnusedInFormals"})
+  public <T> T deserializeWithAdHocMemoizationStrategy(
+      CodedInputStream codedIn, MemoizationStrategy memoizationStrategy)
+      throws IOException, SerializationException {
+    return deserializeInternal(codedIn, memoizationStrategy);
+  }
+
+  @SuppressWarnings({"TypeParameterUnusedInFormals", "unchecked"})
+  private <T> T deserializeInternal(
+      CodedInputStream codedIn, @Nullable MemoizationStrategy customMemoizationStrategy)
+      throws IOException, SerializationException {
+    int tag = codedIn.readSInt32();
+    if (tag == 0) {
+      return null;
+    }
+    if (tag < 0) {
+      // Subtract 1 to undo transformation from SerializationContext to avoid null.
+      return (T) deserializer.getMemoized(-tag - 1); // unchecked cast
+    }
+    T constant = (T) registry.maybeGetConstantByTag(tag);
+    if (constant != null) {
+      return constant;
+    }
+    CodecDescriptor codecDescriptor = registry.getCodecDescriptorByTag(tag);
+    if (deserializer == null) {
+      return (T) codecDescriptor.deserialize(this, codedIn); // unchecked cast
+    } else {
+      @SuppressWarnings("unchecked")
+      ObjectCodec<T> castCodec = (ObjectCodec<T>) codecDescriptor.getCodec();
+      MemoizationStrategy memoizationStrategy =
+          customMemoizationStrategy != null ? customMemoizationStrategy : castCodec.getStrategy();
+      return deserializer.deserialize(this, castCodec, memoizationStrategy, codedIn);
+    }
+  }
+
 
   /**
-   * This is a stub for context where it is less straightforward to thread from the top-level
-   * invocation.
+   * Register an initial value for the currently deserializing value, for use by child objects that
+   * may have references to it.
    *
-   * <p>This is a bug waiting to happen because it is very easy to accidentally modify a codec to
-   * use this context which won't contain any of the expected state.
+   * <p>This is a noop when memoization is disabled.
    */
-  // TODO(bazel-team): delete this and all references to it.
-  public static final DeserializationContext UNTHREADED_PLEASE_FIX =
-      new DeserializationContext(ImmutableMap.of());
+  public <T> void registerInitialValue(T initialValue) {
+    if (deserializer == null) {
+      return;
+    }
+    deserializer.registerInitialValue(initialValue);
+  }
 
-  private final ImmutableMap<Class<?>, Object> dependencies;
+  public <T> T getDependency(Class<T> type) {
+    return checkNotNull(dependencies.getInstance(type), "Missing dependency of type %s", type);
+  }
 
-  public DeserializationContext(ImmutableMap<Class<?>, Object> dependencies) {
-    this.dependencies = dependencies;
+  /**
+   * 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#getMemoizingContext}. The context returned here should
+   * be used instead of the original: memoization may only occur when using the returned context.
+   *
+   * <p>This method is idempotent: calling it on an already memoizing context will return the same
+   * context.
+   */
+  @CheckReturnValue
+  public DeserializationContext getMemoizingContext() {
+    if (deserializer != null) {
+      return this;
+    }
+    return getNewMemoizingContext();
   }
 
-  @SuppressWarnings("unchecked")
-  public <T> T getDependency(Class<T> type) {
-    Preconditions.checkNotNull(type);
-    return (T) dependencies.get(type);
+  /**
+   * Returns a memoizing {@link DeserializationContext}, as getMemoizingContext above. Unlike
+   * getMemoizingContext, this method is not idempotent - the returned context will always be fresh.
+   */
+  public DeserializationContext getNewMemoizingContext() {
+    return new DeserializationContext(this.registry, this.dependencies, new Deserializer());
+  }
+
+  public DeserializationContext getNewNonMemoizingContext() {
+    return new DeserializationContext(this.registry, this.dependencies, null);
   }
 }

dataset/GitHub_Java/bazelbuild.bazel/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java

@@ -26,7 +26,6 @@
  * serving as a layer between the streaming-oriented {@link ObjectCodec} interface and users.
  */
 public class ObjectCodecs {
-  private final ObjectCodecRegistry codecRegistry;
   private final SerializationContext serializationContext;
   private final DeserializationContext deserializationContext;
 
@@ -36,7 +35,6 @@ public class ObjectCodecs {
    */
   public ObjectCodecs(
       ObjectCodecRegistry codecRegistry, ImmutableClassToInstanceMap<Object> dependencies) {
-    this.codecRegistry = codecRegistry;
     serializationContext = new SerializationContext(codecRegistry, dependencies);
     deserializationContext = new DeserializationContext(codecRegistry, dependencies);
   }
@@ -98,10 +96,6 @@ public Object deserializeMemoized(CodedInputStream codedIn) throws Serialization
     return deserializeImpl(codedIn, /*memoize=*/ true);
   }
 
-  public ObjectCodecRegistry getCodecRegistry() {
-    return codecRegistry;
-  }
-
   private static void serializeImpl(
       Object subject, CodedOutputStream codedOut, SerializationContext serializationContext)
       throws SerializationException {