1 files changed, 21 insertions, 3 deletions
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;
     }