Fix unknown merging to be linear in space complexity when referenced in binds

PiperOrigin-RevId: 788995643

Author: l46kok

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

@@ -26,7 +26,7 @@
  * an immutable CelUnknownSet.
  */
 final class AccumulatedUnknowns {
-
+  private static final int MAX_UNKNOWN_ATTRIBUTE_SIZE = 500_000;
   private final List<Long> exprIds;
   private final List<CelAttribute> attributes;
 
@@ -40,6 +40,7 @@ List<CelAttribute> attributes() {
 
   @CanIgnoreReturnValue
   AccumulatedUnknowns merge(AccumulatedUnknowns arg) {
+    enforceMaxAttributeSize(this.attributes, arg.attributes);
     this.exprIds.addAll(arg.exprIds);
     this.attributes.addAll(arg.attributes);
     return this;
@@ -57,6 +58,15 @@ static AccumulatedUnknowns create(Collection<Long> exprIds, Collection<CelAttrib
     return new AccumulatedUnknowns(new ArrayList<>(exprIds), new ArrayList<>(attributes));
   }
 
+  private static void enforceMaxAttributeSize(
+      List<CelAttribute> lhsAttributes, List<CelAttribute> rhsAttributes) {
+    int totalSize = lhsAttributes.size() + rhsAttributes.size();
+    if (totalSize > MAX_UNKNOWN_ATTRIBUTE_SIZE) {
+      throw new IllegalArgumentException(
+          "Exceeded maximum allowed unknown attributes: " + totalSize);
+    }
+  }
+
   private AccumulatedUnknowns(List<Long> exprIds, List<CelAttribute> attributes) {
     this.exprIds = exprIds;
     this.attributes = attributes;

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

@@ -150,7 +150,7 @@ private ScopedResolver(
     DefaultInterpreter.IntermediateResult resolveSimpleName(String name, Long exprId) {
       DefaultInterpreter.IntermediateResult result = lazyEvalResultCache.get(name);
       if (result != null) {
-        return result;
+        return copyIfMutable(result);
       }
       result = shadowedVars.get(name);
       if (result != null) {
@@ -161,7 +161,27 @@ DefaultInterpreter.IntermediateResult resolveSimpleName(String name, Long exprId
 
     @Override
     void cacheLazilyEvaluatedResult(String name, DefaultInterpreter.IntermediateResult result) {
-      lazyEvalResultCache.put(name, result);
+      lazyEvalResultCache.put(name, copyIfMutable(result));
+    }
+
+    /**
+     * Perform a defensive copy of the intermediate result if it is mutable.
+     *
+     * <p>Some internal types are mutable to optimize performance, but this can cause issues when
+     * the result can be reused in multiple subexpressions due to caching.
+     *
+     * <p>Note: this is necessary on both the cache put and get path since the interpreter may use
+     * the same instance that was cached as a return value.
+     */
+    private static DefaultInterpreter.IntermediateResult copyIfMutable(
+        DefaultInterpreter.IntermediateResult result) {
+      if (result.value() instanceof AccumulatedUnknowns) {
+        AccumulatedUnknowns unknowns = (AccumulatedUnknowns) result.value();
+        return DefaultInterpreter.IntermediateResult.create(
+            result.attribute(),
+            AccumulatedUnknowns.create(unknowns.exprIds(), unknowns.attributes()));
+      }
+      return result;
     }
   }
 

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

@@ -52,6 +52,7 @@ java_library(
         "//common/values:proto_message_lite_value_provider",
         "//compiler",
         "//compiler:compiler_builder",
+        "//extensions",
         "//extensions:optional_library",
         "//parser:macro",
         "//parser:unparser",

runtime/src/test/java/dev/cel/runtime/CelRuntimeTest.java

@@ -42,11 +42,13 @@
 import dev.cel.common.ast.CelExpr;
 import dev.cel.common.ast.CelExpr.ExprKind.Kind;
 import dev.cel.common.types.CelV1AlphaTypes;
+import dev.cel.common.types.ListType;
 import dev.cel.common.types.SimpleType;
 import dev.cel.common.types.StructTypeReference;
 import dev.cel.compiler.CelCompiler;
 import dev.cel.compiler.CelCompilerFactory;
 import dev.cel.expr.conformance.proto3.TestAllTypes;
+import dev.cel.extensions.CelExtensions;
 import dev.cel.parser.CelStandardMacro;
 import dev.cel.parser.CelUnparserFactory;
 import java.util.List;
@@ -117,6 +119,50 @@ public void evaluate_v1alpha1CheckedExpr() throws Exception {
     assertThat(evaluatedResult).isEqualTo("Hello world!");
   }
 
+  @Test
+  // Lazy evaluation result cache doesn't allow references to mutate the cached instance.
+  @TestParameters(
+      "{expression: 'cel.bind(x, unknown_attr, (unknown_attr > 0) || [0, 1, 2, 3, 4, 5, 6, 7, 8, 9,"
+          + " 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20].exists(i, x + x > 0))'}")
+  @TestParameters(
+      "{expression: 'cel.bind(x, unknown_attr, x + x + x + x + x + x + x + x + x + x + x + x + x +"
+          + " x + x + x + x + x + x + x + x + x + x + x + x + x + x + x + x + x + x + x + x)'}")
+  // A new unknown is created per 'x' reference.
+  @TestParameters(
+      "{expression: '(my_list.exists(x, (x + x + x + x + x + x + x + x + x + x + x + x + x + x + x"
+          + " + x + x + x + x + x + x + x + x + x + x + x + x + x + x + x + x + x + x) > 100) &&"
+          + " false) || unknown_attr > 0'}")
+  public void advanceEvaluation_withUnknownTracking_noSelfReferenceInMerge(String expression)
+      throws Exception {
+    Cel cel =
+        CelFactory.standardCelBuilder()
+            .setStandardMacros(CelStandardMacro.STANDARD_MACROS)
+            .addCompilerLibraries(CelExtensions.bindings())
+            .setContainer(CelContainer.ofName("cel.expr.conformance.test"))
+            .addVar("unknown_attr", SimpleType.INT)
+            .addVar("my_list", ListType.create(SimpleType.INT))
+            .setOptions(CelOptions.current().enableUnknownTracking(true).build())
+            .build();
+
+    CelUnknownSet result =
+        (CelUnknownSet)
+            cel.createProgram(cel.compile(expression).getAst())
+                .advanceEvaluation(
+                    UnknownContext.create(
+                        (String name) -> {
+                          if (name.equals("my_list")) {
+                            return Optional.of(ImmutableList.of(1));
+                          }
+                          return Optional.empty();
+                        },
+                        ImmutableList.of(
+                            CelAttributePattern.create("unknown_attr"),
+                            CelAttributePattern.create("my_list")
+                                .qualify(CelAttribute.Qualifier.ofInt(0)))));
+
+    assertThat(result.attributes()).containsExactly(CelAttribute.create("unknown_attr"));
+  }
+
   @Test
   public void newWellKnownTypeMessage_withDifferentDescriptorInstance() throws Exception {
     CelCompiler celCompiler =