Add cel.block as an extension library

PiperOrigin-RevId: 602536654

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

@@ -89,7 +89,6 @@ public abstract class CelOptions {
 
   public abstract boolean enableCelValue();
 
-
   public abstract int comprehensionMaxIterations();
 
   public abstract Builder toBuilder();

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

@@ -15,6 +15,7 @@ java_library(
     ],
     deps = [
         ":bindings",
+        ":block",
         ":encoders",
         ":math",
         ":protos",
@@ -77,6 +78,17 @@ java_library(
     ],
 )
 
+java_library(
+    name = "block",
+    srcs = ["CelBlockExtensions.java"],
+    deps = [
+        "//checker:checker_builder",
+        "//common:compiler_common",
+        "//common/types",
+        "//compiler:compiler_builder",
+    ],
+)
+
 java_library(
     name = "bindings",
     srcs = ["CelBindingsExtensions.java"],

extensions/src/main/java/dev/cel/extensions/CelBlockExtensions.java

@@ -0,0 +1,45 @@
+// Copyright 2024 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.extensions;
+
+import dev.cel.checker.CelCheckerBuilder;
+import dev.cel.common.CelFunctionDecl;
+import dev.cel.common.CelOverloadDecl;
+import dev.cel.common.types.ListType;
+import dev.cel.common.types.SimpleType;
+import dev.cel.compiler.CelCompilerLibrary;
+
+/**
+ * Internal implementation of cel.@block extension function.
+ *
+ * <p>Note the extensions only contain function declarations. This requires special support from
+ * interpreters.
+ */
+final class CelBlockExtensions implements CelCompilerLibrary {
+  private static final String CEL_BLOCK_FUNCTION = "cel.@block";
+  private static final String CEL_BLOCK_INDEX_FUNCTION = "cel.@index";
+
+  @Override
+  public void setCheckerOptions(CelCheckerBuilder checkerBuilder) {
+    checkerBuilder.addFunctionDeclarations(
+        CelFunctionDecl.newFunctionDeclaration(
+            CEL_BLOCK_FUNCTION,
+            CelOverloadDecl.newGlobalOverload(
+                "cel_block_list", SimpleType.DYN, ListType.create(SimpleType.DYN), SimpleType.DYN)),
+        CelFunctionDecl.newFunctionDeclaration(
+            CEL_BLOCK_INDEX_FUNCTION,
+            CelOverloadDecl.newGlobalOverload("cel_index_int", SimpleType.DYN, SimpleType.INT)));
+  }
+}

extensions/src/main/java/dev/cel/extensions/CelExtensions.java

@@ -29,6 +29,7 @@ public final class CelExtensions {
   private static final CelStringExtensions STRING_EXTENSIONS_ALL = new CelStringExtensions();
   private static final CelProtoExtensions PROTO_EXTENSIONS = new CelProtoExtensions();
   private static final CelBindingsExtensions BINDINGS_EXTENSIONS = new CelBindingsExtensions();
+  private static final CelBlockExtensions BLOCK_EXTENSIONS = new CelBlockExtensions();
   private static final CelEncoderExtensions ENCODER_EXTENSIONS = new CelEncoderExtensions();
 
   /**
@@ -160,6 +161,23 @@ public static CelBindingsExtensions bindings() {
     return BINDINGS_EXTENSIONS;
   }
 
+  /**
+   * Exposes internal function declarations for cel.@block
+   *
+   * <p>The function can hold a list of expressions that are lazily evaluated then memoized upon its
+   * reference. Conceptually, this can be thought of as cascaded cel.bind calls.
+   *
+   * <p>Example: {@code cel.block([fn_1(), cel.@index(0) && fn_1(), fn_2()], cel.@index(1)]}. fn_1
+   * is invoked once and fn_2 is never invoked.
+   *
+   * <p>Users cannot directly leverage cel.@block in their expressions (hence the @ symbol). The
+   * expected usage is for CEL optimizers to rewrite a long, deeply nested expressions into
+   * cel.@block format to improve performance and safety.
+   */
+  public static CelBlockExtensions block() {
+    return BLOCK_EXTENSIONS;
+  }
+
   /**
    * Extended functions for string, byte and object encodings.
    *

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

@@ -12,6 +12,8 @@ java_library(
         "//common",
         "//common:compiler_common",
         "//common:options",
+        "//common/ast",
+        "//common/navigation",
         "//common/resources/testdata/proto2:messages_extensions_proto2_java_proto",
         "//common/resources/testdata/proto2:messages_proto2_java_proto",
         "//common/resources/testdata/proto2:test_all_types_java_proto",
@@ -24,6 +26,7 @@ java_library(
         "//extensions:math",
         "//extensions:optional_library",
         "//extensions:strings",
+        "//optimizer:mutable_ast",
         "//parser:macro",
         "//runtime",
         "@maven//:com_google_guava_guava",