Implement iterative eval for program planner

PiperOrigin-RevId: 964335064

Author: l46kok

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

@@ -19,6 +19,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
 import org.jspecify.annotations.Nullable;
@@ -36,12 +37,12 @@ public final class AccumulatedUnknowns {
   private final Set<Long> exprIds;
   private final Set<CelAttribute> attributes;
 
-  Set<Long> exprIds() {
-    return exprIds;
+  public Set<Long> exprIds() {
+    return Collections.unmodifiableSet(exprIds);
   }
 
-  Set<CelAttribute> attributes() {
-    return attributes;
+  public Set<CelAttribute> attributes() {
+    return Collections.unmodifiableSet(attributes);
   }
 
   /**

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

@@ -235,7 +235,16 @@ public Object trace(PartialVars partialVars, CelEvaluationListener listener)
 
       @Override
       public Object advanceEvaluation(UnknownContext context) throws CelEvaluationException {
-        throw new UnsupportedOperationException("Unsupported operation.");
+        PlannedProgram plannedProgram = (PlannedProgram) program;
+        return plannedProgram.evalOrThrow(
+            plannedProgram.interpretable(),
+            context.variableResolver(),
+            EMPTY_FUNCTION_RESOLVER,
+            PartialVars.of(
+                (name) -> Optional.ofNullable(context.variableResolver().resolve(name)),
+                context.unresolvedAttributes()),
+            context.createAttributeResolver(),
+            /* listener= */ null);
       }
     };
   }

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

@@ -107,6 +107,16 @@ public GlobalResolver variableResolver() {
     return variableResolver;
   }
 
