Retain tagged extension when it exists in both original and new ASTs

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

@@ -21,6 +21,7 @@
 import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.CheckReturnValue;
 import com.google.errorprone.annotations.Immutable;
@@ -43,7 +44,7 @@ public final class CelSource {
   private final ImmutableList<Integer> lineOffsets;
   private final ImmutableMap<Long, Integer> positions;
   private final ImmutableMap<Long, CelExpr> macroCalls;
-  private final ImmutableList<Extension> extensions;
+  private final ImmutableSet<Extension> extensions;
 
   private CelSource(Builder builder) {
     this.codePoints = checkNotNull(builder.codePoints);
@@ -80,7 +81,7 @@ public ImmutableMap<Long, CelExpr> getMacroCalls() {
     return macroCalls;
   }
 
-  public ImmutableList<Extension> getExtensions() {
+  public ImmutableSet<Extension> getExtensions() {
     return extensions;
   }
 
@@ -209,7 +210,7 @@ public static final class Builder {
     private final List<Integer> lineOffsets;
     private final ImmutableMap.Builder<Long, Integer> positions;
     private final Map<Long, CelExpr> macroCalls;
-    private final ImmutableList.Builder<Extension> extensions;
+    private final ImmutableSet.Builder<Extension> extensions;
 
     private String description;
 
@@ -222,7 +223,7 @@ private Builder(CelCodePointArray codePoints, List<Integer> lineOffsets) {
       this.lineOffsets = checkNotNull(lineOffsets);
       this.positions = ImmutableMap.builder();
       this.macroCalls = new HashMap<>();
-      this.extensions = ImmutableList.builder();
+      this.extensions = ImmutableSet.builder();
       this.description = "";
     }
 
@@ -278,13 +279,21 @@ public Builder clearMacroCall(long exprId) {
       return this;
     }
 
+    /**
+     * Adds one or more {@link Extension}s to the source information. Extensions implement set
+     * semantics and deduped if same ones are provided.
+     */
     @CanIgnoreReturnValue
     public Builder addAllExtensions(Iterable<? extends Extension> extensions) {
       checkNotNull(extensions);
       this.extensions.addAll(extensions);
       return this;
     }
 
+    /**
+     * Adds one or more {@link Extension}s to the source information. Extensions implement set
+     * semantics and deduped if same ones are provided.
+     */
     @CanIgnoreReturnValue
     public Builder addAllExtensions(Extension... extensions) {
       return addAllExtensions(Arrays.asList(extensions));

common/src/test/java/dev/cel/common/CelSourceTest.java

@@ -18,6 +18,7 @@
 import static org.antlr.v4.runtime.IntStream.UNKNOWN_SOURCE_NAME;
 import static org.junit.Assert.assertThrows;
 
+import com.google.common.collect.Iterables;
 import dev.cel.common.CelSource.Extension;
 import dev.cel.common.CelSource.Extension.Component;
 import dev.cel.common.CelSource.Extension.Version;
@@ -172,7 +173,7 @@ public void source_withExtension() {
                     Component.COMPONENT_TYPE_CHECKER))
             .build();
 
-    Extension extension = celSource.getExtensions().get(0);
+    Extension extension = Iterables.getOnlyElement(celSource.getExtensions());
     assertThat(extension.id()).isEqualTo("extension_id");
     assertThat(extension.version().major()).isEqualTo(5L);
     assertThat(extension.version().minor()).isEqualTo(3L);

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

@@ -571,7 +571,11 @@ private static CelSource combine(CelSource celSource1, CelSource celSource2) {
     macroMap.putAll(celSource1.getMacroCalls());
     macroMap.putAll(celSource2.getMacroCalls());
 
-    return CelSource.newBuilder().addAllMacroCalls(macroMap.buildOrThrow()).build();
+    return CelSource.newBuilder()
+        .addAllExtensions(celSource1.getExtensions())
+        .addAllExtensions(celSource2.getExtensions())
+        .addAllMacroCalls(macroMap.buildOrThrow())
+        .build();
   }
 
   /**
@@ -589,7 +593,8 @@ private CelAbstractSyntaxTree stabilizeAst(CelAbstractSyntaxTree ast, long seedE
       return CelAbstractSyntaxTree.newParsedAst(newExprBuilder.build(), ast.getSource());
     }
 
-    CelSource.Builder sourceBuilder = CelSource.newBuilder();
+    CelSource.Builder sourceBuilder =
+        CelSource.newBuilder().addAllExtensions(ast.getSource().getExtensions());
     // Update the macro call IDs and their call IDs
     for (Entry<Long, CelExpr> macroCall : ast.getSource().getMacroCalls().entrySet()) {
       long macroId = macroCall.getKey();

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

@@ -122,7 +122,7 @@ public void mutableAst_macro_sourceMacroCallsPopulated() throws Exception {
   }
 
   @Test
-  public void mutableAst_astContainsTaggedExtension_retained() throws Exception {
+  public void replaceSubtree_astContainsTaggedExtension_retained() throws Exception {
     CelAbstractSyntaxTree ast = CEL.compile("has(TestAllTypes{}.single_int32)").getAst();
     Extension extension = Extension.create("test", Version.of(1, 1));
     CelSource celSource = ast.getSource().toBuilder().addAllExtensions(extension).build();
@@ -137,6 +137,64 @@ public void mutableAst_astContainsTaggedExtension_retained() throws Exception {
     assertThat(mutatedAst.getSource().getExtensions()).containsExactly(extension);
   }
 
+  @Test
+  public void replaceSubtreeWithNewAst_astsContainTaggedExtension_retained() throws Exception {
+    // Setup first AST with a test extension
+    CelAbstractSyntaxTree ast = CEL.compile("has(TestAllTypes{}.single_int32)").getAst();
+    Extension extension = Extension.create("test", Version.of(1, 1));
+    ast =
+        CelAbstractSyntaxTree.newCheckedAst(
+            ast.getExpr(),
+            ast.getSource().toBuilder().addAllExtensions(extension).build(),
+            ast.getReferenceMap(),
+            ast.getTypeMap());
+    // Setup second AST with another test extension
+    CelAbstractSyntaxTree astToReplaceWith = CEL.compile("cel.bind(a, true, a)").getAst();
+    Extension extension2 = Extension.create("test2", Version.of(2, 2));
+    astToReplaceWith =
+        CelAbstractSyntaxTree.newCheckedAst(
+            astToReplaceWith.getExpr(),
+            astToReplaceWith.getSource().toBuilder().addAllExtensions(extension2).build(),
+            astToReplaceWith.getReferenceMap(),
+            astToReplaceWith.getTypeMap());
+
+    // Mutate the original AST with the new AST at the root
+    CelAbstractSyntaxTree mutatedAst =
+        MUTABLE_AST.replaceSubtreeWithNewAst(ast, astToReplaceWith, ast.getExpr().id());
+
+    // Expect that both the extensions are merged
+    assertThat(mutatedAst.getSource().getExtensions()).containsExactly(extension, extension2);
+  }
+
+  @Test
+  public void replaceSubtreeWithNewAst_astsContainSameExtensions_deduped() throws Exception {
+    // Setup first AST with a test extension
+    CelAbstractSyntaxTree ast = CEL.compile("has(TestAllTypes{}.single_int32)").getAst();
+    Extension extension = Extension.create("test", Version.of(1, 1));
+    ast =
+        CelAbstractSyntaxTree.newCheckedAst(
+            ast.getExpr(),
+            ast.getSource().toBuilder().addAllExtensions(extension).build(),
+            ast.getReferenceMap(),
+            ast.getTypeMap());
+    // Setup second AST with the same test extension as above
+    CelAbstractSyntaxTree astToReplaceWith = CEL.compile("cel.bind(a, true, a)").getAst();
+    Extension extension2 = Extension.create("test", Version.of(1, 1));
+    astToReplaceWith =
+        CelAbstractSyntaxTree.newCheckedAst(
+            astToReplaceWith.getExpr(),
+            astToReplaceWith.getSource().toBuilder().addAllExtensions(extension2).build(),
+            astToReplaceWith.getReferenceMap(),
+            astToReplaceWith.getTypeMap());
+
+    // Mutate the original AST with the new AST at the root
+    CelAbstractSyntaxTree mutatedAst =
+        MUTABLE_AST.replaceSubtreeWithNewAst(ast, astToReplaceWith, ast.getExpr().id());
+
+    // Expect that the extension is deduped
+    assertThat(mutatedAst.getSource().getExtensions()).containsExactly(extension);
+  }
+
   @Test
   @TestParameters("{source: '[1].exists(x, x > 0)', expectedMacroCallSize: 1}")
   @TestParameters(