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 694dfc04e..ee7ffeb7e 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel @@ -128,6 +128,7 @@ java_library( "//common:cel_source", "//common:compiler_common", "//common:mutable_ast", + "//common:operator", "//common/ast", "//common/ast:expr_factory", "//common/ast:mutable_expr", diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java index 5d15ad311..2ce52ac31 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java @@ -41,6 +41,7 @@ import dev.cel.common.CelSource.Extension; import dev.cel.common.CelSource.Extension.Component; import dev.cel.common.CelSource.Extension.Version; +import dev.cel.common.Operator; import dev.cel.common.ast.CelConstant; import dev.cel.common.ast.CelExpr.ExprKind.Kind; import dev.cel.common.ast.CelExprIdGeneratorFactory; @@ -59,8 +60,10 @@ import dev.cel.common.navigation.CelNavigableMutableExpr; import dev.cel.common.navigation.TraversalOrder; import dev.cel.common.types.CelKind; +import dev.cel.common.types.CelType; import dev.cel.common.types.CelTypes; import dev.cel.common.types.ListType; +import dev.cel.common.types.OptionalType; import dev.cel.common.types.SimpleType; import dev.cel.common.types.TypeParamType; import dev.cel.common.values.CelByteString; @@ -122,7 +125,7 @@ * *

Map indexing and non-protobuf selects pass through untouched. * - *

Rename Resilience & Dynamic Type Limitations: + *

Limitations: * *

*/ public final class SelectOptimizer implements CelAstOptimizer { @@ -213,6 +221,7 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) { CelMutableAst astToModify = CelMutableAst.fromCelAst(ast); CelNavigableMutableAst navAst = CelNavigableMutableAst.fromAst(astToModify); + navAst.getRoot().allNodes().forEach(node -> checkNotOptionalFieldSelection(navAst, node)); ImmutableList topOfChainSelects = navAst .getRoot() @@ -259,7 +268,6 @@ private void rewriteSelectChain( () -> new IllegalStateException("Expected optimizable field on select node")); fields.add(topField); - // TODO: Support optional field selection (_?._) once integrated with lite runtime. CelMutableExpr currentExpr = topNode.expr().select().operand(); while (currentExpr.getKind() == Kind.SELECT) { CelMutableSelect select = currentExpr.select(); @@ -357,11 +365,43 @@ private Optional getOptimizableField( private Optional getOptimizableFieldForExpr( CelNavigableMutableAst navAst, CelMutableSelect select) { - return navAst - .getType(select.operand().id()) - .filter(type -> type.kind() == CelKind.STRUCT) - .flatMap(type -> descriptorPool.findDescriptor(type.name())) - .map(desc -> desc.findFieldByName(select.field())); + return navAst.getType(select.operand().id()).flatMap(type -> findField(type, select.field())); + } + + private Optional findField(CelType type, String field) { + if (type.kind() != CelKind.STRUCT) { + return Optional.empty(); + } + return descriptorPool.findDescriptor(type.name()).map(desc -> desc.findFieldByName(field)); + } + + // TODO: Support optional field selection once integrated with lite runtime. + private void checkNotOptionalFieldSelection( + CelNavigableMutableAst navAst, CelNavigableMutableExpr node) { + Optional field = Optional.empty(); + if (node.getKind() == Kind.SELECT) { + // A field of an optional message, e.g. `.bb` in `msg.repeated_nested_message[?0].bb` + CelMutableSelect select = node.expr().select(); + field = + navAst + .getType(select.operand().id()) + .filter(OptionalType.class::isInstance) + .flatMap(type -> findField(type.parameters().get(0), select.field())); + } else if (node.getKind() == Kind.CALL + && node.expr().call().function().equals(Operator.OPTIONAL_SELECT.getFunction())) { + // `msg.?field`, where `msg` may itself be optional + List args = node.expr().call().args(); + field = + navAst + .getType(args.get(0).id()) + .map(type -> type instanceof OptionalType ? type.parameters().get(0) : type) + .flatMap(type -> findField(type, args.get(1).constant().stringValue())); + } + if (field.isPresent()) { + throw new UnsupportedOperationException( + "Optimization of optional field selection is currently unimplemented: " + + field.get().getFullName()); + } } private static CelMutableExpr resolveDummyOrDefExpr( diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java index 94567f167..4f85a426e 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java +++ b/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java @@ -367,6 +367,10 @@ private enum RewriteTestCase { + "cel.@attribute(msg, [[227, \"map_string_message\", -1, [9, 11]]]," + " {\"\": cel.expr.conformance.proto3.TestAllTypes.NestedMessage{}}).key, " + "[[1, \"bb\"]])"), + PROTO_MAP_FIELD_SELECT_STOPS_AT_ANY_BOUNDARY( + "msg.map_string_any['k'].value", + "cel.@attribute(msg, [[311, \"map_string_any\", -1, [9, 11]]]," + + " {\"\": google.protobuf.Any{}})[\"k\"].value"), MIXED_BOOLEAN_EXPRESSION( "msg.single_int64 > 0 && has(msg.single_nested_message)", @@ -1469,24 +1473,37 @@ public void optimize_comprehensionOverRepeatedMessageField_preservesStaticElemen } @Test - public void optimize_optionalSelect_passesThroughUntouched() throws Exception { - Cel celWithOptional = - cel.toCelBuilder() - .addCompilerLibraries(CelExtensions.optional()) - .addRuntimeLibraries(CelExtensions.optional()) - .build(); - CelOptimizer optimizer = - CelOptimizerFactory.standardCelOptimizerBuilder(celWithOptional) - .addAstOptimizers( - SelectOptimizer.newInstance( - SelectOptimizerOptions.newBuilder().build(), - TestAllTypes.getDescriptor().getFile())) - .build(); - CelAbstractSyntaxTree ast = celWithOptional.compile("msg.?single_nested_message.bb").getAst(); + public void optimize_optionalFieldSelection_throwsUnsupportedOperationException( + @TestParameter({ + "msg.?single_nested_message", + "msg.repeated_nested_message[?0].?bb", + "msg.repeated_nested_message[?0].bb" + }) + String expression) + throws Exception { + Cel celWithOptional = cel.toCelBuilder().addCompilerLibraries(CelExtensions.optional()).build(); + CelAbstractSyntaxTree ast = celWithOptional.compile(expression).getAst(); + SelectOptimizer optimizer = SelectOptimizer.newInstance(TestAllTypes.getDescriptor().getFile()); - CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast); + UnsupportedOperationException e = + assertThrows( + UnsupportedOperationException.class, () -> optimizer.optimize(ast, celWithOptional)); + + assertThat(e) + .hasMessageThat() + .contains("Optimization of optional field selection is currently unimplemented"); + } - assertThat(optimizedAst.getExpr()).isEqualTo(ast.getExpr()); + @Test + public void optimize_optionalMapKeySelection_rewritesMapField() throws Exception { + Cel celWithOptional = cel.toCelBuilder().addCompilerLibraries(CelExtensions.optional()).build(); + CelAbstractSyntaxTree ast = celWithOptional.compile("msg.map_string_string.?key").getAst(); + + CelAbstractSyntaxTree optimizedAst = newSelectOptimizer(celWithOptional).optimize(ast); + + assertThat(CEL_UNPARSER.unparse(optimizedAst)) + .isEqualTo( + "cel.@attribute(msg, [[61, \"map_string_string\", -1, [9, 9]]], {\"\": \"\"}).?key"); } private static CelOptimizer newSelectOptimizer(Cel cel) {