Clean up dangling source information from macro expansion

PiperOrigin-RevId: 629898444

Author: l46kok

common/src/main/java/dev/cel/common/CelSource.java

@@ -49,7 +49,7 @@ public final class CelSource {
   private CelSource(Builder builder) {
     this.codePoints = checkNotNull(builder.codePoints);
     this.description = checkNotNull(builder.description);
-    this.positions = checkNotNull(builder.positions.buildOrThrow());
+    this.positions = checkNotNull(ImmutableMap.copyOf(builder.positions));
     this.lineOffsets = checkNotNull(ImmutableList.copyOf(builder.lineOffsets));
     this.macroCalls = checkNotNull(ImmutableMap.copyOf(builder.macroCalls));
     this.extensions = checkNotNull(builder.extensions.build());
@@ -208,7 +208,7 @@ public static final class Builder {
 
     private final CelCodePointArray codePoints;
     private final List<Integer> lineOffsets;
-    private final ImmutableMap.Builder<Long, Integer> positions;
+    private final Map<Long, Integer> positions;
     private final Map<Long, CelExpr> macroCalls;
     private final ImmutableSet.Builder<Extension> extensions;
 
@@ -221,7 +221,7 @@ private Builder() {
     private Builder(CelCodePointArray codePoints, List<Integer> lineOffsets) {
       this.codePoints = checkNotNull(codePoints);
       this.lineOffsets = checkNotNull(lineOffsets);
-      this.positions = ImmutableMap.builder();
+      this.positions = new HashMap<>();
       this.macroCalls = new HashMap<>();
       this.extensions = ImmutableSet.builder();
       this.description = "";
@@ -261,6 +261,18 @@ public Builder addPositions(long exprId, int position) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    public Builder clearPositions() {
+      this.positions.clear();
+      return this;
+    }
+
+    @CanIgnoreReturnValue
+    public Builder removePositions(long exprId) {
+      this.positions.remove(exprId);
+      return this;
+    }
+
     @CanIgnoreReturnValue
     public Builder addMacroCalls(long exprId, CelExpr expr) {
       this.macroCalls.put(exprId, expr);
@@ -329,8 +341,8 @@ public Optional<CelSourceLocation> getOffsetLocation(int offset) {
     }
 
     @CheckReturnValue
-    public ImmutableMap<Long, Integer> getPositionsMap() {
-      return this.positions.buildOrThrow();
+    public Map<Long, Integer> getPositionsMap() {
+      return this.positions;
     }
 
     @CheckReturnValue

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

@@ -25,18 +25,12 @@
 /** Factory for generating expression nodes. */
 @Internal
 public class CelExprFactory {
-
-  private final CelExprIdGeneratorFactory.ExprIdGenerator idGenerator;
+  private long exprId = 0L;
 
   public static CelExprFactory newInstance() {
     return new CelExprFactory();
   }
 
-  public static CelExprFactory newInstance(
-      CelExprIdGeneratorFactory.ExprIdGenerator exprIdGenerator) {
-    return new CelExprFactory(exprIdGenerator);
-  }
-
   /** Create a new constant expression. */
   public final CelExpr newConstant(CelConstant constant) {
     return CelExpr.newBuilder().setId(nextExprId()).setConstant(constant).build();
@@ -545,19 +539,15 @@ public final CelExpr newSelect(CelExpr operand, String field, boolean testOnly)
 
   /** Returns the next unique expression ID. */
   protected long nextExprId() {
-    return idGenerator.generate(
-        /* exprId= */ -1); // Unconditionally generate next unique ID (i.e: no renumbering).
+    return ++exprId;
   }
 
-  protected CelExprFactory() {
-    this(CelExprIdGeneratorFactory.newMonotonicIdGenerator(0));
+  /** Attempts to decrement the next expr ID if possible. */
+  protected void maybeDeleteId(long id) {
+    if (id == exprId - 1) {
+      exprId--;
+    }
   }
 
-  private CelExprFactory(CelExprIdGeneratorFactory.MonotonicIdGenerator idGenerator) {
-    this((unused) -> idGenerator.nextExprId());
-  }
-
-  private CelExprFactory(CelExprIdGeneratorFactory.ExprIdGenerator exprIdGenerator) {
-    idGenerator = exprIdGenerator;
-  }
+  protected CelExprFactory() {}
 }

common/src/test/java/dev/cel/common/ast/CelExprFactoryTest.java

@@ -16,7 +16,6 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import dev.cel.common.ast.CelExprIdGeneratorFactory.StableIdGenerator;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -38,15 +37,4 @@ public void nextExprId_startingDefaultIsOne() {
     assertThat(exprFactory.nextExprId()).isEqualTo(1L);
     assertThat(exprFactory.nextExprId()).isEqualTo(2L);
   }
-
-  @Test
-  public void nextExprId_usingStableIdGenerator() {
-    StableIdGenerator stableIdGenerator = CelExprIdGeneratorFactory.newStableIdGenerator(0);
-    CelExprFactory exprFactory = CelExprFactory.newInstance(stableIdGenerator::nextExprId);
-
-    assertThat(exprFactory.nextExprId()).isEqualTo(1L);
-    assertThat(exprFactory.nextExprId()).isEqualTo(2L);
-    assertThat(stableIdGenerator.hasId(-1)).isFalse();
-    assertThat(stableIdGenerator.hasId(0)).isFalse();
-  }
 }

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

@@ -660,6 +660,7 @@ private Optional<CelExpr> visitMacro(
           CelExpr.newBuilder().setCall(callExpr.build()).build());
     }
 
+    exprFactory.maybeDeleteId(expr.id());
     return expandedMacro;
   }
 
@@ -895,6 +896,7 @@ private CelExpr macroOrCall(
     ImmutableList<CelExpr> arguments = visitExprListContext(args);
     Optional<CelExpr> errorArg = arguments.stream().filter(ERROR::equals).findAny();
     if (errorArg.isPresent()) {
+      sourceInfo.clearPositions();
       // Any arguments passed in to the macro may fail parsing.
       // Stop the macro expansion in this case as the result of the macro will be a parse failure.
       return ERROR;
@@ -1109,6 +1111,12 @@ private long nextExprId(int position) {
       return exprId;
     }
 
+    @Override
+    protected void maybeDeleteId(long id) {
+      sourceInfo.removePositions(id);
+      super.maybeDeleteId(id);
+    }
+
     @Override
     public long copyExprId(long id) {
       return nextExprId(getPosition(id));

parser/src/test/java/dev/cel/parser/CelParserParameterizedTest.java

@@ -32,6 +32,7 @@
 import com.google.protobuf.Descriptors.FieldDescriptor;
 import com.google.protobuf.Descriptors.OneofDescriptor;
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
+import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.CelOptions;
 import dev.cel.common.CelProtoAbstractSyntaxTree;
 import dev.cel.common.CelSource;
@@ -248,6 +249,11 @@ public void parser_errors() {
     runTest(parserWithoutOptionalSupport, "[?a, ?b]");
   }
 
+  @Test
+  public void source_info() throws Exception {
+    runSourceInfoTest("[{}, {'field': true}].exists(i, has(i.field))");
+  }
+
   private void runTest(CelParser parser, String expression) {
     runTest(parser, expression, true);
   }
@@ -287,6 +293,15 @@ private void runTest(CelParser parser, String expression, boolean validateParseO
     testOutput().println();
   }
 
+  private void runSourceInfoTest(String expression) throws Exception {
+    CelAbstractSyntaxTree ast = PARSER.parse(expression).getAst();
+    SourceInfo sourceInfo =
+        CelProtoAbstractSyntaxTree.fromCelAst(ast).toParsedExpr().getSourceInfo();
+    testOutput().println("I: " + expression);
+    testOutput().println("=====>");
+    testOutput().println("S: " + sourceInfo);
+  }
+
   private String convertMacroCallsToString(SourceInfo sourceInfo) {
     KindAndIdAdorner macroCallsAdorner = new KindAndIdAdorner(sourceInfo);
     // Sort in ascending order so that nested macro calls are always in the same order for tests

parser/src/test/resources/parser.baseline

@@ -1543,7 +1543,7 @@ L: noop_macro(
 I: get_constant_macro()
 =====>
 P: 10^#1:int64#
-L: 10^#1[1,18]#
+L: 10^#1[NO_POS]#
 M: get_constant_macro()^#0:Expr.Call#
 
 I: a.?b[?0] && a[?c]

parser/src/test/resources/source_info.baseline

@@ -0,0 +1,143 @@
+I: [{}, {'field': true}].exists(i, has(i.field))
+=====>
+S: location: "<input>"
+line_offsets: 46
+positions {
+  key: 1
+  value: 0
+}
+positions {
+  key: 2
+  value: 1
+}
+positions {
+  key: 3
+  value: 5
+}
+positions {
+  key: 4
+  value: 13
+}
+positions {
+  key: 5
+  value: 6
+}
+positions {
+  key: 6
+  value: 15
+}
+positions {
+  key: 8
+  value: 29
+}
+positions {
+  key: 10
+  value: 36
+}
+positions {
+  key: 11
+  value: 37
+}
+positions {
+  key: 12
+  value: 35
+}
+positions {
+  key: 13
+  value: 28
+}
+positions {
+  key: 14
+  value: 28
+}
+positions {
+  key: 15
+  value: 28
+}
+positions {
+  key: 16
+  value: 28
+}
+positions {
+  key: 17
+  value: 28
+}
+positions {
+  key: 18
+  value: 28
+}
+positions {
+  key: 19
+  value: 28
+}
+positions {
+  key: 20
+  value: 28
+}
+macro_calls {
+  key: 12
+  value {
+    call_expr {
+      function: "has"
+      args {
+        id: 11
+        select_expr {
+          operand {
+            id: 10
+            ident_expr {
+              name: "i"
+            }
+          }
+          field: "field"
+        }
+      }
+    }
+  }
+}
+macro_calls {
+  key: 20
+  value {
+    call_expr {
+      target {
+        id: 1
+        list_expr {
+          elements {
+            id: 2
+            struct_expr {
+            }
+          }
+          elements {
+            id: 3
+            struct_expr {
+              entries {
+                id: 4
+                map_key {
+                  id: 5
+                  const_expr {
+                    string_value: "field"
+                  }
+                }
+                value {
+                  id: 6
+                  const_expr {
+                    bool_value: true
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+      function: "exists"
+      args {
+        id: 8
+        ident_expr {
+          name: "i"
+        }
+      }
+      args {
+        id: 12
+      }
+    }
+  }
+}