TypeChecker Refactor

PiperOrigin-RevId: 861155220

Author: l46kok

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

@@ -190,6 +190,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/Env.java

@@ -19,7 +19,6 @@
 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 +42,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 +79,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 +106,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 +119,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 +132,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 +157,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 +198,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 +217,8 @@ public Errors getErrorContext() {
     return errors;
   }
 
-  /** Returns the {@code TypeProvider}. */
-  public TypeProvider getTypeProvider() {
+  /** Returns the {@code CelTypeProvider}. */
+  public CelTypeProvider getTypeProvider() {
     return typeProvider;
   }
 
@@ -491,7 +480,7 @@ 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());
       decls.get(0).putIdent(decl);
@@ -500,7 +489,8 @@ public Env add(String name, Type type) {
 
     // 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);
+    // Integer enumValue = typeProvider.lookupEnumValue(cand);
+    Integer enumValue = 0;
     if (enumValue != null) {
       decl =
           CelIdentDecl.newBuilder()

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

@@ -39,11 +39,15 @@
 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.ListType;
 import dev.cel.common.types.MapType;
 import dev.cel.common.types.OptionalType;
+import dev.cel.common.types.ProtoMessageType;
+import dev.cel.common.types.ProtoMessageType.Extension;
 import dev.cel.common.types.SimpleType;
+import dev.cel.common.types.StructType.Field;
 import dev.cel.common.types.TypeType;
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -143,7 +147,7 @@ public static CelAbstractSyntaxTree typecheck(
   }
 
   private final Env env;
-  private final TypeProvider typeProvider;
+  private final CelTypeProvider typeProvider;
   private final CelContainer container;
   private final Map<Long, Integer> positionMap;
   private final InferenceContext inferenceContext;
@@ -410,7 +414,7 @@ private CelExpr visit(CelExpr expr, CelExpr.CelStruct struct) {
         expr = replaceStructEntryValueSubtree(expr, visitedValueExpr, i);
       }
       CelType fieldType =
-          getFieldType(entry.id(), getPosition(entry), messageType, entry.fieldKey()).celType();
+          getFieldType(entry.id(), getPosition(entry), messageType, entry.fieldKey()).type();
       CelType valueType = env.getType(visitedValueExpr);
       if (entry.optionalEntry()) {
         if (valueType instanceof OptionalType) {
@@ -716,7 +720,7 @@ private OverloadResolution resolveOverload(
   // Return value from visit is not needed as the subtree is not rewritten here.
   @SuppressWarnings("CheckReturnValue")
   private CelType visitSelectField(
-      CelExpr expr, CelExpr operand, String field, boolean isOptional) {
+      CelExpr expr, CelExpr operand, String fieldName, boolean isOptional) {
     CelType operandType = inferenceContext.specialize(env.getType(operand));
     CelType resultType = SimpleType.ERROR;
 
@@ -727,10 +731,9 @@ private CelType visitSelectField(
 
     if (!Types.isDynOrError(operandType)) {
       if (operandType.kind() == CelKind.STRUCT) {
-        TypeProvider.FieldType fieldType =
-            getFieldType(expr.id(), getPosition(expr), operandType, field);
-        // Type of the field
-        resultType = fieldType.celType();
+        Field field =
+            getFieldType(expr.id(), getPosition(expr), operandType, fieldName);
+        resultType = field.type();
       } else if (operandType.kind() == CelKind.MAP) {
         resultType = ((MapType) operandType).valueType();
       } else if (operandType.kind() == CelKind.TYPE_PARAM) {
@@ -805,18 +808,25 @@ private CelExpr visitOptionalCall(CelExpr expr, CelExpr.CelCall call) {
   }
 
   /** Returns the field type give a type instance and field name. */
-  private TypeProvider.FieldType getFieldType(
-      long exprId, int position, CelType type, String fieldName) {
+  private Field getFieldType(long exprId, int position, CelType type, String fieldName) {
     String typeName = type.name();
-    if (typeProvider.lookupCelType(typeName).isPresent()) {
-      TypeProvider.FieldType fieldType = typeProvider.lookupFieldType(type, fieldName);
-      if (fieldType != null) {
-        return fieldType;
+    ProtoMessageType protoMessageType =
+        typeProvider
+            .findType(typeName)
+            .filter(t -> t instanceof ProtoMessageType)
+            .map(ProtoMessageType.class::cast)
+            .orElse(null);
+
+    if (protoMessageType != null) {
+      Field field = protoMessageType.findField(fieldName).orElse(null);
+      if (field != null) {
+        return field;
       }
-      TypeProvider.ExtensionFieldType extensionFieldType =
-          typeProvider.lookupExtensionType(fieldName);
-      if (extensionFieldType != null) {
-        return extensionFieldType.fieldType();
+
+      Extension extensionField = protoMessageType.findExtension(fieldName).orElse(null);
+
+      if (extensionField != null) {
+        return Field.of(extensionField.name(), extensionField.type());
       }
       env.reportError(exprId, position, "undefined field '%s'", fieldName);
     } else {
@@ -831,6 +841,7 @@ private TypeProvider.FieldType getFieldType(
       }
       env.reportError(exprId, position, errorMessage, fieldName, typeName);
     }
+
     return ERROR;
   }
 
@@ -892,8 +903,8 @@ public static OverloadResolution of(CelReference reference, CelType type) {
     }
   }
 
-  /** Helper object to represent a {@link TypeProvider.FieldType} lookup failure. */
-  private static final TypeProvider.FieldType ERROR = TypeProvider.FieldType.of(Types.ERROR);
+  /** Helper object to represent a {@link CelTypeProvider#findType(String)} lookup failure. */
+  private static final Field ERROR = Field.of(SimpleType.ERROR.name(), SimpleType.ERROR);
 
   private static CelExpr replaceIdentSubtree(CelExpr expr, String name) {
     CelExpr.CelIdent newIdent = CelExpr.CelIdent.newBuilder().setName(name).build();