Make FoldContext have equality

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/FoldContext.java

@@ -14,6 +14,8 @@
 import org.elasticsearch.xpack.esql.core.QlClientException;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 
+import java.util.Objects;
+
 /**
  * Context passed to {@link Expression#fold}. This is not thread safe.
  */
@@ -33,6 +35,26 @@ public FoldContext(long allowedBytes) {
         this.allowedBytes = allowedBytes;
     }
 
+    long initialAllowedBytes() {
+        return initialAllowedBytes;
+    }
+
+    long allowedBytes() {
+        return allowedBytes;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (o == null || getClass() != o.getClass()) return false;
+        FoldContext that = (FoldContext) o;
+        return initialAllowedBytes == that.initialAllowedBytes && allowedBytes == that.allowedBytes;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(initialAllowedBytes, allowedBytes);
+    }
+
     /**
      * Track an allocation. Best to call this <strong>before</strong> allocating
      * if possible, but after is ok if the allocation is small.

x-pack/plugin/esql-core/src/test/java/org/elasticsearch/xpack/esql/core/expression/FoldContextTests.java

@@ -9,11 +9,45 @@
 
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.test.EqualsHashCodeTestUtils;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 
 import static org.hamcrest.Matchers.equalTo;
 
 public class FoldContextTests extends ESTestCase {
+    public void testEq() {
+        EqualsHashCodeTestUtils.checkEqualsAndHashCode(randomFoldContext(), this::copy, this::mutate);
+    }
+
+    private FoldContext randomFoldContext() {
+        FoldContext ctx = new FoldContext(randomNonNegativeLong());
+        if (randomBoolean()) {
+            ctx.trackAllocation(Source.EMPTY, randomLongBetween(0, ctx.initialAllowedBytes()));
+        }
+        return ctx;
+    }
+
+    private FoldContext copy(FoldContext ctx) {
+        FoldContext copy = new FoldContext(ctx.initialAllowedBytes());
+        copy.trackAllocation(Source.EMPTY, ctx.initialAllowedBytes() - ctx.allowedBytes());
+        return copy;
+    }
+
+    private FoldContext mutate(FoldContext ctx) {
+        if (randomBoolean()) {
+            FoldContext differentInitial = new FoldContext(ctx.initialAllowedBytes() + 1);
+            differentInitial.trackAllocation(Source.EMPTY, differentInitial.initialAllowedBytes() - ctx.allowedBytes());
+            assertThat(differentInitial.allowedBytes(), equalTo(ctx.allowedBytes()));
+            return differentInitial;
+        } else {
+            FoldContext differentAllowed = new FoldContext(ctx.initialAllowedBytes());
+            long allowed = randomValueOtherThan(ctx.allowedBytes(), () -> randomLongBetween(0, ctx.initialAllowedBytes()));
+            differentAllowed.trackAllocation(Source.EMPTY, ctx.initialAllowedBytes() - allowed);
+            assertThat(differentAllowed.allowedBytes(), equalTo(allowed));
+            return differentAllowed;
+        }
+    }
+
     public void testTrackAllocation() {
         FoldContext ctx = new FoldContext(10);
         ctx.trackAllocation(Source.synthetic("shouldn't break"), 10);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Case.java

@@ -266,7 +266,7 @@ public boolean foldable() {
      * {@code EVAL c=CASE(b, bar, bort)}.
      */
     public Expression partiallyFold(FoldContext ctx) {
-        // TODO don't fold here - look for literal TRUE
+        // TODO don't throw away the results of any `fold`. That might mean looking for literal TRUE on the conditions.
         List<Expression> newChildren = new ArrayList<>(children().size());
         boolean modified = false;
         for (Condition condition : conditions) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/date/DateExtract.java

@@ -113,7 +113,7 @@ public ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) {
         if (children().get(0).foldable()) {
             ChronoField chrono = chronoField(toEvaluator.foldCtx());
             if (chrono == null) {
-                BytesRef field = (BytesRef) children().getFirst().fold(toEvaluator.foldCtx());
+                BytesRef field = (BytesRef) children().get(0).fold(toEvaluator.foldCtx());
                 throw new InvalidArgumentException("invalid date field for [{}]: {}", sourceText(), field.utf8ToString());
             }
             return new DateExtractConstantEvaluator.Factory(source(), fieldEvaluator, chrono, configuration().zoneId());

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalOptimizerContext.java

@@ -35,7 +35,7 @@ public boolean equals(Object obj) {
         if (obj == this) return true;
         if (obj == null || obj.getClass() != this.getClass()) return false;
         var that = (LogicalOptimizerContext) obj;
-        return Objects.equals(this.configuration, that.configuration);
+        return this.configuration.equals(that.configuration) && this.foldCtx.equals(that.foldCtx);
     }
 
     @Override