Adds Class-keyed lookup to ObjectCodecRegistry and populates it using CodecScanner.

Introduces a class, CodecRegisterer, to allow customization of the scan-based
registration process.

PiperOrigin-RevId: 185749655

6 files changed, 326 insertions, 7 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecRegisterer.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecRegisterer.java
new file mode 100644
index 0000000000..ec9ba202f5
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecRegisterer.java
@@ -0,0 +1,44 @@
+// 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;
+
+/**
+ * Custom registration behavior for a codec.
+ *
+ * <p>This class should only be needed for low-level codecs like collections and probably shouldn't
+ * be used in normal application code.
+ *
+ * <p>Instances of this class are discovered automatically by {@link CodecScanner} and used for
+ * custom registration of a codec. This can be useful when a single codec is used for multiple
+ * classes or when the class that is being serialized has a hidden type, e.g., {@link
+ * com.google.common.collect.RegularImmutableList}.
+ *
+ * <p>If a {@link CodecRegisterer} definition exists, the codec will only be registered by the
+ * {@link CodecRegisterer#register} method. Otherwise, an attempt will be made to register the codec
+ * to its generic parameter type.
+ *
+ * <p>Implementations must have a default constructor.
+ *
+ * <p>Multiple {@link CodecRegisterer} implementations for the same {@link ObjectCodec} is illegal.
+ *
+ * <p>Inheriting {@link CodecRegisterer} through a superclass is illegal. It must be directly
+ * implemented. Also, the generic parameter of {@link CodecRegisterer} must be reified.
+ *
+ * <p>Constraint violations will cause exceptions to be raised from {@link CodecScanner}.
+ */
+public interface CodecRegisterer<T extends ObjectCodec<?>> {
+
+  void register(ObjectCodecRegistry.Builder builder);
+}
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 7536c6df1b..ce6db63d63 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
@@ -14,29 +14,188 @@
 
 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 java.io.IOException;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.ParameterizedType;
+import java.lang.reflect.Type;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import java.util.stream.Stream;
+import javax.annotation.Nullable;
 
