Validate cel.@block expressions in a single visitor pass

PiperOrigin-RevId: 996873631

Author: l46kok

common/ast/BUILD.bazel

@@ -48,6 +48,11 @@ java_library(
     exports = ["//common/src/main/java/dev/cel/common/ast:cel_expr_visitor"],
 )
 
+cel_android_library(
+    name = "cel_expr_visitor_android",
+    exports = ["//common/src/main/java/dev/cel/common/ast:cel_expr_visitor_android"],
+)
+
 java_library(
     name = "expr_factory",
     exports = ["//common/src/main/java/dev/cel/common/ast:expr_factory"],

common/src/main/java/dev/cel/common/ast/BUILD.bazel

@@ -7,6 +7,7 @@ package(
     ],
     default_visibility = [
         "//common/ast:__pkg__",
+        "//publish:__pkg__",
     ],
 )
 
@@ -64,9 +65,9 @@ java_library(
     ],
     deps = [
         ":ast",
+        ":cel_expr_visitor",
         "//common:cel_ast",
         "//common/annotations",
-        "//common/navigation",
         "@maven//:com_google_guava_guava",
     ],
 )
@@ -78,9 +79,9 @@ cel_android_library(
     ],
     deps = [
         ":ast_android",
+        ":cel_expr_visitor_android",
         "//common:cel_ast_android",
         "//common/annotations",
-        "//common/navigation:navigation_android",
         "@maven_android//:com_google_guava_guava",
     ],
 )
@@ -143,6 +144,17 @@ java_library(
     ],
 )
 
