Refactor type-checker to use canonical CelTypeProvider

PiperOrigin-RevId: 862321687

Author: l46kok

bundle/src/test/java/dev/cel/bundle/CelImplTest.java

@@ -35,6 +35,7 @@
 import dev.cel.expr.Reference;
 import dev.cel.expr.Type;
 import dev.cel.expr.Type.PrimitiveType;
+import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -63,7 +64,6 @@
 import dev.cel.checker.CelCheckerLegacyImpl;
 import dev.cel.checker.DescriptorTypeProvider;
 import dev.cel.checker.ProtoTypeMask;
-import dev.cel.checker.TypeProvider;
 import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.CelContainer;
 import dev.cel.common.CelDescriptorUtil;
@@ -82,12 +82,14 @@
 import dev.cel.common.types.CelProtoMessageTypes;
 import dev.cel.common.types.CelProtoTypes;
 import dev.cel.common.types.CelType;
+import dev.cel.common.types.CelTypeProvider;
 import dev.cel.common.types.EnumType;
 import dev.cel.common.types.ListType;
 import dev.cel.common.types.MapType;
 import dev.cel.common.types.OptionalType;
 import dev.cel.common.types.ProtoMessageTypeProvider;
 import dev.cel.common.types.SimpleType;
+import dev.cel.common.types.StructType;
 import dev.cel.common.types.StructTypeReference;
 import dev.cel.common.values.CelByteString;
 import dev.cel.common.values.NullValue;
@@ -123,7 +125,6 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.ThreadPoolExecutor;
-import org.jspecify.annotations.Nullable;
 import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -297,13 +298,12 @@ public void compile_customTypeProvider() {
 
   @Test
   public void compile_customTypesWithAliasingCombinedProviders() throws Exception {
-
     // The custom type provider sets up an alias from "Condition" to "google.type.Expr".
     // However, the first type resolution from the alias to the qualified type name won't be
     // sufficient as future checks will expect the resolved alias to also be a type.
-    TypeProvider customTypeProvider =
+    CelTypeProvider customTypeProvider =
         aliasingProvider(
-            ImmutableMap.of("Condition", CelProtoTypes.createMessage("google.type.Expr")));
+            ImmutableMap.of("Condition", StructTypeReference.create("google.type.Expr")));
 
     // The registration of the aliasing TypeProvider and the google.type.Expr descriptor
     // ensures that once the alias is resolved, the additional details about the Expr type
@@ -329,15 +329,19 @@ public void compile_customTypesWithAliasingCombinedProviders() throws Exception
 
   @Test
   public void compile_customTypesWithAliasingSelfContainedProvider() throws Exception {
-
     // The custom type provider sets up an alias from "Condition" to "google.type.Expr".
-    TypeProvider customTypeProvider =
+    StructType exprStruct = StructType.create(
+        "google.type.Expr",
+        ImmutableSet.of("expression"),
+        fieldName -> Optional.of(SimpleType.STRING)
+    );
+    CelTypeProvider customTypeProvider =
         aliasingProvider(
             ImmutableMap.of(
                 "Condition",
-                CelProtoTypes.createMessage("google.type.Expr"),
+                exprStruct,
                 "google.type.Expr",
-                CelProtoTypes.createMessage("google.type.Expr")));
+                exprStruct));
 
     // The registration of the aliasing TypeProvider and the google.type.Expr descriptor
     // ensures that once the alias is resolved, the additional details about the Expr type
@@ -1001,14 +1005,11 @@ public void program_protoActivation() throws Exception {
   }
 
   @Test
