Checker and parser changes to support comprehensionsV2

PiperOrigin-RevId: 788649636

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/CelExprFactory.java

@@ -143,7 +143,7 @@ public final CelExpr.CelStruct.Entry newMessageField(String field, CelExpr value
         .build();
   }
 
-  /** Fold creates a fold comprehension instruction. */
+  /** Fold creates a fold for one variable comprehension instruction. */
   public final CelExpr fold(
       String iterVar,
       CelExpr iterRange,

parser/src/main/java/dev/cel/parser/CelMacroExprFactory.java

@@ -133,6 +133,7 @@ public final CelExpr copy(CelExpr expr) {
         builder.setComprehension(
             CelExpr.CelComprehension.newBuilder()
                 .setIterVar(expr.comprehension().iterVar())
+                .setIterVar2(expr.comprehension().iterVar2())
                 .setIterRange(copy(expr.comprehension().iterRange()))
                 .setAccuVar(expr.comprehension().accuVar())
                 .setAccuInit(copy(expr.comprehension().accuInit()))