diff --git a/optimizer/optimizers/BUILD.bazel b/optimizer/optimizers/BUILD.bazel index 26d98c574..e95d48728 100644 --- a/optimizer/optimizers/BUILD.bazel +++ b/optimizer/optimizers/BUILD.bazel @@ -19,3 +19,8 @@ java_library( name = "inlining", exports = ["//optimizer/src/main/java/dev/cel/optimizer/optimizers:inlining"], ) + +java_library( + name = "select_optimizer", + exports = ["//optimizer/src/main/java/dev/cel/optimizer/optimizers:select_optimizer"], +) 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 0e6509c44..e5201994a 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel @@ -111,6 +111,40 @@ java_library( ], ) +java_library( + name = "select_optimizer", + srcs = [ + "SelectOptimizer.java", + ], + tags = [ + ], + deps = [ + "//:auto_value", + "//bundle:cel", + "//common:cel_ast", + "//common:cel_descriptor_util", + "//common:cel_descriptors", + "//common:cel_source", + "//common:compiler_common", + "//common:mutable_ast", + "//common/ast", + "//common/ast:mutable_expr", + "//common/internal:cel_descriptor_pools", + "//common/navigation:common", + "//common/navigation:mutable_navigation", + "//common/types", + "//common/types:type_providers", + "//common/values", + "//common/values:cel_byte_string", + "//optimizer:ast_optimizer", + "//optimizer:mutable_ast", + "@maven//:com_google_errorprone_error_prone_annotations", + "@maven//:com_google_guava_guava", + "@maven//:com_google_protobuf_protobuf_java", + "@maven//:org_jspecify_jspecify", + ], +) + java_library( name = "default_optimizer_constants", srcs = [ diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java new file mode 100644 index 000000000..0b0dec4ff --- /dev/null +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java @@ -0,0 +1,477 @@ +// Copyright 2026 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.optimizer.optimizers; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.collect.ImmutableList.toImmutableList; + +import com.google.auto.value.AutoValue; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableList; +import com.google.common.primitives.UnsignedLong; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.protobuf.ByteString; +import com.google.protobuf.Descriptors.EnumValueDescriptor; +import com.google.protobuf.Descriptors.FieldDescriptor; +import com.google.protobuf.Descriptors.FileDescriptor; +import dev.cel.bundle.Cel; +import dev.cel.common.CelAbstractSyntaxTree; +import dev.cel.common.CelDescriptorUtil; +import dev.cel.common.CelDescriptors; +import dev.cel.common.CelFunctionDecl; +import dev.cel.common.CelMutableAst; +import dev.cel.common.CelOverloadDecl; +import dev.cel.common.CelSource; +import dev.cel.common.CelSource.Extension; +import dev.cel.common.CelSource.Extension.Component; +import dev.cel.common.CelSource.Extension.Version; +import dev.cel.common.ast.CelConstant; +import dev.cel.common.ast.CelExpr.ExprKind.Kind; +import dev.cel.common.ast.CelMutableExpr; +import dev.cel.common.ast.CelMutableExpr.CelMutableCall; +import dev.cel.common.ast.CelMutableExpr.CelMutableList; +import dev.cel.common.ast.CelMutableExpr.CelMutableMap; +import dev.cel.common.ast.CelMutableExpr.CelMutableSelect; +import dev.cel.common.internal.CelDescriptorPool; +import dev.cel.common.internal.CombinedDescriptorPool; +import dev.cel.common.internal.DefaultDescriptorPool; +// CEL-Internal-1 +import dev.cel.common.navigation.CelNavigableMutableAst; +import dev.cel.common.navigation.CelNavigableMutableExpr; +import dev.cel.common.navigation.TraversalOrder; +import dev.cel.common.types.CelKind; +import dev.cel.common.types.ListType; +import dev.cel.common.types.SimpleType; +import dev.cel.common.values.CelByteString; +import dev.cel.common.values.NullValue; +import dev.cel.optimizer.AstMutator; +import dev.cel.optimizer.CelAstOptimizer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicLong; + +/** + * Performs field selection optimization on protobuf message select chains. + * + *

Embeds protobuf field metadata (field number, field name, type code, default value) directly + * into qualification paths ({@code cel.@attribute} and {@code cel.@hasField}). This accelerates + * nested field evaluation, enables reflection-free field traversal in resource-constrained runtimes + * without descriptor tables, and provides resilience against protobuf field renames. + * + *

Trade-off: Modestly increases serialized AST size over the wire due to the embedded + * metadata tuples. + * + *

AST Rewriting Semantics

+ * + * + * + *

Map indexing and non-protobuf selects pass through untouched. + */ +public final class SelectOptimizer implements CelAstOptimizer { + + private static final String CEL_ATTRIBUTE_FUNCTION_NAME = "cel.@attribute"; + private static final String CEL_HAS_FIELD_FUNCTION_NAME = "cel.@hasField"; + + @VisibleForTesting + static final CelFunctionDecl CEL_ATTRIBUTE_FUNCTION_DECL = + CelFunctionDecl.newFunctionDeclaration( + CEL_ATTRIBUTE_FUNCTION_NAME, + CelOverloadDecl.newGlobalOverload( + "cel_attribute_list", + SimpleType.DYN, + SimpleType.DYN, + ListType.create(SimpleType.DYN))); + + @VisibleForTesting + static final CelFunctionDecl CEL_HAS_FIELD_FUNCTION_DECL = + CelFunctionDecl.newFunctionDeclaration( + CEL_HAS_FIELD_FUNCTION_NAME, + CelOverloadDecl.newGlobalOverload( + "cel_has_field_list", + SimpleType.BOOL, + SimpleType.DYN, + ListType.create(SimpleType.DYN))); + + @VisibleForTesting + static final Extension SELECT_OPTIMIZATION_AST_EXTENSION_TAG = + Extension.create("select_optimization", Version.of(1L, 0L), Component.COMPONENT_RUNTIME); + + private static final SelectOptimizer INSTANCE = + new SelectOptimizer(SelectOptimizerOptions.newBuilder().build()); + + private final SelectOptimizerOptions options; + private final AstMutator astMutator; + + /** Returns a default instance of the select optimizer with preconfigured defaults. */ + public static SelectOptimizer getInstance() { + return INSTANCE; + } + + /** Returns a new select optimizer configured with the provided options. */ + public static SelectOptimizer newInstance(SelectOptimizerOptions options) { + return new SelectOptimizer(options); + } + + /** Returns a new select optimizer configured with the provided options and file descriptors. */ + public static SelectOptimizer newInstance( + SelectOptimizerOptions options, FileDescriptor... fileDescriptors) { + return newInstance(options, Arrays.asList(checkNotNull(fileDescriptors))); + } + + /** Returns a new select optimizer configured with the provided options and file descriptors. */ + public static SelectOptimizer newInstance( + SelectOptimizerOptions options, Iterable fileDescriptors) { + return new SelectOptimizer( + checkNotNull(options).toBuilder().addFileDescriptors(fileDescriptors).build()); + } + + private SelectOptimizer(SelectOptimizerOptions options) { + this.options = checkNotNull(options); + this.astMutator = AstMutator.newInstance(options.iterationLimit()); + } + + @Override + public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) { + checkArgument(ast.isChecked(), "AST must be type-checked."); + + CelMutableAst astToModify = CelMutableAst.fromCelAst(ast); + if (!options.populateMacroCalls()) { + astToModify.source().clearMacroCalls(); + } + + CelNavigableMutableAst navAst = CelNavigableMutableAst.fromAst(astToModify); + ImmutableList topOfChainSelects = + navAst + .getRoot() + .allNodes(TraversalOrder.POST_ORDER) + .filter(node -> isTopOfSelectChain(navAst, node)) + .collect(toImmutableList()); + + if (topOfChainSelects.isEmpty()) { + if (!options.populateMacroCalls() && !ast.getSource().getMacroCalls().isEmpty()) { + return OptimizationResult.create(astToModify.toParsedAst()); + } + return OptimizationResult.create(ast); + } + + long maxId = navAst.getRoot().allNodes().mapToLong(node -> node.expr().id()).max().orElse(0L); + AtomicLong idCounter = new AtomicLong(maxId); + + for (CelNavigableMutableExpr topNode : topOfChainSelects) { + rewriteSelectChain(astToModify, navAst, topNode, idCounter); + } + + astToModify = astMutator.renumberIdsConsecutively(astToModify); + CelAbstractSyntaxTree optimizedAst = tagAstExtension(astToModify.toParsedAst()); + + return OptimizationResult.create( + optimizedAst, + ImmutableList.of(), + ImmutableList.of(CEL_ATTRIBUTE_FUNCTION_DECL, CEL_HAS_FIELD_FUNCTION_DECL)); + } + + private void rewriteSelectChain( + CelMutableAst astToModify, + CelNavigableMutableAst navAst, + CelNavigableMutableExpr topNode, + AtomicLong idCounter) { + boolean isHasField = topNode.expr().select().testOnly(); + astToModify.source().getMacroCalls().remove(topNode.expr().id()); + + List fields = new ArrayList<>(); + FieldDescriptor topField = + getOptimizableField(navAst, topNode) + .orElseThrow( + () -> new IllegalStateException("Expected optimizable field on select node")); + fields.add(topField); + + CelMutableExpr currentExpr = topNode.expr().select().operand(); + while (currentExpr.getKind() == Kind.SELECT) { + CelMutableSelect select = currentExpr.select(); + FieldDescriptor field = getOptimizableFieldForExpr(navAst, select).orElse(null); + if (field == null) { + break; + } + fields.add(field); + currentExpr = select.operand(); + } + + Collections.reverse(fields); + + List qualifierLists = new ArrayList<>(fields.size()); + for (FieldDescriptor field : fields) { + if (isHasField) { + qualifierLists.add( + CelMutableExpr.ofList( + idCounter.incrementAndGet(), + CelMutableList.create( + CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue((long) field.getNumber())), + CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue(field.getName()))))); + } else { + CelMutableExpr defaultValue = resolveDefaultValue(field, idCounter); + qualifierLists.add( + CelMutableExpr.ofList( + idCounter.incrementAndGet(), + CelMutableList.create( + CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue((long) field.getNumber())), + CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue(field.getName())), + CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), + CelConstant.ofValue((long) field.getType().toProto().getNumber())), + defaultValue))); + } + } + + CelMutableExpr qualifiersExpr = + CelMutableExpr.ofList(idCounter.incrementAndGet(), CelMutableList.create(qualifierLists)); + String functionName = isHasField ? CEL_HAS_FIELD_FUNCTION_NAME : CEL_ATTRIBUTE_FUNCTION_NAME; + topNode.expr().setCall(CelMutableCall.create(functionName, currentExpr, qualifiersExpr)); + } + + private boolean isTopOfSelectChain(CelNavigableMutableAst navAst, CelNavigableMutableExpr node) { + return getOptimizableField(navAst, node).isPresent() + && !node.parent().flatMap(parent -> getOptimizableField(navAst, parent)).isPresent(); + } + + private Optional getOptimizableField( + CelNavigableMutableAst navAst, CelNavigableMutableExpr node) { + if (node.getKind() != Kind.SELECT) { + return Optional.empty(); + } + return getOptimizableFieldForExpr(navAst, node.expr().select()); + } + + private Optional getOptimizableFieldForExpr( + CelNavigableMutableAst navAst, CelMutableSelect select) { + return navAst + .getType(select.operand().id()) + .filter(type -> type.kind() == CelKind.STRUCT) + .flatMap(type -> options.descriptorPool().findDescriptor(type.name())) + .map(desc -> desc.findFieldByName(select.field())); + } + + private static CelMutableExpr resolveDefaultValue(FieldDescriptor field, AtomicLong idCounter) { + if (field.isMapField()) { + return CelMutableExpr.ofMap( + idCounter.incrementAndGet(), CelMutableMap.create(ImmutableList.of())); + } + if (field.isRepeated()) { + return CelMutableExpr.ofList(idCounter.incrementAndGet(), CelMutableList.create()); + } + if (field.getType() == FieldDescriptor.Type.MESSAGE + || field.getType() == FieldDescriptor.Type.GROUP) { + String messageFullName = field.getMessageType().getFullName(); + switch (messageFullName) { + case "google.protobuf.Duration": + return CelMutableExpr.ofCall( + idCounter.incrementAndGet(), + CelMutableCall.create( + "duration", + CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue("0s")))); + case "google.protobuf.Timestamp": + return CelMutableExpr.ofCall( + idCounter.incrementAndGet(), + CelMutableCall.create( + "timestamp", + CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue("1970-01-01T00:00:00Z")))); + case "google.protobuf.Struct": + return CelMutableExpr.ofMap( + idCounter.incrementAndGet(), CelMutableMap.create(ImmutableList.of())); + case "google.protobuf.ListValue": + return CelMutableExpr.ofList(idCounter.incrementAndGet(), CelMutableList.create()); + default: + return CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue(NullValue.NULL_VALUE)); + } + } + + Object def = field.getDefaultValue(); + switch (field.getType()) { + case DOUBLE: + return CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue((Double) def)); + case FLOAT: + return CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue(((Float) def).doubleValue())); + case INT64: + case SINT64: + case SFIXED64: + return CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue((Long) def)); + case UINT64: + case FIXED64: + return CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), + CelConstant.ofValue(UnsignedLong.fromLongBits((Long) def))); + case INT32: + case SINT32: + case SFIXED32: + return CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue(((Integer) def).longValue())); + case UINT32: + case FIXED32: + return CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), + CelConstant.ofValue(UnsignedLong.fromLongBits(Integer.toUnsignedLong((Integer) def)))); + case BOOL: + return CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue((Boolean) def)); + case STRING: + return CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue((String) def)); + case BYTES: + ByteString byteString = (ByteString) def; + return CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), + byteString.isEmpty() + ? CelConstant.ofValue(CelByteString.EMPTY) + : CelConstant.ofValue(CelByteString.of(byteString.toByteArray()))); + case ENUM: + EnumValueDescriptor enumValue = (EnumValueDescriptor) def; + return CelMutableExpr.ofConstant( + idCounter.incrementAndGet(), CelConstant.ofValue((long) enumValue.getNumber())); + default: + throw new IllegalArgumentException("Unsupported protobuf field type: " + field.getType()); + } + } + + private static CelAbstractSyntaxTree tagAstExtension(CelAbstractSyntaxTree ast) { + CelSource.Builder celSourceBuilder = + ast.getSource().toBuilder().addAllExtensions(SELECT_OPTIMIZATION_AST_EXTENSION_TAG); + return CelAbstractSyntaxTree.newParsedAst(ast.getExpr(), celSourceBuilder.build()); + } + + /** Options configuring the behavior of {@link SelectOptimizer}. */ + @AutoValue + public abstract static class SelectOptimizerOptions { + + public abstract int iterationLimit(); + + public abstract boolean populateMacroCalls(); + + abstract CelDescriptorPool descriptorPool(); + + /** Builder for configuring {@link SelectOptimizerOptions}. */ + @AutoValue.Builder + public abstract static class Builder { + + public abstract Builder iterationLimit(int value); + + public abstract Builder populateMacroCalls(boolean value); + + abstract Builder descriptorPool(CelDescriptorPool descriptorPool); + + abstract Optional descriptorPool(); + + private final List fileDescriptors; + private boolean linkedMessageTypesEnabled; + + /** + * Sets whether to resolve compiled linked message types in the descriptor pool. + * + *

Note: This setting is only applied when the initial descriptor pool is constructed. It + * has no effect when configuring an options instance from {@link #toBuilder()} whose + * descriptor pool has already been initialized. + */ + @CanIgnoreReturnValue + public Builder enableLinkedMessageTypes(boolean enable) { + this.linkedMessageTypesEnabled = enable; + return this; + } + + /** Adds file descriptors to the descriptor pool. */ + @CanIgnoreReturnValue + public Builder addFileDescriptors(FileDescriptor... fileDescriptors) { + return addFileDescriptors(Arrays.asList(checkNotNull(fileDescriptors))); + } + + /** Adds file descriptors to the descriptor pool. */ + @CanIgnoreReturnValue + public Builder addFileDescriptors(Iterable fileDescriptors) { + checkNotNull(fileDescriptors); + for (FileDescriptor fileDescriptor : fileDescriptors) { + this.fileDescriptors.add(checkNotNull(fileDescriptor)); + } + return this; + } + + abstract SelectOptimizerOptions autoBuild(); + + public SelectOptimizerOptions build() { + CelDescriptorPool pool = + descriptorPool() + .map( + existingPool -> { + if (fileDescriptors.isEmpty()) { + return existingPool; + } + CelDescriptors descriptors = + CelDescriptorUtil.getAllDescriptorsFromFileDescriptor(fileDescriptors); + fileDescriptors.clear(); + return CombinedDescriptorPool.create( + ImmutableList.of( + DefaultDescriptorPool.create(descriptors), existingPool)); + }) + .orElseGet( + () -> { + ImmutableList.Builder pools = ImmutableList.builder(); + if (!fileDescriptors.isEmpty()) { + CelDescriptors descriptors = + CelDescriptorUtil.getAllDescriptorsFromFileDescriptor(fileDescriptors); + pools.add(DefaultDescriptorPool.create(descriptors)); + } + + pools.add(DefaultDescriptorPool.INSTANCE); + fileDescriptors.clear(); + return CombinedDescriptorPool.create(pools.build()); + }); + descriptorPool(pool); + return autoBuild(); + } + + Builder() { + this.fileDescriptors = new ArrayList<>(); + this.linkedMessageTypesEnabled = true; + } + } + + abstract Builder toBuilder(); + + /** Returns a new options builder with recommended defaults. */ + public static Builder newBuilder() { + return new AutoValue_SelectOptimizer_SelectOptimizerOptions.Builder() + .iterationLimit(500) + .populateMacroCalls(true); + } + + // Package-private constructor to prevent external extension, required by @AutoValue. + SelectOptimizerOptions() {} + } +} 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 c912d9570..787012466 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/BUILD.bazel +++ b/optimizer/src/test/java/dev/cel/optimizer/optimizers/BUILD.bazel @@ -18,6 +18,7 @@ java_library( "//common:container", "//common:mutable_ast", "//common:options", + "//common:proto_ast", "//common/ast", "//common/navigation:mutable_navigation", "//common/types", @@ -30,6 +31,7 @@ java_library( "//optimizer/optimizers:common_subexpression_elimination", "//optimizer/optimizers:constant_folding", "//optimizer/optimizers:inlining", + "//optimizer/optimizers:select_optimizer", "//parser:macro", "//parser:unparser", "//runtime", @@ -42,6 +44,8 @@ java_library( "@maven//:junit_junit", "@maven//:com_google_testparameterinjector_test_parameter_injector", "//:java_truth", + "@maven//:com_google_truth_extensions_truth_proto_extension", + "@cel_spec//proto/cel/expr:syntax_java_proto", "@cel_spec//proto/cel/expr/conformance/proto2:test_all_types_java_proto", "@cel_spec//proto/cel/expr/conformance/proto3:test_all_types_java_proto", "@maven//:com_google_guava_guava", diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java new file mode 100644 index 000000000..9e7015486 --- /dev/null +++ b/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java @@ -0,0 +1,765 @@ +// Copyright 2026 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.optimizer.optimizers; + +import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.truth.Truth.assertThat; +import static com.google.common.truth.extensions.proto.ProtoTruth.assertThat; +import static org.junit.Assert.assertThrows; + +import dev.cel.expr.ParsedExpr; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.Descriptors.FileDescriptor; +import com.google.protobuf.TextFormat; +import com.google.testing.junit.testparameterinjector.TestParameter; +import com.google.testing.junit.testparameterinjector.TestParameterInjector; +import dev.cel.bundle.Cel; +import dev.cel.bundle.CelBuilder; +import dev.cel.common.CelAbstractSyntaxTree; +import dev.cel.common.CelFunctionDecl; +import dev.cel.common.CelMutableAst; +import dev.cel.common.CelOptions; +import dev.cel.common.CelProtoAbstractSyntaxTree; +import dev.cel.common.CelValidationException; +import dev.cel.common.navigation.CelNavigableMutableAst; +import dev.cel.common.types.MapType; +import dev.cel.common.types.SimpleType; +import dev.cel.common.types.StructTypeReference; +import dev.cel.expr.conformance.proto2.NestedTestAllTypes; +import dev.cel.expr.conformance.proto2.TestAllTypesProto; +import dev.cel.expr.conformance.proto3.TestAllTypes; +import dev.cel.optimizer.CelOptimizer; +import dev.cel.optimizer.CelOptimizerFactory; +import dev.cel.optimizer.optimizers.SelectOptimizer.SelectOptimizerOptions; +import dev.cel.parser.CelStandardMacro; +import dev.cel.parser.CelUnparser; +import dev.cel.parser.CelUnparserFactory; +import dev.cel.runtime.CelFunctionBinding; +import dev.cel.testing.CelRuntimeFlavor; +import java.util.List; +import java.util.stream.LongStream; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; + +@RunWith(TestParameterInjector.class) +public final class SelectOptimizerTest { + + private static final CelOptions CEL_OPTIONS = + CelOptions.current() + .populateMacroCalls(true) + .enableHeterogeneousNumericComparisons(true) + .build(); + + private static final CelUnparser CEL_UNPARSER = CelUnparserFactory.newUnparser(); + + private static final Descriptor PROTO2_TEST_ALL_TYPES_DESCRIPTOR = + checkNotNull(TestAllTypesProto.getDescriptor().findMessageTypeByName("TestAllTypes")); + + @TestParameter CelRuntimeFlavor runtimeFlavor; + + private Cel cel; + private CelOptimizer celOptimizer; + + @Before + public void setUp() { + cel = setupEnv(runtimeFlavor.builder()); + celOptimizer = + CelOptimizerFactory.standardCelOptimizerBuilder(cel) + .addAstOptimizers( + SelectOptimizer.newInstance( + SelectOptimizerOptions.newBuilder().build(), + TestAllTypes.getDescriptor().getFile(), + PROTO2_TEST_ALL_TYPES_DESCRIPTOR.getFile(), + NestedTestAllTypes.getDescriptor().getFile())) + .build(); + } + + private static Cel setupEnv(CelBuilder celBuilder) { + return celBuilder + .setOptions(CEL_OPTIONS) + .setStandardMacros(CelStandardMacro.STANDARD_MACROS) + .addMessageTypes(TestAllTypes.getDescriptor()) + .addMessageTypes(PROTO2_TEST_ALL_TYPES_DESCRIPTOR) + .addMessageTypes(NestedTestAllTypes.getDescriptor()) + .addVar("msg", StructTypeReference.create(TestAllTypes.getDescriptor().getFullName())) + .addVar( + "proto2_msg", + StructTypeReference.create(PROTO2_TEST_ALL_TYPES_DESCRIPTOR.getFullName())) + .addVar( + "nested_msg", + StructTypeReference.create(NestedTestAllTypes.getDescriptor().getFullName())) + .addVar("map_var", MapType.create(SimpleType.STRING, SimpleType.INT)) + .addVar("x", SimpleType.INT) + .build(); + } + + private enum RewriteTestCase { + // === Selection & Traversal === + PROTO3_SINGLE_FIELD_SELECT( + "msg.single_int64", "cel.@attribute(msg, [[2, \"single_int64\", 3, 0]])"), + PROTO3_CHAINED_FIELD_SELECT( + "msg.single_nested_message.bb", + "cel.@attribute(msg, [[21, \"single_nested_message\", 11, null], [1, \"bb\", 5, 0]])"), + PROTO2_CHAINED_FIELD_SELECT( + "proto2_msg.single_nested_message.bb", + "cel.@attribute(proto2_msg, [[21, \"single_nested_message\", 11, null], [1, \"bb\", 5," + + " 0]])"), + PROTO2_TRIPLE_CHAINED_FIELD_SELECT( + "nested_msg.child.payload.single_int64", + "cel.@attribute(nested_msg, " + + "[[1, \"child\", 11, null], " + + "[2, \"payload\", 11, null], " + + "[2, \"single_int64\", 3, -64]])"), + PROTO2_GROUP_FIELD_SELECT( + "proto2_msg.nestedgroup.single_id", + "cel.@attribute(proto2_msg, [[403, \"nestedgroup\", 10, null], [404, \"single_id\", 5," + + " 0]])"), + + // === Presence Tests: Proto2 (Explicit Presence) vs Proto3 (Implicit/Explicit Presence) === + // In proto2, scalar fields have explicit presence (has-bit). + PROTO2_HAS_SCALAR_INT32( + "has(proto2_msg.single_int32)", "cel.@hasField(proto2_msg, [[1, \"single_int32\"]])"), + PROTO2_HAS_SCALAR_INT64( + "has(proto2_msg.single_int64)", "cel.@hasField(proto2_msg, [[2, \"single_int64\"]])"), + // In proto3, non-optional scalar fields have implicit presence (evaluated as != default). + PROTO3_HAS_SCALAR_INT32("has(msg.single_int32)", "cel.@hasField(msg, [[1, \"single_int32\"]])"), + PROTO3_HAS_SCALAR_INT64("has(msg.single_int64)", "cel.@hasField(msg, [[2, \"single_int64\"]])"), + // In proto3, explicit optional scalars have presence (has-bit). + PROTO3_HAS_OPTIONAL_BOOL( + "has(msg.optional_bool)", "cel.@hasField(msg, [[16, \"optional_bool\"]])"), + PROTO3_HAS_OPTIONAL_STRING( + "has(msg.optional_string)", "cel.@hasField(msg, [[17, \"optional_string\"]])"), + // Messages in both proto2 and proto3 have explicit presence. + PROTO2_HAS_MESSAGE( + "has(proto2_msg.single_nested_message)", + "cel.@hasField(proto2_msg, [[21, \"single_nested_message\"]])"), + PROTO3_HAS_MESSAGE( + "has(msg.single_nested_message)", "cel.@hasField(msg, [[21, \"single_nested_message\"]])"), + PROTO3_HAS_STANDALONE_MESSAGE( + "has(msg.standalone_message)", "cel.@hasField(msg, [[23, \"standalone_message\"]])"), + PROTO3_HAS_ONEOF_ENUM( + "has(msg.single_nested_enum)", "cel.@hasField(msg, [[22, \"single_nested_enum\"]])"), + PROTO2_HAS_CHAINED_MESSAGE( + "has(proto2_msg.single_nested_message.bb)", + "cel.@hasField(proto2_msg, [[21, \"single_nested_message\"], [1, \"bb\"]])"), + PROTO3_HAS_CHAINED_MESSAGE( + "has(msg.single_nested_message.bb)", + "cel.@hasField(msg, [[21, \"single_nested_message\"], [1, \"bb\"]])"), + PROTO2_HAS_TRIPLE_CHAINED_MESSAGE( + "has(nested_msg.child.payload.single_int64)", + "cel.@hasField(nested_msg, [[1, \"child\"], [2, \"payload\"], [2, \"single_int64\"]])"), + + // === Default Value Divergence: Proto2 Custom Defaults vs Proto3 Zero Defaults === + // Int32: proto2 has custom default -32, proto3 has 0 + PROTO2_CUSTOM_INT32( + "proto2_msg.single_int32", "cel.@attribute(proto2_msg, [[1, \"single_int32\", 5, -32]])"), + PROTO3_ZERO_INT32("msg.single_int32", "cel.@attribute(msg, [[1, \"single_int32\", 5, 0]])"), + + // Int64: proto2 has custom default -64, proto3 has 0 + PROTO2_CUSTOM_INT64( + "proto2_msg.single_int64", "cel.@attribute(proto2_msg, [[2, \"single_int64\", 3, -64]])"), + PROTO3_ZERO_INT64("msg.single_int64", "cel.@attribute(msg, [[2, \"single_int64\", 3, 0]])"), + + // Uint32: proto2 has custom default 32, proto3 has 0 + PROTO2_CUSTOM_UINT32( + "proto2_msg.single_uint32", + "cel.@attribute(proto2_msg, [[3, \"single_uint32\", 13, 32u]])"), + PROTO3_ZERO_UINT32( + "msg.single_uint32", "cel.@attribute(msg, [[3, \"single_uint32\", 13, 0u]])"), + + // Uint64: proto2 has custom default 64, proto3 has 0 + PROTO2_CUSTOM_UINT64( + "proto2_msg.single_uint64", "cel.@attribute(proto2_msg, [[4, \"single_uint64\", 4, 64u]])"), + PROTO3_ZERO_UINT64("msg.single_uint64", "cel.@attribute(msg, [[4, \"single_uint64\", 4, 0u]])"), + + // String: proto2 has custom default "empty", proto3 has "" + PROTO2_CUSTOM_STRING( + "proto2_msg.single_string", + "cel.@attribute(proto2_msg, [[14, \"single_string\", 9, \"empty\"]])"), + PROTO3_ZERO_STRING( + "msg.single_string", "cel.@attribute(msg, [[14, \"single_string\", 9, \"\"]])"), + + // Bool: proto2 has custom default true, proto3 has false + PROTO2_CUSTOM_BOOL( + "proto2_msg.single_bool", "cel.@attribute(proto2_msg, [[13, \"single_bool\", 8, true]])"), + PROTO3_ZERO_BOOL("msg.single_bool", "cel.@attribute(msg, [[13, \"single_bool\", 8, false]])"), + + // Float: proto2 has custom default 3.0, proto3 has 0.0 + PROTO2_CUSTOM_FLOAT( + "proto2_msg.single_float", "cel.@attribute(proto2_msg, [[11, \"single_float\", 2, 3.0]])"), + PROTO3_ZERO_FLOAT("msg.single_float", "cel.@attribute(msg, [[11, \"single_float\", 2, 0.0]])"), + + // Double: proto2 has custom default 6.4, proto3 has 0.0 + PROTO2_CUSTOM_DOUBLE( + "proto2_msg.single_double", + "cel.@attribute(proto2_msg, [[12, \"single_double\", 1, 6.4]])"), + PROTO3_ZERO_DOUBLE( + "msg.single_double", "cel.@attribute(msg, [[12, \"single_double\", 1, 0.0]])"), + + // Bytes: proto2 has custom default "none", proto3 has "" + PROTO2_CUSTOM_BYTES( + "proto2_msg.single_bytes", + "cel.@attribute(proto2_msg, [[15, \"single_bytes\", 12, b\"\\156\\157\\156\\145\"]])"), + PROTO3_ZERO_BYTES( + "msg.single_bytes", "cel.@attribute(msg, [[15, \"single_bytes\", 12, b\"\"]])"), + + // Enum: proto2 has custom default 1 (BAR), proto3 has 0 (FOO) + PROTO2_CUSTOM_ENUM( + "proto2_msg.single_nested_enum", + "cel.@attribute(proto2_msg, [[22, \"single_nested_enum\", 14, 1]])"), + PROTO3_ZERO_ENUM( + "msg.single_nested_enum", "cel.@attribute(msg, [[22, \"single_nested_enum\", 14, 0]])"), + + // Fixed / sfixed fields + PROTO3_SFIXED32( + "msg.single_sfixed32", "cel.@attribute(msg, [[9, \"single_sfixed32\", 15, 0]])"), + PROTO3_SFIXED64( + "msg.single_sfixed64", "cel.@attribute(msg, [[10, \"single_sfixed64\", 16, 0]])"), + + // Repeated fields: empty list default + PROTO2_REPEATED_PRIMITIVE( + "proto2_msg.repeated_int64", + "cel.@attribute(proto2_msg, [[32, \"repeated_int64\", 3, []]])"), + PROTO3_REPEATED_PRIMITIVE( + "msg.repeated_int64", "cel.@attribute(msg, [[32, \"repeated_int64\", 3, []]])"), + PROTO3_REPEATED_MESSAGE( + "msg.repeated_nested_message", + "cel.@attribute(msg, [[51, \"repeated_nested_message\", 11, []]])"), + + // Well-known types + PROTO3_TIMESTAMP( + "msg.single_timestamp", + "cel.@attribute(msg, [[102, \"single_timestamp\", 11," + + " timestamp(\"1970-01-01T00:00:00Z\")]])"), + PROTO3_DURATION( + "msg.single_duration", + "cel.@attribute(msg, [[101, \"single_duration\", 11, duration(\"0s\")]])"), + PROTO3_STRUCT("msg.single_struct", "cel.@attribute(msg, [[103, \"single_struct\", 11, {}]])"), + PROTO3_LIST_VALUE("msg.list_value", "cel.@attribute(msg, [[114, \"list_value\", 11, []]])"), + + // Map selects + MAP_FIELD_INDEXING( + "msg.map_int64_message[1].bb", + "cel.@attribute(" + + "cel.@attribute(msg, [[95, \"map_int64_message\", 11, {}]])[1], " + + "[[1, \"bb\", 5, 0]])"), + + // Mixed expressions + MIXED_BOOLEAN_EXPRESSION( + "msg.single_int64 > 0 && has(msg.single_nested_message)", + "cel.@attribute(msg, [[2, \"single_int64\", 3, 0]]) > 0 " + + "&& cel.@hasField(msg, [[21, \"single_nested_message\"]])"); + + private final String expression; + private final String expectedUnparsed; + + RewriteTestCase(String expression, String expectedUnparsed) { + this.expression = expression; + this.expectedUnparsed = expectedUnparsed; + } + } + + @Test + public void optimize_rewritesSelectExpressions(@TestParameter RewriteTestCase testCase) + throws Exception { + CelAbstractSyntaxTree ast = cel.compile(testCase.expression).getAst(); + + CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast); + + assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo(testCase.expectedUnparsed); + assertThat(optimizedAst.getSource().getExtensions()) + .contains(SelectOptimizer.SELECT_OPTIMIZATION_AST_EXTENSION_TAG); + } + + @Test + public void optimize_unoptimizableMapFieldSelect_leavesAstUntouched() throws Exception { + CelAbstractSyntaxTree ast = cel.compile("map_var.key").getAst(); + + CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast); + + assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo("map_var.key"); + assertThat(optimizedAst.getSource().getExtensions()) + .doesNotContain(SelectOptimizer.SELECT_OPTIMIZATION_AST_EXTENSION_TAG); + } + + @Test + public void optimize_unoptimizableMapHasField_leavesAstUntouched() throws Exception { + CelAbstractSyntaxTree ast = cel.compile("has(map_var.key)").getAst(); + + CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast); + + assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo("has(map_var.key)"); + assertThat(optimizedAst.getSource().getExtensions()) + .doesNotContain(SelectOptimizer.SELECT_OPTIMIZATION_AST_EXTENSION_TAG); + } + + @Test + public void optimize_noSelects_returnsOriginalAst() throws Exception { + CelAbstractSyntaxTree ast = cel.compile("1 + 2 == 3").getAst(); + + CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast); + + assertThat(optimizedAst).isEqualTo(ast); + assertThat(optimizedAst.getSource().getExtensions()) + .doesNotContain(SelectOptimizer.SELECT_OPTIMIZATION_AST_EXTENSION_TAG); + } + + @Test + public void optimize_notCheckedAst_throwsIllegalArgumentException() throws Exception { + CelAbstractSyntaxTree parsedAst = cel.parse("msg.single_int64").getAst(); + SelectOptimizer optimizer = SelectOptimizer.getInstance(); + + IllegalArgumentException exception = + assertThrows(IllegalArgumentException.class, () -> optimizer.optimize(parsedAst, cel)); + + assertThat(exception).hasMessageThat().contains("AST must be type-checked."); + } + + @Test + public void optimize_withFileDescriptors_success() throws Exception { + FileDescriptor fd = TestAllTypes.getDescriptor().getFile(); + SelectOptimizer customOptimizer = + SelectOptimizer.newInstance(SelectOptimizerOptions.newBuilder().build(), fd); + CelOptimizer optimizer = + CelOptimizerFactory.standardCelOptimizerBuilder(cel) + .addAstOptimizers(customOptimizer) + .build(); + CelAbstractSyntaxTree ast = cel.compile("msg.single_int64").getAst(); + + CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast); + + assertThat(CEL_UNPARSER.unparse(optimizedAst)) + .isEqualTo("cel.@attribute(msg, [[2, \"single_int64\", 3, 0]])"); + } + + @Test + public void optimize_withFileDescriptorsIterable_success() throws Exception { + FileDescriptor fd = TestAllTypes.getDescriptor().getFile(); + SelectOptimizer customOptimizer = + SelectOptimizer.newInstance( + SelectOptimizerOptions.newBuilder().build(), ImmutableList.of(fd)); + CelOptimizer optimizer = + CelOptimizerFactory.standardCelOptimizerBuilder(cel) + .addAstOptimizers(customOptimizer) + .build(); + CelAbstractSyntaxTree ast = cel.compile("msg.single_int64").getAst(); + + CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast); + + assertThat(CEL_UNPARSER.unparse(optimizedAst)) + .isEqualTo("cel.@attribute(msg, [[2, \"single_int64\", 3, 0]])"); + } + + @Test + public void newInstance_withOptionsAndFileDescriptors_preservesAddedDescriptors() + throws Exception { + FileDescriptor fd = PROTO2_TEST_ALL_TYPES_DESCRIPTOR.getFile(); + SelectOptimizerOptions baseOptions = + SelectOptimizerOptions.newBuilder().enableLinkedMessageTypes(false).build(); + SelectOptimizer optimizer = SelectOptimizer.newInstance(baseOptions, fd); + CelAbstractSyntaxTree ast = cel.compile("proto2_msg.single_int64").getAst(); + + CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast, cel).optimizedAst(); + + assertThat(CEL_UNPARSER.unparse(optimizedAst)) + .isEqualTo("cel.@attribute(proto2_msg, [[2, \"single_int64\", 3, -64]])"); + } + + @Test + public void optimize_defaultOptions_populatesMacroCalls() throws Exception { + CelAbstractSyntaxTree ast = + cel.compile("[1].exists(x, x > 0) && msg.single_int64 > 0").getAst(); + + CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast); + + assertThat(optimizedAst.getSource().getMacroCalls()).isNotEmpty(); + } + + @Test + public void optimize_populateMacroCallsFalse_clearsMacroCalls() throws Exception { + SelectOptimizer optimizerWithoutMacroCalls = + SelectOptimizer.newInstance( + SelectOptimizerOptions.newBuilder().populateMacroCalls(false).build(), + TestAllTypes.getDescriptor().getFile()); + CelOptimizer optimizer = + CelOptimizerFactory.standardCelOptimizerBuilder(cel) + .addAstOptimizers(optimizerWithoutMacroCalls) + .build(); + CelAbstractSyntaxTree ast = + cel.compile("[1].exists(x, x > 0) && msg.single_int64 > 0").getAst(); + + CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast); + + assertThat(optimizedAst.getSource().getMacroCalls()).isEmpty(); + } + + @Test + public void optimize_populateMacroCallsFalse_withoutSelects_clearsMacroCalls() throws Exception { + SelectOptimizer optimizerWithoutMacroCalls = + SelectOptimizer.newInstance( + SelectOptimizerOptions.newBuilder().populateMacroCalls(false).build()); + CelOptimizer optimizer = + CelOptimizerFactory.standardCelOptimizerBuilder(cel) + .addAstOptimizers(optimizerWithoutMacroCalls) + .build(); + CelAbstractSyntaxTree ast = cel.compile("[1].exists(x, x > 0)").getAst(); + + CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast); + + assertThat(optimizedAst.getSource().getMacroCalls()).isEmpty(); + } + + @Test + public void optimize_hasFieldMacroCall_removesHasMacroCallFromSource() throws Exception { + CelAbstractSyntaxTree ast = cel.compile("has(msg.single_int64)").getAst(); + + CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast); + + assertThat(optimizedAst.getSource().getMacroCalls()).isEmpty(); + } + + @Test + public void optimize_renumbersIdsConsecutively() throws Exception { + CelAbstractSyntaxTree ast = cel.compile("msg.single_nested_message.bb").getAst(); + + CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast); + + CelNavigableMutableAst navAst = + CelNavigableMutableAst.fromAst(CelMutableAst.fromCelAst(optimizedAst)); + ImmutableList ids = + navAst + .getRoot() + .allNodes() + .map(node -> node.expr().id()) + .sorted() + .collect(toImmutableList()); + ImmutableList expectedIds = + LongStream.rangeClosed(1, ids.size()).boxed().collect(toImmutableList()); + assertThat(ids).containsExactlyElementsIn(expectedIds).inOrder(); + } + + @Test + public void optimizeAndEvaluate_withAttributeFunctionBinding_evaluatesSuccessfully() + throws Exception { + Cel celWithBinding = + cel.toCelBuilder() + .addFunctionDeclarations(SelectOptimizer.CEL_ATTRIBUTE_FUNCTION_DECL) + .addFunctionBindings( + CelFunctionBinding.from( + "cel_attribute_list", Object.class, List.class, (target, path) -> 42L)) + .build(); + CelOptimizer optimizer = + CelOptimizerFactory.standardCelOptimizerBuilder(celWithBinding) + .addAstOptimizers( + SelectOptimizer.newInstance( + SelectOptimizerOptions.newBuilder().build(), + TestAllTypes.getDescriptor().getFile())) + .build(); + CelAbstractSyntaxTree ast = celWithBinding.compile("msg.single_int64").getAst(); + + CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast); + Object result = + celWithBinding + .createProgram(optimizedAst) + .eval(ImmutableMap.of("msg", TestAllTypes.getDefaultInstance())); + + assertThat(result).isEqualTo(42L); + } + + @Test + public void optimizeAndEvaluate_withHasFieldFunctionBinding_evaluatesSuccessfully() + throws Exception { + Cel celWithBinding = + cel.toCelBuilder() + .addFunctionDeclarations(SelectOptimizer.CEL_HAS_FIELD_FUNCTION_DECL) + .addFunctionBindings( + CelFunctionBinding.from( + "cel_has_field_list", Object.class, List.class, (target, path) -> true)) + .build(); + CelOptimizer optimizer = + CelOptimizerFactory.standardCelOptimizerBuilder(celWithBinding) + .addAstOptimizers( + SelectOptimizer.newInstance( + SelectOptimizerOptions.newBuilder().build(), + TestAllTypes.getDescriptor().getFile())) + .build(); + CelAbstractSyntaxTree ast = celWithBinding.compile("has(msg.single_int64)").getAst(); + + CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast); + Object result = + celWithBinding + .createProgram(optimizedAst) + .eval(ImmutableMap.of("msg", TestAllTypes.getDefaultInstance())); + + assertThat((Boolean) result).isTrue(); + } + + @Test + public void optionsBuilder_toBuilderAddFileDescriptors_combinesPools() { + FileDescriptor fd1 = TestAllTypes.getDescriptor().getFile(); + FileDescriptor fd2 = PROTO2_TEST_ALL_TYPES_DESCRIPTOR.getFile(); + SelectOptimizerOptions baseOptions = + SelectOptimizerOptions.newBuilder() + .enableLinkedMessageTypes(false) + .addFileDescriptors(fd1) + .build(); + + SelectOptimizerOptions options = baseOptions.toBuilder().addFileDescriptors(fd2).build(); + + assertThat(options.descriptorPool().findDescriptor(TestAllTypes.getDescriptor().getFullName())) + .hasValue(TestAllTypes.getDescriptor()); + assertThat( + options.descriptorPool().findDescriptor(PROTO2_TEST_ALL_TYPES_DESCRIPTOR.getFullName())) + .hasValue(PROTO2_TEST_ALL_TYPES_DESCRIPTOR); + } + + @Test + public void optionsBuilder_buildMultipleTimes_isIdempotent() { + FileDescriptor fd = TestAllTypes.getDescriptor().getFile(); + SelectOptimizerOptions.Builder builder = + SelectOptimizerOptions.newBuilder().enableLinkedMessageTypes(false).addFileDescriptors(fd); + + SelectOptimizerOptions options1 = builder.build(); + SelectOptimizerOptions options2 = builder.build(); + + assertThat(options2.descriptorPool()).isSameInstanceAs(options1.descriptorPool()); + } + + @Test + public void optionsBuilder_toBuilderAddFileDescriptorsBuildMultipleTimes_isIdempotent() { + FileDescriptor fd1 = TestAllTypes.getDescriptor().getFile(); + FileDescriptor fd2 = PROTO2_TEST_ALL_TYPES_DESCRIPTOR.getFile(); + SelectOptimizerOptions.Builder builder = + SelectOptimizerOptions.newBuilder() + .enableLinkedMessageTypes(false) + .addFileDescriptors(fd1) + .build() + .toBuilder() + .addFileDescriptors(fd2); + + SelectOptimizerOptions options1 = builder.build(); + SelectOptimizerOptions options2 = builder.build(); + + assertThat(options2.descriptorPool()).isSameInstanceAs(options1.descriptorPool()); + } + + @Test + public void optionsBuilder_toBuilderWithoutFileDescriptors_preservesPool() { + FileDescriptor fd = TestAllTypes.getDescriptor().getFile(); + SelectOptimizerOptions baseOptions = + SelectOptimizerOptions.newBuilder() + .enableLinkedMessageTypes(false) + .addFileDescriptors(fd) + .build(); + + SelectOptimizerOptions options = baseOptions.toBuilder().iterationLimit(100).build(); + + assertThat(options.descriptorPool()).isSameInstanceAs(baseOptions.descriptorPool()); + assertThat(options.iterationLimit()).isEqualTo(100); + } + + @Test + public void optionsBuilder_withLinkedDescriptorsDisabled_containsWellKnownTypes() { + SelectOptimizerOptions options = + SelectOptimizerOptions.newBuilder().enableLinkedMessageTypes(false).build(); + + assertThat(options.descriptorPool().findDescriptor("google.protobuf.Timestamp")).isPresent(); + assertThat(options.descriptorPool().findDescriptor("google.protobuf.Duration")).isPresent(); + assertThat(options.descriptorPool().findDescriptor(TestAllTypes.getDescriptor().getFullName())) + .isEmpty(); + } + + @Test + public void optionsBuilder_addFileDescriptorsIterable_withLinkedDescriptorsDisabled_success() { + FileDescriptor fd = PROTO2_TEST_ALL_TYPES_DESCRIPTOR.getFile(); + + SelectOptimizerOptions options = + SelectOptimizerOptions.newBuilder() + .enableLinkedMessageTypes(false) + .addFileDescriptors(ImmutableList.of(fd)) + .build(); + + assertThat( + options.descriptorPool().findDescriptor(PROTO2_TEST_ALL_TYPES_DESCRIPTOR.getFullName())) + .hasValue(PROTO2_TEST_ALL_TYPES_DESCRIPTOR); + assertThat(options.descriptorPool().findDescriptor("google.protobuf.Timestamp")).isPresent(); + assertThat(options.descriptorPool().findDescriptor("google.protobuf.Duration")).isPresent(); + assertThat(options.descriptorPool().findDescriptor(TestAllTypes.getDescriptor().getFullName())) + .isEmpty(); + } + + @Test + public void optionsBuilder_addFileDescriptorsVarargs_withLinkedDescriptorsDisabled_success() { + FileDescriptor fd = PROTO2_TEST_ALL_TYPES_DESCRIPTOR.getFile(); + + SelectOptimizerOptions options = + SelectOptimizerOptions.newBuilder() + .enableLinkedMessageTypes(false) + .addFileDescriptors(fd) + .build(); + + assertThat( + options.descriptorPool().findDescriptor(PROTO2_TEST_ALL_TYPES_DESCRIPTOR.getFullName())) + .hasValue(PROTO2_TEST_ALL_TYPES_DESCRIPTOR); + assertThat(options.descriptorPool().findDescriptor("google.protobuf.Timestamp")).isPresent(); + assertThat(options.descriptorPool().findDescriptor("google.protobuf.Duration")).isPresent(); + assertThat(options.descriptorPool().findDescriptor(TestAllTypes.getDescriptor().getFullName())) + .isEmpty(); + } + + private enum CompilerRejectionTestCase { + ATTRIBUTE_AT_SIGN( + SelectOptimizer.CEL_ATTRIBUTE_FUNCTION_DECL, + "cel.@attribute(msg, [])", + "token recognition error at: '@'"), + ATTRIBUTE_OVERLOAD( + SelectOptimizer.CEL_ATTRIBUTE_FUNCTION_DECL, + "cel_attribute_list(msg, [])", + "undeclared reference to 'cel_attribute_list'"), + HAS_FIELD_AT_SIGN( + SelectOptimizer.CEL_HAS_FIELD_FUNCTION_DECL, + "cel.@hasField(msg, [])", + "token recognition error at: '@'"), + HAS_FIELD_OVERLOAD( + SelectOptimizer.CEL_HAS_FIELD_FUNCTION_DECL, + "cel_has_field_list(msg, [])", + "undeclared reference to 'cel_has_field_list'"); + + private final CelFunctionDecl functionDecl; + private final String expression; + private final String expectedErrorMessage; + + CompilerRejectionTestCase( + CelFunctionDecl functionDecl, String expression, String expectedErrorMessage) { + this.functionDecl = functionDecl; + this.expression = expression; + this.expectedErrorMessage = expectedErrorMessage; + } + } + + @Test + public void compile_sourceWithInternalFunctionCall_failsCompilation( + @TestParameter CompilerRejectionTestCase testCase) { + Cel celWithDecl = cel.toCelBuilder().addFunctionDeclarations(testCase.functionDecl).build(); + + CelValidationException e = + assertThrows( + CelValidationException.class, () -> celWithDecl.compile(testCase.expression).getAst()); + + assertThat(e).hasMessageThat().contains(testCase.expectedErrorMessage); + } + + @Test + public void optimize_toParsedExpr_matchesExpectedSerializedProto() throws Exception { + CelAbstractSyntaxTree ast = cel.compile("msg.single_nested_message.bb").getAst(); + ParsedExpr expectedParsedExpr = + TextFormat.parse( + "expr {\n" + + " id: 1\n" + + " call_expr {\n" + + " function: \"cel.@attribute\"\n" + + " args {\n" + + " id: 2\n" + + " ident_expr {\n" + + " name: \"msg\"\n" + + " }\n" + + " }\n" + + " args {\n" + + " id: 3\n" + + " list_expr {\n" + + " elements {\n" + + " id: 4\n" + + " list_expr {\n" + + " elements {\n" + + " id: 5\n" + + " const_expr {\n" + + " int64_value: 21\n" + + " }\n" + + " }\n" + + " elements {\n" + + " id: 6\n" + + " const_expr {\n" + + " string_value: \"single_nested_message\"\n" + + " }\n" + + " }\n" + + " elements {\n" + + " id: 7\n" + + " const_expr {\n" + + " int64_value: 11\n" + + " }\n" + + " }\n" + + " elements {\n" + + " id: 8\n" + + " const_expr {\n" + + " null_value: NULL_VALUE\n" + + " }\n" + + " }\n" + + " }\n" + + " }\n" + + " elements {\n" + + " id: 9\n" + + " list_expr {\n" + + " elements {\n" + + " id: 10\n" + + " const_expr {\n" + + " int64_value: 1\n" + + " }\n" + + " }\n" + + " elements {\n" + + " id: 11\n" + + " const_expr {\n" + + " string_value: \"bb\"\n" + + " }\n" + + " }\n" + + " elements {\n" + + " id: 12\n" + + " const_expr {\n" + + " int64_value: 5\n" + + " }\n" + + " }\n" + + " elements {\n" + + " id: 13\n" + + " const_expr {\n" + + " int64_value: 0\n" + + " }\n" + + " }\n" + + " }\n" + + " }\n" + + " }\n" + + " }\n" + + " }\n" + + "}\n" + + "source_info {\n" + + " location: \"\"\n" + + " extensions {\n" + + " id: \"select_optimization\"\n" + + " affected_components: COMPONENT_RUNTIME\n" + + " version {\n" + + " major: 1\n" + + " }\n" + + " }\n" + + "}\n", + ParsedExpr.class); + + CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast); + ParsedExpr parsedExpr = CelProtoAbstractSyntaxTree.fromCelAst(optimizedAst).toParsedExpr(); + + assertThat(parsedExpr).isEqualTo(expectedParsedExpr); + } +}