diff --git a/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java b/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java
index 0f428f75a..e3647de3a 100644
--- a/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java
+++ b/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java
@@ -18,11 +18,13 @@
 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;
@@ -44,11 +46,13 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
 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;
 
@@ -552,6 +556,140 @@ 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>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>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>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,
@@ -983,4 +1121,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 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
+      }
+    }
+  }
 }
diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel b/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel
index 0c4b78826..1012b19c2 100644
--- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel
+++ b/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",
diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java
index 266059426..5aed8acd6 100644
--- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java
+++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java
@@ -55,6 +55,7 @@
 import dev.cel.common.values.StructValue;
 import dev.cel.extensions.CelOptionalLibrary.Function;
 import dev.cel.optimizer.AstMutator;
+import dev.cel.optimizer.AstMutator.SubtreeReplacement;
 import dev.cel.optimizer.CelAstOptimizer;
 import dev.cel.optimizer.CelOptimizationException;
 import dev.cel.runtime.CelAttributePattern;
@@ -134,7 +135,15 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
     Cel optimizerEnv = builder.setResultType(SimpleType.DYN).build();
 
     CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
+    ImmutableMap<String, CelType> identTypes = precomputeIdentTypes(mutableAst);
 
+    mutableAst = foldConstants(optimizerEnv, valueProvider, identTypes, mutableAst);
+    mutableAst = pruneOptionalElements(mutableAst);
+
+    return OptimizationResult.create(astMutator.renumberIdsConsecutively(mutableAst).toParsedAst());
+  }
+
+  private static ImmutableMap<String, CelType> precomputeIdentTypes(CelMutableAst mutableAst) {
     // HACK: The AstMutator strips type metadata during intermediate folds due to ID renumbering.
     // We pre-compute identifier types from the unmutated AST to safely evaluate boolean conditions
     // later.
@@ -151,58 +160,47 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
       Optional<CelType> type = mutableAst.getType(node.id());
       type.ifPresent(celType -> mutableIdentTypes.put(node.expr().ident().name(), celType));
     }
-    ImmutableMap<String, CelType> identTypes = ImmutableMap.copyOf(mutableIdentTypes);
+    return ImmutableMap.copyOf(mutableIdentTypes);
+  }
 
