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..941deba50 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; @@ -276,6 +281,38 @@ private static boolean areChildrenArgConstant(CelNavigableMutableExpr expr) { return false; } + /** + * Checks whether the provided identifier node is a loop variable bound by a surrounding + * comprehension. Identifiers existing strictly within loop initialization bounds (e.g. iterRange) + * are not considered loop variables. + */ + private static boolean isComprehensionBoundLoopVariable(CelNavigableMutableExpr identNode) { + String identName = identNode.expr().ident().name(); + CelNavigableMutableExpr child = identNode; + CelNavigableMutableExpr parent = identNode.parent().orElse(null); + + while (parent != null) { + if (parent.getKind().equals(Kind.COMPREHENSION)) { + CelMutableComprehension comprehension = parent.expr().comprehension(); + long childId = child.expr().id(); + + boolean isInitializationNode = + childId == comprehension.iterRange().id() || childId == comprehension.accuInit().id(); + + if (!isInitializationNode) { + if (identName.equals(comprehension.accuVar()) + || identName.equals(comprehension.iterVar()) + || identName.equals(comprehension.iterVar2())) { + return true; + } + } + } + child = parent; + parent = parent.parent().orElse(null); + } + return false; + } + private static boolean isNestedComprehension(CelNavigableMutableExpr expr) { Optional maybeParent = expr.parent(); while (maybeParent.isPresent()) { @@ -305,10 +342,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 +354,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 +375,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 +406,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 +439,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 +455,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(); @@ -733,6 +802,7 @@ private static Object evaluateExpr(Cel cel, CelNavigableMutableExpr navigableMut navigableMutableExpr .allNodes() .filter(node -> node.getKind().equals(Kind.IDENT)) + .filter(node -> !isComprehensionBoundLoopVariable(node)) .map(node -> node.expr().ident().name()) .map(CelAttributePattern::fromQualifiedIdentifier) .collect(toImmutableList()); 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..7f5a0e102 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,17 @@ 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: '[1, 2, 3].map(item, item * 2)', expected: '[2, 4, 6]'}") + @TestParameters("{source: '[1, 2, 3].filter(item, item > 1)', expected: '[2, 3]'}") + @TestParameters("{source: '[1, 2, 3].exists(item, item > 1)', expected: 'true'}") + @TestParameters("{source: '[1, 2, 3].all(item, item > 1)', expected: 'false'}") + @TestParameters( + "{source: '[1, 2, x].map(item, item * 2)', expected: '[1, 2, x].map(item, item * 2)'}") + @TestParameters( + "{source: '[(x - 1 > 3) ? (x - 1) : 5].exists(x, x - 1 > 3)', expected: '[(x - 1 > 3) ? (x -" + + " 1) : 5].exists(x, x - 1 > 3)'}") + @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 +482,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 +626,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 =