-/** Scans the classpath to find codec instances. */
+/**
+ * Scans the classpath to find {@link ObjectCodec} and {@link CodecRegisterer} instances.
+ *
+ * <p>To avoid loading classes unnecessarily, the scanner filters by class name before loading.
+ * {@link ObjectCodec} implementation class names should end in "Codec" while {@link
+ * CodecRegisterer} implementation class names should end in "CodecRegisterer".
+ *
+ * <p>See {@link CodecRegisterer} for more details.
+ */
 public class CodecScanner {
 
+  private static final Logger log = Logger.getLogger(CodecScanner.class.getName());
+
+  /**
+   * Initializes an {@link ObjectCodecRegistry} builder by scanning a given package prefix.
+   *
+   * @param packagePrefix processes only classes in packages having this prefix
+   * @see CodecRegisterer
+   */
+  @SuppressWarnings("unchecked")
+  public static ObjectCodecRegistry.Builder initializeCodecRegistry(String packagePrefix)
+      throws IOException, ReflectiveOperationException {
+    ArrayList<Class<? extends ObjectCodec<?>>> codecs = new ArrayList<>();
+    ArrayList<Class<? extends CodecRegisterer<?>>> registerers = new ArrayList<>();
+    scanCodecs(packagePrefix)
+        .forEach(
+            type -> {
+              if (!ObjectCodec.class.equals(type) && ObjectCodec.class.isAssignableFrom(type)) {
+                codecs.add((Class<? extends ObjectCodec<?>>) type);
+              }
+              if (!CodecRegisterer.class.equals(type)
+                  && CodecRegisterer.class.isAssignableFrom(type)) {
+                registerers.add((Class<? extends CodecRegisterer<?>>) type);
+              }
+            });
+
+    ObjectCodecRegistry.Builder builder = ObjectCodecRegistry.newBuilder();
+    HashSet<Class<? extends ObjectCodec<?>>> alreadyRegistered =
+        runRegisterers(builder, registerers);
+
+    applyDefaultRegistration(builder, alreadyRegistered, codecs);
+    return builder;
+  }
+
+  @SuppressWarnings("unchecked")
+  private static HashSet<Class<? extends ObjectCodec<?>>> runRegisterers(
+      ObjectCodecRegistry.Builder builder,
+      ArrayList<Class<? extends CodecRegisterer<?>>> registerers)
+      throws ReflectiveOperationException {
+    HashSet<Class<? extends ObjectCodec<?>>> registered = new HashSet<>();
+    for (Class<? extends CodecRegisterer<?>> registererType : registerers) {
+      Class<? extends ObjectCodec<?>> objectCodecType = getObjectCodecType(registererType);
+      Preconditions.checkState(
+          !registered.contains(objectCodecType),
+          "%s has multiple associated CodecRegisterer definitions!",
+          objectCodecType);
+      registered.add(objectCodecType);
+      Constructor<CodecRegisterer<?>> constructor =
+          (Constructor<CodecRegisterer<?>>) registererType.getDeclaredConstructor();
+      constructor.setAccessible(true);
+      constructor.newInstance().register(builder);
+    }
+    return registered;
+  }
+
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  private static void applyDefaultRegistration(
+      ObjectCodecRegistry.Builder builder,
+      HashSet<Class<? extends ObjectCodec<?>>> alreadyRegistered,
+      ArrayList<Class<? extends ObjectCodec<?>>> codecs)
+      throws ReflectiveOperationException {
+    for (Class<? extends ObjectCodec<?>> codecType : codecs) {
+      if (alreadyRegistered.contains(codecType)) {
+        continue;
+      }
+      Class<?> serializedType = getSerializedType(codecType);
+      if (serializedType == null) {
+        log.log(
+            Level.FINE,
+            "Skipping registration of "
+                + codecType
+                + " because it's generic parameter is not reified.");
+        continue;
+      }
+      try {
+        Constructor constructor = codecType.getDeclaredConstructor();
+        constructor.setAccessible(true);
+        builder.add((Class) serializedType, (ObjectCodec) constructor.newInstance());
+      } catch (NoSuchMethodException e) {
+        log.log(
+            Level.FINE,
+            "Skipping registration of " + codecType + " because it had no default constructor.");
+      }
+    }
+  }
+
+  @SuppressWarnings("unchecked")
+  private static Class<? extends ObjectCodec<?>> getObjectCodecType(
+      Class<? extends CodecRegisterer<?>> registererType) throws ReflectiveOperationException {
+    Type typeArg =
+        ((ParameterizedType)
+                registererType.getGenericInterfaces()[getCodecRegistererIndex(registererType)])
+            .getActualTypeArguments()[0];
+    // This occurs when the generic parameter of CodecRegisterer is not reified, for example:
+    //   class MyCodecRegisterer<T> implements CodecRegisterer<T>
+    Preconditions.checkArgument(
+        typeArg instanceof Class,
+        "Illegal CodecRegisterer definition: %s"
+            + "\nCodecRegisterer generic parameter must be reified.",
+        registererType);
+    return (Class<? extends ObjectCodec<?>>) typeArg;
+  }
+
+  private static int getCodecRegistererIndex(Class<? extends CodecRegisterer<?>> registererType) {
+    Class<?>[] interfaces = registererType.getInterfaces();
+    for (int i = 0; i < interfaces.length; ++i) {
+      if (CodecRegisterer.class.equals(interfaces[i])) {
+        return i;
+      }
+    }
+    // The following line is reached when there are multiple layers of inheritance involving
+    // CodecRegisterer, which is prohibited.
+    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 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;
+  }
+
   /**
-   * Returns a stream of likely codec implementations.
+   * Returns a stream of likely codec and registerer implementations.
    *
    * <p>Caller should do additional checks as this method only performs string matching.
    *
    * @param packagePrefix emits only classes in packages having this prefix
    */
-  public static Stream<Class<?>> scanCodecs(String packagePrefix) throws IOException {
+  private static Stream<Class<?>> scanCodecs(String packagePrefix) throws IOException {
     return ClassPath.from(ClassLoader.getSystemClassLoader())
         .getResources()
         .stream()
         .filter(r -> r instanceof ClassInfo)
         .map(r -> (ClassInfo) r)
         .filter(c -> c.getPackageName().startsWith(packagePrefix))
-        .filter(c -> c.getSimpleName().endsWith("Codec"))
+        .filter(c -> c.getName().endsWith("Codec") || c.getName().endsWith("CodecRegisterer"))
         .map(c -> c.load());
   }
 }
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 3fa0e63b5f..20f8708e8b 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
@@ -26,7 +26,7 @@ import javax.annotation.Nullable;
  * classifiers and assigned deterministic numeric identifiers for more compact on-the-wire
  * representation if desired.
  */
-class ObjectCodecRegistry {
+public class ObjectCodecRegistry {
 
   static Builder newBuilder() {
     return new Builder();
@@ -46,6 +46,7 @@ class ObjectCodecRegistry {
       nextTag++;
     }
     this.stringMappedCodecs = codecMappingsBuilder.build();
+
     this.byteStringMappedCodecs = makeByteStringMappedCodecs(stringMappedCodecs);
 
     this.defaultCodecDescriptor = allowDefaultCodec
@@ -82,6 +83,18 @@ class ObjectCodecRegistry {
     }
   }
 
+  public CodecDescriptor getCodecDescriptor(Class<?> type)
+      throws SerializationException.NoCodecException {
+    CodecDescriptor result =
+        stringMappedCodecs.getOrDefault(type.getName(), defaultCodecDescriptor);
+    if (result != null) {
+      return result;
+    } else {
+      throw new SerializationException.NoCodecException(
+          "No codec available for " + type + " and default fallback disabled");
+    }
+  }
+
   /** Returns the {@link CodecDescriptor} associated with the supplied tag. */
   public CodecDescriptor getCodecDescriptorByTag(int tag)
       throws SerializationException.NoCodecException {
@@ -121,7 +134,7 @@ class ObjectCodecRegistry {
   }
 
   /** Builder for {@link ObjectCodecRegistry}. */
-  static class Builder {
+  public static class Builder {
     private final ImmutableMap.Builder<String, ObjectCodec<?>> codecsBuilder =
         ImmutableMap.builder();
     private boolean allowDefaultCodec = true;
@@ -139,6 +152,11 @@ class ObjectCodecRegistry {
       return this;
     }
 
+    public <T> Builder add(Class<? extends T> type, ObjectCodec<T> codec) {
+      add(type.getName(), codec);
+      return this;
+    }
+
     /**
      * Set whether or not we allow fallback to java serialization when no matching codec is found.
      */
@@ -166,7 +184,7 @@ class ObjectCodecRegistry {
     }
 
     public <T> ClassKeyedBuilder add(Class<? extends T> clazz, ObjectCodec<T> codec) {
-      underlying.add(clazz.getName(), codec);
+      underlying.add(clazz, codec);
       return this;
     }
 
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
new file mode 100644
index 0000000000..648b1a5f42
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/UnmodifiableListCodec.java
@@ -0,0 +1,60 @@
+// 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/strings/StringCodecs.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/strings/StringCodecs.java
index 60e5f29f87..aea5a4c67d 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
@@ -14,7 +14,9 @@
 
 package com.google.devtools.build.lib.skyframe.serialization.strings;
 
+import com.google.devtools.build.lib.skyframe.serialization.CodecRegisterer;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecRegistry;
 import java.util.logging.Logger;
 
 /** Utility for accessing (potentially platform-specific) {@link String} {@link ObjectCodec}s. */
@@ -64,4 +66,32 @@ public final class StringCodecs {
   public static ObjectCodec<String> simple() {
     return stringCodec;
   }
+
+  /**
+   * Registers a codec for {@link String}.
+   *
+   * <p>Needed to resolve ambiguity between {@link StringCodec} and {@link FastStringCodec}.
+   */
+  static class StringCodecRegisterer implements CodecRegisterer<StringCodec> {
+    @Override
+    public void register(ObjectCodecRegistry.Builder builder) {
+      if (!supportsOptimizedAscii()) {
+        builder.add(String.class, simple());
+      }
+    }
+  }
+
+  /**
+   * Registers a codec for {@link String}.
+   *
+   * <p>Needed to resolve ambiguity between {@link StringCodec} and {@link FastStringCodec}.
+   */
+  static class FastStringCodecRegisterer implements CodecRegisterer<FastStringCodec> {
+    @Override
+    public void register(ObjectCodecRegistry.Builder builder) {
+      if (supportsOptimizedAscii()) {
+        builder.add(String.class, asciiOptimized());
+      }
+    }
+  }
 }
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 8d914da254..2c729d4671 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
@@ -16,8 +16,10 @@ package com.google.devtools.build.lib.skyframe.serialization.testutils;
 
 import static com.google.common.truth.Truth.assertThat;
 
+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.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;
@@ -90,5 +92,11 @@ public class TestUtils {
       stringCodec.deserialize(context, codedIn);
       return "dummy";
     }
+
+    /** Disables auto-registration of ConstantStringCodec. */
+    static class ConstantStringCodecRegisterer implements CodecRegisterer<ConstantStringCodec> {
+      @Override
+      public void register(ObjectCodecRegistry.Builder unusedBuilder) {}
+    }
   }
 }