Revert "Revert "Short circuit boolean expressions""

2 files changed, 107 insertions, 1 deletions

diff --git a/config-model/src/main/java/com/yahoo/schema/expressiontransforms/BooleanExpressionTransformer.java b/config-model/src/main/java/com/yahoo/schema/expressiontransforms/BooleanExpressionTransformer.java
new file mode 100644
index 00000000000..e1affaf2aac
--- /dev/null
+++ b/config-model/src/main/java/com/yahoo/schema/expressiontransforms/BooleanExpressionTransformer.java
@@ -0,0 +1,105 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.schema.expressiontransforms;
+
+import com.yahoo.searchlib.rankingexpression.evaluation.BooleanValue;
+import com.yahoo.searchlib.rankingexpression.rule.ArithmeticNode;
+import com.yahoo.searchlib.rankingexpression.rule.ArithmeticOperator;
+import com.yahoo.searchlib.rankingexpression.rule.CompositeNode;
+import com.yahoo.searchlib.rankingexpression.rule.ConstantNode;
+import com.yahoo.searchlib.rankingexpression.rule.ExpressionNode;
+import com.yahoo.searchlib.rankingexpression.rule.IfNode;
+import com.yahoo.searchlib.rankingexpression.transform.ExpressionTransformer;
+import com.yahoo.searchlib.rankingexpression.transform.TransformContext;
+
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Deque;
+import java.util.Iterator;
+import java.util.List;
+
+/**
+ * Transforms
+ *        a && b into if(a, b, false)
+ * and
+ *        a || b into if(a, true, b)
+ * to avoid computing b if a is false and true respectively.
+ *
+ * This may increase performance since boolean expressions are short-circuited.
+ *
+ * @author bratseth
+ */
+public class BooleanExpressionTransformer extends ExpressionTransformer<TransformContext> {
+
+    @Override
+    public ExpressionNode transform(ExpressionNode node, TransformContext context) {
+        if (node instanceof CompositeNode composite)
+            node = transformChildren(composite, context);
+
+        if (node instanceof ArithmeticNode arithmetic)
+            node = transformBooleanArithmetics(arithmetic);
+
+        return node;
+    }
+
+    private ExpressionNode transformBooleanArithmetics(ArithmeticNode node) {
+        Iterator<ExpressionNode> child = node.children().iterator();
+
+        // Transform in precedence order:
+        Deque<ChildNode> stack = new ArrayDeque<>();
+        stack.push(new ChildNode(ArithmeticOperator.OR, child.next()));
+        for (Iterator<ArithmeticOperator> it = node.operators().iterator(); it.hasNext() && child.hasNext();) {
+            ArithmeticOperator op = it.next();
+            if ( ! stack.isEmpty()) {
+                while (stack.size() > 1 && ! op.hasPrecedenceOver(stack.peek().op)) {
+                    popStack(stack);
+                }
+            }
+            stack.push(new ChildNode(op, child.next()));
+        }
+        while (stack.size() > 1)
+            popStack(stack);
+        return stack.getFirst().child;
+    }
+
+    private void popStack(Deque<ChildNode> stack) {
+        ChildNode rhs = stack.pop();
+        ChildNode lhs = stack.peek();
+
+        ExpressionNode combination;
+        if (rhs.op == ArithmeticOperator.AND)
+            combination = andByIfNode(lhs.child, rhs.child);
+        else if (rhs.op == ArithmeticOperator.OR)
+            combination = orByIfNode(lhs.child, rhs.child);
+        else
+            combination = new ArithmeticNode(List.of(lhs.child, rhs.child), List.of(rhs.op));
+        lhs.child = combination;
+    }
+
+
+    private IfNode andByIfNode(ExpressionNode a, ExpressionNode b) {
+        return new IfNode(a, b, new ConstantNode(new BooleanValue(false)));
+    }
+
+    private IfNode orByIfNode(ExpressionNode a, ExpressionNode b) {
+        return new IfNode(a, new ConstantNode(new BooleanValue(true)), b);
+    }
+
+    /** A child with the operator to be applied to it when combining it with the previous child. */
+    private static class ChildNode {
+
+        final ArithmeticOperator op;
+        ExpressionNode child;
+
+        public ChildNode(ArithmeticOperator op, ExpressionNode child) {
+            this.op = op;
+            this.child = child;
+        }
+
+        @Override
+        public String toString() {
+            return child.toString();
+        }
+
+    }
+
+}
diff --git a/config-model/src/main/java/com/yahoo/schema/expressiontransforms/ExpressionTransforms.java b/config-model/src/main/java/com/yahoo/schema/expressiontransforms/ExpressionTransforms.java
index 86aedd4332a..132597ee75e 100644
--- a/config-model/src/main/java/com/yahoo/schema/expressiontransforms/ExpressionTransforms.java
+++ b/config-model/src/main/java/com/yahoo/schema/expressiontransforms/ExpressionTransforms.java
@@ -35,7 +35,8 @@ public class ExpressionTransforms {
                                  new FunctionInliner(),
                                  new FunctionShadower(),
                                  new TensorMaxMinTransformer(),
-                                 new Simplifier());
+                                 new Simplifier(),
+                                 new BooleanExpressionTransformer());
     }
 
     public RankingExpression transform(RankingExpression expression, RankProfileTransformContext context) {