Add option to enable updated accumulator variable.

jnthntatum · copybara-github · commit ece46c4e7dd6 · 2024-12-19T11:37:14.000-08:00
PiperOrigin-RevId: 707977249
diff --git a/common/src/main/java/dev/cel/common/CelOptions.java b/common/src/main/java/dev/cel/common/CelOptions.java
@@ -67,6 +67,8 @@ public enum ProtoUnsetFieldOptions {
 
   public abstract boolean retainUnbalancedLogicalExpressions();
 
+  public abstract boolean enableHiddenAccumulatorVar();
+
   // Type-Checker related options
 
   public abstract boolean enableCompileTimeOverloadResolution();
@@ -188,6 +190,7 @@ public static Builder newBuilder() {
         .populateMacroCalls(false)
         .retainRepeatedUnaryOperators(false)
         .retainUnbalancedLogicalExpressions(false)
+        .enableHiddenAccumulatorVar(false)
         // Type-Checker options
         .enableCompileTimeOverloadResolution(false)
         .enableHomogeneousLiterals(false)
@@ -319,6 +322,16 @@ public abstract static class Builder {
      */
     public abstract Builder retainUnbalancedLogicalExpressions(boolean value);
 
+    /**
+     * Enable the use of a hidden accumulator variable name.
+     *
+     * <p>This is a temporary option to transition to using an internal identifier for the
+     * accumulator variable used by builtin comprehension macros. When enabled, parses result in a
+     * semantically equivalent AST, but with a different accumulator variable that can't be directly
+     * referenced in the source expression.
+     */
+    public abstract Builder enableHiddenAccumulatorVar(boolean value);
+
     // Type-Checker related options
 
     /**
diff --git a/parser/src/main/java/dev/cel/parser/CelMacroExprFactory.java b/parser/src/main/java/dev/cel/parser/CelMacroExprFactory.java
@@ -51,6 +51,9 @@ public final CelExpr reportError(String message) {
   /** Reports a {@link CelIssue} and returns a sentinel {@link CelExpr} that indicates an error. */
   public abstract CelExpr reportError(CelIssue error);
 
+  /** Returns the default accumulator variable name used by macros implementing comprehensions. */
+  public abstract String getAccumulatorVarName();
+
   /** Retrieves the source location for the given {@link CelExpr} ID. */
   public final CelSourceLocation getSourceLocation(CelExpr expr) {
     return getSourceLocation(expr.id());
diff --git a/parser/src/main/java/dev/cel/parser/CelStandardMacro.java b/parser/src/main/java/dev/cel/parser/CelStandardMacro.java
@@ -78,8 +78,6 @@ public enum CelStandardMacro {
   public static final ImmutableSet<CelStandardMacro> STANDARD_MACROS =
       ImmutableSet.of(HAS, ALL, EXISTS, EXISTS_ONE, MAP, MAP_FILTER, FILTER);
 
-  private static final String ACCUMULATOR_VAR = "__result__";
-
   private final CelMacro macro;
 
   CelStandardMacro(CelMacro macro) {
@@ -123,14 +121,23 @@ private static Optional<CelExpr> expandAllMacro(
     CelExpr accuInit = exprFactory.newBoolLiteral(true);
     CelExpr condition =
         exprFactory.newGlobalCall(
-            Operator.NOT_STRICTLY_FALSE.getFunction(), exprFactory.newIdentifier(ACCUMULATOR_VAR));
+            Operator.NOT_STRICTLY_FALSE.getFunction(),
+            exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()));
     CelExpr step =
         exprFactory.newGlobalCall(
-            Operator.LOGICAL_AND.getFunction(), exprFactory.newIdentifier(ACCUMULATOR_VAR), arg1);
-    CelExpr result = exprFactory.newIdentifier(ACCUMULATOR_VAR);
+            Operator.LOGICAL_AND.getFunction(),
+            exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()),
+            arg1);
+    CelExpr result = exprFactory.newIdentifier(exprFactory.getAccumulatorVarName());
     return Optional.of(
         exprFactory.fold(
-            arg0.ident().name(), target, ACCUMULATOR_VAR, accuInit, condition, step, result));
+            arg0.ident().name(),
+            target,
+            exprFactory.getAccumulatorVarName(),
+            accuInit,
+            condition,
+            step,
+            result));
   }
 
   // CelMacroExpander implementation for CEL's exists() macro.
