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 35476a792..b7fdbb33a 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel @@ -31,6 +31,7 @@ java_library( "//common/navigation:common", "//common/navigation:mutable_navigation", "//common/types", + "//common/types:type_providers", "//extensions:optional_library", "//optimizer:ast_optimizer", "//optimizer:mutable_ast", @@ -40,6 +41,7 @@ java_library( "//runtime:unknown_attributes", "@maven//:com_google_errorprone_error_prone_annotations", "@maven//:com_google_guava_guava", + "@maven//:com_google_protobuf_protobuf_java", ], ) 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 35d181905..680b369c4 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java @@ -23,6 +23,8 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.protobuf.Descriptors.FieldDescriptor; +import com.google.protobuf.Message; import dev.cel.bundle.Cel; import dev.cel.common.CelAbstractSyntaxTree; import dev.cel.common.CelMutableAst; @@ -42,7 +44,10 @@ import dev.cel.common.navigation.CelNavigableMutableAst; import dev.cel.common.navigation.CelNavigableMutableExpr; import dev.cel.common.navigation.TraversalOrder; +import dev.cel.common.types.CelType; +import dev.cel.common.types.CelTypeProvider; import dev.cel.common.types.SimpleType; +import dev.cel.common.types.StructType; import dev.cel.extensions.CelOptionalLibrary.Function; import dev.cel.optimizer.AstMutator; import dev.cel.optimizer.CelAstOptimizer; @@ -305,10 +310,11 @@ private Optional maybeFold( // ex2: optional.ofNonZeroValue(5) -> optional.of(5) if (result instanceof Optional) { Optional optResult = ((Optional) result); - return maybeRewriteOptional(optResult, mutableAst, node.expr()); + return maybeRewriteOptional(cel.getTypeProvider(), optResult, mutableAst, node.expr()); } - CelMutableExpr adaptedResult = maybeAdaptEvaluatedResult(result).orElse(null); + CelMutableExpr adaptedResult = + maybeAdaptEvaluatedResult(cel.getTypeProvider(), result).orElse(null); if (adaptedResult == null) { return Optional.empty(); } @@ -316,14 +322,16 @@ private Optional maybeFold( return Optional.of(astMutator.replaceSubtree(mutableAst, adaptedResult, node.id())); } - private Optional maybeAdaptEvaluatedResult(Object result) { + private Optional maybeAdaptEvaluatedResult( + CelTypeProvider typeProvider, Object result) { if (CelConstant.isConstantValue(result)) { return Optional.of(CelMutableExpr.ofConstant(CelConstant.ofObjectValue(result))); } else if (result instanceof Collection) { Collection collection = (Collection) result; List listElements = new ArrayList<>(); for (Object evaluatedElement : collection) { - CelMutableExpr adaptedExpr = maybeAdaptEvaluatedResult(evaluatedElement).orElse(null); + CelMutableExpr adaptedExpr = + maybeAdaptEvaluatedResult(typeProvider, evaluatedElement).orElse(null); if (adaptedExpr == null) { return Optional.empty(); } @@ -335,11 +343,13 @@ private Optional maybeAdaptEvaluatedResult(Object result) { Map map = (Map) result; List mapEntries = new ArrayList<>(); for (Map.Entry entry : map.entrySet()) { - CelMutableExpr adaptedKey = maybeAdaptEvaluatedResult(entry.getKey()).orElse(null); + CelMutableExpr adaptedKey = + maybeAdaptEvaluatedResult(typeProvider, entry.getKey()).orElse(null); if (adaptedKey == null) { return Optional.empty(); } - CelMutableExpr adaptedValue = maybeAdaptEvaluatedResult(entry.getValue()).orElse(null); + CelMutableExpr adaptedValue = + maybeAdaptEvaluatedResult(typeProvider, entry.getValue()).orElse(null); if (adaptedValue == null) { return Optional.empty(); } @@ -364,6 +374,32 @@ private Optional maybeAdaptEvaluatedResult(Object result) { CelMutableExpr.ofConstant(CelConstant.ofValue(timestampStrArg))); return Optional.of(CelMutableExpr.ofCall(timestampCall)); + } else if (result instanceof Message) { + Message message = (Message) result; + Optional typeOpt = + typeProvider.findType(message.getDescriptorForType().getFullName()); + if (!typeOpt.isPresent() || !(typeOpt.get() instanceof StructType)) { + return Optional.empty(); + } + StructType structType = (StructType) typeOpt.get(); + List structEntries = new ArrayList<>(); + for (String fieldName : structType.fieldNames()) { + FieldDescriptor fieldDescriptor = message.getDescriptorForType().findFieldByName(fieldName); + if (fieldDescriptor == null) { + continue; + } + if (!message.getAllFields().containsKey(fieldDescriptor)) { + continue; + } + CelMutableExpr adaptedFieldExpr = + maybeAdaptEvaluatedResult(typeProvider, message.getField(fieldDescriptor)).orElse(null); + if (adaptedFieldExpr == null) { + return Optional.empty(); + } + structEntries.add(CelMutableStruct.Entry.create(0, fieldName, adaptedFieldExpr)); + } + return Optional.of( + CelMutableExpr.ofStruct(CelMutableStruct.create(structType.name(), structEntries))); } // Evaluated result cannot be folded (e.g: unknowns) @@ -371,7 +407,7 @@ private Optional maybeAdaptEvaluatedResult(Object result) { } private Optional maybeRewriteOptional( - Optional optResult, CelMutableAst mutableAst, CelMutableExpr expr) { + CelTypeProvider typeProvider, Optional optResult, CelMutableAst mutableAst, CelMutableExpr expr) { Object unwrappedResult = optResult.orElse(null); if (unwrappedResult == null) { if (isCallToFunction(expr, Function.OPTIONAL_NONE.getFunction())) { @@ -387,7 +423,8 @@ private Optional maybeRewriteOptional( return Optional.empty(); } - CelMutableExpr adaptedResult = maybeAdaptEvaluatedResult(unwrappedResult).orElse(null); + CelMutableExpr adaptedResult = + maybeAdaptEvaluatedResult(typeProvider, unwrappedResult).orElse(null); if (adaptedResult == null) { // Evaluated result is not an adaptable constant. Leave the optional as is. return Optional.empty(); diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java index ee62c5ed1..4605517ea 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java +++ b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java @@ -195,6 +195,28 @@ private static Cel setupEnv(CelBuilder celBuilder) { @TestParameters( "{source: 'TestAllTypes{single_nested_message: TestAllTypes.NestedMessage{bb:" + " 42}}.single_nested_message.bb', expected: '42'}") + @TestParameters( + "{source: 'TestAllTypes{single_int32: 3}', " + + " expected: 'cel.expr.conformance.proto3.TestAllTypes{single_int32: 3}'}") + @TestParameters( + "{source: 'TestAllTypes{single_float: 1.5}', " + + " expected: 'cel.expr.conformance.proto3.TestAllTypes{single_float: 1.5}'}") + @TestParameters( + "{source: 'TestAllTypes{single_nested_message: TestAllTypes.NestedMessage{bb: 42}}', " + + " expected: 'cel.expr.conformance.proto3.TestAllTypes{single_nested_message:" + + " cel.expr.conformance.proto3.TestAllTypes.NestedMessage{bb: 42}}'}") + @TestParameters( + "{source: 'TestAllTypes{repeated_int32: [1, 2, 3]}', " + + " expected: 'cel.expr.conformance.proto3.TestAllTypes{repeated_int32: [1, 2, 3]}'}") + @TestParameters( + "{source: 'TestAllTypes{map_int32_int64: {1: 2}}', " + + " expected: 'cel.expr.conformance.proto3.TestAllTypes{map_int32_int64: {1: 2}}'}") + @TestParameters( + "{source: 'TestAllTypes{single_any: google.protobuf.Any{type_url:" + + " \"type.googleapis.com/google.protobuf.Int32Value\", value: b\"\\010\\001\"}}', " + + " expected: 'cel.expr.conformance.proto3.TestAllTypes{single_any:" + + " google.protobuf.Any{type_url:" + + " \"type.googleapis.com/google.protobuf.Int32Value\", value: b\"\\010\\001\"}}'}") @TestParameters("{source: '{\"a\": 1}[\"a\"]', expected: '1'}") @TestParameters("{source: '{\"a\": {\"b\": 2}}[\"a\"][\"b\"]', expected: '2'}") @TestParameters("{source: '{\"hello\": \"world\"}.hello == x', expected: '\"world\" == x'}") @@ -228,6 +250,8 @@ private static Cel setupEnv(CelBuilder celBuilder) { @TestParameters("{source: 'true == false', expected: 'false'}") @TestParameters("{source: 'true == true', expected: 'true'}") @TestParameters("{source: 'false == true', expected: 'false'}") + @TestParameters("{source: 'TestAllTypes{single_int64: 1 + 2 + 3}.single_int64', expected: '6'}") + @TestParameters("{source: 'TestAllTypes{single_int64: 3}.single_int64', expected: '3'}") @TestParameters("{source: 'false == false', expected: 'true'}") @TestParameters("{source: '10 == 42', expected: 'false'}") @TestParameters("{source: '42 == 42', expected: 'true'}") @@ -449,6 +473,38 @@ public void constantFold_addFoldableFunction_success() throws Exception { assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo("true"); } + @Test + public void constantFold_functionReturningUnregisteredMessage_doesNotFold() throws Exception { + Cel customCel = + runtimeFlavor + .builder() + .addVar("x", SimpleType.DYN) + .addFunctionDeclarations( + CelFunctionDecl.newFunctionDeclaration( + "get_unregistered_message", + CelOverloadDecl.newGlobalOverload( + "get_unregistered_message_overload", SimpleType.ANY))) + .addFunctionBindings( + CelFunctionBinding.from( + "get_unregistered_message_overload", + ImmutableList.of(), + unused -> TestAllTypes.getDefaultInstance())) + .build(); + ConstantFoldingOptions options = + ConstantFoldingOptions.newBuilder() + .addFoldableFunctions("get_unregistered_message") + .build(); + CelOptimizer optimizer = + CelOptimizerFactory.standardCelOptimizerBuilder(customCel) + .addAstOptimizers(ConstantFoldingOptimizer.newInstance(options)) + .build(); + CelAbstractSyntaxTree ast = customCel.compile("get_unregistered_message()").getAst(); + + CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast); + + assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo("get_unregistered_message()"); + } + @Test public void constantFold_withExpectedResultTypeSet_success() throws Exception { Cel cel = runtimeFlavor.builder().setResultType(SimpleType.STRING).build(); @@ -561,10 +617,7 @@ public void iterationLimitReached_throws() throws Exception { Cel cel = runtimeFlavor .builder() - .setOptions( - CelOptions.current() - .enableHeterogeneousNumericComparisons(true) - .build()) + .setOptions(CelOptions.current().enableHeterogeneousNumericComparisons(true).build()) .build(); CelAbstractSyntaxTree ast = cel.compile("1 + 1").getAst(); CelOptimizer optimizer =