Retain source description post AST optimization

PiperOrigin-RevId: 631847830

Author: l46kok

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

@@ -35,8 +35,9 @@
  */
 public final class CelMutableSource {
 
-  final Map<Long, CelMutableExpr> macroCalls;
-  final Set<Extension> extensions;
+  private String description;
+  private final Map<Long, CelMutableExpr> macroCalls;
+  private final Set<Extension> extensions;
 
   @CanIgnoreReturnValue
   public CelMutableSource addMacroCalls(long exprId, CelMutableExpr expr) {
@@ -57,6 +58,12 @@ public CelMutableSource addAllExtensions(Collection<? extends Extension> extensi
     return this;
   }
 
+  @CanIgnoreReturnValue
+  public CelMutableSource setDescription(String description) {
+    this.description = checkNotNull(description);
+    return this;
+  }
+
   @CanIgnoreReturnValue
   public CelMutableSource clearMacroCall(long exprId) {
     this.macroCalls.remove(exprId);
@@ -69,6 +76,10 @@ public CelMutableSource clearMacroCalls() {
     return this;
   }
 
+  public String getDescription() {
+    return description;
+  }
+
   public Map<Long, CelMutableExpr> getMacroCalls() {
     return macroCalls;
   }
@@ -79,6 +90,7 @@ public Set<Extension> getExtensions() {
 
   public CelSource toCelSource() {
     return CelSource.newBuilder()
+        .setDescription(description)
         .addAllExtensions(extensions)
         .addAllMacroCalls(
             macroCalls.entrySet().stream()
@@ -89,11 +101,12 @@ public CelSource toCelSource() {
   }
 
   public static CelMutableSource newInstance() {
-    return new CelMutableSource(new HashMap<>(), new HashSet<>());
+    return new CelMutableSource("", new HashMap<>(), new HashSet<>());
   }
 
   public static CelMutableSource fromCelSource(CelSource source) {
     return new CelMutableSource(
+        source.getDescription(),
         source.getMacroCalls().entrySet().stream()
             .collect(
                 Collectors.toMap(
@@ -107,7 +120,9 @@ public static CelMutableSource fromCelSource(CelSource source) {
         source.getExtensions());
   }
 
-  CelMutableSource(Map<Long, CelMutableExpr> macroCalls, Set<Extension> extensions) {
+  CelMutableSource(
+      String description, Map<Long, CelMutableExpr> macroCalls, Set<Extension> extensions) {
+    this.description = checkNotNull(description);
     this.macroCalls = checkNotNull(macroCalls);
     this.extensions = checkNotNull(extensions);
   }

optimizer/src/main/java/dev/cel/optimizer/AstMutator.java

@@ -20,6 +20,7 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Preconditions;
+import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableMap;
 import com.google.errorprone.annotations.Immutable;
 import dev.cel.common.CelAbstractSyntaxTree;
@@ -418,7 +419,8 @@ public CelMutableAst replaceSubtree(
 
     CelMutableExpr mutatedRoot =
         mutateExpr(stableIdGenerator::renumberId, ast.expr(), newAst.expr(), exprIdToReplace);
-    CelMutableSource newAstSource = CelMutableSource.newInstance();
+    CelMutableSource newAstSource =
+        CelMutableSource.newInstance().setDescription(ast.source().getDescription());
     if (!ast.source().getMacroCalls().isEmpty()) {
       newAstSource = combine(newAstSource, ast.source());
     }
@@ -570,6 +572,10 @@ private CelMutableExpr newBindMacroSourceExpr(
   private static CelMutableSource combine(
       CelMutableSource celSource1, CelMutableSource celSource2) {
     return CelMutableSource.newInstance()
+        .setDescription(
+            Strings.isNullOrEmpty(celSource1.getDescription())
+                ? celSource2.getDescription()
+                : celSource1.getDescription())
         .addAllExtensions(celSource1.getExtensions())
         .addAllExtensions(celSource2.getExtensions())
         .addAllMacroCalls(celSource1.getMacroCalls())
@@ -635,7 +641,9 @@ private CelMutableSource normalizeMacroSource(
                     }));
 
     CelMutableSource newMacroSource =
-        CelMutableSource.newInstance().addAllExtensions(source.getExtensions());
+        CelMutableSource.newInstance()
+            .setDescription(source.getDescription())
+            .addAllExtensions(source.getExtensions());
     // Update the macro call IDs and their call references
     for (Entry<Long, CelMutableExpr> existingMacroCall : source.getMacroCalls().entrySet()) {
       long macroId = existingMacroCall.getKey();

optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java

@@ -109,11 +109,11 @@ public void astMutator_nonMacro_sourceCleared() throws Exception {
             .replaceSubtree(mutableAst, newBooleanConst, mutableAst.expr().id())
             .toParsedAst();
 
-    assertThat(mutatedAst.getSource().getDescription()).isEmpty();
     assertThat(mutatedAst.getSource().getLineOffsets()).isEmpty();
     assertThat(mutatedAst.getSource().getPositionsMap()).isEmpty();
     assertThat(mutatedAst.getSource().getExtensions()).isEmpty();
     assertThat(mutatedAst.getSource().getMacroCalls()).isEmpty();
+    assertThat(mutatedAst.getSource().getDescription()).isEqualTo(ast.getSource().getDescription());
   }
 
   @Test
@@ -125,11 +125,11 @@ public void astMutator_macro_sourceMacroCallsPopulated() throws Exception {
     CelAbstractSyntaxTree mutatedAst =
         AST_MUTATOR.replaceSubtree(mutableAst, newBooleanConst, 1).toParsedAst(); // no_op
 
-    assertThat(mutatedAst.getSource().getDescription()).isEmpty();
     assertThat(mutatedAst.getSource().getLineOffsets()).isEmpty();
     assertThat(mutatedAst.getSource().getPositionsMap()).isEmpty();
     assertThat(mutatedAst.getSource().getExtensions()).isEmpty();
     assertThat(mutatedAst.getSource().getMacroCalls()).isNotEmpty();
+    assertThat(mutatedAst.getSource().getDescription()).isEqualTo(ast.getSource().getDescription());
   }
 
   @Test