Skip to content

Commit 30f8e6d

Browse files
l46kokcopybara-github
l46kok
authored and
copybara-github
committed
Add a validation pass for ID uniqueness in optimizers
PiperOrigin-RevId: 963628515
1 parent 94ba8ad commit 30f8e6d

6 files changed

Lines changed: 383 additions & 8 deletions

File tree

‎optimizer/src/main/java/dev/cel/optimizer/BUILD.bazel‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,12 +32,14 @@ java_library(
3232
srcs = [
3333
"CelOptimizer.java",
3434
"CelOptimizerBuilder.java",
35+
"CelOptimizerOptions.java",
3536
],
3637
tags = [
3738
],
3839
deps = [
3940
":ast_optimizer",
4041
":optimization_exception",
42+
"//:auto_value",
4143
"//common:cel_ast",
4244
"@maven//:com_google_errorprone_error_prone_annotations",
4345
],
@@ -57,6 +59,8 @@ java_library(
5759
"//bundle:cel",
5860
"//common:cel_ast",
5961
"//common:compiler_common",
62+
"//common/ast",
63+
"//common/navigation",
6064
"@maven//:com_google_guava_guava",
6165
],
6266
)

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

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,22 +25,48 @@
2525
/** Factory class for constructing an {@link CelOptimizer} instance. */
2626
public final class CelOptimizerFactory {
2727

28+
private static final CelOptimizerOptions DEFAULT_OPTIMIZER_OPTIONS =
29+
CelOptimizerOptions.newBuilder().build();
30+
2831
/** Create a new builder for constructing a {@link CelOptimizer} instance. */
2932
public static CelOptimizerBuilder standardCelOptimizerBuilder(Cel cel) {
30-
return CelOptimizerImpl.newBuilder(cel);
33+
return standardCelOptimizerBuilder(cel, DEFAULT_OPTIMIZER_OPTIONS);
34+
}
35+
36+
/** Create a new builder for constructing a {@link CelOptimizer} instance with custom options. */
37+
public static CelOptimizerBuilder standardCelOptimizerBuilder(
38+
Cel cel, CelOptimizerOptions optimizerOptions) {
39+
return CelOptimizerImpl.newBuilder(cel, optimizerOptions);
3140
}
3241

3342
/** Create a new builder for constructing a {@link CelOptimizer} instance. */
3443
public static CelOptimizerBuilder standardCelOptimizerBuilder(
3544
CelCompiler celCompiler, CelRuntime celRuntime) {
36-
return standardCelOptimizerBuilder(CelFactory.combine(celCompiler, celRuntime));
45+
return standardCelOptimizerBuilder(celCompiler, celRuntime, DEFAULT_OPTIMIZER_OPTIONS);
46+
}
47+
48+
/** Create a new builder for constructing a {@link CelOptimizer} instance with custom options. */
49+
public static CelOptimizerBuilder standardCelOptimizerBuilder(
50+
CelCompiler celCompiler, CelRuntime celRuntime, CelOptimizerOptions optimizerOptions) {
51+
return standardCelOptimizerBuilder(
52+
CelFactory.combine(celCompiler, celRuntime), optimizerOptions);
3753
}
3854

3955
/** Create a new builder for constructing a {@link CelOptimizer} instance. */
4056
public static CelOptimizerBuilder standardCelOptimizerBuilder(
4157
CelParser celParser, CelChecker celChecker, CelRuntime celRuntime) {
4258
return standardCelOptimizerBuilder(
43-
CelCompilerFactory.combine(celParser, celChecker), celRuntime);
59+
celParser, celChecker, celRuntime, DEFAULT_OPTIMIZER_OPTIONS);
60+
}
61+
62+
/** Create a new builder for constructing a {@link CelOptimizer} instance with custom options. */
63+
public static CelOptimizerBuilder standardCelOptimizerBuilder(
64+
CelParser celParser,
65+
CelChecker celChecker,
66+
CelRuntime celRuntime,
67+
CelOptimizerOptions optimizerOptions) {
68+
return standardCelOptimizerBuilder(
69+
CelCompilerFactory.combine(celParser, celChecker), celRuntime, optimizerOptions);
4470
}
4571

4672
private CelOptimizerFactory() {}

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

Lines changed: 76 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,25 @@
2020
import dev.cel.bundle.Cel;
2121
import dev.cel.common.CelAbstractSyntaxTree;
2222
import dev.cel.common.CelValidationException;
23+
import dev.cel.common.ast.CelExpr;
24+
import dev.cel.common.ast.CelExpr.ExprKind.Kind;
25+
import dev.cel.common.navigation.CelNavigableAst;
26+
import dev.cel.common.navigation.CelNavigableExpr;
2327
import dev.cel.optimizer.CelAstOptimizer.OptimizationResult;
2428
import java.util.Arrays;
29+
import java.util.HashMap;
30+
import java.util.Map;
2531

