Decouple DynamicProto by extracting logic for descriptor pool and message factory.

common/src/main/java/dev/cel/common/internal/BUILD.bazel

@@ -91,7 +91,6 @@ java_library(
     ],
 )
 
-# keep sorted
 DYNAMIC_PROTO_SOURCES = [
     "DynamicProto.java",
     "ProtoAdapter.java",
@@ -119,13 +118,12 @@ java_library(
     ],
     deps = [
         ":converter",
-        ":default_instance_message_factory",
+        ":proto_message_factory",
+        ":well_known_proto",
         "//:auto_value",
-        "//common",
         "//common:error_codes",
         "//common:runtime_exception",
         "//common/annotations",
-        "//common/types:cel_types",
         "@cel_spec//proto/cel/expr:expr_java_proto",
         "@maven//:com_google_code_findbugs_annotations",
         "@maven//:com_google_errorprone_error_prone_annotations",

common/src/main/java/dev/cel/common/internal/CombinedDescriptorPool.java

@@ -74,7 +74,8 @@ public ExtensionRegistry getExtensionRegistry() {
   private CombinedDescriptorPool(ImmutableList<CelDescriptorPool> descriptorPools) {
     this.descriptorPools = descriptorPools;
     // TODO: Combine the extension registry. This will become necessary once we accept
-    // ExtensionRegistry through runtime builder.
+    // ExtensionRegistry through runtime builder. Ideally, proto team should open source this
+    // implementation but we may have to create our own.
     this.extensionRegistry =
         descriptorPools.stream()
             .map(CelDescriptorPool::getExtensionRegistry)

common/src/main/java/dev/cel/common/internal/DynamicProto.java

@@ -15,29 +15,16 @@
 package dev.cel.common.internal;
 
 import static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
-import static java.util.Arrays.stream;
 
-import com.google.auto.value.AutoBuilder;
-import com.google.common.collect.ImmutableCollection;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableMultimap;
 import com.google.errorprone.annotations.CheckReturnValue;
 import com.google.errorprone.annotations.Immutable;
 import com.google.protobuf.Any;
 import com.google.protobuf.ByteString;
-import com.google.protobuf.Descriptors.Descriptor;
-import com.google.protobuf.Descriptors.FieldDescriptor;
 import com.google.protobuf.DynamicMessage;
 import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.Message;
-import dev.cel.common.CelDescriptorUtil;
-import dev.cel.common.CelDescriptors;
 import dev.cel.common.annotations.Internal;
-import dev.cel.common.types.CelTypes;
-import java.util.Map.Entry;
 import java.util.Optional;
-import org.jspecify.nullness.Nullable;
 
 /**
  * The {@code DynamicProto} class supports the conversion of {@link Any} values to concrete {@code
@@ -49,58 +36,13 @@
 @CheckReturnValue
 @Internal
 public final class DynamicProto {
-
-  private static final ImmutableMap<String, Descriptor> WELL_KNOWN_DESCRIPTORS =
-      stream(ProtoAdapter.WellKnownProto.values())
-          .collect(toImmutableMap(d -> d.typeName(), d -> d.descriptor()));
-
-  private final ImmutableMap<String, Descriptor> dynamicDescriptors;
-  private final ImmutableMultimap<String, FieldDescriptor> dynamicExtensionDescriptors;
   private final ProtoMessageFactory protoMessageFactory;
 
-  /** {@code ProtoMessageFactory} provides a method to create a protobuf builder objects by name. */
-  @Immutable
-  @FunctionalInterface
-  public interface ProtoMessageFactory {
-    Message.@Nullable Builder newBuilder(String messageName);
-  }
-
-  /** Builder for configuring the {@link DynamicProto}. */
-  @AutoBuilder(ofClass = DynamicProto.class)
-  public abstract static class Builder {
-
-    /** Sets {@link CelDescriptors} to unpack any message types. */
-    public abstract Builder setDynamicDescriptors(CelDescriptors celDescriptors);
-
-    /** Sets a custom type factory to unpack any message types. */
-    public abstract Builder setProtoMessageFactory(ProtoMessageFactory factory);
-
-    /** Builds a new instance of {@link DynamicProto} */
-    @CheckReturnValue
-    public abstract DynamicProto build();
+  public static DynamicProto create(ProtoMessageFactory protoMessageFactory) {
+    return new DynamicProto(protoMessageFactory);
   }
 
-  public static Builder newBuilder() {
-    return new AutoBuilder_DynamicProto_Builder()
-        .setDynamicDescriptors(CelDescriptors.builder().build())
-        .setProtoMessageFactory((typeName) -> null);
-  }
-
-  DynamicProto(
-      CelDescriptors dynamicDescriptors,
-      ProtoMessageFactory protoMessageFactory) {
-    ImmutableMap<String, Descriptor> messageTypeDescriptorMap =
-        CelDescriptorUtil.descriptorCollectionToMap(dynamicDescriptors.messageTypeDescriptors());
-    ImmutableMap<String, Descriptor> filteredDescriptors =
-        messageTypeDescriptorMap.entrySet().stream()
-            .filter(e -> !WELL_KNOWN_DESCRIPTORS.containsKey(e.getKey()))
-            .collect(toImmutableMap(Entry::getKey, Entry::getValue));
-    this.dynamicDescriptors =
-        ImmutableMap.<String, Descriptor>builder()
-            .putAll(WELL_KNOWN_DESCRIPTORS)
-            .putAll(filteredDescriptors)
-            .buildOrThrow();
-    this.dynamicExtensionDescriptors = checkNotNull(dynamicDescriptors.extensionDescriptors());
+  DynamicProto(ProtoMessageFactory protoMessageFactory) {
     this.protoMessageFactory = checkNotNull(protoMessageFactory);
   }
 
@@ -120,7 +62,8 @@ public Message unpack(Any any) throws InvalidProtocolBufferException {
                         String.format("malformed type URL: %s", any.getTypeUrl())));
 
     Message.Builder builder =
-        newMessageBuilder(messageTypeName)
+        protoMessageFactory
+            .newBuilder(messageTypeName)
             .orElseThrow(
                 () ->
                     new InvalidProtocolBufferException(
@@ -136,7 +79,7 @@ public Message unpack(Any any) throws InvalidProtocolBufferException {
    */
   public Message maybeAdaptDynamicMessage(DynamicMessage input) {
     Optional<Message.Builder> maybeBuilder =
-        newMessageBuilder(input.getDescriptorForType().getFullName());
+        protoMessageFactory.newBuilder(input.getDescriptorForType().getFullName());
     if (!maybeBuilder.isPresent() || maybeBuilder.get() instanceof DynamicMessage.Builder) {
       // Just return the same input if:
       // 1. We didn't get a builder back because there's no descriptor (nothing we can do)
@@ -148,61 +91,6 @@ public Message maybeAdaptDynamicMessage(DynamicMessage input) {
     return merge(maybeBuilder.get(), input.toByteString());
   }
 
-  /**
-   * This method instantiates a builder for the given {@code typeName} assuming one is configured
-   * within the descriptor set provided to the {@code DynamicProto} constructor.
-   *
-   * <p>When the {@code useLinkedTypes} flag is set, the {@code Message.Builder} returned will be
-   * the concrete builder instance linked into the binary if it is present; otherwise, the result
-   * will be a {@code DynamicMessageBuilder}.
-   */
-  public Optional<Message.Builder> newMessageBuilder(String typeName) {
-    if (!CelTypes.isWellKnownType(typeName)) {
-      // Check if the message factory can produce a concrete message via custom type factory
-      // first.
-      Message.Builder builder = protoMessageFactory.newBuilder(typeName);
-      if (builder != null) {
-        return Optional.of(builder);
-      }
-    }
-
-    Optional<Descriptor> descriptor = maybeGetDescriptor(typeName);
-    if (!descriptor.isPresent()) {
-      return Optional.empty();
-    }
-    // If the descriptor that's resolved does not match the descriptor instance in the message
-    // factory, the call to fetch the prototype will return null, and a dynamic proto message
-    // should be used as a fallback.
-    Optional<Message> message =
-        DefaultInstanceMessageFactory.getInstance().getPrototype(descriptor.get());
-    if (message.isPresent()) {
-      return Optional.of(message.get().toBuilder());
-    }
-
-    // Fallback to a dynamic proto instance.
-    return Optional.of(DynamicMessage.newBuilder(descriptor.get()));
-  }
-
-  private Optional<Descriptor> maybeGetDescriptor(String typeName) {
-
-    Descriptor descriptor = ProtoRegistryProvider.getTypeRegistry().find(typeName);
-    return Optional.ofNullable(descriptor != null ? descriptor : dynamicDescriptors.get(typeName));
-  }
-
-  /** Gets the corresponding field descriptor for an extension field on a message. */
-  public Optional<FieldDescriptor> maybeGetExtensionDescriptor(
-      Descriptor containingDescriptor, String fieldName) {
-
-    String typeName = containingDescriptor.getFullName();
-    ImmutableCollection<FieldDescriptor> fieldDescriptors =
-        dynamicExtensionDescriptors.get(typeName);
-    if (fieldDescriptors.isEmpty()) {
-      return Optional.empty();
-    }
-
-    return fieldDescriptors.stream().filter(d -> d.getFullName().equals(fieldName)).findFirst();
-  }
-
   /**
    * Merge takes in a Message builder and merges another message bytes into the builder. Some
    * example usages are:
@@ -214,7 +102,9 @@ public Optional<FieldDescriptor> maybeGetExtensionDescriptor(
    */
   private Message merge(Message.Builder builder, ByteString inputBytes) {
     try {
-      return builder.mergeFrom(inputBytes, ProtoRegistryProvider.getExtensionRegistry()).build();
+      return builder
+          .mergeFrom(inputBytes, protoMessageFactory.getDescriptorPool().getExtensionRegistry())
+          .build();
     } catch (InvalidProtocolBufferException e) {
       throw new AssertionError("Failed to merge input message into the message builder", e);
     }

common/src/main/java/dev/cel/common/internal/ProtoAdapter.java

@@ -143,53 +143,6 @@ public final class ProtoAdapter {
       stream(WellKnownProto.values())
           .collect(toImmutableMap(WellKnownProto::typeName, Function.identity()));
 
-  /**
-   * WellKnownProto types used throughout CEL. These types are specially handled to ensure that
-   * bidirectional conversion between CEL native values and these well-known types is performed
-   * consistently across runtimes.
-   */
-  enum WellKnownProto {
-    JSON_VALUE(Value.getDescriptor()),
-    JSON_STRUCT_VALUE(Struct.getDescriptor()),
-    JSON_LIST_VALUE(ListValue.getDescriptor()),
-    ANY_VALUE(Any.getDescriptor()),
-    BOOL_VALUE(BoolValue.getDescriptor(), true),
-    BYTES_VALUE(BytesValue.getDescriptor(), true),
-    DOUBLE_VALUE(DoubleValue.getDescriptor(), true),
-    FLOAT_VALUE(FloatValue.getDescriptor(), true),
-    INT32_VALUE(Int32Value.getDescriptor(), true),
-    INT64_VALUE(Int64Value.getDescriptor(), true),
-    STRING_VALUE(StringValue.getDescriptor(), true),
-    UINT32_VALUE(UInt32Value.getDescriptor(), true),
-    UINT64_VALUE(UInt64Value.getDescriptor(), true),
-    DURATION_VALUE(Duration.getDescriptor()),
-    TIMESTAMP_VALUE(Timestamp.getDescriptor());
-
-    private final Descriptor descriptor;
-    private final boolean isWrapperType;
-
-    WellKnownProto(Descriptor descriptor) {
-      this(descriptor, /* isWrapperType= */ false);
-    }
-
-    WellKnownProto(Descriptor descriptor, boolean isWrapperType) {
-      this.descriptor = descriptor;
-      this.isWrapperType = isWrapperType;
-    }
-
-    Descriptor descriptor() {
-      return descriptor;
-    }
-
-    String typeName() {
-      return descriptor.getFullName();
-    }
-
-    boolean isWrapperType() {
-      return isWrapperType;
-    }
-  }
-
   private final DynamicProto dynamicProto;
   private final boolean enableUnsignedLongs;