Add evaluation listener

PiperOrigin-RevId: 569010652

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

@@ -36,6 +36,7 @@ java_library(
         "//parser:macro",
         "//parser:parser_builder",
         "//runtime",
+        "//runtime:evaluation_listener",
         "@cel_spec//proto/cel/expr:expr_java_proto",
         "@maven//:com_google_code_findbugs_annotations",
         "@maven//:com_google_errorprone_error_prone_annotations",

bundle/src/main/java/dev/cel/bundle/CelImpl.java

@@ -49,6 +49,7 @@
 import dev.cel.parser.CelParserBuilder;
 import dev.cel.parser.CelStandardMacro;
 import dev.cel.runtime.CelEvaluationException;
+import dev.cel.runtime.CelEvaluationListener;
 import dev.cel.runtime.CelRuntime;
 import dev.cel.runtime.CelRuntimeBuilder;
 import dev.cel.runtime.CelRuntimeLegacyImpl;
@@ -92,6 +93,12 @@ public CelRuntime.Program createProgram(CelAbstractSyntaxTree ast) throws CelEva
     return runtime.get().createProgram(ast);
   }
 
+  @Override
+  public Program createProgram(
+      CelAbstractSyntaxTree ast, CelEvaluationListener evaluationListener) {
+    return runtime.get().createProgram(ast, evaluationListener);
+  }
+
   @Override
   public void accept(EnvVisitor envVisitor) {
     CelCompiler celCompiler = compiler.get();

common/src/main/java/dev/cel/common/CelDescriptorUtil.java

@@ -14,7 +14,6 @@
 
 package dev.cel.common;
 
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
@@ -63,7 +62,6 @@ public static ImmutableSet<FileDescriptor> getFileDescriptorsForDescriptors(
    * <p>Warning: This will produce unique FileDescriptor instances. Use with care especially in
    * hermetic environments.
    */
-  @VisibleForTesting
   public static ImmutableSet<FileDescriptor> getFileDescriptorsFromFileDescriptorSet(
       FileDescriptorSet fileDescriptorSet) {
     return FileDescriptorSetConverter.convert(fileDescriptorSet);

runtime/BUILD.bazel

@@ -40,3 +40,8 @@ java_library(
     name = "interpreter_util",
     exports = ["//runtime/src/main/java/dev/cel/runtime:interpreter_util"],
 )
+
+java_library(
+    name = "evaluation_listener",
+    exports = ["//runtime/src/main/java/dev/cel/runtime:evaluation_listener"],
+)

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

@@ -73,6 +73,7 @@ java_library(
     exports = [":base"],
     deps = [
         ":base",
+        ":evaluation_listener",
         ":runtime_helper",
         ":unknown_attributes",
         "//:auto_value",
@@ -85,7 +86,6 @@ java_library(
         "//common/ast",
         "//common/internal:dynamic_proto",
         "//common/types:type_providers",
-        "//runtime:unknown_attributes",
         "@cel_spec//proto/cel/expr:expr_java_proto",
         "@maven//:com_google_code_findbugs_annotations",
         "@maven//:com_google_errorprone_error_prone_annotations",
@@ -142,6 +142,7 @@ java_library(
     ],
     deps = [
         ":base",
+        ":evaluation_listener",
         ":interpreter",
         ":unknown_attributes",
         "//:auto_value",
@@ -150,7 +151,6 @@ java_library(
         "//common:options",
         "//common/annotations",
         "//common/internal:dynamic_proto",
-        "//runtime:unknown_attributes",
         "@maven//:com_google_code_findbugs_annotations",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",
@@ -178,12 +178,12 @@ java_library(
     tags = [
     ],
     deps = [
+        ":unknown_attributes",
         "//common",
         "//common:compiler_common",
         "//parser",
         "//parser:operator",
         "//parser:parser_builder",
-        "//runtime:unknown_attributes",
         "@cel_spec//proto/cel/expr:expr_java_proto",
         "@maven//:com_google_guava_guava",
     ],
@@ -216,3 +216,15 @@ java_library(
         "@maven//:org_jspecify_jspecify",
     ],
 )
+
+java_library(
+    name = "evaluation_listener",
+    srcs = ["CelEvaluationListener.java"],
+    tags = [
+    ],
+    deps = [
+        "//common/ast",
+        "@maven//:com_google_code_findbugs_annotations",
+        "@maven//:com_google_errorprone_error_prone_annotations",
+    ],
+)

runtime/src/main/java/dev/cel/runtime/CelEvaluationListener.java

@@ -0,0 +1,40 @@
+// Copyright 2023 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.runtime;
+
+import javax.annotation.concurrent.ThreadSafe;
+import dev.cel.common.ast.CelExpr;
+
+/**
+ * Functional interface for a callback method invoked by the runtime. Implementations must ensure
+ * that its instances are unconditionally thread-safe.
+ */
+@FunctionalInterface
+@ThreadSafe
+public interface CelEvaluationListener {
+
+  /**
+   * Callback method invoked by the CEL runtime as evaluation progresses through the AST.
+   *
+   * @param expr CelExpr that was evaluated to produce the evaluated result.
+   * @param evaluatedResult Evaluated result.
+   */
+  void callback(CelExpr expr, Object evaluatedResult);
+
+  /** Construct a listener that does nothing. */
+  static CelEvaluationListener noOpListener() {
+    return (arg1, arg2) -> {};
+  }
+}

runtime/src/main/java/dev/cel/runtime/CelRuntime.java

@@ -37,6 +37,16 @@ public interface CelRuntime {
   @CanIgnoreReturnValue
   Program createProgram(CelAbstractSyntaxTree ast) throws CelEvaluationException;
 
+  /**
+   * Creates a {@link Program} from the input {@code ast}.
+   *
+   * @param evaluationListener A listener that gets invoked as evaluation progresses through the
+   *     AST. Implementations must ensure that an instance of a listener is unconditionally
+   *     thread-safe.
+   */
+  @CanIgnoreReturnValue
+  Program createProgram(CelAbstractSyntaxTree ast, CelEvaluationListener evaluationListener);
+
   /** Creates an evaluable {@code Program} instance which is thread-safe and immutable. */
   @AutoValue
   @Immutable

runtime/src/main/java/dev/cel/runtime/CelRuntimeLegacyImpl.java

@@ -56,6 +56,14 @@ public CelRuntime.Program createProgram(CelAbstractSyntaxTree ast) {
     return CelRuntime.Program.from(interpreter.createInterpretable(ast), options);
   }
 
+  @Override
+  public Program createProgram(
+      CelAbstractSyntaxTree ast, CelEvaluationListener evaluationListener) {
+    checkState(ast.isChecked(), "programs must be created from checked expressions");
+    return CelRuntime.Program.from(
+        interpreter.createInterpretable(ast, evaluationListener), options);
+  }
+
   /** Create a new builder for constructing a {@code CelRuntime} instance. */
   public static Builder newBuilder() {
     return new Builder();
@@ -75,91 +83,79 @@ public static final class Builder implements CelRuntimeBuilder {
     private boolean standardEnvironmentEnabled;
     private Function<String, Message.Builder> customTypeFactory;
 
-    /** {@inheritDoc} */
     @Override
     @CanIgnoreReturnValue
     public Builder setOptions(CelOptions options) {
       this.options = options;
       return this;
     }
 
-    /** {@inheritDoc} */
     @Override
     @CanIgnoreReturnValue
     public Builder addFunctionBindings(CelFunctionBinding... bindings) {
       return addFunctionBindings(Arrays.asList(bindings));
     }
 
-    /** {@inheritDoc} */
     @Override
     @CanIgnoreReturnValue
     public Builder addFunctionBindings(Iterable<CelFunctionBinding> bindings) {
       bindings.forEach(o -> functionBindings.put(o.getOverloadId(), o));
       return this;
     }
 
-    /** {@inheritDoc} */
     @Override
     @CanIgnoreReturnValue
     public Builder addMessageTypes(Descriptor... descriptors) {
       return addMessageTypes(Arrays.asList(descriptors));
     }
 
-    /** {@inheritDoc} */
     @Override
     @CanIgnoreReturnValue
     public Builder addMessageTypes(Iterable<Descriptor> descriptors) {
       return addFileTypes(CelDescriptorUtil.getFileDescriptorsForDescriptors(descriptors));
     }
 
-    /** {@inheritDoc} */
     @Override
     @CanIgnoreReturnValue
     public Builder addFileTypes(FileDescriptor... fileDescriptors) {
       return addFileTypes(Arrays.asList(fileDescriptors));
     }
 
-    /** {@inheritDoc} */
     @Override
     @CanIgnoreReturnValue
     public Builder addFileTypes(Iterable<FileDescriptor> fileDescriptors) {
       this.fileTypes.addAll(checkNotNull(fileDescriptors));
       return this;
     }
 
-    /** {@inheritDoc} */
     @Override
     @CanIgnoreReturnValue
     public Builder addFileTypes(FileDescriptorSet fileDescriptorSet) {
       return addFileTypes(
           CelDescriptorUtil.getFileDescriptorsFromFileDescriptorSet(fileDescriptorSet));
     }
 
-    /** {@inheritDoc} */
     @Override
     @CanIgnoreReturnValue
     public Builder setTypeFactory(Function<String, Message.Builder> typeFactory) {
       this.customTypeFactory = typeFactory;
       return this;
     }
 
-    /** {@inheritDoc} */
     @Override
     @CanIgnoreReturnValue
     public Builder setStandardEnvironmentEnabled(boolean value) {
       standardEnvironmentEnabled = value;
       return this;
     }
 
-    /** {@inheritDoc} */
     @Override
     @CanIgnoreReturnValue
     public Builder addLibraries(CelRuntimeLibrary... libraries) {
       checkNotNull(libraries);
       return this.addLibraries(Arrays.asList(libraries));
     }
 
-    /** {@inheritDoc} */
     @Override
     @CanIgnoreReturnValue
     public Builder addLibraries(Iterable<? extends CelRuntimeLibrary> libraries) {

runtime/src/main/java/dev/cel/runtime/DefaultInterpreter.java

@@ -127,17 +127,22 @@ public DefaultInterpreter(
     this.celOptions = celOptions;
   }
 
-  /** {@inheritDoc} */
   @Override
   @Deprecated
   public Interpretable createInterpretable(CheckedExpr checkedExpr) {
     return createInterpretable(CelProtoAbstractSyntaxTree.fromCheckedExpr(checkedExpr).getAst());
   }
 
-  /** {@inheritDoc} */
   @Override
   public Interpretable createInterpretable(CelAbstractSyntaxTree ast) {
-    return new DefaultInterpretable(typeProvider, dispatcher, ast, celOptions);
+    return new DefaultInterpretable(
+        typeProvider, dispatcher, ast, celOptions, CelEvaluationListener.noOpListener());
+  }
+
+  @Override
+  public Interpretable createInterpretable(
+      CelAbstractSyntaxTree ast, CelEvaluationListener listener) {
+    return new DefaultInterpretable(typeProvider, dispatcher, ast, celOptions, listener);
   }
 
   @Immutable
@@ -149,16 +154,23 @@ private static final class DefaultInterpretable
     private final CelAbstractSyntaxTree ast;
     private final CelOptions celOptions;
 
+    @SuppressWarnings(
+        "Immutable") // Used for callbacks only. A listener instance being mutable won't affect the
+    // correctness here.
+    private final CelEvaluationListener evaluationListener;
+
     DefaultInterpretable(
         RuntimeTypeProvider typeProvider,
         Dispatcher dispatcher,
         CelAbstractSyntaxTree ast,
-        CelOptions celOptions) {
+        CelOptions celOptions,
+        CelEvaluationListener evaluationListener) {
       this.typeProvider = Preconditions.checkNotNull(typeProvider);
       this.dispatcher = Preconditions.checkNotNull(dispatcher).immutableCopy();
       this.ast = Preconditions.checkNotNull(ast);
       this.metadata = new DefaultMetadata(ast);
       this.celOptions = Preconditions.checkNotNull(celOptions);
+      this.evaluationListener = Preconditions.checkNotNull(evaluationListener);
     }
 
     @Override
@@ -182,27 +194,38 @@ private IntermediateResult evalInternal(ExecutionFrame frame, CelExpr expr)
         throws InterpreterException {
       try {
         ExprKind.Kind exprKind = expr.exprKind().getKind();
+        IntermediateResult result;
         switch (exprKind) {
           case CONSTANT:
-            return IntermediateResult.create(evalConstant(frame, expr, expr.constant()));
+            result = IntermediateResult.create(evalConstant(frame, expr, expr.constant()));
+            break;
           case IDENT:
-            return evalIdent(frame, expr, expr.ident());
+            result = evalIdent(frame, expr, expr.ident());
+            break;
           case SELECT:
-            return evalSelect(frame, expr, expr.select());
+            result = evalSelect(frame, expr, expr.select());
+            break;
           case CALL:
-            return evalCall(frame, expr, expr.call());
+            result = evalCall(frame, expr, expr.call());
+            break;
           case CREATE_LIST:
-            return evalList(frame, expr, expr.createList());
+            result = evalList(frame, expr, expr.createList());
+            break;
           case CREATE_STRUCT:
-            return evalStruct(frame, expr, expr.createStruct());
+            result = evalStruct(frame, expr, expr.createStruct());
+            break;
           case CREATE_MAP:
-            return evalMap(frame, expr.createMap());
+            result = evalMap(frame, expr.createMap());
+            break;
           case COMPREHENSION:
-            return evalComprehension(frame, expr, expr.comprehension());
+            result = evalComprehension(frame, expr, expr.comprehension());
+            break;
           default:
             throw new IllegalStateException(
                 "unexpected expression kind: " + expr.exprKind().getKind());
         }
+        evaluationListener.callback(expr, result.value());
+        return result;
       } catch (RuntimeException e) {
         throw new InterpreterException.Builder(e, e.getMessage())
             .setLocation(metadata, expr.id())

runtime/src/main/java/dev/cel/runtime/Interpreter.java

@@ -44,4 +44,11 @@ public interface Interpreter {
    * <p>This method may run pre-processing and partial evaluation of the expression it gets passed.
    */
   Interpretable createInterpretable(CelAbstractSyntaxTree ast);
+
+  /**
+   * Creates an interpretable for the given expression. This accepts a listener that gets invoked as
+   * evaluation progresses through the AST.
+   */
+  Interpretable createInterpretable(
+      CelAbstractSyntaxTree ast, CelEvaluationListener evaluationListener);
 }