Add optimizer listener

PiperOrigin-RevId: 974424751

Author: Cristina Borza

optimizer/BUILD.bazel

@@ -15,6 +15,11 @@ java_library(
     exports = ["//optimizer/src/main/java/dev/cel/optimizer:optimizer_builder"],
 )
 
+java_library(
+    name = "optimizer_listener",
+    exports = ["//optimizer/src/main/java/dev/cel/optimizer:optimizer_listener"],
+)
+
 java_library(
     name = "ast_optimizer",
     exports = ["//optimizer/src/main/java/dev/cel/optimizer:ast_optimizer"],

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

@@ -27,6 +27,19 @@ java_library(
     ],
 )
 
+java_library(
+    name = "optimizer_listener",
+    srcs = ["CelOptimizerListener.java"],
+    tags = [
+    ],
+    deps = [
+        ":ast_optimizer",
+        "//common:cel_ast",
+        "@maven//:com_google_code_findbugs_annotations",
+        "@maven//:com_google_errorprone_error_prone_annotations",
+    ],
+)
+
 java_library(
     name = "optimizer_builder",
     srcs = [
@@ -38,6 +51,7 @@ java_library(
     deps = [
         ":ast_optimizer",
         ":optimization_exception",
+        ":optimizer_listener",
         "//common:cel_ast",
         "@maven//:com_google_errorprone_error_prone_annotations",
     ],
@@ -54,6 +68,7 @@ java_library(
         ":ast_optimizer",
         ":optimization_exception",
         ":optimizer_builder",
+        ":optimizer_listener",
         "//bundle:cel",
         "//common:cel_ast",
         "//common:compiler_common",

optimizer/src/main/java/dev/cel/optimizer/CelOptimizerBuilder.java

@@ -28,6 +28,14 @@ public interface CelOptimizerBuilder {
   @CanIgnoreReturnValue
   CelOptimizerBuilder addAstOptimizers(Iterable<CelAstOptimizer> astOptimizers);
 
+  /** Adds one or more listeners to observe optimization lifecycle. */
+  @CanIgnoreReturnValue
+  CelOptimizerBuilder addOptimizerListeners(CelOptimizerListener... listeners);
+
+  /** Adds one or more listeners to observe optimization lifecycle. */
+  @CanIgnoreReturnValue
+  CelOptimizerBuilder addOptimizerListeners(Iterable<CelOptimizerListener> listeners);
+
   /** Build a new instance of the {@link CelOptimizer}. */
   @CheckReturnValue
   CelOptimizer build();

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

@@ -32,10 +32,15 @@
 final class CelOptimizerImpl implements CelOptimizer {
   private final Cel cel;
   private final ImmutableSet<CelAstOptimizer> astOptimizers;
+  private final ImmutableSet<CelOptimizerListener> listeners;
 
-  CelOptimizerImpl(Cel cel, ImmutableSet<CelAstOptimizer> astOptimizers) {
+  CelOptimizerImpl(
+      Cel cel,
+      ImmutableSet<CelAstOptimizer> astOptimizers,
+      ImmutableSet<CelOptimizerListener> listeners) {
     this.cel = cel;
     this.astOptimizers = astOptimizers;
+    this.listeners = listeners;
   }
 
   @Override
@@ -44,27 +49,51 @@ public CelAbstractSyntaxTree optimize(CelAbstractSyntaxTree ast) throws CelOptim
       throw new IllegalArgumentException("AST must be type-checked.");
     }
 
+    listeners.forEach(listener -> listener.onOptimizationStart(ast));
+
     Cel celOptimizerEnv = cel;
     CelAbstractSyntaxTree optimizedAst = ast;
+
     try {
       for (CelAstOptimizer optimizer : astOptimizers) {
-        OptimizationResult result = optimizer.optimize(optimizedAst, celOptimizerEnv);
-        if (!result.newFunctionDecls().isEmpty() || !result.newVarDecls().isEmpty()) {
-          celOptimizerEnv =
-              celOptimizerEnv
-                  .toCelBuilder()
-                  .addVarDeclarations(result.newVarDecls())
-                  .addFunctionDeclarations(result.newFunctionDecls())
-                  .build();
+        CelAbstractSyntaxTree preAst = optimizedAst;
+        try {
+          for (CelOptimizerListener listener : listeners) {
+            listener.onPassStart(optimizer, preAst);
+          }
+
+          OptimizationResult result = optimizer.optimize(optimizedAst, celOptimizerEnv);
+
+          if (!result.newFunctionDecls().isEmpty() || !result.newVarDecls().isEmpty()) {
+            celOptimizerEnv =
+                celOptimizerEnv
+                    .toCelBuilder()
+                    .addVarDeclarations(result.newVarDecls())
+                    .addFunctionDeclarations(result.newFunctionDecls())
+                    .build();
+          }
+          optimizedAst = celOptimizerEnv.check(result.optimizedAst()).getAst();
+          assertAstIdCorrectness(optimizedAst);
+
+          for (CelOptimizerListener listener : listeners) {
+            listener.onPassEnd(optimizer, preAst, optimizedAst);
+          }
+        } catch (CelValidationException e) {
+          notifyPassFailure(optimizer, preAst, e);
+          throw new CelOptimizationException(
+              "Optimized AST failed to type-check: " + e.getMessage(), e);
+        } catch (CelOptimizationException e) {
+          notifyPassFailure(optimizer, preAst, e);
+          throw e;
+        } catch (RuntimeException e) {
+          notifyPassFailure(optimizer, preAst, e);
+          throw new CelOptimizationException("Optimization failure: " + e.getMessage(), e);
         }
-        optimizedAst = celOptimizerEnv.check(result.optimizedAst()).getAst();
-        assertAstIdCorrectness(optimizedAst);
       }
-    } catch (CelValidationException e) {
-      throw new CelOptimizationException(
-          "Optimized AST failed to type-check: " + e.getMessage(), e);
-    } catch (RuntimeException e) {
-      throw new CelOptimizationException("Optimization failure: " + e.getMessage(), e);
+    } finally {
+      for (CelOptimizerListener listener : listeners) {
+        listener.onOptimizationEnd(ast, optimizedAst);
+      }
     }
 
     return optimizedAst;
@@ -123,6 +152,13 @@ private static void assertAstIdCorrectness(CelAbstractSyntaxTree ast) {
     }
   }
 
+  private void notifyPassFailure(
+      CelAstOptimizer optimizer, CelAbstractSyntaxTree ast, Exception failure) {
+    for (CelOptimizerListener listener : listeners) {
+      listener.onPassFailure(optimizer, ast, failure);
+    }
+  }
+
   /** Create a new builder for constructing a {@link CelOptimizer} instance. */
   static CelOptimizerImpl.Builder newBuilder(Cel cel) {
     return new CelOptimizerImpl.Builder(cel);
@@ -132,10 +168,12 @@ static CelOptimizerImpl.Builder newBuilder(Cel cel) {
   static final class Builder implements CelOptimizerBuilder {
     private final Cel cel;
     private final ImmutableSet.Builder<CelAstOptimizer> astOptimizers;
+    private final ImmutableSet.Builder<CelOptimizerListener> listeners;
 
     private Builder(Cel cel) {
       this.cel = cel;
       this.astOptimizers = ImmutableSet.builder();
+      this.listeners = ImmutableSet.builder();
     }
 
     @Override
@@ -151,9 +189,22 @@ public CelOptimizerBuilder addAstOptimizers(Iterable<CelAstOptimizer> astOptimiz
       return this;
     }
 
+    @Override
+    public CelOptimizerBuilder addOptimizerListeners(CelOptimizerListener... listeners) {
+      checkNotNull(listeners);
+      return addOptimizerListeners(Arrays.asList(listeners));
+    }
+
+    @Override
+    public CelOptimizerBuilder addOptimizerListeners(Iterable<CelOptimizerListener> listeners) {
+      checkNotNull(listeners);
+      this.listeners.addAll(listeners);
+      return this;
+    }
+
     @Override
     public CelOptimizer build() {
-      return new CelOptimizerImpl(cel, astOptimizers.build());
+      return new CelOptimizerImpl(cel, astOptimizers.build(), listeners.build());
     }
   }
 }

optimizer/src/main/java/dev/cel/optimizer/CelOptimizerListener.java

@@ -0,0 +1,72 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dev.cel.optimizer;
+
+import javax.annotation.concurrent.ThreadSafe;
+import dev.cel.common.CelAbstractSyntaxTree;
+
+/**
+ * Listener interface for observing the execution lifecycle of {@link CelOptimizer}.
+ *
+ * <p>Implementations must be thread-safe.
+ */
+@ThreadSafe
+public interface CelOptimizerListener {
+  /**
+   * Invoked before the optimization pipeline begins.
+   *
+   * @param ast the initial AST to be optimized.
+   */
+  default void onOptimizationStart(CelAbstractSyntaxTree ast) {}
+
+  /**
+   * Invoked before a specific {@link CelAstOptimizer} pass executes.
+   *
+   * @param optimizer the optimizer pass that is about to execute.
+   * @param ast the initial AST that is about to be optimized.
+   */
+  default void onPassStart(CelAstOptimizer optimizer, CelAbstractSyntaxTree ast) {}
+
+  /**
+   * Invoked after a specific {@link CelAstOptimizer} pass completes successfully.
+   *
+   * @param optimizer the optimizer pass that just completed.
+   * @param preAst the initial AST that was passed to the optimizer pass.
+   * @param optimizedAst the AST after the optimizer pass completed.
+   */
+  default void onPassEnd(
+      CelAstOptimizer optimizer,
+      CelAbstractSyntaxTree preAst,
+      CelAbstractSyntaxTree optimizedAst) {}
+
+  /**
+   * Invoked if an optimizer pass throws an unhandled exception.
+   *
+   * @param optimizer the optimizer pass that threw the exception.
+   * @param ast the initial AST that was passed to the optimizer pass.
+   * @param failure the exception that was thrown by the optimizer pass.
+   */
+  default void onPassFailure(
+      CelAstOptimizer optimizer, CelAbstractSyntaxTree ast, Exception failure) {}
+
+  /**
+   * Invoked after all optimization passes and final type-checks complete.
+   *
+   * @param initialAst the initial AST that was passed to the optimizer.
+   * @param finalAst the final AST after all optimization passes.
+   */
+  default void onOptimizationEnd(
+      CelAbstractSyntaxTree initialAst, CelAbstractSyntaxTree finalAst) {}
+}

optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel

@@ -32,6 +32,7 @@ java_library(
         "//optimizer:optimization_exception",
         "//optimizer:optimizer_builder",
         "//optimizer:optimizer_impl",
+        "//optimizer:optimizer_listener",
         "//parser:macro",
         "//parser:parser_factory",
         "//parser:unparser",

optimizer/src/test/java/dev/cel/optimizer/CelOptimizerImplTest.java

@@ -44,6 +44,41 @@ public class CelOptimizerImplTest {
           .setStandardMacros(CelStandardMacro.STANDARD_MACROS)
           .build();
 
+  private final List<String> events = new ArrayList<>();
+
+  private final CelOptimizerListener listener =
+      new CelOptimizerListener() {
+        @Override
+        public void onOptimizationStart(CelAbstractSyntaxTree ast) {
+          events.add("start");
+        }
+
+        @Override
+        public void onPassStart(CelAstOptimizer optimizer, CelAbstractSyntaxTree ast) {
+          events.add("pass_start");
+        }
+
+        @Override
+        public void onPassEnd(
+            CelAstOptimizer optimizer,
+            CelAbstractSyntaxTree preAst,
+            CelAbstractSyntaxTree optimizedAst) {
+          events.add("pass_end");
+        }
+
+        @Override
+        public void onOptimizationEnd(
+            CelAbstractSyntaxTree initialAst, CelAbstractSyntaxTree finalAst) {
+          events.add("end");
+        }
+
+        @Override
+        public void onPassFailure(
+            CelAstOptimizer optimizer, CelAbstractSyntaxTree ast, Exception failure) {
+          events.add("pass_failure");
+        }
+      };
+
   @Test
   public void constructCelOptimizer_success() {
     CelOptimizer celOptimizer =
@@ -313,4 +348,35 @@ public void optimize_validMacroCalls_success() throws Exception {
     assertThat(optimizedAst).isNotNull();
     assertThat(optimizedAst.getSource().getMacroCalls()).hasSize(1);
   }
+
+  @Test
+  public void optimize_withListener_invokesListenerMethods() throws Exception {
+    CelOptimizer celOptimizer =
+        CelOptimizerImpl.newBuilder(CEL)
+            .addAstOptimizers((navigableAst, cel) -> OptimizationResult.create(navigableAst))
+            .addOptimizerListeners(listener)
+            .build();
+
+    CelAbstractSyntaxTree ast = CEL.compile("'hello world'").getAst();
+    CelAbstractSyntaxTree unused = celOptimizer.optimize(ast);
+
+    assertThat(events).containsExactly("start", "pass_start", "pass_end", "end").inOrder();
+  }
+
+  @Test
+  public void optimize_withListener_onPassFailure_invokesListenerMethods() throws Exception {
+    CelOptimizer celOptimizer =
+        CelOptimizerImpl.newBuilder(CEL)
+            .addAstOptimizers(
+                (navigableAst, cel) -> {
+                  throw new RuntimeException("Test failure");
+                })
+            .addOptimizerListeners(listener)
+            .build();
+
+    CelAbstractSyntaxTree ast = CEL.compile("'hello world'").getAst();
+    assertThrows(CelOptimizationException.class, () -> celOptimizer.optimize(ast));
+
+    assertThat(events).containsExactly("start", "pass_start", "pass_failure", "end").inOrder();
+  }
 }