2632
final class CelOptimizerImpl implements CelOptimizer {
2733
private final Cel cel;
2834
private final ImmutableSet<CelAstOptimizer> astOptimizers;
35+
private final CelOptimizerOptions optimizerOptions;
2936

30-
CelOptimizerImpl(Cel cel, ImmutableSet<CelAstOptimizer> astOptimizers) {
37+
CelOptimizerImpl(
38+
Cel cel, ImmutableSet<CelAstOptimizer> astOptimizers, CelOptimizerOptions optimizerOptions) {
3139
this.cel = cel;
3240
this.astOptimizers = astOptimizers;
41+
this.optimizerOptions = optimizerOptions;
3342
}
3443

3544
@Override
@@ -52,6 +61,9 @@ public CelAbstractSyntaxTree optimize(CelAbstractSyntaxTree ast) throws CelOptim
5261
.build();
5362
}
5463
optimizedAst = celOptimizerEnv.check(result.optimizedAst()).getAst();
64+
if (optimizerOptions.enableAstValidation()) {
65+
assertAstIdCorrectness(optimizedAst);
66+
}
5567
}
5668
} catch (CelValidationException e) {
5769
throw new CelOptimizationException(
@@ -63,18 +75,78 @@ public CelAbstractSyntaxTree optimize(CelAbstractSyntaxTree ast) throws CelOptim
6375
return optimizedAst;
6476
}
6577

78+
private static void assertAstIdCorrectness(CelAbstractSyntaxTree ast) {
79+
Map<Long, CelExpr> allExprs = new HashMap<>();
80+
CelNavigableAst.fromAst(ast)
81+
.getRoot()
82+
.allNodes()
83+
.forEach(
84+
navExpr -> {
85+
CelExpr expr = navExpr.expr();
86+
CelExpr existing = allExprs.put(expr.id(), expr);
87+
if (existing != null) {
88+
throw new IllegalStateException(
89+
String.format("Duplicate expr ID %d detected in the AST.", expr.id()));
90+
}
91+
});
92+
93+
for (CelExpr macroCall : ast.getSource().getMacroCalls().values()) {
94+
if (macroCall.id() != 0) {
95+
throw new IllegalStateException(
96+
String.format("Expected macro call root ID to be 0, but was %d.", macroCall.id()));
97+
}
98+
CelNavigableExpr.fromExpr(macroCall)
99+
.descendants()
100+
.forEach(
101+
navExpr -> {
102+
CelExpr macroExpr = navExpr.expr();
103+
CelExpr astExpr = allExprs.get(macroExpr.id());
104+
// A node may not exist in the AST if it is a synthetic macro node or was eliminated
105+
// during optimization passes.
106+
if (astExpr == null) {
107+
return;
108+
}
109+
110+
if (astExpr.exprKind().getKind().equals(Kind.COMPREHENSION)) {
111+
if (!macroExpr.exprKind().getKind().equals(Kind.NOT_SET)) {
112+
throw new IllegalStateException(
113+
String.format(
114+
"Expected macro call node %d to be NOT_SET for comprehension, but"
115+
+ " was %s.",
116+
macroExpr.id(), macroExpr.exprKind().getKind()));
117+
}
118+
} else if (!macroExpr.exprKind().getKind().equals(astExpr.exprKind().getKind())) {
119+
throw new IllegalStateException(
120+
String.format(
121+
"Macro call node %d kind mismatch: expected %s (from AST), but was %s"
122+
+ " (in macro call).",
123+
macroExpr.id(),
124+
astExpr.exprKind().getKind(),
125+
macroExpr.exprKind().getKind()));
126+
}
127+
});
128+
}
129+
}
130+
66131
/** Create a new builder for constructing a {@link CelOptimizer} instance. */
67132
static CelOptimizerImpl.Builder newBuilder(Cel cel) {
68-
return new CelOptimizerImpl.Builder(cel);
133+
return newBuilder(cel, CelOptimizerOptions.newBuilder().build());
134+
}
135+
136+
/** Create a new builder for constructing a {@link CelOptimizer} instance with custom options. */
137+
static CelOptimizerImpl.Builder newBuilder(Cel cel, CelOptimizerOptions optimizerOptions) {
138+
return new CelOptimizerImpl.Builder(cel, optimizerOptions);
69139
}
70140

71141
/** Builder class for {@link CelOptimizerImpl}. */
72142
static final class Builder implements CelOptimizerBuilder {
73143
private final Cel cel;
144+
private final CelOptimizerOptions optimizerOptions;
74145
private final ImmutableSet.Builder<CelAstOptimizer> astOptimizers;
75146

76-
private Builder(Cel cel) {
147+
private Builder(Cel cel, CelOptimizerOptions optimizerOptions) {
77148
this.cel = cel;
149+
this.optimizerOptions = checkNotNull(optimizerOptions);
78150
this.astOptimizers = ImmutableSet.builder();
79151
}
80152

@@ -93,7 +165,7 @@ public CelOptimizerBuilder addAstOptimizers(Iterable<CelAstOptimizer> astOptimiz
93165

94166
@Override
95167
public CelOptimizer build() {
96-
return new CelOptimizerImpl(cel, astOptimizers.build());
168+
return new CelOptimizerImpl(cel, astOptimizers.build(), optimizerOptions);
97169
}
98170
}
99171
}

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

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package dev.cel.optimizer;
16+
17+
import com.google.auto.value.AutoValue;
18+
19+
/** Options to configure how {@link CelOptimizer} behaves. */
20+
@AutoValue
21+
public abstract class CelOptimizerOptions {
22+
23+
/**
24+
* Returns true if AST validation is enabled. When enabled, each optimizer pass verifies AST
25+
* invariants (such as expression ID uniqueness and macro source consistency) after type-checking.
26+
*/
27+
public abstract boolean enableAstValidation();
28+
29+
/** Builder for configuring the {@link CelOptimizerOptions}. */
30+
@AutoValue.Builder
31+
public abstract static class Builder {
32+
33+
/**
34+
* Enables or disables post-pass AST validation. When enabled, each optimizer pass verifies that
35+
* expression IDs are unique and macro calls in the AST source are consistent with the
36+
* expression nodes.
37+
*/
38+
public abstract Builder enableAstValidation(boolean value);
39+
40+
public abstract CelOptimizerOptions build();
41+
42+
Builder() {}
43+
}
44+
45+
/** Returns a new options builder with recommended defaults pre-configured. */
46+
public static Builder newBuilder() {
47+
return new AutoValue_CelOptimizerOptions.Builder().enableAstValidation(false);
48+
}
49+
50+
CelOptimizerOptions() {}
51+
}

‎optimizer/src/test/java/dev/cel/optimizer/CelOptimizerFactoryTest.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,19 @@ public void standardCelOptimizerBuilder_withParserCheckerAndRuntime() {
3939
assertThat(builder.build()).isNotNull();
4040
}
4141

42+
@Test
43+
public void standardCelOptimizerBuilder_withParserCheckerRuntimeAndOptions() {
44+
CelOptimizerBuilder builder =
45+
CelOptimizerFactory.standardCelOptimizerBuilder(
46+
CelParserFactory.standardCelParserBuilder().build(),
47+
CelCompilerFactory.standardCelCheckerBuilder().build(),
48+
CelRuntimeFactory.standardCelRuntimeBuilder().build(),
49+
CelOptimizerOptions.newBuilder().enableAstValidation(true).build());
50+
51+
assertThat(builder).isNotNull();
52+
assertThat(builder.build()).isNotNull();
53+
}
54+
4255
@Test
4356
public void standardCelOptimizerBuilder_withCompilerAndRuntime() {
4457
CelOptimizerBuilder builder =
@@ -50,6 +63,18 @@ public void standardCelOptimizerBuilder_withCompilerAndRuntime() {
5063
assertThat(builder.build()).isNotNull();
5164
}
5265

66+
@Test
67+
public void standardCelOptimizerBuilder_withCompilerRuntimeAndOptions() {
68+
CelOptimizerBuilder builder =
69+
CelOptimizerFactory.standardCelOptimizerBuilder(
70+
CelCompilerFactory.standardCelCompilerBuilder().build(),
71+
CelRuntimeFactory.standardCelRuntimeBuilder().build(),
72+
CelOptimizerOptions.newBuilder().enableAstValidation(true).build());
73+
74+
assertThat(builder).isNotNull();
75+
assertThat(builder.build()).isNotNull();
76+
}
77+
5378
@Test
5479
public void standardCelOptimizerBuilder_withCel() {
5580
CelOptimizerBuilder builder =
@@ -58,4 +83,15 @@ public void standardCelOptimizerBuilder_withCel() {
5883
assertThat(builder).isNotNull();
5984
assertThat(builder.build()).isNotNull();
6085
}
86+
87+
@Test
88+
public void standardCelOptimizerBuilder_withCelAndOptions() {
89+
CelOptimizerBuilder builder =
90+
CelOptimizerFactory.standardCelOptimizerBuilder(
91+
CelFactory.standardCelBuilder().build(),
92+
CelOptimizerOptions.newBuilder().enableAstValidation(true).build());
93+
94+
assertThat(builder).isNotNull();
95+
assertThat(builder.build()).isNotNull();
96+
}
6197
}

0 commit comments

Comments
 (0)