Allow resolution of proto messages from dyn-typed functions in lite runtime

PiperOrigin-RevId: 761247740

Author: l46kok

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

@@ -314,6 +314,7 @@ java_library(
     deps = [
         "//protobuf:cel_lite_descriptor",
         "@maven//:com_google_errorprone_error_prone_annotations",
+        "@maven_android//:com_google_protobuf_protobuf_javalite",
     ],
 )
 
@@ -325,6 +326,7 @@ cel_android_library(
     deps = [
         "//protobuf:cel_lite_descriptor",
         "@maven//:com_google_errorprone_error_prone_annotations",
+        "@maven_android//:com_google_protobuf_protobuf_javalite",
     ],
 )
 

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

@@ -15,6 +15,7 @@
 package dev.cel.common.internal;
 
 import com.google.errorprone.annotations.Immutable;
+import com.google.protobuf.MessageLite;
 import dev.cel.protobuf.CelLiteDescriptor.MessageLiteDescriptor;
 import java.util.Optional;
 
@@ -25,5 +26,7 @@
 public interface CelLiteDescriptorPool {
   Optional<MessageLiteDescriptor> findDescriptor(String protoTypeName);
 
+  Optional<MessageLiteDescriptor> findDescriptor(MessageLite messageLite);
+
   MessageLiteDescriptor getDescriptorOrThrow(String protoTypeName);
 }

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

@@ -50,6 +50,7 @@
 @Internal
 public final class DefaultLiteDescriptorPool implements CelLiteDescriptorPool {
   private final ImmutableMap<String, MessageLiteDescriptor> protoFqnToMessageInfo;
+  private final ImmutableMap<Class<?>, MessageLiteDescriptor> classToMessageInfo;
 
   public static DefaultLiteDescriptorPool newInstance(CelLiteDescriptor... descriptors) {
     return newInstance(ImmutableSet.copyOf(descriptors));
@@ -59,6 +60,11 @@ public static DefaultLiteDescriptorPool newInstance(ImmutableSet<CelLiteDescript
     return new DefaultLiteDescriptorPool(descriptors);
   }
 
+  @Override
+  public Optional<MessageLiteDescriptor> findDescriptor(MessageLite messageLite) {
+    return Optional.ofNullable(classToMessageInfo.get(messageLite.getClass()));
+  }
+
   @Override
   public Optional<MessageLiteDescriptor> findDescriptor(String protoTypeName) {
     return Optional.ofNullable(protoFqnToMessageInfo.get(protoTypeName));
@@ -292,15 +298,28 @@ private static FieldLiteDescriptor newPrimitiveFieldDescriptor(
 
   private DefaultLiteDescriptorPool(ImmutableSet<CelLiteDescriptor> descriptors) {
     ImmutableMap.Builder<String, MessageLiteDescriptor> protoFqnMapBuilder = ImmutableMap.builder();
+    ImmutableMap.Builder<Class<?>, MessageLiteDescriptor> classMapBuilder = ImmutableMap.builder();
     for (WellKnownProto wellKnownProto : WellKnownProto.values()) {
       MessageLiteDescriptor wktMessageInfo = newMessageInfo(wellKnownProto);
       protoFqnMapBuilder.put(wellKnownProto.typeName(), wktMessageInfo);
+      classMapBuilder.put(wellKnownProto.messageClass(), wktMessageInfo);
     }
 
     for (CelLiteDescriptor descriptor : descriptors) {
       protoFqnMapBuilder.putAll(descriptor.getProtoTypeNamesToDescriptors());
+
+      for (MessageLiteDescriptor messageLiteDescriptor :
+          descriptor.getProtoTypeNamesToDescriptors().values()) {
+        // Note: message builder is null for proto maps.
+        Optional.ofNullable(messageLiteDescriptor.newMessageBuilder())
+            .ifPresent(
+                builder ->
+                    classMapBuilder.put(
+                        builder.getDefaultInstanceForType().getClass(), messageLiteDescriptor));
+      }
     }
 
     this.protoFqnToMessageInfo = protoFqnMapBuilder.buildOrThrow();
+    this.classToMessageInfo = classMapBuilder.buildOrThrow();
   }
 }

common/src/main/java/dev/cel/common/values/ProtoLiteCelValueConverter.java

@@ -44,6 +44,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.TreeMap;
 
 /**
@@ -360,7 +361,13 @@ public CelValue fromProtoMessageToCelValue(String protoTypeName, MessageLite msg
     checkNotNull(msg);
     checkNotNull(protoTypeName);
 
-    MessageLiteDescriptor descriptor = descriptorPool.getDescriptorOrThrow(protoTypeName);
+    MessageLiteDescriptor descriptor =
+        descriptorPool
+            .findDescriptor(msg)
+            .orElseThrow(
+                () ->
+                    new NoSuchElementException(
+                        "Could not find a descriptor for: " + protoTypeName));
     WellKnownProto wellKnownProto =
         WellKnownProto.getByTypeName(descriptor.getProtoTypeName()).orElse(null);
 

runtime/src/test/java/dev/cel/runtime/CelLiteRuntimeTest.java

@@ -633,4 +633,31 @@ public void eval_withLateBoundFunction() throws Exception {
 
     assertThat(result).isEqualTo("hello world");
   }
+
+  @Test
+  public void eval_dynFunctionReturnsProto() throws Exception {
+    CelCompiler celCompiler =
+        CelCompilerFactory.standardCelCompilerBuilder()
+            .addFunctionDeclarations(
+                CelFunctionDecl.newFunctionDeclaration(
+                    "func", CelOverloadDecl.newGlobalOverload("func_identity", SimpleType.DYN)))
+            .build();
+    CelLiteRuntime celRuntime =
+        CelLiteRuntimeFactory.newLiteRuntimeBuilder()
+            .setValueProvider(
+                ProtoMessageLiteValueProvider.newInstance(
+                    TestAllTypesCelDescriptor.getDescriptor()))
+            .addFunctionBindings(
+                CelFunctionBinding.from(
+                    "func_identity",
+                    ImmutableList.of(),
+                    unused -> TestAllTypes.getDefaultInstance()))
+            .build();
+
+    CelAbstractSyntaxTree ast = celCompiler.compile("func()").getAst();
+
+    TestAllTypes result = (TestAllTypes) celRuntime.createProgram(ast).eval();
+
+    assertThat(result).isEqualToDefaultInstance();
+  }
 }