From 0437f08f4e7a41fe27fa54e1710a2840cfd607ca Mon Sep 17 00:00:00 2001
From: Ievgen Degtiarenko <ievgen.degtiarenko@elastic.co>
Date: Fri, 21 Mar 2025 09:08:41 +0100
Subject: [PATCH] Attempt to make some operations cheaper (#125228)

---
 .../xpack/esql/core/tree/Node.java            | 12 ++++++++--
 .../xpack/esql/plan/QueryPlan.java            | 22 +++++++++++++++++--
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/tree/Node.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/tree/Node.java
index 667720536d884..82c31c0dbdd7e 100644
--- a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/tree/Node.java
+++ b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/tree/Node.java
@@ -67,7 +67,11 @@ public List<T> children() {
     @SuppressWarnings("unchecked")
     public void forEachDown(Consumer<? super T> action) {
         action.accept((T) this);
-        children().forEach(c -> c.forEachDown(action));
+        // please do not refactor it to a for-each loop to avoid
+        // allocating iterator that performs concurrent modification checks and extra stack frames
+        for (int c = 0, size = children.size(); c < size; c++) {
+            children.get(c).forEachDown(action);
+        }
     }
 
     @SuppressWarnings("unchecked")
@@ -81,7 +85,11 @@ public <E extends T> void forEachDown(Class<E> typeToken, Consumer<? super E> ac
 
     @SuppressWarnings("unchecked")
     public void forEachUp(Consumer<? super T> action) {
-        children().forEach(c -> c.forEachUp(action));
+        // please do not refactor it to a for-each loop to avoid
+        // allocating iterator that performs concurrent modification checks and extra stack frames
+        for (int c = 0, size = children.size(); c < size; c++) {
+            children.get(c).forEachUp(action);
+        }
         action.accept((T) this);
     }
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/QueryPlan.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/QueryPlan.java
index d70177508d847..5ebf3a211a57c 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/QueryPlan.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/QueryPlan.java
@@ -140,7 +140,24 @@ private static Object doTransformExpression(Object arg, Function<Expression, ? e
         // preserving the type information is hacky and weird (a lot of context needs to be passed around and the lambda itself
         // has no type info so it's difficult to have automatic checking without having base classes).
 
-        if (arg instanceof Collection<?> c) {
+        if (arg instanceof List<?> c) {
+            List<Object> transformed = null;
+            boolean hasChanged = false;
+            // please do not refactor it to a for-each loop to avoid
+            // allocating iterator that performs concurrent modification checks and extra stack frames
+            for (int i = 0, size = c.size(); i < size; i++) {
+                var e = c.get(i);
+                Object next = doTransformExpression(e, traversal);
+                if (e.equals(next) == false) {
+                    if (hasChanged == false) {
+                        hasChanged = true;
+                        transformed = new ArrayList<>(c);
+                    }
+                    transformed.set(i, next);
+                }
+            }
+            return hasChanged ? transformed : arg;
+        } else if (arg instanceof Collection<?> c) {
             List<Object> transformed = null;
             boolean hasChanged = false;
             int i = 0;
@@ -149,13 +166,14 @@ private static Object doTransformExpression(Object arg, Function<Expression, ? e
                 if (e.equals(next) == false) {
                     if (hasChanged == false) {
                         hasChanged = true;
+                        // TODO if collection is a set then this silently changes its semantics by allowing duplicates
+                        // We should fix it or confirm this branch is never needed
                         transformed = new ArrayList<>(c);
                     }
                     transformed.set(i, next);
                 }
                 i++;
             }
-
             return hasChanged ? transformed : arg;
         }