Add a runtime CEL type resolution for an unknown value

l46kok · copybara-github · commit af008243288f · 2024-10-15T19:45:55.000-07:00
PiperOrigin-RevId: 686325313
diff --git a/bundle/src/test/java/dev/cel/bundle/CelImplTest.java b/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 dev.cel.expr.Value;
 import com.google.api.expr.test.v1.proto2.Proto2ExtensionScopedMessage;
 import com.google.api.expr.test.v1.proto3.TestAllTypesProto.TestAllTypes;
 import com.google.common.collect.ImmutableList;
@@ -1926,6 +1927,54 @@ public void program_functionParamWithWellKnownType() throws Exception {
     assertThat(result).isTrue();
   }
 
+  @Test
+  public void program_nativeTypeUnknownsEnabled_asIdentifiers() throws Exception {
+    Cel cel =
+        CelFactory.standardCelBuilder()
+            .addVar("x", SimpleType.BOOL)
+            .addVar("y", SimpleType.BOOL)
+            .setOptions(CelOptions.current().adaptUnknownValueSetToNativeType(true).build())
+            .build();
+    CelAbstractSyntaxTree ast = cel.compile("x || y").getAst();
+
+    CelUnknownSet result = (CelUnknownSet) cel.createProgram(ast).eval();
+
+    assertThat(result.unknownExprIds()).containsExactly(1L, 3L);
+    assertThat(result.attributes()).isEmpty();
+  }
+
+  @Test
+  public void program_nativeTypeUnknownsEnabled_asCallArguments() throws Exception {
+    Cel cel =
+        CelFactory.standardCelBuilder()
+            .addVar("x", SimpleType.BOOL)
+            .addFunctionDeclarations(
+                newFunctionDeclaration(
+                    "foo", newGlobalOverload("foo_bool", SimpleType.BOOL, SimpleType.BOOL)))
+            .setOptions(CelOptions.current().adaptUnknownValueSetToNativeType(true).build())
+            .build();
+    CelAbstractSyntaxTree ast = cel.compile("foo(x)").getAst();
+
+    CelUnknownSet result = (CelUnknownSet) cel.createProgram(ast).eval();
+
+    assertThat(result.unknownExprIds()).containsExactly(2L);
+    assertThat(result.attributes()).isEmpty();
+  }
+
+  @Test
+  public void program_nativeTypeUnknownsEnabled_resolveUnknownType() throws Exception {
+    Cel cel =
+        CelFactory.standardCelBuilder()
+            .addVar("x", SimpleType.BOOL)
+            .setOptions(CelOptions.current().adaptUnknownValueSetToNativeType(true).build())
+            .build();
+    CelAbstractSyntaxTree ast = cel.compile("type(x)").getAst();
+
+    Value result = (Value) cel.createProgram(ast).eval();
+
+    assertThat(result.getTypeValue()).isEqualTo("unknown");
+  }
+
   @Test
   public void toBuilder_isImmutable() {
     CelBuilder celBuilder = CelFactory.standardCelBuilder();
diff --git a/common/src/main/java/dev/cel/common/CelOptions.java b/common/src/main/java/dev/cel/common/CelOptions.java
@@ -109,6 +109,8 @@ public enum ProtoUnsetFieldOptions {
 
   public abstract ProtoUnsetFieldOptions fromProtoUnsetFieldOption();
 
+  public abstract boolean adaptUnknownValueSetToNativeType();
+
   public abstract Builder toBuilder();
 
   public ImmutableSet<ExprFeatures> toExprFeatures() {
@@ -200,6 +202,7 @@ public static Builder newBuilder() {
         .enableCelValue(false)
         .comprehensionMaxIterations(-1)
         .unwrapWellKnownTypesOnFunctionDispatch(true)
+        .adaptUnknownValueSetToNativeType(false)
         .fromProtoUnsetFieldOption(ProtoUnsetFieldOptions.BIND_DEFAULT);
   }
 
@@ -504,6 +507,16 @@ public abstract static class Builder {
      */
     public abstract Builder fromProtoUnsetFieldOption(ProtoUnsetFieldOptions value);
 
+    /**
+     * If enabled, when the result of an evaluation is a set of unknown values, a native-type
+     * equivalent {@code CelUnknownSet} is returned as a result. Otherwise, ExprValue from {@code
+     * eval.proto} is returned instead.
+     *
+     * <p>This is a temporary flag. It will be removed once the consumers have migrated to the
+     * native-type representation.
+     */
+    public abstract Builder adaptUnknownValueSetToNativeType(boolean value);
+
     public abstract CelOptions build();
   }
 }
diff --git a/runtime/src/main/java/dev/cel/runtime/BUILD.bazel b/runtime/src/main/java/dev/cel/runtime/BUILD.bazel
@@ -19,6 +19,7 @@ BASE_SOURCES = [
     "TypeResolver.java",
 ]
 
+# keep sorted
 INTERPRETER_SOURCES = [
     "Activation.java",
     "CallArgumentChecker.java",
@@ -44,6 +45,7 @@ java_library(
     ],
     deps = [
         ":runtime_helper",
+        ":unknown_attributes",
         "//:auto_value",
         "//common",
         "//common:error_codes",
@@ -249,6 +251,7 @@ java_library(
     ],
     deps = [
         ":base",
+        ":unknown_attributes",
         "//common/annotations",
         "@cel_spec//proto/cel/expr:expr_java_proto",
         "@maven//:com_google_errorprone_error_prone_annotations",
diff --git a/runtime/src/main/java/dev/cel/runtime/CallArgumentChecker.java b/runtime/src/main/java/dev/cel/runtime/CallArgumentChecker.java
@@ -32,13 +32,13 @@
 @Internal
 class CallArgumentChecker {
   private final ArrayList<Long> exprIds;
-  private Optional<CelUnknownSet> unknowns;
   private final RuntimeUnknownResolver resolver;
   private final boolean acceptPartial;
+  private Optional<CelUnknownSet> unknowns;
 
   private CallArgumentChecker(RuntimeUnknownResolver resolver, boolean acceptPartial) {
-    exprIds = new ArrayList<>();
-    unknowns = Optional.empty();
+    this.exprIds = new ArrayList<>();
+    this.unknowns = Optional.empty();
     this.resolver = resolver;
     this.acceptPartial = acceptPartial;
   }
@@ -76,8 +76,13 @@ void checkArg(DefaultInterpreter.IntermediateResult arg) {
 
     // support for ExprValue unknowns.
     if (InterpreterUtil.isUnknown(arg.value())) {
-      ExprValue exprValue = (ExprValue) arg.value();
-      exprIds.addAll(exprValue.getUnknown().getExprsList());
+      if (InterpreterUtil.isExprValueUnknown(arg.value())) {
+        ExprValue exprValue = (ExprValue) arg.value();
+        exprIds.addAll(exprValue.getUnknown().getExprsList());
+      } else if (resolver.getAdaptUnknownValueSetOption()) {
+        CelUnknownSet unknownSet = (CelUnknownSet) arg.value();
+        exprIds.addAll(unknownSet.unknownExprIds());
+      }
     }
   }
 
@@ -98,8 +103,14 @@ Optional<Object> maybeUnknowns() {
     }
 
     if (!exprIds.isEmpty()) {
-      return Optional.of(
-          ExprValue.newBuilder().setUnknown(UnknownSet.newBuilder().addAllExprs(exprIds)).build());
+      if (resolver.getAdaptUnknownValueSetOption()) {
+        return Optional.of(CelUnknownSet.create(exprIds));
+      } else {
+        return Optional.of(
+            ExprValue.newBuilder()
+                .setUnknown(UnknownSet.newBuilder().addAllExprs(exprIds))
+                .build());
+      }
     }
 
     return Optional.empty();
diff --git a/runtime/src/main/java/dev/cel/runtime/CelUnknownSet.java b/runtime/src/main/java/dev/cel/runtime/CelUnknownSet.java
@@ -28,21 +28,46 @@
  */
 @AutoValue
 public abstract class CelUnknownSet {
+
+  /**
+   * Set of attributes with a series of selection or index operations marked unknown. This set is
+   * always empty if enableUnknownTracking is disabled in {@code CelOptions}.
+   */
   public abstract ImmutableSet<CelAttribute> attributes();
 
-  public static CelUnknownSet create(ImmutableSet<CelAttribute> attributes) {
-    return new AutoValue_CelUnknownSet(attributes);
-  }
+  /** Set of subexpression IDs that were decided to be unknown and in the critical path. */
+  public abstract ImmutableSet<Long> unknownExprIds();
 
   public static CelUnknownSet create(CelAttribute attribute) {
     return create(ImmutableSet.of(attribute));
   }
 
+  public static CelUnknownSet create(ImmutableSet<CelAttribute> attributes) {
+    return create(attributes, ImmutableSet.of());
+  }
+
+  public static CelUnknownSet create(Long... unknownExprIds) {
+    return create(ImmutableSet.copyOf(unknownExprIds));
+  }
+
+  static CelUnknownSet create(Iterable<Long> unknownExprIds) {
+    return create(ImmutableSet.of(), ImmutableSet.copyOf(unknownExprIds));
+  }
+
+  private static CelUnknownSet create(
+      ImmutableSet<CelAttribute> attributes, ImmutableSet<Long> unknownExprIds) {
+    return new AutoValue_CelUnknownSet(attributes, unknownExprIds);
+  }
+
   public static CelUnknownSet union(CelUnknownSet lhs, CelUnknownSet rhs) {
     return create(
         ImmutableSet.<CelAttribute>builder()
             .addAll(lhs.attributes())
             .addAll(rhs.attributes())
+            .build(),
+        ImmutableSet.<Long>builder()
+            .addAll(lhs.unknownExprIds())
+            .addAll(rhs.unknownExprIds())
             .build());
   }
 
diff --git a/runtime/src/main/java/dev/cel/runtime/DefaultInterpreter.java b/runtime/src/main/java/dev/cel/runtime/DefaultInterpreter.java
@@ -108,10 +108,6 @@ static IntermediateResult create(Object value) {
     }
   }
 
-  public DefaultInterpreter(RuntimeTypeProvider typeProvider, Dispatcher dispatcher) {
-    this(typeProvider, dispatcher, CelOptions.LEGACY);
-  }
-
   /**
    * Creates a new interpreter
    *
@@ -167,7 +163,10 @@ public Object eval(GlobalResolver resolver) throws InterpreterException {
     @Override
     public Object eval(GlobalResolver resolver, CelEvaluationListener listener)
         throws InterpreterException {
-      return evalTrackingUnknowns(RuntimeUnknownResolver.fromResolver(resolver), listener);
+      return evalTrackingUnknowns(
+          RuntimeUnknownResolver.fromResolver(
+              resolver, celOptions.adaptUnknownValueSetToNativeType()),
+          listener);
     }
 
     @Override
@@ -333,7 +332,9 @@ private IntermediateResult evalFieldSelect(
       Object fieldValue = typeProvider.selectField(operand, field);
 
       return IntermediateResult.create(
-          attribute, InterpreterUtil.valueOrUnknown(fieldValue, expr.id()));
+          attribute,
+          InterpreterUtil.valueOrUnknown(
+              fieldValue, expr.id(), celOptions.adaptUnknownValueSetToNativeType()));
     }
 
     private IntermediateResult evalCall(ExecutionFrame frame, CelExpr expr, CelCall callExpr)
@@ -486,7 +487,8 @@ private IntermediateResult mergeBooleanUnknowns(IntermediateResult lhs, Intermed
 
       // Otherwise fallback to normal impl
       return IntermediateResult.create(
-          InterpreterUtil.shortcircuitUnknownOrThrowable(lhs.value(), rhs.value()));
+          InterpreterUtil.shortcircuitUnknownOrThrowable(
+              lhs.value(), rhs.value(), celOptions.adaptUnknownValueSetToNativeType()));
     }
 
     private enum ShortCircuitableOperators {
diff --git a/runtime/src/main/java/dev/cel/runtime/InterpreterUtil.java b/runtime/src/main/java/dev/cel/runtime/InterpreterUtil.java
@@ -17,9 +17,7 @@
 import dev.cel.expr.ExprValue;
 import dev.cel.expr.UnknownSet;
 import dev.cel.common.annotations.Internal;
-import java.util.Arrays;
 import java.util.LinkedHashSet;
-import java.util.List;
 import java.util.Set;
 import org.jspecify.annotations.Nullable;
 
@@ -56,8 +54,19 @@ public static Object strict(Object valueOrThrowable) throws InterpreterException
    * @return boolean value if object is unknown.
    */
   public static boolean isUnknown(Object obj) {
-    return obj instanceof ExprValue
-        && ((ExprValue) obj).getKindCase() == ExprValue.KindCase.UNKNOWN;
+    if (isExprValueUnknown(obj)) {
+      return true;
+    }
+    return obj instanceof CelUnknownSet;
+  }
+
+  /** TODO: Remove once clients have been migrated. */
+  static boolean isExprValueUnknown(Object obj) {
+    if (obj instanceof ExprValue) {
+      return ((ExprValue) obj).getKindCase() == ExprValue.KindCase.UNKNOWN;
+    }
+
+    return false;
   }
 
   /**
@@ -67,7 +76,7 @@ public static boolean isUnknown(Object obj) {
    * @return A new ExprValue object which has all unknown expr ids from input objects, without
    *     duplication.
    */
-  public static ExprValue combineUnknownExprValue(Object... objs) {
+  static ExprValue combineUnknownProtoExprValue(Object... objs) {
     UnknownSet.Builder unknownsetBuilder = UnknownSet.newBuilder();
     Set<Long> ids = new LinkedHashSet<>();
     for (Object object : objs) {
@@ -79,21 +88,25 @@ public static ExprValue combineUnknownExprValue(Object... objs) {
     return ExprValue.newBuilder().setUnknown(unknownsetBuilder).build();
   }
 
-  /** Create a {@code ExprValue} for one or more {@code ids} representing an unknown set. */
-  public static ExprValue createUnknownExprValue(Long... ids) {
-    return createUnknownExprValue(Arrays.asList(ids));
+  static CelUnknownSet combineUnknownExprValue(Object... objs) {
+    Set<Long> ids = new LinkedHashSet<>();
+    for (Object object : objs) {
+      if (isUnknown(object)) {
+        ids.addAll(((CelUnknownSet) object).unknownExprIds());
+      }
+    }
+
+    return CelUnknownSet.create(ids);
   }
 
   /**
    * Create an ExprValue object has UnknownSet, from a list of unknown expr ids
    *
-   * @param ids List of unknown expr ids
+   * @param id unknown expr id
    * @return A new ExprValue object which has all unknown expr ids from input list
    */
-  public static ExprValue createUnknownExprValue(List<Long> ids) {
-    ExprValue.Builder exprValueBuilder = ExprValue.newBuilder();
-    exprValueBuilder.setUnknown(UnknownSet.newBuilder().addAllExprs(ids));
-    return exprValueBuilder.build();
+  private static ExprValue createUnknownExprValue(long id) {
+    return ExprValue.newBuilder().setUnknown(UnknownSet.newBuilder().addExprs(id)).build();
   }
 
   /**
@@ -104,11 +117,14 @@ public static ExprValue createUnknownExprValue(List<Long> ids) {
    * from any boolean arguments alone. This allows us to consolidate the error/unknown handling for
    * both of these operators.
    */
-  public static Object shortcircuitUnknownOrThrowable(Object left, Object right)
+  public static Object shortcircuitUnknownOrThrowable(
+      Object left, Object right, boolean adaptUnknownValueSetToNativeType)
       throws InterpreterException {
     // unknown <op> unknown ==> unknown combined
     if (InterpreterUtil.isUnknown(left) && InterpreterUtil.isUnknown(right)) {
-      return InterpreterUtil.combineUnknownExprValue(left, right);
+      return adaptUnknownValueSetToNativeType
+          ? InterpreterUtil.combineUnknownExprValue(left, right)
+          : InterpreterUtil.combineUnknownProtoExprValue(left, right);
     }
     // unknown <op> <error> ==> unknown
     // unknown <op> t|f ==> unknown
@@ -132,18 +148,24 @@ public static Object shortcircuitUnknownOrThrowable(Object left, Object right)
         "Left or/and right object is neither bool, unknown nor error, unexpected behavior.");
   }
 
-  public static Object valueOrUnknown(@Nullable Object valueOrThrowable, Long id) {
+  public static Object valueOrUnknown(
+      @Nullable Object valueOrThrowable, Long id, boolean adaptUnknownValueSet) {
     // Handle the unknown value case.
     if (isUnknown(valueOrThrowable)) {
-      ExprValue value = (ExprValue) valueOrThrowable;
-      if (value.getUnknown().getExprsCount() != 0) {
-        return valueOrThrowable;
+      if (adaptUnknownValueSet) {
+        return CelUnknownSet.create(id);
+      } else {
+        // TODO: Remove once clients have been migrated.
+        ExprValue value = (ExprValue) valueOrThrowable;
+        if (value.getUnknown().getExprsCount() != 0) {
+          return valueOrThrowable;
+        }
+        return createUnknownExprValue(id);
       }
-      return createUnknownExprValue(id);
     }
     // Handle the null value case.
     if (valueOrThrowable == null) {
-      return createUnknownExprValue(id);
+      return adaptUnknownValueSet ? CelUnknownSet.create(id) : createUnknownExprValue(id);
     }
     return valueOrThrowable;
   }
diff --git a/runtime/src/main/java/dev/cel/runtime/RuntimeUnknownResolver.java b/runtime/src/main/java/dev/cel/runtime/RuntimeUnknownResolver.java
diff --git a/runtime/src/main/java/dev/cel/runtime/StandardTypeResolver.java b/runtime/src/main/java/dev/cel/runtime/StandardTypeResolver.java