+  /** Accessor for unresolved attribute patterns. */
+  ImmutableList<CelAttributePattern> unresolvedAttributes() {
+    return unresolvedAttributes;
+  }
+
+  /** Accessor for resolved attribute values. */
+  ImmutableMap<CelAttribute, Object> resolvedAttributes() {
+    return resolvedAttributes;
+  }
+
   /**
    * Creates a new unknown context that is a copy of the current context with the provided
    * additional attribute values.
@@ -168,10 +178,27 @@ public Optional<Object> resolve(CelAttribute attribute) {
 
     @Override
     public Optional<CelUnknownSet> maybePartialUnknown(CelAttribute attribute) {
-      return unresolvedAttributes.stream()
-          .filter(p -> p.isPartialMatch(attribute))
-          .findFirst()
-          .map(p -> CelUnknownSet.create(p.simplify(attribute)));
+      if (attribute.equals(CelAttribute.EMPTY) || attribute.qualifiers().isEmpty()) {
+        return Optional.empty();
+      }
+      Optional<CelUnknownSet> fromUnresolved =
+          unresolvedAttributes.stream()
+              .filter(p -> p.isPartialMatch(attribute))
+              .findFirst()
+              .map(p -> CelUnknownSet.create(p.simplify(attribute)));
+      if (fromUnresolved.isPresent()) {
+        return fromUnresolved;
+      }
+      for (CelAttribute resolved : resolvedAttributes.keySet()) {
+        if (resolved.qualifiers().size() > attribute.qualifiers().size()
+            && resolved
+                .qualifiers()
+                .subList(0, attribute.qualifiers().size())
+                .equals(attribute.qualifiers())) {
+          return Optional.of(CelUnknownSet.create(attribute));
+        }
+      }
+      return Optional.empty();
     }
   }
 }

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

@@ -89,6 +89,7 @@ java_library(
         "//runtime:partial_vars",
         "//runtime:program",
         "//runtime:resolved_overload",
+        "//runtime:unknown_attributes",
         "//runtime:variable_resolver",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:org_jspecify_jspecify",
@@ -131,6 +132,7 @@ java_library(
         "//common/types:type_providers",
         "//common/values",
         "//runtime:accumulated_unknowns",
+        "//runtime:evaluation_exception",
         "//runtime:interpretable",
         "//runtime:interpreter_util",
         "//runtime:partial_vars",
@@ -221,12 +223,17 @@ java_library(
     deps = [
         ":eval_helpers",
         ":planned_interpretable",
+        "//common:operator",
         "//common/ast",
         "//common/values",
         "//runtime:accumulated_unknowns",
         "//runtime:evaluation_exception",
         "//runtime:interpretable",
+        "//runtime:interpreter_util",
+        "//runtime:partial_vars",
         "//runtime:resolved_overload",
+        "//runtime:unknown_attributes",
+        "@maven//:com_google_guava_guava",
     ],
 )
 
@@ -518,6 +525,7 @@ java_library(
         "//runtime:interpreter_util",
         "//runtime:partial_vars",
         "//runtime:resolved_overload",
+        "//runtime:unknown_attributes",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:org_jspecify_jspecify",
     ],
@@ -596,6 +604,7 @@ cel_android_library(
         "//runtime:evaluation_exception_builder",
         "//runtime:interpretable_android",
         "//runtime:resolved_overload_android",
+        "//runtime:unknown_attributes_android",
         "//runtime:variable_resolver",
         "//runtime/src/main/java/dev/cel/runtime:evaluation_listener_android",
         "//runtime/src/main/java/dev/cel/runtime:function_resolver_android",
@@ -642,6 +651,7 @@ cel_android_library(
         "//common/types:type_providers_android",
         "//common/types:types_android",
         "//common/values:values_android",
+        "//runtime:evaluation_exception",
         "//runtime:interpretable_android",
         "//runtime:unknown_attributes_android",
         "//runtime/src/main/java/dev/cel/runtime:accumulated_unknowns_android",
@@ -733,12 +743,17 @@ cel_android_library(
     deps = [
         ":eval_helpers_android",
         ":planned_interpretable_android",
+        "//common:operator_android",
         "//common/ast:ast_android",
         "//common/values:values_android",
         "//runtime:evaluation_exception",
         "//runtime:interpretable_android",
+        "//runtime:interpreter_util_android",
+        "//runtime:partial_vars_android",
         "//runtime:resolved_overload_android",
+        "//runtime:unknown_attributes_android",
         "//runtime/src/main/java/dev/cel/runtime:accumulated_unknowns_android",
+        "@maven_android//:com_google_guava_guava",
     ],
 )
 
@@ -1023,6 +1038,7 @@ cel_android_library(
         "//runtime:interpretable_android",
         "//runtime:interpreter_util_android",
         "//runtime:resolved_overload_android",
+        "//runtime:unknown_attributes_android",
         "//runtime/src/main/java/dev/cel/runtime:function_resolver_android",
         "//runtime/src/main/java/dev/cel/runtime:partial_vars_android",
         "@maven//:com_google_errorprone_error_prone_annotations",

runtime/src/main/java/dev/cel/runtime/planner/EvalBinary.java

@@ -17,12 +17,23 @@
 import static dev.cel.runtime.planner.EvalHelpers.evalNonstrictly;
 import static dev.cel.runtime.planner.EvalHelpers.evalStrictly;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.primitives.UnsignedLong;
+import dev.cel.common.Operator;
 import dev.cel.common.ast.CelExpr;
 import dev.cel.common.values.CelValueConverter;
+import dev.cel.common.values.ErrorValue;
 import dev.cel.runtime.AccumulatedUnknowns;
+import dev.cel.runtime.CelAttribute;
+import dev.cel.runtime.CelAttributePattern;
+import dev.cel.runtime.CelAttributeResolver;
 import dev.cel.runtime.CelEvaluationException;
 import dev.cel.runtime.CelResolvedOverload;
 import dev.cel.runtime.GlobalResolver;
+import dev.cel.runtime.InterpreterUtil;
+import dev.cel.runtime.PartialVars;
+import java.util.Optional;
+import org.jspecify.annotations.Nullable;
 
 final class EvalBinary extends PlannedInterpretable {
 
@@ -39,7 +50,13 @@ Object evalInternal(GlobalResolver resolver, ExecutionFrame frame) throws CelEva
         isStrict ? evalStrictly(arg1, resolver, frame) : evalNonstrictly(arg1, resolver, frame);
     Object argVal2 =
         isStrict ? evalStrictly(arg2, resolver, frame) : evalNonstrictly(arg2, resolver, frame);
+
     if (isStrict) {
+      Object indexUnknownResult = maybeEvaluateIndexUnknown(argVal1, argVal2, frame);
+      if (indexUnknownResult != null) {
+        return indexUnknownResult;
+      }
+
       AccumulatedUnknowns unknowns = AccumulatedUnknowns.maybeMerge(null, argVal1);
       unknowns = AccumulatedUnknowns.maybeMerge(unknowns, argVal2);
       if (unknowns != null) {
@@ -51,6 +68,78 @@ Object evalInternal(GlobalResolver resolver, ExecutionFrame frame) throws CelEva
         functionName, resolvedOverload, celValueConverter, argVal1, argVal2);
   }
 
+  private @Nullable Object maybeEvaluateIndexUnknown(
+      Object target, Object index, ExecutionFrame frame) throws CelEvaluationException {
+    if (!Operator.INDEX.getFunction().equals(functionName)
+        || !(target instanceof AccumulatedUnknowns)
+        || target instanceof ErrorValue
+        || index instanceof AccumulatedUnknowns
+        || index instanceof ErrorValue) {
+      return null;
+    }
+
+    Optional<CelAttribute.Qualifier> optionalQualifier = toQualifier(index);
+    if (!optionalQualifier.isPresent()) {
+      return null;
+    }
+    CelAttribute.Qualifier qualifier = optionalQualifier.get();
+
+    AccumulatedUnknowns targetUnknowns = (AccumulatedUnknowns) target;
+    CelAttributeResolver attributeResolver = frame.attributeResolver().orElse(null);
+    PartialVars partialVars = frame.partialVars().orElse(null);
+
+    ImmutableList.Builder<CelAttribute> qualifiedAttrs = ImmutableList.builder();
+    for (CelAttribute attr : targetUnknowns.attributes()) {
+      CelAttribute qualifiedAttr = attr.qualify(qualifier);
+      if (attributeResolver != null) {
+        Optional<Object> resolved = attributeResolver.resolve(qualifiedAttr);
+        if (resolved.isPresent()) {
+          return adaptResolvedValue(resolved.get());
+        }
+      }
+      qualifiedAttrs.add(simplifyAttribute(qualifiedAttr, attr, partialVars));
+    }
+
+    return AccumulatedUnknowns.create(targetUnknowns.exprIds(), qualifiedAttrs.build());
+  }
+
+  private static CelAttribute simplifyAttribute(
+      CelAttribute qualifiedAttr, CelAttribute fallbackAttr, @Nullable PartialVars partialVars) {
+    if (partialVars == null) {
+      return qualifiedAttr;
+    }
+    for (CelAttributePattern pattern : partialVars.unknowns()) {
+      if (pattern.isPartialMatch(qualifiedAttr)) {
+        return pattern.simplify(qualifiedAttr);
+      }
+    }
+    return fallbackAttr;
+  }
+
+  private static Optional<CelAttribute.Qualifier> toQualifier(Object value) {
+    if (value instanceof UnsignedLong) {
+      return Optional.of(CelAttribute.Qualifier.ofUint((UnsignedLong) value));
+    }
+    if (value instanceof Long) {
+      return Optional.of(CelAttribute.Qualifier.ofInt((Long) value));
+    }
+    if (value instanceof Integer) {
+      return Optional.of(CelAttribute.Qualifier.ofInt(((Integer) value).longValue()));
+    }
+    if (value instanceof Boolean) {
+      return Optional.of(CelAttribute.Qualifier.ofBool((Boolean) value));
+    }
+    if (value instanceof String) {
+      return Optional.of(CelAttribute.Qualifier.ofString((String) value));
+    }
+    return Optional.empty();
+  }
+
+  private static Object adaptResolvedValue(Object resolvedVal) throws CelEvaluationException {
+    resolvedVal = InterpreterUtil.strict(resolvedVal);
+    return InterpreterUtil.maybeAdaptToAccumulatedUnknowns(resolvedVal);
+  }
+
   static EvalBinary create(
       CelExpr expr,
       String functionName,

runtime/src/main/java/dev/cel/runtime/planner/ExecutionFrame.java

@@ -16,6 +16,7 @@
 
 import dev.cel.common.CelOptions;
 import dev.cel.common.exceptions.CelIterationLimitExceededException;
+import dev.cel.runtime.CelAttributeResolver;
 import dev.cel.runtime.CelEvaluationException;
 import dev.cel.runtime.CelEvaluationListener;
 import dev.cel.runtime.CelFunctionResolver;
@@ -31,6 +32,7 @@ final class ExecutionFrame {
   private final int comprehensionIterationLimit;
   private final CelFunctionResolver functionResolver;
   private final PartialVars partialVars;
+  private final @Nullable CelAttributeResolver attributeResolver;
   private final @Nullable CelEvaluationListener listener;
   private int iterationCount;
   private BlockMemoizer blockMemoizer;
@@ -68,15 +70,33 @@ static ExecutionFrame create(
       CelFunctionResolver functionResolver,
       CelOptions celOptions,
       @Nullable PartialVars partialVars,
+      @Nullable CelAttributeResolver attributeResolver,
       @Nullable CelEvaluationListener listener) {
     return new ExecutionFrame(
-        functionResolver, celOptions.comprehensionMaxIterations(), partialVars, listener);
+        functionResolver,
+        celOptions.comprehensionMaxIterations(),
+        partialVars,
+        attributeResolver,
+        listener);
+  }
+
+  static ExecutionFrame create(
+      CelFunctionResolver functionResolver,
+      CelOptions celOptions,
+      @Nullable PartialVars partialVars,
+      @Nullable CelEvaluationListener listener) {
+    return create(
+        functionResolver, celOptions, partialVars, /* attributeResolver= */ null, listener);
   }
 
   Optional<PartialVars> partialVars() {
     return Optional.ofNullable(partialVars);
   }
 
+  Optional<CelAttributeResolver> attributeResolver() {
+    return Optional.ofNullable(attributeResolver);
+  }
+
   @Nullable CelEvaluationListener getListener() {
     return listener;
   }
@@ -85,10 +105,12 @@ private ExecutionFrame(
       CelFunctionResolver functionResolver,
       int limit,
       @Nullable PartialVars partialVars,
+      @Nullable CelAttributeResolver attributeResolver,
       @Nullable CelEvaluationListener listener) {
     this.comprehensionIterationLimit = limit;
     this.functionResolver = functionResolver;
     this.partialVars = partialVars;
+    this.attributeResolver = attributeResolver;
     this.listener = listener;
   }
 }