AST and checker changes to support comprehensionsV2

PiperOrigin-RevId: 786471983

Author: maskri17

checker/src/main/java/dev/cel/checker/ExprChecker.java

@@ -22,6 +22,7 @@
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Joiner;
 import com.google.common.base.Optional;
+import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Maps;
 import com.google.errorprone.annotations.CheckReturnValue;
@@ -510,13 +511,21 @@ private CelExpr visit(CelExpr expr, CelExpr.CelComprehension compre) {
     CelType accuType = env.getType(visitedInit);
     CelType rangeType = inferenceContext.specialize(env.getType(visitedRange));
     CelType varType;
+    CelType varType2 = null;
     switch (rangeType.kind()) {
       case LIST:
         varType = ((ListType) rangeType).elemType();
+        if (!Strings.isNullOrEmpty(compre.iterVar2())) {
+          varType2 = varType;
+          varType = SimpleType.INT;
+        }
         break;
       case MAP:
         // Ranges over the keys.
         varType = ((MapType) rangeType).keyType();
+        if (!Strings.isNullOrEmpty(compre.iterVar2())) {
+          varType2 = ((MapType) rangeType).valueType();
+        }
         break;
       case DYN:
       case ERROR:
@@ -547,6 +556,9 @@ private CelExpr visit(CelExpr expr, CelExpr.CelComprehension compre) {
     // Declare iteration variable on inner scope.
     env.enterScope();
     env.add(CelIdentDecl.newIdentDeclaration(compre.iterVar(), varType));
+    if (!Strings.isNullOrEmpty(compre.iterVar2())) {
+      env.add(CelIdentDecl.newIdentDeclaration(compre.iterVar2(), varType2));
+    }
     CelExpr condition = visit(compre.loopCondition());
     assertType(condition, SimpleType.BOOL);
     CelExpr visitedStep = visit(compre.loopStep());

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

@@ -52,7 +52,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",
         "@maven_android//:com_google_protobuf_protobuf_javalite",
     ],
 )
@@ -139,7 +138,6 @@ cel_android_library(
         "//:auto_value",
         "//common/annotations",
         "@maven//:com_google_errorprone_error_prone_annotations",
-        "@maven//:org_jspecify_jspecify",
         "@maven_android//:com_google_guava_guava",
         "@maven_android//:com_google_protobuf_protobuf_javalite",
     ],

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

@@ -23,7 +23,6 @@
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.CheckReturnValue;
 import com.google.errorprone.annotations.Immutable;
-import dev.cel.common.ast.CelExpr.ExprKind.Kind;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Optional;
@@ -902,6 +901,9 @@ public abstract static class CelComprehension implements Expression.Comprehensio
     @Override
     public abstract String iterVar();
 
+    @Override
+    public abstract String iterVar2();
+
     @Override
     public abstract CelExpr iterRange();
 
@@ -937,6 +939,8 @@ public abstract static class Builder {
 
       public abstract Builder setIterVar(String value);
 
+      public abstract Builder setIterVar2(String value);
+
       public abstract Builder setIterRange(CelExpr value);
 
       public abstract Builder setAccuVar(String value);
@@ -958,6 +962,7 @@ public abstract static class Builder {
     public static Builder newBuilder() {
       return new AutoValue_CelExpr_CelComprehension.Builder()
           .setIterVar("")
+          .setIterVar2("")
           .setIterRange(CelExpr.newBuilder().build())
           .setAccuVar("")
           .setAccuInit(CelExpr.newBuilder().build())
@@ -1073,12 +1078,28 @@ public static CelExpr ofComprehension(
       CelExpr loopCondition,
       CelExpr loopStep,
       CelExpr result) {
+    return ofComprehension(
+        id, iterVar, "", iterRange, accuVar, accuInit, loopCondition, loopStep, result);
+  }
+
+  @SuppressWarnings("TooManyParameters")
+  public static CelExpr ofComprehension(
+      long id,
+      String iterVar,
+      String iterVar2,
+      CelExpr iterRange,
+      String accuVar,
+      CelExpr accuInit,
+      CelExpr loopCondition,
+      CelExpr loopStep,
+      CelExpr result) {
     return newBuilder()
         .setId(id)
         .setExprKind(
             AutoOneOf_CelExpr_ExprKind.comprehension(
                 CelComprehension.newBuilder()
                     .setIterVar(iterVar)
+                    .setIterVar2(iterVar2)
                     .setIterRange(iterRange)
                     .setAccuVar(accuVar)
                     .setAccuInit(accuInit)

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

@@ -84,6 +84,7 @@ public static Expr fromCelExpr(CelExpr celExpr) {
         return expr.setComprehensionExpr(
                 Comprehension.newBuilder()
                     .setIterVar(celComprehension.iterVar())
+                    .setIterVar2(celComprehension.iterVar2())
                     .setIterRange(fromCelExpr(celComprehension.iterRange()))
                     .setAccuVar(celComprehension.accuVar())
                     .setAccuInit(fromCelExpr(celComprehension.accuInit()))
@@ -132,6 +133,7 @@ public static CelExpr fromExpr(Expr expr) {
         return CelExpr.ofComprehension(
             expr.getId(),
             comprehensionExpr.getIterVar(),
+            comprehensionExpr.getIterVar2(),
             fromExpr(comprehensionExpr.getIterRange()),
             comprehensionExpr.getAccuVar(),
             fromExpr(comprehensionExpr.getAccuInit()),