Allow configuration of MutableAst by making it instantiable

optimizer/BUILD.bazel

@@ -25,8 +25,6 @@ java_library(
 
 java_library(
     name = "mutable_ast",
-    testonly = 1,
-    visibility = ["//optimizer/src/test/java/dev/cel/optimizer:__pkg__"],
     exports = ["//optimizer/src/main/java/dev/cel/optimizer:mutable_ast"],
 )
 

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

@@ -66,18 +66,19 @@ java_library(
     tags = [
     ],
     deps = [
-        ":mutable_ast",
         ":optimization_exception",
         "//bundle:cel",
         "//common",
-        "//common/ast",
         "//common/navigation",
     ],
 )
 
 java_library(
     name = "mutable_ast",
-    srcs = ["MutableAst.java"],
+    srcs = [
+        "MutableAst.java",
+        "MutableExprVisitor.java",
+    ],
     tags = [
     ],
     deps = [
@@ -87,6 +88,7 @@ java_library(
         "//common/ast",
         "//common/ast:expr_factory",
         "//common/navigation",
+        "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",
     ],
 )

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

@@ -16,7 +16,6 @@
 
 import dev.cel.bundle.Cel;
 import dev.cel.common.CelAbstractSyntaxTree;
-import dev.cel.common.ast.CelExpr;
 import dev.cel.common.navigation.CelNavigableAst;
 
 /** Public interface for performing a single, custom optimization on an AST. */
@@ -25,106 +24,4 @@ public interface CelAstOptimizer {
   /** Optimizes a single AST. */
   CelAbstractSyntaxTree optimize(CelNavigableAst navigableAst, Cel cel)
       throws CelOptimizationException;
-
-  /**
-   * Replaces a subtree in the given expression node. This operation is intended for AST
-   * optimization purposes.
-   *
-   * <p>This is a very dangerous operation. Callers should re-typecheck the mutated AST and
-   * additionally verify that the resulting AST is semantically valid.
-   *
-   * <p>All expression IDs will be renumbered in a stable manner to ensure there's no ID collision
-   * between the nodes. The renumbering occurs even if the subtree was not replaced.
-   *
-   * @param celExpr Original expression node to rewrite.
-   * @param newExpr New CelExpr to replace the subtree with.
-   * @param exprIdToReplace Expression id of the subtree that is getting replaced.
-   */
-  default CelExpr replaceSubtree(CelExpr celExpr, CelExpr newExpr, long exprIdToReplace) {
-    return MutableAst.replaceSubtree(celExpr, newExpr, exprIdToReplace);
-  }
-
-  /**
-   * Replaces a subtree in the given AST. This operation is intended for AST optimization purposes.
-   *
-   * <p>This is a very dangerous operation. Callers should re-typecheck the mutated AST and
-   * additionally verify that the resulting AST is semantically valid.
-   *
-   * <p>All expression IDs will be renumbered in a stable manner to ensure there's no ID collision
-   * between the nodes. The renumbering occurs even if the subtree was not replaced.
-   *
-   * <p>This will scrub out the description, positions and line offsets from {@code CelSource}. If
-   * the source contains macro calls, its call IDs will be to be consistent with the renumbered IDs
-   * in the AST.
-   *
-   * @param ast Original ast to mutate.
-   * @param newExpr New CelExpr to replace the subtree with.
-   * @param exprIdToReplace Expression id of the subtree that is getting replaced.
-   */
-  default CelAbstractSyntaxTree replaceSubtree(
-      CelAbstractSyntaxTree ast, CelExpr newExpr, long exprIdToReplace) {
-    return MutableAst.replaceSubtree(ast, newExpr, exprIdToReplace);
-  }
-
-  /**
-   * Generates a new bind macro using the provided initialization and result expression, then
-   * replaces the subtree using the new bind expr at the designated expr ID.
-   *
-   * <p>The bind call takes the format of: {@code cel.bind(varInit, varName, resultExpr)}
-   *
-   * @param ast Original ast to mutate.
-   * @param varName New variable name for the bind macro call.
-   * @param varInit Initialization expression to bind to the local variable.
-   * @param resultExpr Result expression
-   * @param exprIdToReplace Expression ID of the subtree that is getting replaced.
-   */
-  default CelAbstractSyntaxTree replaceSubtreeWithNewBindMacro(
-      CelAbstractSyntaxTree ast,
-      String varName,
-      CelExpr varInit,
-      CelExpr resultExpr,
-      long exprIdToReplace) {
-    return MutableAst.replaceSubtreeWithNewBindMacro(
-        ast, varName, varInit, resultExpr, exprIdToReplace);
-  }
-
-  /**
-   * Replaces all comprehension identifier names with a unique name based on the given prefix.
-   *
-   * <p>The purpose of this is to avoid errors that can be caused by shadowed variables while
-   * augmenting an AST. As an example: {@code [2, 3].exists(x, x - 1 > 3) || x - 1 > 3}. Note that
-   * the scoping of `x - 1` is different between th two LOGICAL_OR branches. Iteration variable `x`
-   * in `exists` will be mangled to {@code [2, 3].exists(@c0, @c0 - 1 > 3) || x - 1 > 3} to avoid
-   * erroneously extracting x - 1 as common subexpression.
-   *
-   * <p>The expression IDs are not modified when the identifier names are changed.
-   *
-   * <p>Iteration variables in comprehensions are numbered based on their comprehension nesting
-   * levels. Examples:
-   *
-   * <ul>
-   *   <li>{@code [true].exists(i, i) && [true].exists(j, j)} -> {@code [true].exists(@c0, @c0) &&
-   *       [true].exists(@c0, @c0)} // Note that i,j gets replaced to the same @c0 in this example
-   *   <li>{@code [true].exists(i, i && [true].exists(j, j))} -> {@code [true].exists(@c0, @c0 &&
-   *       [true].exists(@c1, @c1))}
-   * </ul>
-   *
-   * @param ast AST to mutate
-   * @param newIdentPrefix Prefix to use for new identifier names. For example, providing @c will
-   *     produce @c0, @c1, @c2... as new names.
-   */
-  default CelAbstractSyntaxTree mangleComprehensionIdentifierNames(
-      CelAbstractSyntaxTree ast, String newIdentPrefix) {
-    return MutableAst.mangleComprehensionIdentifierNames(ast, newIdentPrefix);
-  }
-
-  /** Sets all expr IDs in the expression tree to 0. */
-  default CelExpr clearExprIds(CelExpr celExpr) {
-    return MutableAst.clearExprIds(celExpr);
-  }
-
-  /** Renumbers all the expr IDs in the given AST in a consecutive manner starting from 1. */
-  default CelAbstractSyntaxTree renumberIdsConsecutively(CelAbstractSyntaxTree ast) {
-    return MutableAst.renumberIdsConsecutively(ast);
-  }
 }