-    int iterCount = 0;
-    boolean continueFolding = true;
-    while (continueFolding) {
-      if (iterCount >= constantFoldingOptions.maxIterationLimit()) {
-        throw new IllegalStateException("Max iteration count reached.");
+  private CelMutableAst foldConstants(
+      Cel optimizerEnv,
+      @Nullable CelValueProvider valueProvider,
+      ImmutableMap<String, CelType> identTypes,
+      CelMutableAst mutableAst)
+      throws CelOptimizationException {
+    for (int iterCount = 0; iterCount < constantFoldingOptions.maxIterationLimit(); iterCount++) {
+      Optional<SubtreeReplacement> replacement =
+          findNextFoldableSubtree(optimizerEnv, valueProvider, identTypes, mutableAst);
+      if (!replacement.isPresent()) {
+        return mutableAst;
       }
-      iterCount++;
-      continueFolding = false;
-      ImmutableList<CelNavigableMutableExpr> foldableExprs =
-          CelNavigableMutableAst.fromAst(mutableAst)
-              .getRoot()
-              .allNodes(TraversalOrder.PRE_ORDER)
-              .filter(this::canFold)
-              .collect(toImmutableList());
-      for (CelNavigableMutableExpr foldableExpr : foldableExprs) {
-        iterCount++;
-
-        Optional<CelMutableAst> mutatedResult;
-        // Attempt to prune if it is a non-strict call
-        mutatedResult = maybePruneBranches(mutableAst, identTypes, foldableExpr.expr());
-        if (!mutatedResult.isPresent()) {
-          // Evaluate the call then fold
-          try {
-            mutatedResult = maybeFold(optimizerEnv, valueProvider, mutableAst, foldableExpr);
-          } catch (CelEvaluationException e) {
-            throw new CelOptimizationException(
-                "Constant folding failure. Failed to evaluate subtree due to: " + e.getMessage(),
-                e);
-          }
-        }
-
-        if (!mutatedResult.isPresent()) {
-          // Skip this expr. It's neither prune-able nor foldable.
-          continue;
-        }
+      mutableAst = astMutator.replaceSubtree(mutableAst, replacement.get());
+    }
+    throw new IllegalStateException("Max iteration count reached.");
+  }
 
-        continueFolding = true;
-        mutableAst = mutatedResult.get();
-        // Break the loop because we mutated the AST. Since we traverse in PRE_ORDER (top-down),
-        // mutating a parent node means its children are now obsolete or folded.
-        // We restart the traversal to gather a fresh list of foldable expressions.
-        break;
+  private Optional<SubtreeReplacement> findNextFoldableSubtree(
+      Cel optimizerEnv,
+      @Nullable CelValueProvider valueProvider,
+      ImmutableMap<String, CelType> identTypes,
+      CelMutableAst mutableAst)
+      throws CelOptimizationException {
+    CelNavigableMutableAst navAst = CelNavigableMutableAst.fromAst(mutableAst);
+    Iterator<CelNavigableMutableExpr> foldableExprs =
+        navAst.getRoot().allNodes(TraversalOrder.PRE_ORDER).filter(this::canFold).iterator();
+    while (foldableExprs.hasNext()) {
+      CelNavigableMutableExpr foldableExpr = foldableExprs.next();
+      Optional<CelMutableExpr> pruned = maybePruneBranches(identTypes, foldableExpr.expr());
+      if (pruned.isPresent()) {
+        return Optional.of(SubtreeReplacement.of(foldableExpr.id(), pruned.get()));
+      }
+      Optional<CelMutableExpr> folded = maybeFold(optimizerEnv, valueProvider, foldableExpr);
+      if (folded.isPresent()) {
+        return Optional.of(SubtreeReplacement.of(foldableExpr.id(), folded.get()));
       }
     }
-
-    // If the output is a list, map, or struct which contains optional entries, then prune it
-    // to make sure that the optionals, if resolved, do not surface in the output literal.
-    mutableAst = pruneOptionalElements(mutableAst);
-
-    return OptimizationResult.create(astMutator.renumberIdsConsecutively(mutableAst).toParsedAst());
+    return Optional.empty();
   }
 
   private boolean canFold(CelNavigableMutableExpr navigableExpr) {
@@ -317,12 +315,9 @@ private static boolean isNestedComprehension(CelNavigableMutableExpr expr) {
     return false;
   }
 
-  private Optional<CelMutableAst> maybeFold(
-      Cel cel,
-      CelValueProvider valueProvider,
-      CelMutableAst mutableAst,
-      CelNavigableMutableExpr node)
-      throws CelOptimizationException, CelEvaluationException {
+  private Optional<CelMutableExpr> maybeFold(
+      Cel cel, @Nullable CelValueProvider valueProvider, CelNavigableMutableExpr node)
+      throws CelOptimizationException {
     if (!node.getKind().equals(Kind.COMPREHENSION)
         && CelNavigableExprUtil.hasComprehensionVariable(node)) {
       return Optional.empty();
@@ -330,7 +325,7 @@ private Optional<CelMutableAst> maybeFold(
     Object result;
     try {
       result = evaluateExpr(cel, node);
-    } catch (CelValidationException e) {
+    } catch (CelEvaluationException | CelValidationException e) {
       throw new CelOptimizationException(
           "Constant folding failure. Failed to evaluate subtree due to: " + e.getMessage(), e);
     }
@@ -340,17 +335,10 @@ private Optional<CelMutableAst> maybeFold(
     // ex2: optional.ofNonZeroValue(5) -> optional.of(5)
     if (result instanceof Optional<?>) {
       Optional<?> optResult = ((Optional<?>) result);
-      return maybeRewriteOptional(
-          cel.getTypeProvider(), valueProvider, optResult, mutableAst, node.expr());
+      return maybeRewriteOptional(cel.getTypeProvider(), valueProvider, optResult, node.expr());
     }
 
-    CelMutableExpr adaptedResult =
-        maybeAdaptEvaluatedResult(cel.getTypeProvider(), valueProvider, result).orElse(null);
-    if (adaptedResult == null) {
-      return Optional.empty();
-    }
-
-    return Optional.of(astMutator.replaceSubtree(mutableAst, adaptedResult, node.id()));
+    return maybeAdaptEvaluatedResult(cel.getTypeProvider(), valueProvider, result);
   }
 
   private Optional<CelMutableExpr> maybeAdaptEvaluatedResult(
@@ -440,11 +428,10 @@ private Optional<CelMutableExpr> maybeAdaptEvaluatedResult(
     return Optional.empty();
   }
 
-  private Optional<CelMutableAst> maybeRewriteOptional(
+  private Optional<CelMutableExpr> maybeRewriteOptional(
       CelTypeProvider typeProvider,
       CelValueProvider valueProvider,
       Optional<?> optResult,
-      CelMutableAst mutableAst,
       CelMutableExpr expr) {
     Object unwrappedResult = optResult.orElse(null);
     if (unwrappedResult == null) {
@@ -454,7 +441,7 @@ private Optional<CelMutableAst> maybeRewriteOptional(
       // An empty optional value was encountered. Rewrite the tree with optional.none call.
       // This is to account for other optional functions returning an empty optional value
       // e.g: optional.ofNonZeroValue(0)
-      return Optional.of(astMutator.replaceSubtree(mutableAst, newOptionalNoneExpr(), expr.id()));
+      return Optional.of(newOptionalNoneExpr());
     }
 
     if (isCallToFunction(expr, Function.OPTIONAL_OF.getFunction())) {
@@ -472,7 +459,7 @@ private Optional<CelMutableAst> maybeRewriteOptional(
         CelMutableExpr.ofCall(
             CelMutableCall.create(Function.OPTIONAL_OF.getFunction(), adaptedResult));
 
-    return Optional.of(astMutator.replaceSubtree(mutableAst, newOptionalOfCall, expr.id()));
+    return Optional.of(newOptionalOfCall);
   }
 
   private static boolean isCallToFunction(CelMutableExpr expr, String functionName) {
@@ -480,8 +467,8 @@ private static boolean isCallToFunction(CelMutableExpr expr, String functionName
   }
 
   /** Inspects the non-strict calls to determine whether a branch can be removed. */
-  private Optional<CelMutableAst> maybePruneBranches(
-      CelMutableAst mutableAst, Map<String, CelType> identTypes, CelMutableExpr expr) {
+  private Optional<CelMutableExpr> maybePruneBranches(
+      Map<String, CelType> identTypes, CelMutableExpr expr) {
     if (!expr.getKind().equals(Kind.CALL)) {
       return Optional.empty();
     }
@@ -490,7 +477,7 @@ private Optional<CelMutableAst> maybePruneBranches(
     String function = call.function();
     if (function.equals(Operator.LOGICAL_AND.getFunction())
         || function.equals(Operator.LOGICAL_OR.getFunction())) {
-      return maybeShortCircuitCall(mutableAst, identTypes, expr);
+      return maybeShortCircuitCall(identTypes, expr);
     } else if (function.equals(Operator.CONDITIONAL.getFunction())) {
       CelMutableExpr cond = call.args().get(0);
       CelMutableExpr truthy = call.args().get(1);
@@ -501,7 +488,7 @@ private Optional<CelMutableAst> maybePruneBranches(
       }
       CelMutableExpr result = cond.constant().booleanValue() ? truthy : falsy;
 
-      return Optional.of(astMutator.replaceSubtree(mutableAst, result, expr.id()));
+      return Optional.of(result);
     } else if (function.equals(Operator.IN.getFunction())) {
       CelMutableExpr callArg = call.args().get(1);
       if (!callArg.getKind().equals(Kind.LIST)) {
@@ -510,9 +497,7 @@ private Optional<CelMutableAst> maybePruneBranches(
 
       CelMutableList haystack = callArg.list();
       if (haystack.elements().isEmpty()) {
-        return Optional.of(
-            astMutator.replaceSubtree(
-                mutableAst, CelMutableExpr.ofConstant(CelConstant.ofValue(false)), expr.id()));
+        return Optional.of(CelMutableExpr.ofConstant(CelConstant.ofValue(false)));
       }
 
       CelMutableExpr needle = call.args().get(0);
@@ -529,24 +514,13 @@ private Optional<CelMutableAst> maybePruneBranches(
                   && Double.isNaN(needle.constant().doubleValue())) {
                 continue;
               }
-              return Optional.of(
-                  astMutator.replaceSubtree(
-                      mutableAst.expr(),
-                      CelMutableExpr.ofConstant(CelConstant.ofValue(true)),
-                      expr.id()));
+              return Optional.of(CelMutableExpr.ofConstant(CelConstant.ofValue(true)));
             }
 
-            CelType needleType =
-                mutableAst
-                    .getType(needle.id())
-                    .orElseGet(() -> identTypes.get(needle.ident().name()));
+            CelType needleType = identTypes.get(needle.ident().name());
 
             if (needleType != null && isSafeForExactEquality(needleType)) {
-              return Optional.of(
-                  astMutator.replaceSubtree(
-                      mutableAst.expr(),
-                      CelMutableExpr.ofConstant(CelConstant.ofValue(true)),
-                      expr.id()));
+              return Optional.of(CelMutableExpr.ofConstant(CelConstant.ofValue(true)));
             }
           }
         }
@@ -566,7 +540,7 @@ private Optional<CelMutableAst> maybePruneBranches(
         return Optional.empty();
       } else if (lhsIsBooleanConstant
           && (!constantFoldingOptions.enableSafeLogicalOptimization()
-              || evaluatesToBoolean(mutableAst, identTypes, rhs))) {
+              || evaluatesToBoolean(identTypes, rhs))) {
         boolean cond = invertCondition != lhs.constant().booleanValue();
         replacementExpr =
             Optional.of(
@@ -576,7 +550,7 @@ private Optional<CelMutableAst> maybePruneBranches(
                         CelMutableCall.create(Operator.LOGICAL_NOT.getFunction(), rhs)));
       } else if (rhsIsBooleanConstant
           && (!constantFoldingOptions.enableSafeLogicalOptimization()
-              || evaluatesToBoolean(mutableAst, identTypes, lhs))) {
+              || evaluatesToBoolean(identTypes, lhs))) {
         boolean cond = invertCondition != rhs.constant().booleanValue();
         replacementExpr =
             Optional.of(
@@ -586,14 +560,14 @@ private Optional<CelMutableAst> maybePruneBranches(
                         CelMutableCall.create(Operator.LOGICAL_NOT.getFunction(), lhs)));
       }
 
-      return replacementExpr.map(node -> astMutator.replaceSubtree(mutableAst, node, expr.id()));
+      return replacementExpr;
     }
 
     return Optional.empty();
   }
 
-  private Optional<CelMutableAst> maybeShortCircuitCall(
-      CelMutableAst mutableAst, Map<String, CelType> identTypes, CelMutableExpr expr) {
+  private Optional<CelMutableExpr> maybeShortCircuitCall(
+      Map<String, CelType> identTypes, CelMutableExpr expr) {
     CelMutableCall call = expr.call();
     boolean shortCircuit = false;
     boolean skip = true;
@@ -613,7 +587,7 @@ private Optional<CelMutableAst> maybeShortCircuitCall(
       }
 
       if (arg.constant().booleanValue() == shortCircuit) {
-        return Optional.of(astMutator.replaceSubtree(mutableAst, arg, expr.id()));
+        return Optional.of(arg);
       }
     }
 
@@ -621,13 +595,13 @@ private Optional<CelMutableAst> maybeShortCircuitCall(
     if (newArgs.isEmpty()) {
       CelMutableExpr shortCircuitTarget =
           call.args().get(0); // either args(0) or args(1) would work here
-      return Optional.of(astMutator.replaceSubtree(mutableAst, shortCircuitTarget, expr.id()));
+      return Optional.of(shortCircuitTarget);
     }
     if (newArgs.size() == 1) {
       CelMutableExpr remainingArg = newArgs.get(0);
       if (!constantFoldingOptions.enableSafeLogicalOptimization()
-          || evaluatesToBoolean(mutableAst, identTypes, remainingArg)) {
-        return Optional.of(astMutator.replaceSubtree(mutableAst, remainingArg, expr.id()));
+          || evaluatesToBoolean(identTypes, remainingArg)) {
+        return Optional.of(remainingArg);
       }
       return Optional.empty();
     }
@@ -637,8 +611,8 @@ private Optional<CelMutableAst> maybeShortCircuitCall(
         "Folding variadic logical operator is not supported yet.");
   }
 
-  private boolean evaluatesToBoolean(
-      CelMutableAst mutableAst, Map<String, CelType> identTypes, CelMutableExpr expr) {
+  private static boolean evaluatesToBoolean(
+      Map<String, CelType> identTypes, CelMutableExpr expr) {
     if (isExprConstantOfKind(expr, CelConstant.Kind.BOOLEAN_VALUE)) {
       return true;
     }
@@ -654,7 +628,7 @@ private boolean evaluatesToBoolean(
     if (expr.getKind().equals(Kind.IDENT)) {
       return Objects.equals(identTypes.get(expr.ident().name()), SimpleType.BOOL);
     }
-    return mutableAst.getType(expr.id()).map(SimpleType.BOOL::equals).orElse(false);
+    return false;
   }
 
   private boolean isFoldedAggregateLiteral(CelMutableExpr expr) {
diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/InliningOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/InliningOptimizer.java
index 147673e47..61fd19347 100644
--- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/InliningOptimizer.java
+++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/InliningOptimizer.java
@@ -14,8 +14,6 @@
 
 package dev.cel.optimizer.optimizers;
 
-import static com.google.common.collect.ImmutableList.toImmutableList;
-
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.common.primitives.UnsignedLong;
@@ -29,13 +27,14 @@
 import dev.cel.common.ast.CelMutableExpr.CelMutableCall;
 import dev.cel.common.ast.CelMutableExpr.CelMutableStruct;
 import dev.cel.common.navigation.CelNavigableExprUtil;
-import dev.cel.common.navigation.CelNavigableMutableAst;
 import dev.cel.common.navigation.CelNavigableMutableExpr;
+import dev.cel.common.navigation.TraversalOrder;
 import dev.cel.common.types.CelKind;
 import dev.cel.common.types.CelType;
 import dev.cel.common.types.SimpleType;
 import dev.cel.common.values.NullValue;
 import dev.cel.optimizer.AstMutator;
+import dev.cel.optimizer.AstMutator.SubtreeReplacement;
 import dev.cel.optimizer.CelAstOptimizer;
 import java.util.ArrayList;
 import java.util.List;
@@ -105,28 +104,25 @@ public static InliningOptimizer newInstance(
   public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) {
     CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
     for (InlineVariable inlineVariable : inlineVariables) {
-      ImmutableList<CelNavigableMutableExpr> inlinableExprs =
-          CelNavigableMutableAst.fromAst(mutableAst)
-              .getRoot()
-              .allNodes()
-              .filter(node -> canInline(node, inlineVariable.name()))
-              .collect(toImmutableList());
-
-      for (CelNavigableMutableExpr inlinableExpr : inlinableExprs) {
-        CelMutableAst inlineVariableAst = CelMutableAst.fromCelAst(inlineVariable.ast());
-        CelMutableExpr replacementExpr = inlineVariableAst.expr();
-
-        if (inlinableExpr.getKind().equals(Kind.SELECT)
-            && inlinableExpr.expr().select().testOnly()) {
-          replacementExpr = rewritePresenceExpr(inlineVariable, replacementExpr);
-        }
-
-        mutableAst =
-            astMutator.replaceSubtree(
-                mutableAst,
-                CelMutableAst.of(replacementExpr, inlineVariableAst.source()),
-                inlinableExpr.id());
-      }
+      mutableAst =
+          astMutator.mutateUntilFixedPoint(
+              mutableAst,
+              TraversalOrder.POST_ORDER,
+              node -> {
+                if (!canInline(node, inlineVariable.name())) {
+                  return Optional.empty();
+                }
+                CelMutableAst inlineVariableAst = CelMutableAst.fromCelAst(inlineVariable.ast());
+                CelMutableExpr replacementExpr = inlineVariableAst.expr();
+
+                if (node.getKind().equals(Kind.SELECT) && node.expr().select().testOnly()) {
+                  replacementExpr = rewritePresenceExpr(inlineVariable, replacementExpr);
+                }
+
+                return Optional.of(
+                    SubtreeReplacement.of(
+                        node.id(), CelMutableAst.of(replacementExpr, inlineVariableAst.source())));
+              });
     }
 
     return OptimizationResult.create(astMutator.renumberIdsConsecutively(mutableAst).toParsedAst());
diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java
index 5eebb1c54..6a9860750 100644
--- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java
+++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java
@@ -65,6 +65,7 @@
 import java.util.Comparator;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Stream;
 
@@ -161,27 +162,23 @@ private OptimizationResult optimizeUsingCelBlock(CelAbstractSyntaxTree ast, Cel
         break;
       }
 
+      CelMutableExpr targetCseShape = normalizeForEquality(cseCandidates.get(0));
       subexpressions.add(cseCandidates.get(0));
 
       String blockIdentifier = BLOCK_INDEX_PREFIX + blockIdentifierIndex++;
 
       // Replace all CSE candidates with new block index identifier
-      for (CelMutableExpr cseCandidate : cseCandidates) {
-        iterCount++;
-
-        astToModify =
-            astMutator.replaceSubtree(
-                navAst,
-                CelNavigableMutableAst.fromAst(
-                    CelMutableAst.of(
-                        CelMutableExpr.ofIdent(blockIdentifier), navAst.getAst().source())),
-                cseCandidate.id());
-
-        // Retain the existing macro calls in case if the block identifiers are replacing a subtree
-        // that contains a comprehension.
-        sourceToModify.addAllMacroCalls(astToModify.source().getMacroCalls());
-        astToModify = CelMutableAst.of(astToModify.expr(), sourceToModify);
-      }
+      astToModify =
+          astMutator.mutateUntilFixedPoint(
+              astToModify,
+              TraversalOrder.POST_ORDER,
+              node -> normalizeForEquality(node.expr()).equals(targetCseShape),
+              node -> Optional.of(CelMutableExpr.ofIdent(blockIdentifier)));
+
+      // Retain the existing macro calls in case if the block identifiers are replacing a subtree
+      // that contains a comprehension.
+      sourceToModify.addAllMacroCalls(astToModify.source().getMacroCalls());
+      astToModify = CelMutableAst.of(astToModify.expr(), sourceToModify);
     }
 
     if (iterCount >= cseOptions.iterationLimit()) {
diff --git a/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java b/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java
index fa896ebca..f0c3a7045 100644
--- a/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java
+++ b/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java
@@ -41,24 +41,31 @@
 import dev.cel.common.ast.CelMutableExprConverter;
 import dev.cel.common.navigation.CelNavigableAst;
 import dev.cel.common.navigation.CelNavigableExpr;
+import dev.cel.common.navigation.TraversalOrder;
 import dev.cel.common.types.SimpleType;
 import dev.cel.common.types.StructTypeReference;
 import dev.cel.expr.conformance.proto3.TestAllTypes;
 import dev.cel.extensions.CelExtensions;
 import dev.cel.extensions.CelOptionalLibrary;
+import dev.cel.optimizer.AstMutator.SubtreeReplacement;
 import dev.cel.parser.CelStandardMacro;
 import dev.cel.parser.CelUnparser;
 import dev.cel.parser.CelUnparserFactory;
 import java.util.List;
+import java.util.Optional;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
 @RunWith(TestParameterInjector.class)
 public class AstMutatorTest {
   private static final Cel CEL =
-      CelFactory.standardCelBuilder()
+      CelFactory.plannerCelBuilder()
           .setStandardMacros(CelStandardMacro.STANDARD_MACROS)
-          .setOptions(CelOptions.current().populateMacroCalls(true).build())
+          .setOptions(
+              CelOptions.current()
+                  .populateMacroCalls(true)
+                  .enableHeterogeneousNumericComparisons(true)
+                  .build())
           .addMessageTypes(TestAllTypes.getDescriptor())
           .addCompilerLibraries(
               CelOptionalLibrary.INSTANCE, CelExtensions.bindings(), CelExtensions.comprehensions())
@@ -66,6 +73,7 @@ public class AstMutatorTest {
           .setContainer(CelContainer.ofName("cel.expr.conformance.proto3"))
           .addVar("msg", StructTypeReference.create(TestAllTypes.getDescriptor().getFullName()))
           .addVar("x", SimpleType.INT)
+          .addVar("b", SimpleType.BOOL)
           .build();
 
   private static final CelUnparser CEL_UNPARSER = CelUnparserFactory.newUnparser();
@@ -440,7 +448,7 @@ public void memberCallExpr_replaceLeafTarget() throws Exception {
     //      10 [1]       func [4]
     //                4 [3]       5 [5]
     Cel cel =
-        CelFactory.standardCelBuilder()
+        CelFactory.plannerCelBuilder()
             .addFunctionDeclarations(
                 CelFunctionDecl.newFunctionDeclaration(
                     "func",
@@ -463,7 +471,7 @@ public void memberCallExpr_replaceLeafArgument() throws Exception {
     //      10 [1]       func [4]
     //                4 [3]       5 [5]
     Cel cel =
-        CelFactory.standardCelBuilder()
+        CelFactory.plannerCelBuilder()
             .addFunctionDeclarations(
                 CelFunctionDecl.newFunctionDeclaration(
                     "func",
@@ -486,7 +494,7 @@ public void memberCallExpr_replaceMiddleBranchTarget() throws Exception {
     //      10 [1]       func [4]
     //                4 [3]       5 [5]
     Cel cel =
-        CelFactory.standardCelBuilder()
+        CelFactory.plannerCelBuilder()
             .addFunctionDeclarations(
                 CelFunctionDecl.newFunctionDeclaration(
                     "func",
@@ -509,7 +517,7 @@ public void memberCallExpr_replaceMiddleBranchArgument() throws Exception {
     //      10 [1]       func [4]
     //                4 [3]       5 [5]
     Cel cel =
-        CelFactory.standardCelBuilder()
+        CelFactory.plannerCelBuilder()
             .addFunctionDeclarations(
                 CelFunctionDecl.newFunctionDeclaration(
                     "func",
@@ -864,9 +872,13 @@ public void mangleComprehensionVariable_adjacentMacros_differentIterVarTypes() t
   public void mangleComprehensionVariable_macroSourceDisabled_macroCallMapIsEmpty()
       throws Exception {
     Cel cel =
-        CelFactory.standardCelBuilder()
+        CelFactory.plannerCelBuilder()
             .setStandardMacros(CelStandardMacro.STANDARD_MACROS)
-            .setOptions(CelOptions.current().populateMacroCalls(false).build())
+            .setOptions(
+                CelOptions.current()
+                    .populateMacroCalls(false)
+                    .enableHeterogeneousNumericComparisons(true)
+                    .build())
             .build();
     CelAbstractSyntaxTree ast = cel.compile("[false].exists(i, i)").getAst();
 
@@ -1044,6 +1056,159 @@ public void newGlobalCallAst_success() throws Exception {
         .isEqualTo("func([1].exists(x, x >= 1), \"hello\")");
   }
 
+  @Test
+  public void replaceSubtree_withSubtreeReplacement_expr() throws Exception {
+    CelAbstractSyntaxTree ast = CEL.compile("1 + 2").getAst();
+    CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
+    SubtreeReplacement replacement =
+        SubtreeReplacement.of(1, CelMutableExpr.ofConstant(CelConstant.ofValue(10)));
+
+    CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, replacement);
+
+    assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("10 + 2");
+  }
+
+  @Test
+  public void replaceSubtree_withSubtreeReplacement_ast() throws Exception {
+    CelAbstractSyntaxTree ast = CEL.compile("true && false").getAst();
+    CelAbstractSyntaxTree macroAst = CEL.compile("[1].exists(x, x > 0)").getAst();
+    CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
+    SubtreeReplacement replacement = SubtreeReplacement.of(3, CelMutableAst.fromCelAst(macroAst));
+
+    CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, replacement);
+
+    assertThat(CEL_UNPARSER.unparse(result.toParsedAst()))
+        .isEqualTo("true && [1].exists(x, x > 0)");
+    assertThat(result.source().getMacroCalls()).hasSize(1);
+  }
+
+  @Test
+  public void mutateUntilFixedPoint_astRewriter_success() throws Exception {
+    // Repeatedly simplifies addition with 0: "1 + 0 + 0" -> "1"
+    CelAbstractSyntaxTree ast = CEL.compile("1 + 0 + 0").getAst();
+    CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
+
+    CelMutableAst result =
+        AST_MUTATOR.mutateUntilFixedPoint(
+            mutableAst,
+            navAst ->
+                navAst
+                    .getRoot()
+                    .allNodes()
+                    .filter(
+                        node ->
+                            node.getKind().equals(Kind.CALL)
+                                && node.expr().call().function().equals("_+_"))
+                    .filter(
+                        node -> {
+                          List<CelMutableExpr> args = node.expr().call().args();
+                          return args.get(1).getKind().equals(Kind.CONSTANT)
+                              && args.get(1).constant().int64Value() == 0;
+                        })
+                    .map(node -> SubtreeReplacement.of(node.id(), node.expr().call().args().get(0)))
+                    .findFirst());
+
+    assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("1");
+  }
+
+  @Test
+  public void mutateUntilFixedPoint_nodeRewriter_success() throws Exception {
+    // Rewrites nested calls: "func(func(1))" -> "1"
+    Cel cel =
+        CelFactory.plannerCelBuilder()
+            .addFunctionDeclarations(
+                CelFunctionDecl.newFunctionDeclaration(
+                    "func",
+                    CelOverloadDecl.newGlobalOverload(
+                        "func_overload", SimpleType.INT, SimpleType.INT)))
+            .build();
+    CelAbstractSyntaxTree ast = cel.compile("func(func(10))").getAst();
+    CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
+
+    CelMutableAst result =
+        AST_MUTATOR.mutateUntilFixedPoint(
+            mutableAst,
+            TraversalOrder.POST_ORDER,
+            node -> {
+              if (node.getKind().equals(Kind.CALL)
+                  && node.expr().call().function().equals("func")) {
+                return Optional.of(
+                    SubtreeReplacement.of(node.id(), node.expr().call().args().get(0)));
+              }
+              return Optional.empty();
+            });
+
+    assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("10");
+  }
+
+  @Test
+  public void mutateUntilFixedPoint_matcherAndNodeRewriter_success() throws Exception {
+    // Replaces all variables named 'x' with constant 5 in "x + x + x" -> "5 + 5 + 5"
+    CelAbstractSyntaxTree ast = CEL.compile("x + x + x").getAst();
+    CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
+
+    CelMutableAst result =
+        AST_MUTATOR.mutateUntilFixedPoint(
+            mutableAst,
+            TraversalOrder.POST_ORDER,
+            node -> node.getKind().equals(Kind.IDENT) && node.expr().ident().name().equals("x"),
+            node -> Optional.of(CelMutableExpr.ofConstant(CelConstant.ofValue(5))));
+
+    assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("5 + 5 + 5");
+  }
+
+  @Test
+  public void mutateUntilFixedPoint_withReplacementAst_preservesMacroSource() throws Exception {
+    // Replaces identifier 'b' with macro AST in "b && true"
+    CelAbstractSyntaxTree ast = CEL.compile("b && true").getAst();
+    CelAbstractSyntaxTree macroAst = CEL.compile("[1].exists(i, i > 0)").getAst();
+    CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
+
+    CelMutableAst result =
+        AST_MUTATOR.mutateUntilFixedPoint(
+            mutableAst,
+            navAst ->
+                navAst
+                    .getRoot()
+                    .allNodes()
+                    .filter(
+                        node ->
+                            node.getKind().equals(Kind.IDENT)
+                                && node.expr().ident().name().equals("b"))
+                    .map(
+                        node ->
+                            SubtreeReplacement.of(node.id(), CelMutableAst.fromCelAst(macroAst)))
+                    .findFirst());
+
+    assertThat(CEL_UNPARSER.unparse(result.toParsedAst()))
+        .isEqualTo("[1].exists(i, i > 0) && true");
+    assertThat(result.source().getMacroCalls()).hasSize(1);
+    assertThat(CEL.createProgram(CEL.check(result.toParsedAst()).getAst()).eval()).isEqualTo(true);
+  }
+
+  @Test
+  public void mutateUntilFixedPoint_exceedsIterationLimit_throws() throws Exception {
+    // Circular rewrite rule that alternates between 1 and 2
+    CelAbstractSyntaxTree ast = CEL.compile("1").getAst();
+    CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
+
+    IllegalStateException e =
+        assertThrows(
+            IllegalStateException.class,
+            () ->
+                AST_MUTATOR.mutateUntilFixedPoint(
+                    mutableAst,
+                    TraversalOrder.POST_ORDER,
+                    node -> node.getKind().equals(Kind.CONSTANT),
+                    node -> {
+                      long currentVal = node.expr().constant().int64Value();
+                      long nextVal = currentVal == 1L ? 2L : 1L;
+                      return Optional.of(CelMutableExpr.ofConstant(CelConstant.ofValue(nextVal)));
+                    }));
+
+    assertThat(e).hasMessageThat().isEqualTo("Max iteration count reached.");
+  }
+
   @Test
   public void newMemberCallAst_success() throws Exception {
     CelMutableAst targetAst = CelMutableAst.fromCelAst(CEL.compile("'hello'").getAst());
diff --git a/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel b/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel
index 8ea72a261..702fe23f3 100644
--- a/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel
+++ b/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel
@@ -21,6 +21,7 @@ java_library(
         "//common/ast",
         "//common/ast:mutable_expr",
         "//common/navigation",
+        "//common/navigation:common",
         "//common/types",
         "//compiler",
         "//extensions",
diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java
index 613b53ea3..d00b7beac 100644
--- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java
+++ b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java
@@ -128,13 +128,13 @@ private static Cel setupEnv(CelBuilder celBuilder) {
             CelExtensions.comprehensions(),
             CelExtensions.bindings(),
             CelOptionalLibrary.INSTANCE,
-            CelExtensions.math(CEL_OPTIONS),
+            CelExtensions.math(),
             CelExtensions.strings(),
             CelExtensions.sets(CEL_OPTIONS),
             CelExtensions.encoders(CEL_OPTIONS))
         .addRuntimeLibraries(
             CelOptionalLibrary.INSTANCE,
-            CelExtensions.math(CEL_OPTIONS),
+            CelExtensions.math(),
             CelExtensions.strings(),
             CelExtensions.sets(CEL_OPTIONS),
             CelExtensions.encoders(CEL_OPTIONS))