diff --git a/common/src/main/java/dev/cel/common/CelSource.java b/common/src/main/java/dev/cel/common/CelSource.java index 62a74873a..6adaf8773 100644 --- a/common/src/main/java/dev/cel/common/CelSource.java +++ b/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 lineOffsets; private final ImmutableMap positions; private final ImmutableMap macroCalls; - private final ImmutableList extensions; + private final ImmutableSet extensions; private CelSource(Builder builder) { this.codePoints = checkNotNull(builder.codePoints); @@ -80,7 +81,7 @@ public ImmutableMap getMacroCalls() { return macroCalls; } - public ImmutableList getExtensions() { + public ImmutableSet getExtensions() { return extensions; } @@ -209,7 +210,7 @@ public static final class Builder { private final List lineOffsets; private final ImmutableMap.Builder positions; private final Map macroCalls; - private final ImmutableList.Builder extensions; + private final ImmutableSet.Builder extensions; private String description; @@ -222,7 +223,7 @@ private Builder(CelCodePointArray codePoints, List lineOffsets) { this.lineOffsets = checkNotNull(lineOffsets); this.positions = ImmutableMap.builder(); this.macroCalls = new HashMap<>(); - this.extensions = ImmutableList.builder(); + this.extensions = ImmutableSet.builder(); this.description = ""; } @@ -278,6 +279,10 @@ 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 extensions) { checkNotNull(extensions); @@ -285,6 +290,10 @@ public Builder addAllExtensions(Iterable 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)); diff --git a/common/src/test/java/dev/cel/common/CelSourceTest.java b/common/src/test/java/dev/cel/common/CelSourceTest.java index 74e350a35..a97eeb853 100644 --- a/common/src/test/java/dev/cel/common/CelSourceTest.java +++ b/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); diff --git a/optimizer/src/main/java/dev/cel/optimizer/MutableAst.java b/optimizer/src/main/java/dev/cel/optimizer/MutableAst.java index 1a466df67..181dda25c 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/MutableAst.java +++ b/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 macroCall : ast.getSource().getMacroCalls().entrySet()) { long macroId = macroCall.getKey(); diff --git a/optimizer/src/test/java/dev/cel/optimizer/MutableAstTest.java b/optimizer/src/test/java/dev/cel/optimizer/MutableAstTest.java index da8ecfb7b..af0e1241e 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/MutableAstTest.java +++ b/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(