diff --git a/common/BUILD.bazel b/common/BUILD.bazel index 032a61150..7ad38b830 100644 --- a/common/BUILD.bazel +++ b/common/BUILD.bazel @@ -39,6 +39,16 @@ java_library( exports = ["//common/src/main/java/dev/cel/common:error_codes"], ) +java_library( + name = "mutable_ast", + exports = ["//common/src/main/java/dev/cel/common:mutable_ast"], +) + +java_library( + name = "mutable_source", + exports = ["//common/src/main/java/dev/cel/common:mutable_source"], +) + java_library( name = "runtime_exception", visibility = ["//visibility:public"], diff --git a/common/ast/BUILD.bazel b/common/ast/BUILD.bazel index cd18dca31..5400e6a8e 100644 --- a/common/ast/BUILD.bazel +++ b/common/ast/BUILD.bazel @@ -34,6 +34,6 @@ java_library( ) java_library( - name = "mutable_ast", - exports = ["//common/src/main/java/dev/cel/common/ast:mutable_ast"], + name = "mutable_expr", + exports = ["//common/src/main/java/dev/cel/common/ast:mutable_expr"], ) diff --git a/common/src/main/java/dev/cel/common/BUILD.bazel b/common/src/main/java/dev/cel/common/BUILD.bazel index 31b593708..de5995c9b 100644 --- a/common/src/main/java/dev/cel/common/BUILD.bazel +++ b/common/src/main/java/dev/cel/common/BUILD.bazel @@ -134,3 +134,31 @@ java_library( "//common/annotations", ], ) + +java_library( + name = "mutable_ast", + srcs = ["CelMutableAst.java"], + tags = [ + ], + deps = [ + ":mutable_source", + "//common", + "//common/ast", + "//common/ast:mutable_expr", + "//common/types:type_providers", + ], +) + +java_library( + name = "mutable_source", + srcs = ["CelMutableSource.java"], + tags = [ + ], + deps = [ + ":common", + "//:auto_value", + "//common/ast:mutable_expr", + "@maven//:com_google_errorprone_error_prone_annotations", + "@maven//:com_google_guava_guava", + ], +) diff --git a/common/src/main/java/dev/cel/common/ast/CelMutableAst.java b/common/src/main/java/dev/cel/common/CelMutableAst.java similarity index 80% rename from common/src/main/java/dev/cel/common/ast/CelMutableAst.java rename to common/src/main/java/dev/cel/common/CelMutableAst.java index f529e2e8d..fd3fe28f8 100644 --- a/common/src/main/java/dev/cel/common/ast/CelMutableAst.java +++ b/common/src/main/java/dev/cel/common/CelMutableAst.java @@ -12,10 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -package dev.cel.common.ast; +package dev.cel.common; -import dev.cel.common.CelAbstractSyntaxTree; -import dev.cel.common.CelSource; +import dev.cel.common.ast.CelMutableExpr; +import dev.cel.common.ast.CelMutableExprConverter; +import dev.cel.common.ast.CelReference; import dev.cel.common.types.CelType; import java.util.HashMap; import java.util.Map; @@ -30,7 +31,7 @@ */ public final class CelMutableAst { private final CelMutableExpr mutatedExpr; - private final CelSource.Builder source; + private final CelMutableSource source; private final Map references; private final Map types; @@ -40,9 +41,10 @@ public CelMutableExpr expr() { } /** - * Returns the {@link CelSource} that was used during construction of the abstract syntax tree. + * Returns the {@link CelMutableSource} that was used during construction of the abstract syntax + * tree. */ - public CelSource.Builder source() { + public CelMutableSource source() { return source; } @@ -69,7 +71,7 @@ public Optional getType(long exprId) { /** Converts this mutable AST into a parsed {@link CelAbstractSyntaxTree}. */ public CelAbstractSyntaxTree toParsedAst() { return CelAbstractSyntaxTree.newParsedAst( - CelMutableExprConverter.fromMutableExpr(mutatedExpr), source.build()); + CelMutableExprConverter.fromMutableExpr(mutatedExpr), source.toCelSource()); } /** @@ -79,7 +81,7 @@ public CelAbstractSyntaxTree toParsedAst() { public static CelMutableAst fromCelAst(CelAbstractSyntaxTree ast) { return new CelMutableAst( CelMutableExprConverter.fromCelExpr(ast.getExpr()), - ast.getSource().toBuilder(), + CelMutableSource.fromCelSource(ast.getSource()), ast.getReferenceMap(), ast.getTypeMap()); } @@ -88,21 +90,21 @@ public static CelMutableAst fromCelAst(CelAbstractSyntaxTree ast) { * Constructs an instance of {@link CelMutableAst} with the mutable expression and its source * builder. */ - public static CelMutableAst of(CelMutableExpr mutableExpr, CelSource.Builder sourceBuilder) { - return new CelMutableAst(mutableExpr, sourceBuilder); + public static CelMutableAst of(CelMutableExpr mutableExpr, CelMutableSource mutableSource) { + return new CelMutableAst(mutableExpr, mutableSource); } - private CelMutableAst(CelMutableExpr mutatedExpr, CelSource.Builder source) { - this(mutatedExpr, source, new HashMap<>(), new HashMap<>()); + private CelMutableAst(CelMutableExpr mutatedExpr, CelMutableSource mutableSource) { + this(mutatedExpr, mutableSource, new HashMap<>(), new HashMap<>()); } private CelMutableAst( CelMutableExpr mutatedExpr, - CelSource.Builder source, + CelMutableSource mutableSource, Map references, Map types) { this.mutatedExpr = mutatedExpr; - this.source = source; + this.source = mutableSource; this.references = new HashMap<>(references); this.types = new HashMap<>(types); } diff --git a/common/src/main/java/dev/cel/common/CelMutableSource.java b/common/src/main/java/dev/cel/common/CelMutableSource.java new file mode 100644 index 000000000..250963b9a --- /dev/null +++ b/common/src/main/java/dev/cel/common/CelMutableSource.java @@ -0,0 +1,114 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dev.cel.common; + +import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.collect.ImmutableMap.toImmutableMap; + +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import dev.cel.common.CelSource.Extension; +import dev.cel.common.ast.CelMutableExpr; +import dev.cel.common.ast.CelMutableExprConverter; +import java.util.Collection; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Map.Entry; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Represents the mutable portion of the {@link CelSource}. This is intended for the purposes of + * augmenting an AST through CEL optimizers. + */ +public final class CelMutableSource { + + final Map macroCalls; + final Set extensions; + + @CanIgnoreReturnValue + public CelMutableSource addMacroCalls(long exprId, CelMutableExpr expr) { + this.macroCalls.put(exprId, checkNotNull(CelMutableExpr.newInstance(expr))); + return this; + } + + @CanIgnoreReturnValue + public CelMutableSource addAllMacroCalls(Map macroCalls) { + this.macroCalls.putAll(macroCalls); + return this; + } + + @CanIgnoreReturnValue + public CelMutableSource addAllExtensions(Collection extensions) { + checkNotNull(extensions); + this.extensions.addAll(extensions); + return this; + } + + @CanIgnoreReturnValue + public CelMutableSource clearMacroCall(long exprId) { + this.macroCalls.remove(exprId); + return this; + } + + @CanIgnoreReturnValue + public CelMutableSource clearMacroCalls() { + this.macroCalls.clear(); + return this; + } + + public Map getMacroCalls() { + return macroCalls; + } + + public Set getExtensions() { + return extensions; + } + + public CelSource toCelSource() { + return CelSource.newBuilder() + .addAllExtensions(extensions) + .addAllMacroCalls( + macroCalls.entrySet().stream() + .collect( + toImmutableMap( + Entry::getKey, v -> CelMutableExprConverter.fromMutableExpr(v.getValue())))) + .build(); + } + + public static CelMutableSource newInstance() { + return new CelMutableSource(new HashMap<>(), new HashSet<>()); + } + + public static CelMutableSource fromCelSource(CelSource source) { + return new CelMutableSource( + source.getMacroCalls().entrySet().stream() + .collect( + Collectors.toMap( + Entry::getKey, + v -> CelMutableExprConverter.fromCelExpr(v.getValue()), + (prev, next) -> { + throw new IllegalStateException( + "Unexpected source collision at ID: " + prev.id()); + }, + HashMap::new)), + source.getExtensions()); + } + + CelMutableSource(Map macroCalls, Set extensions) { + this.macroCalls = checkNotNull(macroCalls); + this.extensions = checkNotNull(extensions); + } +} diff --git a/common/src/main/java/dev/cel/common/CelSource.java b/common/src/main/java/dev/cel/common/CelSource.java index 82b835180..2422f7b20 100644 --- a/common/src/main/java/dev/cel/common/CelSource.java +++ b/common/src/main/java/dev/cel/common/CelSource.java @@ -279,12 +279,6 @@ public Builder clearMacroCall(long exprId) { return this; } - @CanIgnoreReturnValue - public Builder clearMacroCalls() { - this.macroCalls.clear(); - return this; - } - public ImmutableSet getExtensions() { return extensions.build(); } diff --git a/common/src/main/java/dev/cel/common/ast/BUILD.bazel b/common/src/main/java/dev/cel/common/ast/BUILD.bazel index 9ed6c80fc..c59178bcd 100644 --- a/common/src/main/java/dev/cel/common/ast/BUILD.bazel +++ b/common/src/main/java/dev/cel/common/ast/BUILD.bazel @@ -33,8 +33,7 @@ EXPR_FACTORY_SOURCES = [ ] # keep sorted -MUTABLE_AST_SOURCES = [ - "CelMutableAst.java", +MUTABLE_EXPR_SOURCES = [ "CelMutableExpr.java", "CelMutableExprConverter.java", ] @@ -119,14 +118,12 @@ java_library( ) java_library( - name = "mutable_ast", - srcs = MUTABLE_AST_SOURCES, + name = "mutable_expr", + srcs = MUTABLE_EXPR_SOURCES, tags = [ ], deps = [ ":ast", - "//common", - "//common/types:type_providers", "@maven//:com_google_guava_guava", ], ) diff --git a/common/src/main/java/dev/cel/common/ast/CelMutableExpr.java b/common/src/main/java/dev/cel/common/ast/CelMutableExpr.java index f83021b0b..ff2c83285 100644 --- a/common/src/main/java/dev/cel/common/ast/CelMutableExpr.java +++ b/common/src/main/java/dev/cel/common/ast/CelMutableExpr.java @@ -1064,6 +1064,9 @@ private CelMutableExpr(CelMutableExpr other) { this.id = other.id; this.exprKind = other.exprKind; switch (other.getKind()) { + case NOT_SET: + this.exprValue = CelExpr.newBuilder().build().exprKind().notSet(); + break; case CONSTANT: this.exprValue = other.exprValue; // Constant is immutable. break; diff --git a/common/src/main/java/dev/cel/common/ast/CelMutableExprConverter.java b/common/src/main/java/dev/cel/common/ast/CelMutableExprConverter.java index 64b9cdb10..2e6ffdadd 100644 --- a/common/src/main/java/dev/cel/common/ast/CelMutableExprConverter.java +++ b/common/src/main/java/dev/cel/common/ast/CelMutableExprConverter.java @@ -134,8 +134,6 @@ private static CelMutableCreateMap fromCelMapToMutableMap(CelCreateMap celCreate return CelMutableCreateMap.create(entries); } - /////////////////////// - public static CelExpr fromMutableExpr(CelMutableExpr mutableExpr) { long id = mutableExpr.id(); switch (mutableExpr.getKind()) { diff --git a/common/src/main/java/dev/cel/common/navigation/BUILD.bazel b/common/src/main/java/dev/cel/common/navigation/BUILD.bazel index 903ae3f0f..b3bccf67b 100644 --- a/common/src/main/java/dev/cel/common/navigation/BUILD.bazel +++ b/common/src/main/java/dev/cel/common/navigation/BUILD.bazel @@ -56,7 +56,8 @@ java_library( deps = [ ":common", "//:auto_value", - "//common/ast:mutable_ast", + "//common:mutable_ast", + "//common/ast:mutable_expr", "//common/types:type_providers", "@maven//:com_google_errorprone_error_prone_annotations", "@maven//:com_google_guava_guava", diff --git a/common/src/main/java/dev/cel/common/navigation/CelNavigableMutableAst.java b/common/src/main/java/dev/cel/common/navigation/CelNavigableMutableAst.java index bda18dc03..70d33ab41 100644 --- a/common/src/main/java/dev/cel/common/navigation/CelNavigableMutableAst.java +++ b/common/src/main/java/dev/cel/common/navigation/CelNavigableMutableAst.java @@ -14,7 +14,7 @@ package dev.cel.common.navigation; -import dev.cel.common.ast.CelMutableAst; +import dev.cel.common.CelMutableAst; import dev.cel.common.types.CelType; import java.util.Optional; diff --git a/common/src/test/java/dev/cel/common/ast/BUILD.bazel b/common/src/test/java/dev/cel/common/ast/BUILD.bazel index 103f4804a..61e6beb11 100644 --- a/common/src/test/java/dev/cel/common/ast/BUILD.bazel +++ b/common/src/test/java/dev/cel/common/ast/BUILD.bazel @@ -13,13 +13,15 @@ java_library( "//:java_truth", "//common", "//common:compiler_common", + "//common:mutable_ast", + "//common:mutable_source", "//common:options", "//common/ast", "//common/ast:cel_expr_visitor", "//common/ast:expr_converter", "//common/ast:expr_factory", "//common/ast:expr_v1alpha1_converter", - "//common/ast:mutable_ast", + "//common/ast:mutable_expr", "//common/resources/testdata/proto3:test_all_types_java_proto", "//common/types", "//compiler", diff --git a/common/src/test/java/dev/cel/common/ast/CelMutableAstTest.java b/common/src/test/java/dev/cel/common/ast/CelMutableAstTest.java index 1b249484d..6d2c07b17 100644 --- a/common/src/test/java/dev/cel/common/ast/CelMutableAstTest.java +++ b/common/src/test/java/dev/cel/common/ast/CelMutableAstTest.java @@ -17,8 +17,9 @@ import static com.google.common.truth.Truth.assertThat; import dev.cel.common.CelAbstractSyntaxTree; +import dev.cel.common.CelMutableAst; +import dev.cel.common.CelMutableSource; import dev.cel.common.CelOptions; -import dev.cel.common.CelSource; import dev.cel.common.types.SimpleType; import dev.cel.compiler.CelCompiler; import dev.cel.compiler.CelCompilerFactory; @@ -33,12 +34,12 @@ public final class CelMutableAstTest { @Test public void constructMutableAst() { CelMutableExpr mutableExpr = CelMutableExpr.ofConstant(1L, CelConstant.ofValue("hello world")); - CelSource.Builder sourceBuilder = CelSource.newBuilder(); + CelMutableSource mutableSource = CelMutableSource.newInstance(); - CelMutableAst celMutableAst = CelMutableAst.of(mutableExpr, sourceBuilder); + CelMutableAst celMutableAst = CelMutableAst.of(mutableExpr, mutableSource); assertThat(celMutableAst.expr()).isEqualTo(mutableExpr); - assertThat(celMutableAst.source()).isSameInstanceAs(sourceBuilder); + assertThat(celMutableAst.source()).isSameInstanceAs(mutableSource); } @Test diff --git a/common/src/test/java/dev/cel/common/navigation/BUILD.bazel b/common/src/test/java/dev/cel/common/navigation/BUILD.bazel index 40e04a70d..f79a59482 100644 --- a/common/src/test/java/dev/cel/common/navigation/BUILD.bazel +++ b/common/src/test/java/dev/cel/common/navigation/BUILD.bazel @@ -10,9 +10,10 @@ java_library( "//:java_truth", "//common", "//common:compiler_common", + "//common:mutable_ast", "//common:options", "//common/ast", - "//common/ast:mutable_ast", + "//common/ast:mutable_expr", "//common/navigation", "//common/navigation:common", "//common/navigation:mutable_navigation", diff --git a/common/src/test/java/dev/cel/common/navigation/CelNavigableMutableAstTest.java b/common/src/test/java/dev/cel/common/navigation/CelNavigableMutableAstTest.java index 3d7f8bfcd..6b78b5f23 100644 --- a/common/src/test/java/dev/cel/common/navigation/CelNavigableMutableAstTest.java +++ b/common/src/test/java/dev/cel/common/navigation/CelNavigableMutableAstTest.java @@ -16,8 +16,8 @@ import static com.google.common.truth.Truth.assertThat; +import dev.cel.common.CelMutableAst; import dev.cel.common.ast.CelConstant; -import dev.cel.common.ast.CelMutableAst; import dev.cel.common.ast.CelMutableExpr; import dev.cel.compiler.CelCompiler; import dev.cel.compiler.CelCompilerFactory; diff --git a/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java b/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java index 87942d2bd..7a8b636fb 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java +++ b/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java @@ -25,20 +25,17 @@ import com.google.common.collect.Table; import com.google.errorprone.annotations.Immutable; import dev.cel.common.CelAbstractSyntaxTree; -import dev.cel.common.CelSource; +import dev.cel.common.CelMutableAst; +import dev.cel.common.CelMutableSource; import dev.cel.common.ast.CelConstant; -import dev.cel.common.ast.CelExpr; import dev.cel.common.ast.CelExpr.ExprKind.Kind; import dev.cel.common.ast.CelExprIdGeneratorFactory; import dev.cel.common.ast.CelExprIdGeneratorFactory.ExprIdGenerator; import dev.cel.common.ast.CelExprIdGeneratorFactory.StableIdGenerator; -import dev.cel.common.ast.CelMutableAst; import dev.cel.common.ast.CelMutableExpr; import dev.cel.common.ast.CelMutableExpr.CelMutableCall; import dev.cel.common.ast.CelMutableExpr.CelMutableComprehension; import dev.cel.common.ast.CelMutableExpr.CelMutableCreateList; -import dev.cel.common.ast.CelMutableExprConverter; -import dev.cel.common.navigation.CelNavigableExpr; import dev.cel.common.navigation.CelNavigableMutableAst; import dev.cel.common.navigation.CelNavigableMutableExpr; import dev.cel.common.navigation.TraversalOrder; @@ -125,7 +122,7 @@ public CelMutableAst replaceSubtreeWithNewBindMacro( CelMutableExpr newBindMacroExpr = newBindMacroExpr( varName, varInit, CelMutableExpr.newInstance(resultExpr), stableIdGenerator); - CelSource.Builder celSource = CelSource.newBuilder(); + CelMutableSource celSource = CelMutableSource.newInstance(); if (populateMacroSource) { CelMutableExpr newBindMacroSourceExpr = newBindMacroSourceExpr(newBindMacroExpr, varName, stableIdGenerator); @@ -134,13 +131,11 @@ public CelMutableAst replaceSubtreeWithNewBindMacro( // macro. celSource = normalizeMacroSource( - ast.source(), - -1, // Do not replace any of the subexpr in the macro map. - newBindMacroSourceExpr, - stableIdGenerator::renumberId) - .addMacroCalls( - newBindMacroExpr.id(), - CelMutableExprConverter.fromMutableExpr(newBindMacroSourceExpr)); + ast.source(), + -1, // Do not replace any of the subexpr in the macro map. + newBindMacroSourceExpr, + stableIdGenerator::renumberId); + celSource.addMacroCalls(newBindMacroExpr.id(), newBindMacroSourceExpr); } CelMutableAst newBindAst = CelMutableAst.of(newBindMacroExpr, celSource); @@ -152,7 +147,7 @@ public CelMutableAst replaceSubtreeWithNewBindMacro( public CelMutableAst renumberIdsConsecutively(CelMutableAst mutableAst) { StableIdGenerator stableIdGenerator = CelExprIdGeneratorFactory.newStableIdGenerator(0); CelMutableExpr mutableExpr = renumberExprIds(stableIdGenerator::renumberId, mutableAst.expr()); - CelSource.Builder newSource = + CelMutableSource newSource = normalizeMacroSource( mutableAst.source(), Integer.MIN_VALUE, mutableExpr, stableIdGenerator::renumberId); @@ -275,7 +270,7 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames( Table comprehensionLevelToType = HashBasedTable.create(); CelMutableExpr mutatedComprehensionExpr = navigableMutableAst.getAst().expr(); - CelSource.Builder newSource = navigableMutableAst.getAst().source(); + CelMutableSource newSource = navigableMutableAst.getAst().source(); for (Entry comprehensionEntry : comprehensionsToMangle.entrySet()) { iterCount++; @@ -353,7 +348,8 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames( */ public CelMutableAst replaceSubtree( CelMutableExpr root, CelMutableExpr newExpr, long exprIdToReplace) { - return replaceSubtree(CelMutableAst.of(root, CelSource.newBuilder()), newExpr, exprIdToReplace); + return replaceSubtree( + CelMutableAst.of(root, CelMutableSource.newInstance()), newExpr, exprIdToReplace); } /** @@ -381,7 +377,7 @@ public CelMutableAst replaceSubtree( newExpr, // Copy the macro call information to the new AST such that macro call map can be // normalized post-replacement. - CelSource.newBuilder().addAllMacroCalls(ast.source().getMacroCalls())), + ast.source()), exprIdToReplace); } @@ -444,7 +440,7 @@ public CelMutableAst replaceSubtree( CelMutableExpr mutatedRoot = mutateExpr(stableIdGenerator::renumberId, ast.expr(), newAst.expr(), exprIdToReplace); - CelSource.Builder newAstSource = CelSource.newBuilder(); + CelMutableSource newAstSource = CelMutableSource.newInstance(); if (!ast.source().getMacroCalls().isEmpty()) { newAstSource = combine(newAstSource, ast.source()); } @@ -512,8 +508,8 @@ private void replaceIdentName( } } - private CelSource.Builder mangleIdentsInMacroSource( - CelSource.Builder sourceBuilder, + private CelMutableSource mangleIdentsInMacroSource( + CelMutableSource sourceBuilder, CelMutableExpr mutatedComprehensionExpr, String originalIterVar, MangledComprehensionName mangledComprehensionName, @@ -524,7 +520,7 @@ private CelSource.Builder mangleIdentsInMacroSource( // First, normalize the macro source. // ex: [x].exists(x, [x].exists(x, x == 1)) -> [x].exists(x, [@c1].exists(x, @c0 == 1)). - CelSource.Builder newSource = + CelMutableSource newSource = normalizeMacroSource(sourceBuilder, -1, mutatedComprehensionExpr, (id) -> id); // Note that in the above example, the iteration variable is not replaced after normalization. @@ -534,8 +530,7 @@ private CelSource.Builder mangleIdentsInMacroSource( // variable actually exists in the main AST thus, this step isn't needed. // ex: [1].map(x, [2].filter(y, x == y). Here, the variable declaration `x` exists in the AST // but not `y`. - CelMutableExpr macroExpr = - CelMutableExprConverter.fromCelExpr(newSource.getMacroCalls().get(originalComprehensionId)); + CelMutableExpr macroExpr = newSource.getMacroCalls().get(originalComprehensionId); // By convention, the iteration variable is always the first argument of the // macro call expression. CelMutableExpr identToMangle = macroExpr.call().args().get(0); @@ -549,8 +544,7 @@ private CelSource.Builder mangleIdentsInMacroSource( identToMangle.id()); } - newSource.addMacroCalls( - originalComprehensionId, CelMutableExprConverter.fromMutableExpr(macroExpr)); + newSource.addMacroCalls(originalComprehensionId, macroExpr); return newSource; } @@ -595,16 +589,13 @@ private CelMutableExpr newBindMacroSourceExpr( bindMacroExpr.comprehension().result())); } - private static CelSource.Builder combine( - CelSource.Builder celSource1, CelSource.Builder celSource2) { - ImmutableMap.Builder macroMap = ImmutableMap.builder(); - macroMap.putAll(celSource1.getMacroCalls()); - macroMap.putAll(celSource2.getMacroCalls()); - - return CelSource.newBuilder() + private static CelMutableSource combine( + CelMutableSource celSource1, CelMutableSource celSource2) { + return CelMutableSource.newInstance() .addAllExtensions(celSource1.getExtensions()) .addAllExtensions(celSource2.getExtensions()) - .addAllMacroCalls(macroMap.buildOrThrow()); + .addAllMacroCalls(celSource1.getMacroCalls()) + .addAllMacroCalls(celSource2.getMacroCalls()); } /** @@ -614,40 +605,37 @@ private static CelSource.Builder combine( */ private CelMutableAst stabilizeAst(CelMutableAst mutableAst, long seedExprId) { CelMutableExpr mutableExpr = mutableAst.expr(); - CelSource.Builder source = mutableAst.source(); + CelMutableSource source = mutableAst.source(); StableIdGenerator stableIdGenerator = CelExprIdGeneratorFactory.newStableIdGenerator(seedExprId); CelMutableExpr mutatedExpr = renumberExprIds(stableIdGenerator::nextExprId, mutableExpr); - CelSource.Builder sourceBuilder = - CelSource.newBuilder().addAllExtensions(source.getExtensions()); + CelMutableSource sourceBuilder = + CelMutableSource.newInstance().addAllExtensions(source.getExtensions()); // Update the macro call IDs and their call IDs - for (Entry macroCall : source.getMacroCalls().entrySet()) { + for (Entry macroCall : source.getMacroCalls().entrySet()) { long macroId = macroCall.getKey(); long newCallId = stableIdGenerator.renumberId(macroId); - CelMutableExpr existingMacroCallExpr = - CelMutableExprConverter.fromCelExpr(macroCall.getValue()); + CelMutableExpr existingMacroCallExpr = CelMutableExpr.newInstance(macroCall.getValue()); CelMutableExpr newCall = renumberExprIds(stableIdGenerator::renumberId, existingMacroCallExpr); - sourceBuilder.addMacroCalls(newCallId, CelMutableExprConverter.fromMutableExpr(newCall)); + sourceBuilder.addMacroCalls(newCallId, newCall); } return CelMutableAst.of(mutatedExpr, sourceBuilder); } - private CelSource.Builder normalizeMacroSource( - CelSource.Builder celSource, + private CelMutableSource normalizeMacroSource( + CelMutableSource source, long exprIdToReplace, CelMutableExpr mutatedRoot, ExprIdGenerator idGenerator) { // Remove the macro metadata that no longer exists in the AST due to being replaced. - celSource.clearMacroCall(exprIdToReplace); - CelSource.Builder sourceBuilder = - CelSource.newBuilder().addAllExtensions(celSource.getExtensions()); - if (celSource.getMacroCalls().isEmpty()) { - return sourceBuilder; + source.clearMacroCall(exprIdToReplace); + if (source.getMacroCalls().isEmpty()) { + return source; } ImmutableMap allExprs = @@ -668,8 +656,10 @@ private CelSource.Builder normalizeMacroSource( "Expected expressions to be the same for id: " + expr1.id()); })); + CelMutableSource newMacroSource = + CelMutableSource.newInstance().addAllExtensions(source.getExtensions()); // Update the macro call IDs and their call references - for (Entry existingMacroCall : celSource.getMacroCalls().entrySet()) { + for (Entry existingMacroCall : source.getMacroCalls().entrySet()) { long macroId = existingMacroCall.getKey(); long callId = idGenerator.generate(macroId); @@ -678,7 +668,7 @@ private CelSource.Builder normalizeMacroSource( } CelMutableExpr existingMacroCallExpr = - CelMutableExprConverter.fromCelExpr(existingMacroCall.getValue()); + CelMutableExpr.newInstance(existingMacroCall.getValue()); CelMutableExpr newMacroCallExpr = renumberExprIds(idGenerator, existingMacroCallExpr); CelNavigableMutableExpr callNav = CelNavigableMutableExpr.fromExpr(newMacroCallExpr); @@ -711,13 +701,12 @@ private CelSource.Builder normalizeMacroSource( } } - sourceBuilder.addMacroCalls( - callId, CelMutableExprConverter.fromMutableExpr(newMacroCallExpr)); + newMacroSource.addMacroCalls(callId, newMacroCallExpr); } // Replace comprehension nodes with a NOT_SET reference to reduce AST size. - for (Entry macroCall : sourceBuilder.getMacroCalls().entrySet()) { - CelMutableExpr macroCallExpr = CelMutableExprConverter.fromCelExpr(macroCall.getValue()); + for (Entry macroCall : newMacroSource.getMacroCalls().entrySet()) { + CelMutableExpr macroCallExpr = macroCall.getValue(); CelNavigableMutableExpr.fromExpr(macroCallExpr) .allNodes() .filter(node -> node.getKind().equals(Kind.COMPREHENSION)) @@ -730,7 +719,7 @@ private CelSource.Builder normalizeMacroSource( macroCallExpr, CelMutableExpr.ofNotSet(node.id()), node.id()); - macroCall.setValue(CelMutableExprConverter.fromMutableExpr(mutatedNode)); + macroCall.setValue(mutatedNode); }); // Prune any NOT_SET (comprehension) nodes that no longer exist in the main AST @@ -748,11 +737,10 @@ private CelSource.Builder normalizeMacroSource( .collect(toCollection(ArrayList::new)); CelMutableCall call = macroCallExpr.call(); call.setArgs(newCallArgs); - macroCall.setValue(CelMutableExprConverter.fromMutableExpr(macroCallExpr)); }); } - return sourceBuilder; + return newMacroSource; } /** @@ -813,21 +801,14 @@ private CelMutableExpr renumberExprIds(ExprIdGenerator idGenerator, CelMutableEx return mutableAst.visit(root); } - private static long getMaxId(CelExpr newExpr) { - return CelNavigableExpr.fromExpr(newExpr) - .allNodes() - .mapToLong(CelNavigableExpr::id) - .max() - .orElseThrow(NoSuchElementException::new); - } - private static long getMaxId(CelMutableAst mutableAst) { return getMaxId(CelNavigableMutableAst.fromAst(mutableAst)); } private static long getMaxId(CelNavigableMutableAst navAst) { long maxId = navAst.getRoot().maxId(); - for (Entry macroCall : navAst.getAst().source().getMacroCalls().entrySet()) { + for (Entry macroCall : + navAst.getAst().source().getMacroCalls().entrySet()) { maxId = max(maxId, getMaxId(macroCall.getValue())); } diff --git a/optimizer/src/main/java/dev/cel/optimizer/BUILD.bazel b/optimizer/src/main/java/dev/cel/optimizer/BUILD.bazel index 9ddbf2f8a..fcd85ce72 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/BUILD.bazel +++ b/optimizer/src/main/java/dev/cel/optimizer/BUILD.bazel @@ -85,11 +85,12 @@ java_library( deps = [ "//:auto_value", "//common", + "//common:mutable_ast", + "//common:mutable_source", "//common/annotations", "//common/ast", "//common/ast:expr_factory", - "//common/ast:mutable_ast", - "//common/navigation", + "//common/ast:mutable_expr", "//common/navigation:common", "//common/navigation:mutable_navigation", "//common/types:type_providers", diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel b/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel index 409421e08..2b0cc563b 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel @@ -18,9 +18,10 @@ java_library( "//bundle:cel", "//common", "//common:compiler_common", + "//common:mutable_ast", "//common/ast", "//common/ast:expr_util", - "//common/ast:mutable_ast", + "//common/ast:mutable_expr", "//common/navigation:mutable_navigation", "//extensions:optional_library", "//optimizer:ast_optimizer", @@ -45,8 +46,10 @@ java_library( "//checker:checker_legacy_environment", "//common", "//common:compiler_common", + "//common:mutable_ast", + "//common:mutable_source", "//common/ast", - "//common/ast:mutable_ast", + "//common/ast:mutable_expr", "//common/navigation", "//common/navigation:common", "//common/navigation:mutable_navigation", diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java index fd17df9d6..cd7e7fcd0 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java @@ -20,11 +20,11 @@ import com.google.common.collect.ImmutableList; import dev.cel.bundle.Cel; import dev.cel.common.CelAbstractSyntaxTree; +import dev.cel.common.CelMutableAst; import dev.cel.common.CelValidationException; import dev.cel.common.ast.CelConstant; import dev.cel.common.ast.CelExpr.ExprKind.Kind; import dev.cel.common.ast.CelExprUtil; -import dev.cel.common.ast.CelMutableAst; import dev.cel.common.ast.CelMutableExpr; import dev.cel.common.ast.CelMutableExpr.CelMutableCall; import dev.cel.common.ast.CelMutableExpr.CelMutableCreateList; diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java index 311ce34de..7057eb4ba 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java @@ -33,6 +33,8 @@ import dev.cel.checker.Standard; import dev.cel.common.CelAbstractSyntaxTree; import dev.cel.common.CelFunctionDecl; +import dev.cel.common.CelMutableAst; +import dev.cel.common.CelMutableSource; import dev.cel.common.CelOverloadDecl; import dev.cel.common.CelSource; import dev.cel.common.CelSource.Extension; @@ -43,7 +45,6 @@ import dev.cel.common.ast.CelExpr; import dev.cel.common.ast.CelExpr.CelCall; import dev.cel.common.ast.CelExpr.ExprKind.Kind; -import dev.cel.common.ast.CelMutableAst; import dev.cel.common.ast.CelMutableExpr; import dev.cel.common.ast.CelMutableExprConverter; import dev.cel.common.navigation.CelNavigableExpr; @@ -151,7 +152,7 @@ private OptimizationResult optimizeUsingCelBlock(CelAbstractSyntaxTree ast, Cel MANGLED_COMPREHENSION_IDENTIFIER_PREFIX, MANGLED_COMPREHENSION_RESULT_PREFIX); astToModify = mangledComprehensionAst.mutableAst(); - CelSource.Builder sourceToModify = astToModify.source(); + CelMutableSource sourceToModify = astToModify.source(); int blockIdentifierIndex = 0; int iterCount; @@ -177,9 +178,7 @@ private OptimizationResult optimizeUsingCelBlock(CelAbstractSyntaxTree ast, Cel navAst, CelNavigableMutableAst.fromAst( CelMutableAst.of( - CelMutableExpr.ofIdent(blockIdentifier), - CelSource.newBuilder() - .addAllMacroCalls(navAst.getAst().source().getMacroCalls()))), + CelMutableExpr.ofIdent(blockIdentifier), navAst.getAst().source())), cseCandidate.id()); // Retain the existing macro calls in case if the block identifiers are replacing a subtree @@ -354,7 +353,7 @@ private OptimizationResult optimizeUsingCelBind(CelAbstractSyntaxTree ast) { MANGLED_COMPREHENSION_IDENTIFIER_PREFIX, MANGLED_COMPREHENSION_RESULT_PREFIX) .mutableAst(); - CelSource.Builder sourceToModify = astToModify.source(); + CelMutableSource sourceToModify = astToModify.source(); int bindIdentifierIndex = 0; int iterCount; diff --git a/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java b/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java index 8947ad34a..b3c6313da 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java +++ b/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java @@ -25,6 +25,7 @@ import dev.cel.bundle.CelFactory; import dev.cel.common.CelAbstractSyntaxTree; import dev.cel.common.CelFunctionDecl; +import dev.cel.common.CelMutableAst; import dev.cel.common.CelOptions; import dev.cel.common.CelOverloadDecl; import dev.cel.common.CelSource; @@ -33,7 +34,6 @@ import dev.cel.common.ast.CelConstant; import dev.cel.common.ast.CelExpr; import dev.cel.common.ast.CelExpr.ExprKind.Kind; -import dev.cel.common.ast.CelMutableAst; import dev.cel.common.ast.CelMutableExpr; import dev.cel.common.ast.CelMutableExpr.CelMutableCall; import dev.cel.common.ast.CelMutableExpr.CelMutableCreateList; diff --git a/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel b/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel index 4d751e4f9..83ccd8a3c 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel +++ b/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel @@ -11,9 +11,10 @@ java_library( "//bundle:cel", "//common", "//common:compiler_common", + "//common:mutable_ast", "//common:options", "//common/ast", - "//common/ast:mutable_ast", + "//common/ast:mutable_expr", "//common/navigation", "//common/navigation:mutable_navigation", "//common/resources/testdata/proto3:test_all_types_java_proto", diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/BUILD.bazel b/optimizer/src/test/java/dev/cel/optimizer/optimizers/BUILD.bazel index 17389a5a5..e1ac3bd52 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/BUILD.bazel +++ b/optimizer/src/test/java/dev/cel/optimizer/optimizers/BUILD.bazel @@ -12,9 +12,9 @@ java_library( "//bundle:cel", "//common", "//common:compiler_common", + "//common:mutable_ast", "//common:options", "//common/ast", - "//common/ast:mutable_ast", "//common/navigation:mutable_navigation", "//common/resources/testdata/proto3:test_all_types_java_proto", "//common/types", diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerTest.java index a174238d9..8d657cd0a 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerTest.java +++ b/optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerTest.java @@ -29,6 +29,7 @@ import dev.cel.bundle.CelFactory; import dev.cel.common.CelAbstractSyntaxTree; import dev.cel.common.CelFunctionDecl; +import dev.cel.common.CelMutableAst; import dev.cel.common.CelOptions; import dev.cel.common.CelOverloadDecl; import dev.cel.common.CelSource.Extension; @@ -39,7 +40,6 @@ import dev.cel.common.ast.CelConstant; import dev.cel.common.ast.CelExpr; import dev.cel.common.ast.CelExpr.ExprKind.Kind; -import dev.cel.common.ast.CelMutableAst; import dev.cel.common.navigation.CelNavigableMutableAst; import dev.cel.common.navigation.CelNavigableMutableExpr; import dev.cel.common.types.ListType;