QL: Introduce infrastructure for surrogate functions (#54795)

Some functions act as shortcuts for more verbose declarations (sometimes
with certain constraints). This PR removes the boilerplate around
declaring such functions as well as a dedicated rule for the optimizer
to perform the actual substitution.

Fix #54334

(cherry picked from commit 3231d01b0c583deb89252fafe84db48878da3246)

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/expression/function/scalar/string/Wildcard.java

@@ -6,14 +6,11 @@
 
 package org.elasticsearch.xpack.eql.expression.function.scalar.string;
 
-import org.elasticsearch.xpack.eql.EqlIllegalArgumentException;
 import org.elasticsearch.xpack.eql.util.StringUtils;
 import org.elasticsearch.xpack.ql.expression.Expression;
-import org.elasticsearch.xpack.ql.expression.Expressions;
 import org.elasticsearch.xpack.ql.expression.Expressions.ParamOrdinal;
 import org.elasticsearch.xpack.ql.expression.function.scalar.ScalarFunction;
-import org.elasticsearch.xpack.ql.expression.gen.pipeline.Pipe;
+import org.elasticsearch.xpack.ql.expression.function.scalar.BaseSurrogateFunction;
-import org.elasticsearch.xpack.ql.expression.gen.script.ScriptTemplate;
 import org.elasticsearch.xpack.ql.expression.predicate.logical.Or;
 import org.elasticsearch.xpack.ql.expression.predicate.regex.Like;
 import org.elasticsearch.xpack.ql.tree.NodeInfo;
@@ -33,7 +30,7 @@ import static org.elasticsearch.xpack.ql.expression.TypeResolutions.isStringAndE
  * EQL wildcard function. Matches the form:
  *     wildcard(field, "*wildcard*pattern*", ...)
  */
-public class Wildcard extends ScalarFunction {
+public class Wildcard extends BaseSurrogateFunction {
 
     private final Expression field;
     private final List<Expression> patterns;
@@ -95,26 +92,7 @@ public class Wildcard extends ScalarFunction {
     }
 
     @Override
-    public boolean foldable() {
+    public ScalarFunction makeSubstitute() {
-        return Expressions.foldable(children()) && asLikes().foldable();
-    }
-
-    @Override
-    public Object fold() {
-        return asLikes().fold();
-    }
-
-    @Override
-    protected Pipe makePipe() {
-        throw new EqlIllegalArgumentException("Wildcard.makePipe() should not be called directly");
-    }
-
-    @Override
-    public ScriptTemplate asScript() {
-        throw new EqlIllegalArgumentException("Wildcard.asScript() should not be called directly");
-    }
-
-    public ScalarFunction asLikes() {
         ScalarFunction result = null;
 
         for (Expression pattern: patterns) {
@@ -125,4 +103,4 @@ public class Wildcard extends ScalarFunction {
 
         return result;
     }
-}
+}

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/optimizer/Optimizer.java

@@ -6,7 +6,6 @@
 
 package org.elasticsearch.xpack.eql.optimizer;
 
-import org.elasticsearch.xpack.eql.expression.function.scalar.string.Wildcard;
 import org.elasticsearch.xpack.eql.util.StringUtils;
 import org.elasticsearch.xpack.ql.expression.Expression;
 import org.elasticsearch.xpack.ql.expression.predicate.logical.Not;
@@ -24,6 +23,7 @@ import org.elasticsearch.xpack.ql.optimizer.OptimizerRules.OptimizerRule;
 import org.elasticsearch.xpack.ql.optimizer.OptimizerRules.PropagateEquals;
 import org.elasticsearch.xpack.ql.optimizer.OptimizerRules.PruneFilters;
 import org.elasticsearch.xpack.ql.optimizer.OptimizerRules.PruneLiteralsInOrderBy;
+import org.elasticsearch.xpack.ql.optimizer.OptimizerRules.ReplaceSurrogateFunction;
 import org.elasticsearch.xpack.ql.optimizer.OptimizerRules.SetAsOptimized;
 import org.elasticsearch.xpack.ql.plan.logical.Filter;
 import org.elasticsearch.xpack.ql.plan.logical.LogicalPlan;
@@ -39,6 +39,9 @@ public class Optimizer extends RuleExecutor<LogicalPlan> {
 
     @Override
     protected Iterable<RuleExecutor<LogicalPlan>.Batch> batches() {
+        Batch substitutions = new Batch("Operator Replacement", Limiter.ONCE,
+                new ReplaceSurrogateFunction());
+                
         Batch operators = new Batch("Operator Optimization",
                 new ConstantFolding(),
                 // boolean
@@ -49,7 +52,6 @@ public class Optimizer extends RuleExecutor<LogicalPlan> {
                 new ReplaceNullChecks(),
                 new PropagateEquals(),
                 new CombineBinaryComparisons(),
-                new ReplaceWildcardFunction(),
                 // prune/elimination
                 new PruneFilters(),
                 new PruneLiteralsInOrderBy()
@@ -58,16 +60,7 @@ public class Optimizer extends RuleExecutor<LogicalPlan> {
         Batch label = new Batch("Set as Optimized", Limiter.ONCE,
                 new SetAsOptimized());
 
-        return Arrays.asList(operators, label);
+        return Arrays.asList(substitutions, operators, label);
-    }
-
-
-    private static class ReplaceWildcardFunction extends OptimizerRule<Filter> {
-
-        @Override
-        protected LogicalPlan rule(Filter filter) {
-            return filter.transformExpressionsUp(e -> e instanceof Wildcard ? ((Wildcard) e).asLikes() : e);
-        }
     }
 
     private static class ReplaceWildcards extends OptimizerRule<Filter> {

x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/function/scalar/BaseSurrogateFunction.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.xpack.ql.expression.function.scalar;
+
+import org.elasticsearch.xpack.ql.expression.Expression;
+import org.elasticsearch.xpack.ql.expression.gen.pipeline.Pipe;
+import org.elasticsearch.xpack.ql.expression.gen.script.ScriptTemplate;
+import org.elasticsearch.xpack.ql.tree.Source;
+
+import java.util.List;
+
+public abstract class BaseSurrogateFunction extends ScalarFunction implements SurrogateFunction {
+
+    private ScalarFunction lazySubstitute;
+
+    public BaseSurrogateFunction(Source source) {
+        super(source);
+    }
+
+    public BaseSurrogateFunction(Source source, List<Expression> fields) {
+        super(source, fields);
+    }
+
+    @Override
+    public ScalarFunction substitute() {
+        if (lazySubstitute == null) {
+            lazySubstitute = makeSubstitute();
+        }
+        return lazySubstitute;
+    }
+
+    protected abstract ScalarFunction makeSubstitute();
+
+    @Override
+    public boolean foldable() {
+        return substitute().foldable();
+    }
+
+    @Override
+    public Object fold() {
+        return substitute().fold();
+    }
+
+    @Override
+    protected Pipe makePipe() {
+        return substitute().asPipe();
+    }
+
+    @Override
+    public ScriptTemplate asScript() {
+        return substitute().asScript();
+    }
+}

x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/function/scalar/SurrogateFunction.java

@@ -0,0 +1,12 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.xpack.ql.expression.function.scalar;
+
+public interface SurrogateFunction {
+
+    ScalarFunction substitute();
+}

x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/optimizer/OptimizerRules.java

@@ -9,6 +9,7 @@ import org.elasticsearch.xpack.ql.expression.Expression;
 import org.elasticsearch.xpack.ql.expression.Expressions;
 import org.elasticsearch.xpack.ql.expression.Literal;
 import org.elasticsearch.xpack.ql.expression.Order;
+import org.elasticsearch.xpack.ql.expression.function.scalar.SurrogateFunction;
 import org.elasticsearch.xpack.ql.expression.predicate.BinaryOperator;
 import org.elasticsearch.xpack.ql.expression.predicate.BinaryPredicate;
 import org.elasticsearch.xpack.ql.expression.predicate.Negatable;
@@ -1000,6 +1001,21 @@ public final class OptimizerRules {
         }
 
     }
+
+    public static class ReplaceSurrogateFunction extends OptimizerExpressionRule {
+
+        public ReplaceSurrogateFunction() {
+            super(TransformDirection.DOWN);
+        }
+
+        @Override
+        protected Expression rule(Expression e) {
+            if (e instanceof SurrogateFunction) {
+                e = ((SurrogateFunction) e).substitute();
+            }
+            return e;
+        }
+    }
     
     public static final class PruneFilters extends OptimizerRule<Filter> {