+cel_android_library(
+    name = "cel_expr_visitor_android",
+    srcs = ["CelExprVisitor.java"],
+    tags = [
+    ],
+    deps = [
+        ":ast_android",
+        "//common:cel_ast_android",
+    ],
+)
+
 java_library(
     name = "expr_factory",
     srcs = EXPR_FACTORY_SOURCES,

common/src/main/java/dev/cel/common/ast/CelBlock.java

@@ -14,13 +14,13 @@
 
 package dev.cel.common.ast;
 
-import static com.google.common.collect.ImmutableList.toImmutableList;
-
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.annotations.Internal;
-import dev.cel.common.navigation.CelNavigableExpr;
+import dev.cel.common.ast.CelExpr.CelCall;
+import dev.cel.common.ast.CelExpr.CelIdent;
+import dev.cel.common.ast.CelExpr.ExprKind.Kind;
 import java.util.Optional;
 
 /**
@@ -61,27 +61,17 @@ public CelExpr expr() {
    * @throws IllegalArgumentException if the block is malformed or its indices are invalid.
    */
   public static Optional<CelBlock> extract(CelAbstractSyntaxTree ast) {
-    CelNavigableExpr celNavigableExpr = CelNavigableExpr.fromExpr(ast.getExpr());
-
-    ImmutableList<CelExpr> allCelBlocks =
-        celNavigableExpr
-            .allNodes()
-            .map(CelNavigableExpr::expr)
-            .filter(expr -> expr.callOrDefault().function().equals(FUNCTION_NAME))
-            .collect(toImmutableList());
-    if (allCelBlocks.isEmpty()) {
+    CelExpr root = ast.getExpr();
+    if (!isBlockCall(root)) {
+      new BlockValidator(root).validate(root, 0);
       return Optional.empty();
     }
 
-    Preconditions.checkArgument(
-        allCelBlocks.size() == 1,
-        "Expected 1 cel.block function to be present but found %s",
-        allCelBlocks.size());
-    Preconditions.checkArgument(
-        celNavigableExpr.expr().equals(allCelBlocks.get(0)),
-        "Expected cel.block to be present at root");
+    return Optional.of(fromExpr(root));
+  }
 
-    return Optional.of(fromExpr(allCelBlocks.get(0)));
+  private static boolean isBlockCall(CelExpr expr) {
+    return expr.getKind().equals(Kind.CALL) && expr.call().function().equals(FUNCTION_NAME);
   }
 
   /**
@@ -100,45 +90,71 @@ private static CelBlock fromExpr(CelExpr expr) {
     Preconditions.checkArgument(
         expr.call().args().get(0).exprKind().getKind() == CelExpr.ExprKind.Kind.LIST,
         "Expected first argument of cel.@block to be a list");
+    Preconditions.checkArgument(
+        !expr.call().target().isPresent(), "Expected cel.@block to be a global call");
 
     CelBlock block = new CelBlock(expr);
+    BlockValidator validator = new BlockValidator(expr);
 
     // Assert correctness on block indices used in subexpressions
     ImmutableList<CelExpr> subexprs = block.indices();
     for (int i = 0; i < subexprs.size(); i++) {
-      verifyBlockIndex(subexprs.get(i), i, expr);
+      validator.validate(subexprs.get(i), i);
     }
 
     // Assert correctness on block indices used in block result
-    CelExpr blockResult = block.result();
-    verifyBlockIndex(blockResult, subexprs.size(), expr);
-    boolean resultHasAtLeastOneBlockIndex =
-        CelNavigableExpr.fromExpr(blockResult)
-            .allNodes()
-            .map(CelNavigableExpr::expr)
-            .anyMatch(e -> e.identOrDefault().name().startsWith(INDEX_PREFIX));
     Preconditions.checkArgument(
-        resultHasAtLeastOneBlockIndex,
+        validator.validate(block.result(), subexprs.size()),
         "Expected at least one reference of index in cel.block result");
 
     return block;
   }
 
-  private static void verifyBlockIndex(CelExpr celExpr, int maxIndexValue, CelExpr rootBlock) {
-    boolean areAllIndicesValid =
-        CelNavigableExpr.fromExpr(celExpr)
-            .allNodes()
-            .map(CelNavigableExpr::expr)
-            .filter(expr -> expr.identOrDefault().name().startsWith(INDEX_PREFIX))
-            .map(CelExpr::ident)
-            .allMatch(
-                blockIdent ->
-                    Integer.parseInt(blockIdent.name().substring(INDEX_PREFIX.length()))
-                        < maxIndexValue);
-    Preconditions.checkArgument(
-        areAllIndicesValid,
-        "Illegal block index found. The index value must be less than %s. Expr: %s",
-        maxIndexValue,
-        rootBlock);
+  private static final class BlockValidator extends CelExprVisitor {
+    private final CelExpr root;
+    private int maxIndexValue;
+    private boolean hasBlockIndex;
+
+    private boolean validate(CelExpr expr, int maxIndexValue) {
+      this.maxIndexValue = maxIndexValue;
+      this.hasBlockIndex = false;
+      visit(expr);
+      return hasBlockIndex;
+    }
+
+    @Override
+    public void visit(CelExpr expr) {
+      if (!expr.getKind().equals(Kind.NOT_SET)) {
+        super.visit(expr);
+      }
+    }
+
+    @Override
+    protected void visit(CelExpr expr, CelCall call) {
+      if (call.function().equals(FUNCTION_NAME)) {
+        throw new IllegalArgumentException(
+            isBlockCall(root)
+                ? "Expected 1 cel.block function to be present but found 2"
+                : "Expected cel.block to be present at root");
+      }
+      super.visit(expr, call);
+    }
+
+    @Override
+    protected void visit(CelExpr expr, CelIdent ident) {
+      if (ident.name().startsWith(INDEX_PREFIX)) {
+        hasBlockIndex = true;
+        int indexValue = Integer.parseInt(ident.name().substring(INDEX_PREFIX.length()));
+        Preconditions.checkArgument(
+            indexValue >= 0 && indexValue < maxIndexValue,
+            "Illegal block index found. The index value must be less than %s. Expr: %s",
+            maxIndexValue,
+            root);
+      }
+    }
+
+    private BlockValidator(CelExpr root) {
+      this.root = root;
+    }
   }
 }

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

@@ -16,6 +16,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static dev.cel.common.CelOverloadDecl.newGlobalOverload;
+import static dev.cel.optimizer.optimizers.SubexpressionOptimizer.verifyOptimizedAstCorrectness;
 import static org.junit.Assert.assertThrows;
 
 import com.google.common.collect.ImmutableList;
@@ -641,6 +642,32 @@ public void verifyOptimizedAstCorrectness_indexOutOfBounds_throws(String source)
         .contains("Illegal block index found. The index value must be less than");
   }
 
+  @Test
+  public void verifyOptimizedAstCorrectness_negativeIndex_throws() throws Exception {
+    CelMutableAst mutableAst =
+        CelMutableAst.fromCelAst(compileUsingInternalFunctions("cel.block([1, 2], index0)", true));
+    mutableAst.expr().call().args().get(1).ident().setName("@index-1");
+    CelAbstractSyntaxTree ast = mutableAst.toParsedAst();
+
+    IllegalArgumentException e =
+        assertThrows(IllegalArgumentException.class, () -> verifyOptimizedAstCorrectness(ast));
+    assertThat(e)
+        .hasMessageThat()
+        .contains("Illegal block index found. The index value must be less than");
+  }
+
+  @Test
+  public void verifyOptimizedAstCorrectness_celBlockHasTarget_throws() throws Exception {
+    CelMutableAst mutableAst =
+        CelMutableAst.fromCelAst(compileUsingInternalFunctions("cel.block([1, 2], index0)", true));
+    mutableAst.expr().call().setTarget(mutableAst.expr().call().args().get(1));
+    CelAbstractSyntaxTree ast = mutableAst.toParsedAst();
+
+    IllegalArgumentException e =
+        assertThrows(IllegalArgumentException.class, () -> verifyOptimizedAstCorrectness(ast));
+    assertThat(e).hasMessageThat().isEqualTo("Expected cel.@block to be a global call");
+  }
+
   @Test
   @TestParameters("{source: 'cel.block([index0], index0)'}")
   @TestParameters("{source: 'cel.block([1, index1, 2], index2)'}")
@@ -678,9 +705,9 @@ public void block_lazyEvaluationContainsError_cleansUpCycleState() throws Except
     CelAbstractSyntaxTree ast =
         compileUsingInternalFunctions(
             "cel.block([1/0 > 0], (index0 && false) || (index0 && true))");
+    Program program = cel.createProgram(ast);
 
-    CelEvaluationException e =
-        assertThrows(CelEvaluationException.class, () -> cel.createProgram(ast).eval());
+    CelEvaluationException e = assertThrows(CelEvaluationException.class, program::eval);
 
     assertThat(e).hasMessageThat().contains("/ by zero");
     assertThat(e).hasMessageThat().doesNotContain("Cycle detected");