Add cel.@block test coverage for parsed-only mode

PiperOrigin-RevId: 899267021

Author: l46kok

extensions/src/main/java/dev/cel/extensions/CelComprehensionsExtensions.java

@@ -118,29 +118,18 @@ public void setRuntimeOptions(CelRuntimeBuilder runtimeBuilder) {
   @Override
   public void setRuntimeOptions(
       CelRuntimeBuilder runtimeBuilder, RuntimeEquality runtimeEquality, CelOptions celOptions) {
-    for (Function function : functions) {
-      for (CelOverloadDecl overload : function.functionDecl.overloads()) {
-        switch (overload.overloadId()) {
-          case MAP_INSERT_OVERLOAD_MAP_MAP:
-            runtimeBuilder.addFunctionBindings(
-                CelFunctionBinding.from(
-                    MAP_INSERT_OVERLOAD_MAP_MAP,
-                    Map.class,
-                    Map.class,
-                    (map1, map2) -> mapInsertMap(map1, map2, runtimeEquality)));
-            break;
-          case MAP_INSERT_OVERLOAD_KEY_VALUE:
-            runtimeBuilder.addFunctionBindings(
-                CelFunctionBinding.from(
-                    MAP_INSERT_OVERLOAD_KEY_VALUE,
-                    ImmutableList.of(Map.class, Object.class, Object.class),
-                    args -> mapInsertKeyValue(args, runtimeEquality)));
-            break;
-          default:
-            // Nothing to add.
-        }
-      }
-    }
+    runtimeBuilder.addFunctionBindings(
+        CelFunctionBinding.fromOverloads(
+            MAP_INSERT_FUNCTION,
+            CelFunctionBinding.from(
+                MAP_INSERT_OVERLOAD_MAP_MAP,
+                Map.class,
+                Map.class,
+                (map1, map2) -> mapInsertMap(map1, map2, runtimeEquality)),
+            CelFunctionBinding.from(
+                MAP_INSERT_OVERLOAD_KEY_VALUE,
+                ImmutableList.of(Map.class, Object.class, Object.class),
+                args -> mapInsertKeyValue(args, runtimeEquality))));
   }
 
   @Override

optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerBaselineTest.java

@@ -83,8 +83,13 @@ private static Cel setupCelEnv(CelBuilder celBuilder) {
           .addFunctionBindings(
               // This is pure, but for the purposes of excluding it as a CSE candidate, pretend that
               // it isn't.
-              CelFunctionBinding.from("non_pure_custom_func_overload", Long.class, val -> val),
-              CelFunctionBinding.from("pure_custom_func_overload", Long.class, val -> val))
+              CelFunctionBinding.fromOverloads(
+                  "non_pure_custom_func",
+                  CelFunctionBinding.from("non_pure_custom_func_overload", Long.class, val -> val)))
+          .addFunctionBindings(
+              CelFunctionBinding.fromOverloads(
+                  "pure_custom_func",
+                  CelFunctionBinding.from("pure_custom_func_overload", Long.class, val -> val)))
           .addVar("x", SimpleType.DYN)
           .addVar("y", SimpleType.DYN)
           .addVar("opt_x", OptionalType.create(SimpleType.DYN))
@@ -153,6 +158,28 @@ public void allOptimizers_producesSameEvaluationResult(
     assertThat(optimizedEvalResult).isEqualTo(expectedEvalResult);
   }
 
+  @Test
+  public void allOptimizers_producesSameEvaluationResult_parsedOnly(
+      @TestParameter CseTestCase cseTestCase, @TestParameter CseTestOptimizer cseTestOptimizer)
+      throws Exception {
+    skipBaselineVerification();
+    if (runtimeEnv == RuntimeEnv.LEGACY) {
+      return;
+    }
+    CelAbstractSyntaxTree ast = runtimeEnv.cel.compile(cseTestCase.source).getAst();
+    ImmutableMap<String, Object> inputMap =
+        ImmutableMap.of("msg", TEST_ALL_TYPES_INPUT, "x", 5L, "y", 6L, "opt_x", Optional.of(5L));
+    Object expectedEvalResult = runtimeEnv.cel.createProgram(ast).eval(inputMap);
+
+    CelAbstractSyntaxTree optimizedAst = cseTestOptimizer.newCseOptimizer(runtimeEnv).optimize(ast);
+    CelAbstractSyntaxTree parsedOnlyOptimizedAst =
+        CelAbstractSyntaxTree.newParsedAst(optimizedAst.getExpr(), optimizedAst.getSource());
+
+    Object optimizedEvalResult =
+        runtimeEnv.cel.createProgram(parsedOnlyOptimizedAst).eval(inputMap);
+    assertThat(optimizedEvalResult).isEqualTo(expectedEvalResult);
+  }
+
   @Test
   public void subexpression_unparsed() throws Exception {
     for (CseTestCase cseTestCase : EnumSet.allOf(CseTestCase.class)) {

optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerTest.java

@@ -416,6 +416,19 @@ public void block_success(@TestParameter BlockTestCase testCase) throws Exceptio
     assertThat(evaluatedResult).isNotNull();
   }
 
+  @Test
+  public void block_success_parsedOnly(@TestParameter BlockTestCase testCase) throws Exception {
+    if (runtimeEnv == RuntimeEnv.LEGACY) {
+      return;
+    }
+    CelAbstractSyntaxTree ast =
+        compileUsingInternalFunctions(testCase.source, /* parsedOnly= */ true);
+
+    Object evaluatedResult = runtimeEnv.celForEvaluatingBlock.createProgram(ast).eval();
+
+    assertThat(evaluatedResult).isNotNull();
+  }
+
   @Test
   @SuppressWarnings("Immutable") // Test only
   public void lazyEval_blockIndexNeverReferenced() throws Exception {
@@ -694,7 +707,7 @@ public void block_lazyEvaluationContainsError_cleansUpCycleState() throws Except
    * Converts AST containing cel.block related test functions to internal functions (e.g: cel.block
    * -> cel.@block)
    */
-  private CelAbstractSyntaxTree compileUsingInternalFunctions(String expression)
+  private CelAbstractSyntaxTree compileUsingInternalFunctions(String expression, boolean parsedOnly)
       throws CelValidationException {
     CelAbstractSyntaxTree astToModify =
         runtimeEnv.celForEvaluatingBlock.compile(expression).getAst();
@@ -719,6 +732,14 @@ private CelAbstractSyntaxTree compileUsingInternalFunctions(String expression)
               indexExpr.ident().setName(internalIdentName);
             });
 
+    if (parsedOnly) {
+      return mutableAst.toParsedAst();
+    }
     return runtimeEnv.celForEvaluatingBlock.check(mutableAst.toParsedAst()).getAst();
   }
+
+  private CelAbstractSyntaxTree compileUsingInternalFunctions(String expression)
+      throws CelValidationException {
+    return compileUsingInternalFunctions(expression, false);
+  }
 }