TypeChecker Refactor

Author: l46kok

.bazelversion

@@ -0,0 +1 @@
+9.0.0

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

@@ -1001,14 +1001,12 @@ 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())
+                CelOptions.current().resolveTypeDependencies(true).build())
             .setContainer(
                 CelContainer.ofName("dev.cel.testing.testdata.proto3.StandaloneGlobalEnum"))
             .setResultType(SimpleType.BOOL)

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",
 ]
 
@@ -190,6 +191,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,11 @@ 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 +496,7 @@ public CelCheckerLegacyImpl build() {
           identDeclarationSet,
           functionDeclarations.build(),
           Optional.fromNullable(expectedResultType),
-          legacyProvider,
+          customTypeProvider,
           messageTypeProvider,
           standardEnvironmentEnabled,
           standardDeclarations,

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

@@ -14,6 +14,11 @@
 
 package dev.cel.checker;
 
+import com.google.common.collect.ImmutableCollection;
+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 dev.cel.expr.Type;
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Ascii;
@@ -40,20 +45,19 @@
 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.
  */
 @Immutable
 @Internal
+@Deprecated
 public class DescriptorTypeProvider implements TypeProvider {
 
   @SuppressWarnings("Immutable")
@@ -86,6 +90,43 @@ 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") // Legacy type defs
+    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. */

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

@@ -14,12 +14,13 @@
 
 package dev.cel.checker;
 
+import dev.cel.common.types.EnumType;
+import dev.cel.common.types.TypeType;
 import dev.cel.expr.Constant;
 import dev.cel.expr.Decl;
 import dev.cel.expr.Decl.FunctionDecl.Overload;
 import dev.cel.expr.Expr;
 import dev.cel.expr.Type;
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -43,7 +44,9 @@
 import dev.cel.common.types.CelKind;
 import dev.cel.common.types.CelProtoTypes;
 import dev.cel.common.types.CelType;
+import dev.cel.common.types.CelTypeProvider;
 import dev.cel.common.types.CelTypes;
+import dev.cel.common.types.ProtoMessageTypeProvider;
 import dev.cel.common.types.SimpleType;
 import dev.cel.parser.CelStandardMacro;
 import java.util.ArrayList;
@@ -78,7 +81,7 @@ public class Env {
       CelFunctionDecl.newBuilder().setName("*error*").build();
 
   /** Type provider responsible for resolving CEL message references to strong types. */
-  private final TypeProvider typeProvider;
+  private final CelTypeProvider typeProvider;
 
   /**
    * Stack of declaration groups where each entry in stack represents a scope capable of hinding
@@ -105,7 +108,7 @@ public class Env {
           .build();
 
   private Env(
-      Errors errors, TypeProvider typeProvider, DeclGroup declGroup, CelOptions celOptions) {
+      Errors errors, CelTypeProvider typeProvider, DeclGroup declGroup, CelOptions celOptions) {
     this.celOptions = celOptions;
     this.errors = Preconditions.checkNotNull(errors);
     this.typeProvider = Preconditions.checkNotNull(typeProvider);
@@ -118,27 +121,10 @@ private Env(
    */
   @Deprecated
   public static Env unconfigured(Errors errors) {
-    return unconfigured(errors, LEGACY_TYPE_CHECKER_OPTIONS);
+    return unconfigured(errors, new ProtoMessageTypeProvider(), LEGACY_TYPE_CHECKER_OPTIONS);
   }
 
-  /**
-   * Creates an unconfigured {@code Env} value without the standard CEL types, functions, and
-   * operators with a reference to the configured {@code celOptions}.
-   */
-  @VisibleForTesting
-  static Env unconfigured(Errors errors, CelOptions celOptions) {
-    return unconfigured(errors, new DescriptorTypeProvider(), celOptions);
-  }
-
-  /**
-   * Creates an unconfigured {@code Env} value without the standard CEL types, functions, and
-   * operators using a custom {@code typeProvider}.
-   *
-   * @deprecated Do not use. This exists for compatibility reasons. Migrate to CEL-Java fluent APIs.
-   *     See {@code CelCompilerFactory}.
-   */
-  @Deprecated
-  public static Env unconfigured(Errors errors, TypeProvider typeProvider, CelOptions celOptions) {
+  static Env unconfigured(Errors errors, CelTypeProvider typeProvider, CelOptions celOptions) {
     return new Env(errors, typeProvider, new DeclGroup(), celOptions);
   }
 
@@ -148,7 +134,7 @@ public static Env unconfigured(Errors errors, TypeProvider typeProvider, CelOpti
    */
   @Deprecated
   public static Env standard(Errors errors) {
-    return standard(errors, new DescriptorTypeProvider());
+    return standard(errors, new ProtoMessageTypeProvider(), LEGACY_TYPE_CHECKER_OPTIONS);
   }
 
   /**
@@ -173,6 +159,11 @@ public static Env standard(Errors errors, TypeProvider typeProvider) {
    */
   @Deprecated
   public static Env standard(Errors errors, TypeProvider typeProvider, CelOptions celOptions) {
+    CelTypeProvider adapted = new TypeProviderLegacyImpl(typeProvider);
+    return standard(errors, adapted, celOptions);
+  }
+
+  static Env standard(Errors errors, CelTypeProvider typeProvider, CelOptions celOptions) {
     CelStandardDeclarations celStandardDeclaration =
         CelStandardDeclarations.newBuilder()
             .filterFunctions(
@@ -209,10 +200,10 @@ public static Env standard(Errors errors, TypeProvider typeProvider, CelOptions
     return standard(celStandardDeclaration, errors, typeProvider, celOptions);
   }
 
-  public static Env standard(
+  static Env standard(
       CelStandardDeclarations celStandardDeclaration,
       Errors errors,
-      TypeProvider typeProvider,
+      CelTypeProvider typeProvider,
       CelOptions celOptions) {
     Env env = Env.unconfigured(errors, typeProvider, celOptions);
     // Isolate the standard declarations into their own scope for forward compatibility.
@@ -228,8 +219,8 @@ public Errors getErrorContext() {
     return errors;
   }
 
-  /** Returns the {@code TypeProvider}. */
-  public TypeProvider getTypeProvider() {
+  /** Returns the {@code CelTypeProvider}. */
+  public CelTypeProvider getTypeProvider() {
     return typeProvider;
   }
 
@@ -491,30 +482,47 @@ public Env add(String name, Type type) {
 
     // Next try to import the name as a reference to a message type.
     // This is done via the type provider.
-    Optional<CelType> type = typeProvider.lookupCelType(cand);
+    Optional<CelType> type = typeProvider.findType(cand);
     if (type.isPresent()) {
-      decl = CelIdentDecl.newIdentDeclaration(cand, type.get());
+      decl = CelIdentDecl.newIdentDeclaration(cand, TypeType.create(type.get()));
       decls.get(0).putIdent(decl);
       return decl;
     }
 
     // Next try to import this as an enum value by splitting the name in a type prefix and
     // the enum inside.
-    Integer enumValue = typeProvider.lookupEnumValue(cand);
-    if (enumValue != null) {
+    Optional<Integer> enumValue = findEnumValue(cand);
+    if (enumValue.isPresent()) {
       decl =
           CelIdentDecl.newBuilder()
               .setName(cand)
               .setType(SimpleType.INT)
-              .setConstant(CelConstant.ofValue(enumValue))
-              .build();
+              .setConstant(CelConstant.ofValue(enumValue.get())).build();
 
       decls.get(0).putIdent(decl);
       return decl;
     }
+
     return null;
   }
 
+  private Optional<Integer> findEnumValue(String fullyQualifiedEnumName) {
+    int dot = fullyQualifiedEnumName.lastIndexOf('.');
+    if (dot <= 0) {
+      return Optional.empty();
+    }
+
+    String enumTypeName = fullyQualifiedEnumName.substring(0, dot);
+    EnumType enumType = typeProvider.findType(enumTypeName).filter(t -> t instanceof EnumType)
+        .map(EnumType.class::cast).orElse(null);
+    if (enumType == null) {
+      return Optional.empty();
+    }
+
+    String enumValueName = fullyQualifiedEnumName.substring(dot + 1);
+    return enumType.findNumberByName(enumValueName);
+  }
+
   /**
    * Lookup a local identifier by name. This searches only comprehension scopes, bypassing standard
    * environment or user-defined environment.