Support partial evaluation via unknowns in planner

PiperOrigin-RevId: 877616092

Author: l46kok

extensions/src/test/java/dev/cel/extensions/BUILD.bazel

@@ -38,6 +38,8 @@ java_library(
         "//runtime:interpreter_util",
         "//runtime:lite_runtime",
         "//runtime:lite_runtime_factory",
+        "//runtime:partial_vars",
+        "//runtime:unknown_attributes",
         "@cel_spec//proto/cel/expr/conformance/proto2:test_all_types_java_proto",
         "@cel_spec//proto/cel/expr/conformance/proto3:test_all_types_java_proto",
         "@cel_spec//proto/cel/expr/conformance/test:simple_java_proto",

extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java

@@ -49,10 +49,12 @@
 import dev.cel.expr.conformance.proto3.TestAllTypes.NestedMessage;
 import dev.cel.parser.CelMacro;
 import dev.cel.parser.CelStandardMacro;
+import dev.cel.runtime.CelAttributePattern;
 import dev.cel.runtime.CelEvaluationException;
 import dev.cel.runtime.CelFunctionBinding;
 import dev.cel.runtime.CelRuntime;
 import dev.cel.runtime.InterpreterUtil;
+import dev.cel.runtime.PartialVars;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.List;
@@ -897,14 +899,12 @@ public void optionalIndex_onMap_returnsOptionalValue() throws Exception {
   @TestParameters("{source: '{?x: x}'}")
   public void optionalIndex_onMapWithUnknownInput_returnsUnknownResult(String source)
       throws Exception {
-    if (testMode.equals(TestMode.PLANNER_CHECKED) || testMode.equals(TestMode.PLANNER_PARSE_ONLY)) {
-      // TODO: Uncomment once unknowns is implemented
-      return;
-    }
     Cel cel = newCelBuilder().addVar("x", OptionalType.create(SimpleType.INT)).build();
     CelAbstractSyntaxTree ast = compile(cel, source);
 
-    Object result = cel.createProgram(ast).eval();
+    Object result =
+        cel.createProgram(ast)
+            .eval(PartialVars.of(CelAttributePattern.fromQualifiedIdentifier("x")));
 
     assertThat(InterpreterUtil.isUnknown(result)).isTrue();
   }
@@ -987,18 +987,16 @@ public void optionalIndex_onOptionalList_returnsOptionalValue() throws Exception
 
   @Test
   public void optionalIndex_onListWithUnknownInput_returnsUnknownResult() throws Exception {
-    if (testMode.equals(TestMode.PLANNER_CHECKED) || testMode.equals(TestMode.PLANNER_PARSE_ONLY)) {
-      // TODO: Uncomment once unknowns is implemented
-      return;
-    }
     Cel cel =
         newCelBuilder()
             .addVar("x", OptionalType.create(SimpleType.INT))
             .setResultType(ListType.create(SimpleType.INT))
             .build();
     CelAbstractSyntaxTree ast = compile(cel, "[?x]");
 
-    Object result = cel.createProgram(ast).eval();
+    Object result =
+        cel.createProgram(ast)
+            .eval(PartialVars.of(CelAttributePattern.fromQualifiedIdentifier("x")));
 
     assertThat(InterpreterUtil.isUnknown(result)).isTrue();
   }
@@ -1017,6 +1015,29 @@ public void traditionalIndex_onOptionalList_returnsOptionalEmpty() throws Except
     assertThat(result).isEqualTo(Optional.empty());
   }
 
+  @Test
+  public void optionalFieldSelect_fieldMarkedUnknown_returnsUnknownSet() throws Exception {
+    if (testMode.equals(TestMode.LEGACY_CHECKED)) {
+      // This case is not possible to setup for legacy runtime
+      return;
+    }
+
+    Cel cel =
+        newCelBuilder()
+            .addVar("msg", StructTypeReference.create(TestAllTypes.getDescriptor().getFullName()))
+            .build();
+    CelAbstractSyntaxTree ast = compile(cel, "msg.?single_int32");
+
+    Object result =
+        cel.createProgram(ast)
+            .eval(
+                PartialVars.of(
+                    ImmutableMap.of("msg", TestAllTypes.newBuilder().setSingleInt32(42).build()),
+                    CelAttributePattern.fromQualifiedIdentifier("msg.single_int32")));
+
+    assertThat(InterpreterUtil.isUnknown(result)).isTrue();
+  }
+
   @Test
   // LHS
   @TestParameters("{expression: 'optx.or(optional.of(1))'}")
@@ -1026,18 +1047,16 @@ public void traditionalIndex_onOptionalList_returnsOptionalEmpty() throws Except
   @TestParameters("{expression: 'optional.none().orValue(optx)'}")
   public void optionalChainedFunctions_lhsIsUnknown_returnsUnknown(String expression)
       throws Exception {
-    if (testMode.equals(TestMode.PLANNER_CHECKED) || testMode.equals(TestMode.PLANNER_PARSE_ONLY)) {
-      // TODO: Uncomment once unknowns is implemented
-      return;
-    }
     Cel cel =
         newCelBuilder()
             .addVar("optx", OptionalType.create(SimpleType.INT))
             .addVar("x", SimpleType.INT)
             .build();
     CelAbstractSyntaxTree ast = compile(cel, expression);
 
-    Object result = cel.createProgram(ast).eval();
+    Object result =
+        cel.createProgram(ast)
+            .eval(PartialVars.of(CelAttributePattern.fromQualifiedIdentifier("optx")));
 
     assertThat(InterpreterUtil.isUnknown(result)).isTrue();
   }

runtime/BUILD.bazel

@@ -351,3 +351,17 @@ java_library(
         "//runtime/src/main/java/dev/cel/runtime:runtime_planner_impl",
     ],
 )
