Add helpers for performing fixed point optimization

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

@@ -18,17 +18,19 @@
 import static java.lang.Math.max;
 import static java.util.stream.Collectors.toCollection;
 
+import com.google.auto.value.AutoOneOf;
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import com.google.common.collect.HashBasedTable;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Streams;
 import com.google.common.collect.Table;
 import com.google.errorprone.annotations.Immutable;
 import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.CelMutableAst;
 import dev.cel.common.CelMutableSource;
-import dev.cel.common.ast.CelExpr.ExprKind.Kind;
+import dev.cel.common.ast.CelExpr.ExprKind;
 import dev.cel.common.ast.CelExprIdGeneratorFactory;
 import dev.cel.common.ast.CelExprIdGeneratorFactory.ExprIdGenerator;
 import dev.cel.common.ast.CelExprIdGeneratorFactory.StableIdGenerator;
@@ -49,6 +51,7 @@
 import java.util.Map.Entry;
 import java.util.NoSuchElementException;
 import java.util.Optional;
+import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
@@ -213,7 +216,7 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames(
     Predicate<CelNavigableMutableExpr> comprehensionIdentifierPredicate = x -> true;
     comprehensionIdentifierPredicate =
         comprehensionIdentifierPredicate
-            .and(node -> node.getKind().equals(Kind.COMPREHENSION))
+            .and(node -> node.getKind().equals(ExprKind.Kind.COMPREHENSION))
             .and(node -> !node.expr().comprehension().iterVar().startsWith(newIterVarPrefix + ":"))
             .and(node -> !node.expr().comprehension().accuVar().startsWith(newAccuVarPrefix + ":"))
             .and(
@@ -235,7 +238,7 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames(
                   String result = node.expr().comprehension().result().ident().name();
                   return CelNavigableMutableExpr.fromExpr(node.expr().comprehension().loopStep())
                       .allNodes()
-                      .filter(subNode -> subNode.getKind().equals(Kind.IDENT))
+                      .filter(subNode -> subNode.getKind().equals(ExprKind.Kind.IDENT))
                       .map(subNode -> subNode.expr().ident())
                       .anyMatch(
                           ident ->
@@ -259,7 +262,7 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames(
                               .allNodes()
                               .filter(
                                   loopStepNode ->
-                                      loopStepNode.getKind().equals(Kind.IDENT)
+                                      loopStepNode.getKind().equals(ExprKind.Kind.IDENT)
                                           && loopStepNode.expr().ident().name().equals(iterVar))
                               .map(CelNavigableMutableExpr::id)
                               .findAny();
@@ -269,7 +272,7 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames(
                               .filter(
                                   loopStepNode ->
                                       !iterVar2.isEmpty()
-                                          && loopStepNode.getKind().equals(Kind.IDENT)
+                                          && loopStepNode.getKind().equals(ExprKind.Kind.IDENT)
                                           && loopStepNode.expr().ident().name().equals(iterVar2))
                               .map(CelNavigableMutableExpr::id)
                               .findAny();
@@ -406,13 +409,13 @@ private static MangledComprehensionName getMangledComprehensionName(
   private static int countComprehensionNestingLevel(CelNavigableMutableExpr comprehensionExpr) {
     return comprehensionExpr
         .descendants()
-        .filter(node -> node.getKind().equals(Kind.COMPREHENSION))
+        .filter(node -> node.getKind().equals(ExprKind.Kind.COMPREHENSION))
         .mapToInt(
             node -> {
               int nestedLevel = 1;
               CelNavigableMutableExpr maybeParent = node.parent().orElse(null);
               while (maybeParent != null && maybeParent.id() != comprehensionExpr.id()) {
-                if (maybeParent.getKind().equals(Kind.COMPREHENSION)) {
+                if (maybeParent.getKind().equals(ExprKind.Kind.COMPREHENSION)) {
                   nestedLevel++;
                 }
                 maybeParent = maybeParent.parent().orElse(null);
@@ -552,6 +555,151 @@ public CelMutableAst replaceSubtree(
     return CelMutableAst.of(mutatedRoot, newAstSource);
   }
 
+  /**
+   * Replaces a subtree in the given AST with the specified {@link SubtreeReplacement}.
+   *
+   * <p>This operation is intended for AST optimization purposes.
+   *
+   * <p>This is a very dangerous operation. Callers must 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 ast Original AST to mutate.
+   * @param replacement Subtree replacement containing the target node ID and the new expression or
+   *     AST.
+   */
+  public CelMutableAst replaceSubtree(CelMutableAst ast, SubtreeReplacement replacement) {
+    Preconditions.checkNotNull(ast);
+    Preconditions.checkNotNull(replacement);
+    switch (replacement.replacement().kind()) {
+      case EXPR:
+        return replaceSubtree(ast, replacement.replacement().expr(), replacement.exprIdToReplace());
+      case AST:
+        return replaceSubtree(ast, replacement.replacement().ast(), replacement.exprIdToReplace());
+    }
+    throw new IllegalArgumentException(
+        "Unsupported replacement kind: " + replacement.replacement().kind());
+  }
+
+  /**
+   * Repeatedly applies AST mutations using the provided AST-level rewriter until no further
+   * replacements match (fixed point reached) or the mutator's iteration limit is exhausted.
+   *
+   * <p>Per iteration pass, the AST is traversed to find and perform at most one matching subtree
+   * substitution before restarting traversal on the newly mutated AST.
+   *
+   * <p>This operation is intended for AST optimization purposes.
+   *
+   * <p>This is a very dangerous operation. Callers must 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.
+   *
+   * @param ast Initial mutable AST to mutate.
+   * @param astRewriter Function returning a {@link SubtreeReplacement} or {@code Optional.empty()}
+   *     when no further rewrites are possible.
+   * @return Mutated {@link CelMutableAst} at fixed point.
+   * @throws IllegalStateException If the iteration count exceeds {@code iterationLimit}.
+   */
+  public CelMutableAst mutateUntilFixedPoint(
+      CelMutableAst ast,
+      Function<CelNavigableMutableAst, Optional<SubtreeReplacement>> astRewriter) {
+    Preconditions.checkNotNull(ast);
+    Preconditions.checkNotNull(astRewriter);
+    CelMutableAst mutableAst = ast;
+    for (long i = 0; i < iterationLimit; i++) {
+      CelNavigableMutableAst navAst = CelNavigableMutableAst.fromAst(mutableAst);
+      Optional<SubtreeReplacement> replacement = astRewriter.apply(navAst);
+      if (!replacement.isPresent()) {
+        return mutableAst;
+      }
+      mutableAst = replaceSubtree(mutableAst, replacement.get());
+    }
+    throw new IllegalStateException("Max iteration count reached.");
+  }
+
+  /**
+   * Traverses nodes using the specified {@link TraversalOrder} and repeatedly rewrites matching
+   * subtrees until a fixed point is reached.
+   *
+   * <p>Per iteration pass, the AST is walked in the given order to find and perform the first
+   * matching substitution. The traversal then restarts on the freshly mutated AST until no nodes
+   * match or the mutator's iteration limit is exhausted.
+   *
+   * <p>This operation is intended for AST optimization purposes.
+   *
+   * <p>This is a very dangerous operation. Callers must 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.
+   *
+   * @param ast Initial mutable AST to mutate.
+   * @param traversalOrder Order in which nodes are visited per iteration pass.
+   * @param nodeRewriter Function returning a {@link SubtreeReplacement} or {@code
+   *     Optional.empty()}.
+   * @return Mutated {@link CelMutableAst} at fixed point.
+   * @throws IllegalStateException If the iteration count exceeds {@code iterationLimit}.
+   */
+  public CelMutableAst mutateUntilFixedPoint(
+      CelMutableAst ast,
+      TraversalOrder traversalOrder,
+      Function<CelNavigableMutableExpr, Optional<SubtreeReplacement>> nodeRewriter) {
+    Preconditions.checkNotNull(traversalOrder);
+    Preconditions.checkNotNull(nodeRewriter);
+    return mutateUntilFixedPoint(
+        ast,
+        navAst ->
+            navAst
+                .getRoot()
+                .allNodes(traversalOrder)
+                .flatMap(node -> Streams.stream(nodeRewriter.apply(node)))
+                .findFirst());
+  }
+
+  /**
+   * Traverses nodes using the specified {@link TraversalOrder}, applies the node matcher, and
+   * substitutes matching nodes with the returned replacement expression (targeting {@code
+   * node.id()}).
+   *
+   * <p>Per iteration pass, the AST is walked in the given order to find and perform the first
+   * matching substitution. The traversal then restarts on the freshly mutated AST until no nodes
+   * match or the mutator's iteration limit is exhausted.
+   *
+   * <p>This operation is intended for AST optimization purposes.
+   *
+   * <p>This is a very dangerous operation. Callers must 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.
+   *
+   * @param ast Initial mutable AST to mutate.
+   * @param traversalOrder Order in which nodes are visited per iteration pass.
+   * @param nodeMatcher Predicate to filter candidate nodes.
+   * @param nodeRewriter Function producing the new {@link CelMutableExpr} for matched nodes.
+   * @return Mutated {@link CelMutableAst} at fixed point.
+   * @throws IllegalStateException If the iteration count exceeds {@code iterationLimit}.
+   */
+  public CelMutableAst mutateUntilFixedPoint(
+      CelMutableAst ast,
+      TraversalOrder traversalOrder,
+      Predicate<CelNavigableMutableExpr> nodeMatcher,
+      Function<CelNavigableMutableExpr, Optional<CelMutableExpr>> nodeRewriter) {
+    Preconditions.checkNotNull(nodeMatcher);
+    Preconditions.checkNotNull(nodeRewriter);
+    return mutateUntilFixedPoint(
+        ast,
+        traversalOrder,
+        node ->
+            nodeMatcher.test(node)
+                ? nodeRewriter.apply(node).map(newExpr -> SubtreeReplacement.of(node.id(), newExpr))
+                : Optional.empty());
+  }
+
   private CelMutableExpr mangleIdentsInComprehensionExpr(
       CelMutableExpr root,
       CelMutableExpr comprehensionExpr,
@@ -590,7 +738,7 @@ private void replaceIdentName(
               .map(CelNavigableMutableExpr::expr)
               .filter(
                   node ->
-                      node.getKind().equals(Kind.IDENT)
+                      node.getKind().equals(ExprKind.Kind.IDENT)
                           && node.ident().name().equals(originalIdentName))
               .findAny()
               .orElse(null);
@@ -776,7 +924,7 @@ private CelMutableSource normalizeMacroSource(
         long replacedId = idGenerator.generate(exprIdToReplace);
         boolean isListExprBeingReplaced =
             allExprs.containsKey(replacedId)
-                && allExprs.get(replacedId).getKind().equals(Kind.LIST);
+                && allExprs.get(replacedId).getKind().equals(ExprKind.Kind.LIST);
         if (isListExprBeingReplaced) {
           unwrapListArgumentsInMacroCallExpr(
               allExprs.get(callId).comprehension(), newMacroCallExpr);
@@ -791,7 +939,7 @@ private CelMutableSource normalizeMacroSource(
       CelMutableExpr macroCallExpr = macroCall.getValue();
       CelNavigableMutableExpr.fromExpr(macroCallExpr)
           .allNodes()
-          .filter(node -> node.getKind().equals(Kind.COMPREHENSION))
+          .filter(node -> node.getKind().equals(ExprKind.Kind.COMPREHENSION))
           .map(CelNavigableMutableExpr::expr)
           .forEach(
               node -> {
@@ -808,7 +956,7 @@ private CelMutableSource normalizeMacroSource(
       // This can occur from pulling out a nested comprehension into a separate cel.block index
       CelNavigableMutableExpr.fromExpr(macroCallExpr)
           .allNodes()
-          .filter(node -> node.getKind().equals(Kind.NOT_SET))
+          .filter(node -> node.getKind().equals(ExprKind.Kind.NOT_SET))
           .map(CelNavigableMutableExpr::id)
           .filter(id -> !allExprs.containsKey(id))
           .forEach(
@@ -840,7 +988,7 @@ private CelMutableSource normalizeMacroSource(
   private static void unwrapListArgumentsInMacroCallExpr(
       CelMutableComprehension comprehension, CelMutableExpr newMacroCallExpr) {
     CelMutableExpr accuInit = comprehension.accuInit();
-    if (!accuInit.getKind().equals(Kind.LIST) || !accuInit.list().elements().isEmpty()) {
+    if (!accuInit.getKind().equals(ExprKind.Kind.LIST) || !accuInit.list().elements().isEmpty()) {
       // Does not contain an extraneous list.
       return;
     }
@@ -983,4 +1131,53 @@ private static MangledComprehensionName of(
           iterVarName, iterVar2Name, resultName);
     }
   }
+
+  /**
+   * Represents a planned subtree replacement containing the target node ID to replace and either a
+   * {@link CelMutableExpr} or {@link CelMutableAst}.
+   */
+  @AutoValue
+  public abstract static class SubtreeReplacement {
+
+    public abstract long exprIdToReplace();
+
+    public abstract Replacement replacement();
+
+    public static SubtreeReplacement of(long exprIdToReplace, CelMutableExpr replacementExpr) {
+      return new AutoValue_AstMutator_SubtreeReplacement(
+          exprIdToReplace, Replacement.ofExpr(replacementExpr));
+    }
+
+    public static SubtreeReplacement of(long exprIdToReplace, CelMutableAst replacementAst) {
+      return new AutoValue_AstMutator_SubtreeReplacement(
+          exprIdToReplace, Replacement.ofAst(replacementAst));
+    }
+
+    /** Discriminated union of either a {@link CelMutableExpr} or a {@link CelMutableAst}. */
+    @AutoOneOf(Replacement.Kind.class)
+    public abstract static class Replacement {
+
+      public abstract CelMutableExpr expr();
+
+      public abstract CelMutableAst ast();
+
+      public abstract Replacement.Kind kind();
+
+      public static Replacement ofExpr(CelMutableExpr expr) {
+        return AutoOneOf_AstMutator_SubtreeReplacement_Replacement.expr(
+            Preconditions.checkNotNull(expr));
+      }
+
+      public static Replacement ofAst(CelMutableAst ast) {
+        return AutoOneOf_AstMutator_SubtreeReplacement_Replacement.ast(
+            Preconditions.checkNotNull(ast));
+      }
+
+      /** Kind of {@link Replacement}. */
+      public enum Kind {
+        EXPR,
+        AST
+      }
+    }
+  }
 }

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

@@ -96,6 +96,7 @@ java_library(
         "//common:operator",
         "//common/ast",
         "//common/ast:mutable_expr",
+        "//common/navigation:common",
         "//common/navigation:expr_util",
         "//common/navigation:mutable_navigation",
         "//common/types",