Skip to content

Commit 67851ab

Browse files
Cristina Borzacopybara-github
Cristina Borza
authored and
copybara-github
committed
Migrate CEL optimizer metrics collection to CelOptimizerListener.
PiperOrigin-RevId: 977757083
1 parent d374a0d commit 67851ab

4 files changed

Lines changed: 119 additions & 25 deletions

File tree

‎optimizer/src/main/java/dev/cel/optimizer/CelOptimizerImpl.java‎

Lines changed: 30 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
package dev.cel.optimizer;
1616

1717
import static com.google.common.base.Preconditions.checkNotNull;
18+
import static com.google.common.base.Preconditions.checkState;
1819

1920
import com.google.common.collect.ImmutableSet;
2021
import dev.cel.bundle.Cel;
@@ -44,6 +45,7 @@ final class CelOptimizerImpl implements CelOptimizer {
4445
}
4546

4647
@Override
48+
@SuppressWarnings("ReferenceEquality")
4749
public CelAbstractSyntaxTree optimize(CelAbstractSyntaxTree ast) throws CelOptimizationException {
4850
if (!ast.isChecked()) {
4951
throw new IllegalArgumentException("AST must be type-checked.");
@@ -64,16 +66,18 @@ public CelAbstractSyntaxTree optimize(CelAbstractSyntaxTree ast) throws CelOptim
6466

6567
OptimizationResult result = optimizer.optimize(optimizedAst, celOptimizerEnv);
6668

67-
if (!result.newFunctionDecls().isEmpty() || !result.newVarDecls().isEmpty()) {
68-
celOptimizerEnv =
69-
celOptimizerEnv
70-
.toCelBuilder()
71-
.addVarDeclarations(result.newVarDecls())
72-
.addFunctionDeclarations(result.newFunctionDecls())
73-
.build();
69+
if (result.optimizedAst() != optimizedAst) {
70+
if (!result.newFunctionDecls().isEmpty() || !result.newVarDecls().isEmpty()) {
71+
celOptimizerEnv =
72+
celOptimizerEnv
73+
.toCelBuilder()
74+
.addVarDeclarations(result.newVarDecls())
75+
.addFunctionDeclarations(result.newFunctionDecls())
76+
.build();
77+
}
78+
optimizedAst = celOptimizerEnv.check(result.optimizedAst()).getAst();
79+
assertAstIdCorrectness(optimizedAst);
7480
}
75-
optimizedAst = celOptimizerEnv.check(result.optimizedAst()).getAst();
76-
assertAstIdCorrectness(optimizedAst);
7781

7882
for (CelOptimizerListener listener : listeners) {
7983
listener.onPassEnd(optimizer, preAst, optimizedAst);
@@ -131,19 +135,26 @@ private static void assertAstIdCorrectness(CelAbstractSyntaxTree ast) {
131135
return;
132136
}
133137

134-
if (astExpr.exprKind().getKind().equals(Kind.COMPREHENSION)) {
135-
if (!macroExpr.exprKind().getKind().equals(Kind.NOT_SET)) {
136-
throw new IllegalStateException(
137-
String.format(
138-
"Expected macro call node %d to be NOT_SET for comprehension, but"
139-
+ " was %s.",
140-
macroExpr.id(), macroExpr.exprKind().getKind()));
141-
}
138+
if (macroExpr.exprKind().getKind().equals(Kind.NOT_SET)) {
139+
// If a macro node is NOT_SET, its ID must be present in the main AST.
140+
checkState(
141+
ast.getSource().getMacroCalls().containsKey(macroExpr.id()),
142+
"Expected macro call node %s to be present in macro calls map, but was not.",
143+
macroExpr.id());
144+
} else if (astExpr.exprKind().getKind().equals(Kind.COMPREHENSION)) {
145+
// We encountered something other than NOT_SET in macro source for comprehension
146+
// node. This is an error.
147+
throw new IllegalStateException(
148+
String.format(
149+
"Expected macro call node %d to be NOT_SET for comprehension, but was"
150+
+ " %s.",
151+
macroExpr.id(), macroExpr.exprKind().getKind()));
142152
} else if (!macroExpr.exprKind().getKind().equals(astExpr.exprKind().getKind())) {
153+
// Otherwise for all cases, the AST node should match exactly.
143154
throw new IllegalStateException(
144155
String.format(
145-
"Macro call node %d kind mismatch: expected %s (from AST), but was %s"
146-
+ " (in macro call).",
156+
"Macro call node %d kind mismatch: expected %s (from AST), but was %s (in"
157+
+ " macro call).",
147158
macroExpr.id(),
148159
astExpr.exprKind().getKind(),
149160
macroExpr.exprKind().getKind()));

‎optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java‎

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@ private static CelMutableExpr newOptionalNoneExpr() {
121121
}
122122

123123
@Override
124+
@SuppressWarnings("ReferenceEquality")
124125
public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
125126
throws CelOptimizationException {
126127
CelBuilder builder = cel.toCelBuilder();
@@ -134,12 +135,17 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
134135
// Override the environment's expected type to generally allow all subtrees to be folded.
135136
Cel optimizerEnv = builder.setResultType(SimpleType.DYN).build();
136137

137-
CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
138-
ImmutableMap<String, CelType> identTypes = precomputeIdentTypes(mutableAst);
138+
CelMutableAst initialMutableAst = CelMutableAst.fromCelAst(ast);
139+
ImmutableMap<String, CelType> identTypes = precomputeIdentTypes(initialMutableAst);
139140

140-
mutableAst = foldConstants(optimizerEnv, valueProvider, identTypes, mutableAst);
141+
CelMutableAst mutableAst =
142+
foldConstants(optimizerEnv, valueProvider, identTypes, initialMutableAst);
141143
mutableAst = pruneOptionalElements(mutableAst);
142144

145+
if (mutableAst == initialMutableAst) {
146+
return OptimizationResult.create(ast);
147+
}
148+
143149
return OptimizationResult.create(astMutator.renumberIdsConsecutively(mutableAst).toParsedAst());
144150
}
145151

@@ -735,10 +741,21 @@ private CelMutableAst pruneOptionalListElements(CelMutableAst mutableAst, CelMut
735741
updatedIndicesBuilder.add(newOptIndex);
736742
}
737743

744+
// An optional list is modified if:
745+
// 1. An optional.none() was dropped - it this case, the updatedElements.size() decreases.
746+
// 2. An optional.of(literal) was unwrapped into a regular element - in this case,
747+
// updatedIndices.size() decreases.
748+
// If both counts are unchanged, neither case occurred, and we can return the original AST.
749+
ImmutableList<CelMutableExpr> updatedElements = updatedElemBuilder.build();
750+
ImmutableList<Integer> updatedIndices = updatedIndicesBuilder.build();
751+
if (updatedElements.size() == list.elements().size()
752+
&& updatedIndices.size() == list.optionalIndices().size()) {
753+
return mutableAst;
754+
}
755+
738756
return astMutator.replaceSubtree(
739757
mutableAst,
740-
CelMutableExpr.ofList(
741-
CelMutableList.create(updatedElemBuilder.build(), updatedIndicesBuilder.build())),
758+
CelMutableExpr.ofList(CelMutableList.create(updatedElements, updatedIndices)),
742759
expr.id());
743760
}
744761

‎optimizer/src/main/java/dev/cel/optimizer/optimizers/InliningOptimizer.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,8 +101,10 @@ public static InliningOptimizer newInstance(
101101
}
102102

103103
@Override
104+
@SuppressWarnings("ReferenceEquality")
104105
public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) {
105-
CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
106+
CelMutableAst initialMutableAst = CelMutableAst.fromCelAst(ast);
107+
CelMutableAst mutableAst = initialMutableAst;
106108
for (InlineVariable inlineVariable : inlineVariables) {
107109
mutableAst =
108110
astMutator.mutateUntilFixedPoint(
@@ -125,6 +127,10 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) {
125127
});
126128
}
127129

130+
if (mutableAst == initialMutableAst) {
131+
return OptimizationResult.create(ast);
132+
}
133+
128134
return OptimizationResult.create(astMutator.renumberIdsConsecutively(mutableAst).toParsedAst());
129135
}
130136

‎optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerTest.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -702,6 +702,66 @@ public void block_lazyEvaluationContainsError_cleansUpCycleState() throws Except
702702
assertThat(e).hasMessageThat().doesNotContain("Cycle detected");
703703
}
704704

705+
@Test
706+
public void cse_nestedMacro_noOp_assertAstIdCorrectness() throws Exception {
707+
Cel cel =
708+
runtimeFlavor
709+
.builder()
710+
.addVar("x", SimpleType.DYN)
711+
.setStandardMacros(CelStandardMacro.STANDARD_MACROS)
712+
.setOptions(CelOptions.current().populateMacroCalls(true).build())
713+
.addCompilerLibraries(CelExtensions.comprehensions())
714+
.addRuntimeLibraries(CelExtensions.comprehensions())
715+
.build();
716+
CelOptimizer celOptimizer =
717+
CelOptimizerFactory.standardCelOptimizerBuilder(cel)
718+
.addAstOptimizers(SubexpressionOptimizer.getInstance())
719+
.build();
720+
CelAbstractSyntaxTree ast =
721+
cel.compile("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))").getAst();
722+
723+
CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
724+
725+
assertThat(CEL_UNPARSER.unparse(optimizedAst))
726+
.isEqualTo("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))");
727+
assertThat(optimizedAst).isSameInstanceAs(ast);
728+
}
729+
730+
@Test
731+
public void cse_nestedMacro_withOptimization_assertAstIdCorrectness() throws Exception {
732+
Cel cel =
733+
runtimeFlavor
734+
.builder()
735+
.addVar("x", SimpleType.DYN)
736+
.setStandardMacros(CelStandardMacro.STANDARD_MACROS)
737+
.setOptions(
738+
CelOptions.current()
739+
.populateMacroCalls(true)
740+
.enableHeterogeneousNumericComparisons(true)
741+
.build())
742+
.addCompilerLibraries(CelExtensions.comprehensions())
743+
.addRuntimeLibraries(CelExtensions.comprehensions())
744+
.build();
745+
CelOptimizer celOptimizer =
746+
CelOptimizerFactory.standardCelOptimizerBuilder(cel)
747+
.addAstOptimizers(
748+
SubexpressionOptimizer.newInstance(
749+
SubexpressionOptimizerOptions.newBuilder().populateMacroCalls(true).build()))
750+
.build();
751+
CelAbstractSyntaxTree ast =
752+
cel.compile(
753+
"[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a)) == [{}, {\"a\": 1}, {\"b\":"
754+
+ " 2}].filter(m, has(x.a))")
755+
.getAst();
756+
757+
CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
758+
759+
assertThat(CEL_UNPARSER.unparse(optimizedAst))
760+
.isEqualTo(
761+
"cel.@block([[{}, {\"a\": 1}, {\"b\": 2}].filter(@it:0:0, has(x.a))], @index0 =="
762+
+ " @index0)");
763+
}
764+
705765
/**
706766
* Converts AST containing cel.block related test functions to internal functions (e.g: cel.block
707767
* -> cel.@block)

0 commit comments

Comments
 (0)