From bae08999cbb449bffaa0f550cce0a5dcc11e1815 Mon Sep 17 00:00:00 2001
From: Adam Retter <adam.retter@googlemail.com>
Date: Sat, 2 Nov 2024 15:19:11 +0100
Subject: [PATCH] [optimise] Rewrite predicates that contain a statically
 computable position() expression with a number, e.g. [position() eq 1]
 becomes [1]. This means that the evaluation of the predicate will now instead
 have POSITIONAL execution mode instead of BOOLEAN, which in turn means that
 the predicate can be evaluated once instead of single-stepping over each
 context item.

---
 .../main/java/org/exist/xquery/Optimizer.java | 62 ++++++++++++++++++-
 1 file changed, 61 insertions(+), 1 deletion(-)

diff --git a/exist-core/src/main/java/org/exist/xquery/Optimizer.java b/exist-core/src/main/java/org/exist/xquery/Optimizer.java
index 53b4f334eae..54adfc34b9e 100644
--- a/exist-core/src/main/java/org/exist/xquery/Optimizer.java
+++ b/exist-core/src/main/java/org/exist/xquery/Optimizer.java
@@ -30,16 +30,20 @@
  */
 package org.exist.xquery;
 
+import com.evolvedbinary.j8fu.tuple.Tuple2;
 import org.exist.storage.DBBroker;
 import org.exist.xquery.functions.array.ArrayConstructor;
+import org.exist.xquery.functions.fn.FunPosition;
 import org.exist.xquery.pragmas.Optimize;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.exist.xquery.util.ExpressionDumper;
+import org.exist.xquery.value.Type;
 
 import javax.annotation.Nullable;
 import java.util.*;
 
+import static com.evolvedbinary.j8fu.tuple.Tuple.Tuple;
 import static org.apache.commons.lang3.ArrayUtils.isNotEmpty;
 
 /**
@@ -165,7 +169,39 @@ public void visitLocationStep(final LocationStep locationStep) {
         }
     }
 
-    public void visitFilteredExpr(FilteredExpression filtered) {
+    /**
+     * Test whether we can optimise position() withing a predicate to a POSITIONAL version of a predicate.
+     *
+     * @param predicate the predicate to test
+     *
+     * @return the inner expr to be replaced and the literal value to replace it with if it can be optimised, null otherwise.
+     */
+    private @Nullable Tuple2<Expression, LiteralValue> canOptimizeToPositionalPredicate(final Predicate predicate) {
+        if (predicate.getSubExpressionCount() == 1) {
+            final Expression predInnerExpr = predicate.getSubExpression(0);
+            if (predInnerExpr instanceof final ValueComparison valueComparison && valueComparison.getRelation() == Constants.Comparison.EQ) {
+                final Expression left = valueComparison.getLeft();
+                final Expression right = valueComparison.getRight();
+
+                if (left instanceof final InternalFunctionCall leftInternalFunctionCall && leftInternalFunctionCall.getFunction() instanceof FunPosition
+                        && right instanceof final LiteralValue rightLiteralValue && rightLiteralValue.returnsType() == Type.INTEGER) {
+
+                    // NOTE(AR) if the predicate is like [position() eq 1] or [position() = 1] then we could rewrite it to [1]
+                    return Tuple(predInnerExpr, rightLiteralValue);
+
+                } else if (left instanceof final LiteralValue leftLiteralValue && leftLiteralValue.returnsType() == Type.INTEGER
+                        && right instanceof final InternalFunctionCall rightInternalFunctionCall && rightInternalFunctionCall.getFunction() instanceof FunPosition) {
+
+                    // NOTE(AR) if the predicate is like [1 eq position()]  or [1 = position()] then we could rewrite it to [1]
+                    return Tuple(predInnerExpr, leftLiteralValue);
+                }
+            }
+        }
+        return null;
+    }
+
+    @Override
+    public void visitFilteredExpr(final FilteredExpression filtered) {
         super.visitFilteredExpr(filtered);
 
         // check if filtered expression can be simplified:
@@ -309,6 +345,30 @@ public void visitGeneralComparison(GeneralComparison comparison) {
 
     public void visitPredicate(Predicate predicate) {
         ++predicates;
+
+        // NOTE(AR) try and optimise position() withing a predicate to a POSITIONAL version of a predicate
+        @Nullable final Tuple2<Expression, LiteralValue> predOptimizableInnerExpr = canOptimizeToPositionalPredicate(predicate);
+        if (predOptimizableInnerExpr != null) {
+            try {
+                // Replace the predicates expression with a literal number
+
+                predicate.replace(predOptimizableInnerExpr._1, predOptimizableInnerExpr._2);
+
+                // traverse to outermost predicate and then call analyze! to make sure execution mode is correct
+                Expression parentExpr = predicate;
+                Predicate outermostPredicate = predicate;
+                while ((parentExpr = parentExpr.getParent()) != null) {
+                    if (parentExpr instanceof final Predicate parentPredicate) {
+                        outermostPredicate = parentPredicate;
+                    }
+                }
+                outermostPredicate.analyze(new AnalyzeContextInfo(outermostPredicate.getParent(), 0));
+
+            } catch (final XPathException e) {
+                LOG.warn("Failed to optimize expression within predicate: {}: {}", predOptimizableInnerExpr, e.getMessage(), e);
+            }
+        }
+
         super.visitPredicate(predicate);
         --predicates;
     }