@@ -149,14 +156,23 @@ private static Optional<CelExpr> expandExistsMacro(
         exprFactory.newGlobalCall(
             Operator.NOT_STRICTLY_FALSE.getFunction(),
             exprFactory.newGlobalCall(
-                Operator.LOGICAL_NOT.getFunction(), exprFactory.newIdentifier(ACCUMULATOR_VAR)));
+                Operator.LOGICAL_NOT.getFunction(),
+                exprFactory.newIdentifier(exprFactory.getAccumulatorVarName())));
     CelExpr step =
         exprFactory.newGlobalCall(
-            Operator.LOGICAL_OR.getFunction(), exprFactory.newIdentifier(ACCUMULATOR_VAR), arg1);
-    CelExpr result = exprFactory.newIdentifier(ACCUMULATOR_VAR);
+            Operator.LOGICAL_OR.getFunction(),
+            exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()),
+            arg1);
+    CelExpr result = exprFactory.newIdentifier(exprFactory.getAccumulatorVarName());
     return Optional.of(
         exprFactory.fold(
-            arg0.ident().name(), target, ACCUMULATOR_VAR, accuInit, condition, step, result));
+            arg0.ident().name(),
+            target,
+            exprFactory.getAccumulatorVarName(),
+            accuInit,
+            condition,
+            step,
+            result));
   }
 
   // CelMacroExpander implementation for CEL's exists_one() macro.
@@ -178,17 +194,23 @@ private static Optional<CelExpr> expandExistsOneMacro(
             arg1,
             exprFactory.newGlobalCall(
                 Operator.ADD.getFunction(),
-                exprFactory.newIdentifier(ACCUMULATOR_VAR),
+                exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()),
                 exprFactory.newIntLiteral(1)),
-            exprFactory.newIdentifier(ACCUMULATOR_VAR));
+            exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()));
     CelExpr result =
         exprFactory.newGlobalCall(
             Operator.EQUALS.getFunction(),
-            exprFactory.newIdentifier(ACCUMULATOR_VAR),
+            exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()),
             exprFactory.newIntLiteral(1));
     return Optional.of(
         exprFactory.fold(
-            arg0.ident().name(), target, ACCUMULATOR_VAR, accuInit, condition, step, result));
+            arg0.ident().name(),
+            target,
+            exprFactory.getAccumulatorVarName(),
+            accuInit,
+            condition,
+            step,
+            result));
   }
 
   // CelMacroExpander implementation for CEL's map() macro.
@@ -218,25 +240,25 @@ private static Optional<CelExpr> expandMapMacro(
     CelExpr step =
         exprFactory.newGlobalCall(
             Operator.ADD.getFunction(),
-            exprFactory.newIdentifier(ACCUMULATOR_VAR),
+            exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()),
             exprFactory.newList(arg1));
     if (arg2 != null) {
       step =
           exprFactory.newGlobalCall(
               Operator.CONDITIONAL.getFunction(),
               arg2,
               step,
-              exprFactory.newIdentifier(ACCUMULATOR_VAR));
+              exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()));
     }
     return Optional.of(
         exprFactory.fold(
             arg0.ident().name(),
             target,
-            ACCUMULATOR_VAR,
+            exprFactory.getAccumulatorVarName(),
             accuInit,
             condition,
             step,
-            exprFactory.newIdentifier(ACCUMULATOR_VAR)));
+            exprFactory.newIdentifier(exprFactory.getAccumulatorVarName())));
   }
 
   // CelMacroExpander implementation for CEL's filter() macro.