-  @TestParameters("{resolveTypeDependencies: false}")
-  @TestParameters("{resolveTypeDependencies: true}")
-  public void program_enumTypeDirectResolution(boolean resolveTypeDependencies) throws Exception {
+  public void program_enumTypeDirectResolution() throws Exception {
     Cel cel =
         standardCelBuilderWithMacros()
             .addFileTypes(StandaloneGlobalEnum.getDescriptor().getFile())
-            .setOptions(
-                CelOptions.current().resolveTypeDependencies(resolveTypeDependencies).build())
+            .setOptions(CelOptions.current().resolveTypeDependencies(true).build())
             .setContainer(
                 CelContainer.ofName("dev.cel.testing.testdata.proto3.StandaloneGlobalEnum"))
             .setResultType(SimpleType.BOOL)
@@ -2193,28 +2194,16 @@ public void toBuilder_isImmutable() {
     assertThat(newRuntimeBuilder).isNotEqualTo(celImpl.toRuntimeBuilder());
   }
 
-  private static TypeProvider aliasingProvider(ImmutableMap<String, Type> typeAliases) {
-    return new TypeProvider() {
-      @Override
-      public @Nullable Type lookupType(String typeName) {
-        Type alias = typeAliases.get(typeName);
-        if (alias != null) {
-          return CelProtoTypes.create(alias);
-        }
-        return null;
-      }
-
+  private static CelTypeProvider aliasingProvider(ImmutableMap<String, CelType> typeAliases) {
+    return new CelTypeProvider() {
       @Override
-      public @Nullable Integer lookupEnumValue(String enumName) {
-        return null;
+      public ImmutableCollection<CelType> types() {
+        return typeAliases.values();
       }
 
       @Override
-      public TypeProvider.@Nullable FieldType lookupFieldType(Type type, String fieldName) {
-        if (typeAliases.containsKey(type.getMessageType())) {
-          return TypeProvider.FieldType.of(CelProtoTypes.STRING);
-        }
-        return null;
+      public Optional<CelType> findType(String typeName) {
+        return Optional.ofNullable(typeAliases.get(typeName));
       }
     };
   }

checker/src/main/java/dev/cel/checker/BUILD.bazel

@@ -32,6 +32,7 @@ CHECKER_LEGACY_ENV_SOURCES = [
     "InferenceContext.java",
     "TypeFormatter.java",
     "TypeProvider.java",
+    "TypeProviderLegacyImpl.java",
     "Types.java",
 ]
 
@@ -70,7 +71,6 @@ java_library(
         ":checker_legacy_environment",
         ":proto_type_mask",
         ":standard_decl",
-        ":type_provider_legacy_impl",
         "//:auto_value",
         "//common:cel_ast",
         "//common:cel_descriptor_util",
@@ -158,12 +158,9 @@ java_library(
         "//:auto_value",
         "//common/annotations",
         "//common/types",
-        "//common/types:cel_proto_types",
         "//common/types:type_providers",
-        "@cel_spec//proto/cel/expr:checked_java_proto",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",
-        "@maven//:org_jspecify_jspecify",
     ],
 )
 
@@ -190,6 +187,7 @@ java_library(
         "//common/types",
         "//common/types:cel_proto_types",
         "//common/types:cel_types",
+        "//common/types:message_type_provider",
         "//common/types:type_providers",
         "//parser:macro",
         "@cel_spec//proto/cel/expr:checked_java_proto",

checker/src/main/java/dev/cel/checker/CelCheckerLegacyImpl.java

@@ -163,11 +163,11 @@ public void accept(EnvVisitor envVisitor) {
   private Env getEnv(Errors errors) {
     Env env;
     if (standardEnvironmentEnabled) {
-      env = Env.standard(errors, typeProvider, celOptions);
+      env = Env.standard(errors, celTypeProvider, celOptions);
     } else if (overriddenStandardDeclarations != null) {
-      env = Env.standard(overriddenStandardDeclarations, errors, typeProvider, celOptions);
+      env = Env.standard(overriddenStandardDeclarations, errors, celTypeProvider, celOptions);
     } else {
-      env = Env.unconfigured(errors, typeProvider, celOptions);
+      env = Env.unconfigured(errors, celTypeProvider, celOptions);
     }
     identDeclarations.forEach(env::add);
     functionDeclarations.forEach(env::add);
@@ -483,11 +483,10 @@ public CelCheckerLegacyImpl build() {
         messageTypeProvider = protoTypeMaskTypeProvider;
       }
 
-      TypeProvider legacyProvider = new TypeProviderLegacyImpl(messageTypeProvider);
       if (customTypeProvider != null) {
-        legacyProvider =
-            new TypeProvider.CombinedTypeProvider(
-                ImmutableList.of(customTypeProvider, legacyProvider));
+        messageTypeProvider =
+            new CelTypeProvider.CombinedCelTypeProvider(
+                messageTypeProvider, new TypeProviderLegacyImpl(customTypeProvider));
       }
 
       return new CelCheckerLegacyImpl(
@@ -496,7 +495,7 @@ public CelCheckerLegacyImpl build() {
           identDeclarationSet,
           functionDeclarations.build(),
           Optional.fromNullable(expectedResultType),
-          legacyProvider,
+          customTypeProvider,
           messageTypeProvider,
           standardEnvironmentEnabled,
           standardDeclarations,

checker/src/main/java/dev/cel/checker/DescriptorTypeProvider.java

@@ -18,6 +18,7 @@
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Ascii;
 import com.google.common.base.Joiner;
+import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -32,28 +33,29 @@
 import com.google.protobuf.Descriptors.FieldDescriptor;
 import com.google.protobuf.Descriptors.FileDescriptor;
 import com.google.protobuf.Descriptors.OneofDescriptor;
-import dev.cel.common.annotations.Internal;
 import dev.cel.common.internal.FileDescriptorSetConverter;
+import dev.cel.common.types.CelProtoTypes;
+import dev.cel.common.types.CelType;
+import dev.cel.common.types.ProtoMessageType;
+import dev.cel.common.types.TypeType;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import org.jspecify.annotations.Nullable;
 
 /**
  * The {@code DescriptorTypeProvider} provides type information for one or more {@link Descriptor}
  * instances of proto messages.
  *
- * <p>TODO: Unify implementation across the runtime (i.e: DescriptorMessageProvider)
- * and the compilation. This class can likely be eliminated as part of the work.
- *
- * <p>CEL Library Internals. Do Not Use.
+ * @deprecated Do not use. Migrate to {@code ProtoMessageTypeProvider).
  */
 @Immutable
-@Internal
+@Deprecated
 public class DescriptorTypeProvider implements TypeProvider {
 
   @SuppressWarnings("Immutable")
@@ -86,6 +88,45 @@ public DescriptorTypeProvider(Iterable<Descriptor> descriptors) {
     return typeDef != null ? Types.create(Types.createMessage(typeDef.name())) : null;
   }
 
+  @Override
+  public Optional<TypeType> lookupCelType(String typeName) {
+    TypeDef typeDef = lookupMessageTypeDef(typeName);
+    if (typeDef == null) {
+      return Optional.empty();
+    }
+
+    ImmutableSet.Builder<String> fieldsBuilder = ImmutableSet.builder();
+    for (FieldDef fieldDef : typeDef.fields()) {
+      fieldsBuilder.add(fieldDef.name());
+    }
+
+    @SuppressWarnings("Immutable")
+    ProtoMessageType protoMessageType =
+        ProtoMessageType.create(
+            typeName,
+            fieldsBuilder.build(),
+            fieldName -> {
+              FieldDef fieldDef = typeDef.lookupField(fieldName);
+              if (fieldDef == null) {
+                return Optional.empty();
+              }
+
+              Type type = fieldDefToType(fieldDef);
+              return Optional.of(CelProtoTypes.typeToCelType(type));
+            },
+            extensionFieldName -> {
+              ExtensionFieldType extensionFieldType =
+                  symbolTable.lookupExtension(extensionFieldName);
+              if (extensionFieldType == null) {
+                return Optional.empty();
+              }
+
+              return Optional.of(extensionFieldType.fieldType().celType());
+            });
+
+    return Optional.of(TypeType.create(protoMessageType));
+  }
+
   @Override
   public @Nullable Integer lookupEnumValue(String enumName) {
     int dot = enumName.lastIndexOf('.');
@@ -339,6 +380,8 @@ private TypeDef buildTypeDef(EnumDescriptor descriptor, Map<String, TypeDef> typ
 
   /** Value object for a proto-based primitive, message, or enum definition. */
   @AutoValue
+  @AutoValue.CopyAnnotations
+  @SuppressWarnings("Immutable")
   protected abstract static class TypeDef {
 
     /** The qualified name of the message or enum. */
@@ -434,12 +477,28 @@ static TypeDef ofEnum(String name, Iterable<EnumValueDef> enumValues) {
     }
   }
 
+  @Override
+  public ImmutableCollection<CelType> types() {
+    ImmutableList.Builder<CelType> typesBuilder = ImmutableList.builder();
+    for (TypeDef typeDef : symbolTable.typeMap.values()) {
+      TypeType typeType = lookupCelType(typeDef.name()).orElse(null);
+      if (typeType == null) {
+        continue;
+      }
+
+      typesBuilder.add(typeType.type());
+    }
+
+    return typesBuilder.build();
+  }
+
   /**
    * Value object for a proto-based field definition.
    *
    * <p>Only one of the {@link #type} or {@link #mapEntryType} may be set.
    */
   @AutoValue
+  @AutoValue.CopyAnnotations
   protected abstract static class FieldDef {
 
     /** The field name. */