+
+java_library(
+    name = "accumulated_unknowns",
+    visibility = ["//:internal"],
+    exports = [
+        "//runtime/src/main/java/dev/cel/runtime:accumulated_unknowns",
+    ],
+)
+
+java_library(
+    name = "partial_vars",
+    visibility = ["//:internal"],
+    exports = ["//runtime/src/main/java/dev/cel/runtime:partial_vars"],
+)

runtime/src/main/java/dev/cel/runtime/AccumulatedUnknowns.java

@@ -15,18 +15,23 @@
 package dev.cel.runtime;
 
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import dev.cel.common.annotations.Internal;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.Set;
+import org.jspecify.annotations.Nullable;
 
 /**
  * An internal representation used for fast accumulation of unknown expr IDs and attributes. For
  * safety, this object should never be returned as an evaluated result and instead be adapted into
  * an immutable CelUnknownSet.
+ *
+ * <p>CEL Library Internals. Do Not Use.
  */
-final class AccumulatedUnknowns {
+@Internal
+public final class AccumulatedUnknowns {
   private static final int MAX_UNKNOWN_ATTRIBUTE_SIZE = 500_000;
   private final Set<Long> exprIds;
   private final Set<CelAttribute> attributes;
@@ -39,8 +44,21 @@ Set<CelAttribute> attributes() {
     return attributes;
   }
 
+  /**
+   * Evaluates if the right hand side is an accumulated unknown, and if so, merges it into the accumulator.
+   */
+  @Nullable
+  public static AccumulatedUnknowns maybeMerge(
+      @Nullable AccumulatedUnknowns accumulator, Object newValue) {
+    if (newValue instanceof AccumulatedUnknowns) {
+      AccumulatedUnknowns newUnknowns = (AccumulatedUnknowns) newValue;
+      return accumulator == null ? newUnknowns : accumulator.merge(newUnknowns);
+    }
+    return accumulator;
+  }
+
   @CanIgnoreReturnValue
-  AccumulatedUnknowns merge(AccumulatedUnknowns arg) {
+  public AccumulatedUnknowns merge(AccumulatedUnknowns arg) {
     enforceMaxAttributeSize(this.attributes, arg.attributes);
     this.exprIds.addAll(arg.exprIds);
     this.attributes.addAll(arg.attributes);
@@ -55,7 +73,8 @@ static AccumulatedUnknowns create(Collection<Long> ids) {
     return create(ids, new ArrayList<>());
   }
 
-  static AccumulatedUnknowns create(Collection<Long> exprIds, Collection<CelAttribute> attributes) {
+  public static AccumulatedUnknowns create(
+      Collection<Long> exprIds, Collection<CelAttribute> attributes) {
     return new AccumulatedUnknowns(new HashSet<>(exprIds), new HashSet<>(attributes));
   }
 

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

@@ -826,6 +826,7 @@ java_library(
         ":evaluation_listener",
         ":function_binding",
         ":function_resolver",
+        ":partial_vars",
         ":program",
         ":proto_message_runtime_equality",
         ":runtime",
@@ -938,6 +939,7 @@ java_library(
         ":function_resolver",
         ":interpretable",
         ":interpreter",
+        ":partial_vars",
         ":program",
         ":proto_message_activation_factory",
         ":runtime_equality",
@@ -955,7 +957,6 @@ java_library(
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",
         "@maven//:com_google_protobuf_protobuf_java",
-        "@maven//:org_jspecify_jspecify",
     ],
 )
 
@@ -1014,6 +1015,7 @@ java_library(
         ":evaluation_exception",
         ":function_resolver",
         ":interpretable",
+        ":partial_vars",
         ":program",
         ":variable_resolver",
         "//:auto_value",
@@ -1029,6 +1031,7 @@ cel_android_library(
         ":evaluation_exception",
         ":function_resolver_android",
         ":interpretable_android",
+        ":partial_vars_android",
         ":program_android",
         ":variable_resolver",
         "//:auto_value",
@@ -1199,6 +1202,7 @@ java_library(
         ":unknown_attributes",
         "//common/annotations",
         "@maven//:com_google_errorprone_error_prone_annotations",
+        "@maven//:com_google_guava_guava",
         "@maven//:org_jspecify_jspecify",
     ],
 )
@@ -1214,6 +1218,7 @@ cel_android_library(
         "//common/annotations",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:org_jspecify_jspecify",
+        "@maven_android//:com_google_guava_guava",
     ],
 )
 
@@ -1273,9 +1278,11 @@ java_library(
 java_library(
     name = "accumulated_unknowns",
     srcs = ["AccumulatedUnknowns.java"],
-    visibility = ["//visibility:private"],
+    tags = [
+    ],
     deps = [
         ":unknown_attributes",
+        "//common/annotations",
         "@maven//:com_google_errorprone_error_prone_annotations",
     ],
 )
@@ -1286,6 +1293,7 @@ cel_android_library(
     visibility = ["//visibility:private"],
     deps = [
         ":unknown_attributes_android",
+        "//common/annotations",
         "@maven//:com_google_errorprone_error_prone_annotations",
     ],
 )
@@ -1318,6 +1326,34 @@ cel_android_library(
     ],
 )
 
+java_library(
+    name = "partial_vars",
+    srcs = ["PartialVars.java"],
+    tags = [
+    ],
+    deps = [
+        ":variable_resolver",
+        "//:auto_value",
+        "//runtime:unknown_attributes",
+        "@maven//:com_google_errorprone_error_prone_annotations",
+        "@maven//:com_google_guava_guava",
+    ],
+)
+
+cel_android_library(
+    name = "partial_vars_android",
+    srcs = ["PartialVars.java"],
+    tags = [
+    ],
+    deps = [
+        ":variable_resolver",
+        "//:auto_value",
+        "//runtime:unknown_attributes_android",
+        "@maven//:com_google_errorprone_error_prone_annotations",
+        "@maven_android//:com_google_guava_guava",
+    ],
+)
+
 java_library(
     name = "program",
     srcs = ["Program.java"],
@@ -1326,6 +1362,7 @@ java_library(
     deps = [
         ":evaluation_exception",
         ":function_resolver",
+        ":partial_vars",
         ":variable_resolver",
         "@maven//:com_google_errorprone_error_prone_annotations",
     ],
@@ -1339,8 +1376,8 @@ cel_android_library(
     deps = [
         ":evaluation_exception",
         ":function_resolver_android",
+        ":partial_vars_android",
         ":variable_resolver",
-        "//:auto_value",
         "@maven//:com_google_errorprone_error_prone_annotations",
     ],
 )

runtime/src/main/java/dev/cel/runtime/CelRuntimeImpl.java

@@ -134,6 +134,11 @@ public Object eval(
         return program.eval(resolver, lateBoundFunctionResolver);
       }
 
+      @Override
+      public Object eval(PartialVars partialVars) throws CelEvaluationException {
+        return program.eval(partialVars);
+      }
+
       @Override
       public Object trace(CelEvaluationListener listener) throws CelEvaluationException {
         throw new UnsupportedOperationException("Trace is not yet supported.");

runtime/src/main/java/dev/cel/runtime/CelUnknownSet.java

@@ -59,7 +59,7 @@ static CelUnknownSet create(Iterable<Long> unknownExprIds) {
     return create(ImmutableSet.of(), ImmutableSet.copyOf(unknownExprIds));
   }
 
-  static CelUnknownSet create(
+  public static CelUnknownSet create(
       ImmutableSet<CelAttribute> attributes, ImmutableSet<Long> unknownExprIds) {
     return new AutoValue_CelUnknownSet(attributes, unknownExprIds);
   }

runtime/src/main/java/dev/cel/runtime/InterpreterUtil.java

@@ -14,6 +14,7 @@
 
 package dev.cel.runtime;
 
+import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CheckReturnValue;
 import dev.cel.common.annotations.Internal;
 import org.jspecify.annotations.Nullable;
@@ -55,23 +56,33 @@ public static boolean isUnknown(Object obj) {
     return obj instanceof CelUnknownSet;
   }
 
-  static boolean isAccumulatedUnknowns(Object obj) {
+  public static boolean isAccumulatedUnknowns(Object obj) {
     return obj instanceof AccumulatedUnknowns;
   }
 
   /** If the argument is {@link CelUnknownSet}, adapts it into {@link AccumulatedUnknowns} */
-  static Object maybeAdaptToAccumulatedUnknowns(Object val) {
+  public static Object maybeAdaptToAccumulatedUnknowns(Object val) {
     if (!(val instanceof CelUnknownSet)) {
       return val;
     }
 
     return adaptToAccumulatedUnknowns((CelUnknownSet) val);
   }
 
-  static AccumulatedUnknowns adaptToAccumulatedUnknowns(CelUnknownSet unknowns) {
+  public static AccumulatedUnknowns adaptToAccumulatedUnknowns(CelUnknownSet unknowns) {
     return AccumulatedUnknowns.create(unknowns.unknownExprIds(), unknowns.attributes());
   }
 
+  public static Object maybeAdaptToCelUnknownSet(Object val) {
+    if (!(val instanceof AccumulatedUnknowns)) {
+      return val;
+    }
+
+    AccumulatedUnknowns unknowns = (AccumulatedUnknowns) val;
+    return CelUnknownSet.create(
+        ImmutableSet.copyOf(unknowns.attributes()), ImmutableSet.copyOf(unknowns.exprIds()));
+  }
+
   /**
    * Enforces strictness on both lhs/rhs arguments from logical operators (i.e: intentionally throws
    * an appropriate exception when {@link Throwable} is encountered as part of evaluated result.

runtime/src/main/java/dev/cel/runtime/LiteProgramImpl.java

@@ -52,6 +52,12 @@ public Object eval(CelVariableResolver resolver) throws CelEvaluationException {
     throw new UnsupportedOperationException("To be implemented");
   }
 
+  @Override
+  public Object eval(PartialVars partialVars) throws CelEvaluationException {
+    // TODO: Wire in program planner
+    throw new UnsupportedOperationException("To be implemented");
+  }
+
   static Program plan(Interpretable interpretable) {
     return new AutoValue_LiteProgramImpl(interpretable);
   }