@@ -255,23 +277,23 @@ private static Optional<CelExpr> expandFilterMacro(
     CelExpr step =
         exprFactory.newGlobalCall(
             Operator.ADD.getFunction(),
-            exprFactory.newIdentifier(ACCUMULATOR_VAR),
+            exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()),
             exprFactory.newList(arg0));
     step =
         exprFactory.newGlobalCall(
             Operator.CONDITIONAL.getFunction(),
             arg1,
             step,
-            exprFactory.newIdentifier(ACCUMULATOR_VAR));
+            exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()));
     return Optional.of(
         exprFactory.fold(
             arg0.ident().name(),
             target,
-            ACCUMULATOR_VAR,
+            exprFactory.getAccumulatorVarName(),
             accuInit,
             condition,
             step,
-            exprFactory.newIdentifier(ACCUMULATOR_VAR)));
+            exprFactory.newIdentifier(exprFactory.getAccumulatorVarName())));
   }
 
   private static CelExpr reportArgumentError(CelMacroExprFactory exprFactory, CelExpr argument) {
diff --git a/parser/src/main/java/dev/cel/parser/Parser.java b/parser/src/main/java/dev/cel/parser/Parser.java
@@ -125,6 +125,8 @@ final class Parser extends CELBaseVisitor<CelExpr> {
           "var",
           "void",
           "while");
+  private static final String ACCUMULATOR_NAME = "__result__";
+  private static final String HIDDEN_ACCUMULATOR_NAME = "@result";
 
   static CelValidationResult parse(CelParserImpl parser, CelSource source, CelOptions options) {
     if (source.getContent().size() > options.maxExpressionCodePointSize()) {
@@ -142,7 +144,11 @@ static CelValidationResult parse(CelParserImpl parser, CelSource source, CelOpti
     CELParser antlrParser = new CELParser(new CommonTokenStream(antlrLexer));
     CelSource.Builder sourceInfo = source.toBuilder();
     sourceInfo.setDescription(source.getDescription());
-    ExprFactory exprFactory = new ExprFactory(antlrParser, sourceInfo);
+    ExprFactory exprFactory =
+        new ExprFactory(
+            antlrParser,
+            sourceInfo,
+            options.enableHiddenAccumulatorVar() ? HIDDEN_ACCUMULATOR_NAME : ACCUMULATOR_NAME);
     Parser parserImpl = new Parser(parser, options, sourceInfo, exprFactory);
     ErrorListener errorListener = new ErrorListener(exprFactory);
     antlrLexer.removeErrorListeners();
@@ -1033,12 +1039,17 @@ private static final class ExprFactory extends CelMacroExprFactory {
     private final CelSource.Builder sourceInfo;
     private final ArrayList<CelIssue> issues;
     private final ArrayDeque<Integer> positions;
+    private final String accumulatorVarName;
 
-    private ExprFactory(org.antlr.v4.runtime.Parser recognizer, CelSource.Builder sourceInfo) {
+    private ExprFactory(
+        org.antlr.v4.runtime.Parser recognizer,
+        CelSource.Builder sourceInfo,
+        String accumulatorVarName) {
       this.recognizer = recognizer;
       this.sourceInfo = sourceInfo;
       this.issues = new ArrayList<>();
       this.positions = new ArrayDeque<>(1); // Currently this usually contains at most 1 position.
+      this.accumulatorVarName = accumulatorVarName;
     }
 
     // Implementation of CelExprFactory.
@@ -1062,6 +1073,11 @@ public CelExpr reportError(CelIssue error) {
       return ERROR;
     }
 
+    @Override
+    public String getAccumulatorVarName() {
+      return accumulatorVarName;
+    }
+
     // Internal methods used by the parser but not part of the public API.
     @FormatMethod
     @CanIgnoreReturnValue
diff --git a/parser/src/test/java/dev/cel/parser/CelMacroExprFactoryTest.java b/parser/src/test/java/dev/cel/parser/CelMacroExprFactoryTest.java
@@ -61,6 +61,11 @@ public CelExpr reportError(CelIssue issue) {
       return CelExpr.newBuilder().setId(nextExprId()).setConstant(Constants.ERROR).build();
     }
 
+    @Override
+    public String getAccumulatorVarName() {
+      return "__result__";
+    }
+
     @Override
     protected CelSourceLocation getSourceLocation(long exprId) {
       return CelSourceLocation.NONE;
diff --git a/parser/src/test/java/dev/cel/parser/CelParserParameterizedTest.java b/parser/src/test/java/dev/cel/parser/CelParserParameterizedTest.java
@@ -72,6 +72,16 @@ public final class CelParserParameterizedTest extends BaselineTestCase {
           .setOptions(CelOptions.current().populateMacroCalls(true).build())
           .build();
 
+  private static final CelParser PARSER_WITH_UPDATED_ACCU_VAR =
+      PARSER
+          .toParserBuilder()
+          .setOptions(
+              CelOptions.current()
+                  .populateMacroCalls(true)
+                  .enableHiddenAccumulatorVar(true)
+                  .build())
+          .build();
+
   @Test
   public void parser() {
     runTest(PARSER, "x * 2");
@@ -193,6 +203,17 @@ public void parser() {
         "while");
   }
 
+  @Test
+  public void parser_updatedAccuVar() {
+    runTest(PARSER_WITH_UPDATED_ACCU_VAR, "x * 2");
+    runTest(PARSER_WITH_UPDATED_ACCU_VAR, "has(m.f)");
+    runTest(PARSER_WITH_UPDATED_ACCU_VAR, "m.exists_one(v, f)");
+    runTest(PARSER_WITH_UPDATED_ACCU_VAR, "m.all(v, f)");
+    runTest(PARSER_WITH_UPDATED_ACCU_VAR, "m.map(v, f)");
+    runTest(PARSER_WITH_UPDATED_ACCU_VAR, "m.map(v, p, f)");
+    runTest(PARSER_WITH_UPDATED_ACCU_VAR, "m.filter(v, p)");
+  }
+
   @Test
   public void parser_errors() {
     runTest(PARSER, "*@a | b");
diff --git a/parser/src/test/resources/parser_updatedAccuVar.baseline b/parser/src/test/resources/parser_updatedAccuVar.baseline
