Only optimize boolean expressions on primitives

Only transform a && b to if(a, b, false) etc.0
if both a and b are primitives, not tensors,
as if requires both branches to return the same type.

1 files changed, 42 insertions, 26 deletions

diff --git a/config-model/src/test/java/com/yahoo/schema/expressiontransforms/BooleanExpressionTransformerTestCase.java b/config-model/src/test/java/com/yahoo/schema/expressiontransforms/BooleanExpressionTransformerTestCase.java
index d692b69d3c8..d06573f7bae 100644
--- a/config-model/src/test/java/com/yahoo/schema/expressiontransforms/BooleanExpressionTransformerTestCase.java
+++ b/config-model/src/test/java/com/yahoo/schema/expressiontransforms/BooleanExpressionTransformerTestCase.java
@@ -2,10 +2,13 @@
 package com.yahoo.schema.expressiontransforms;
 
 import com.yahoo.searchlib.rankingexpression.RankingExpression;
+import com.yahoo.searchlib.rankingexpression.Reference;
 import com.yahoo.searchlib.rankingexpression.evaluation.MapContext;
 import com.yahoo.searchlib.rankingexpression.evaluation.MapTypeContext;
 import com.yahoo.searchlib.rankingexpression.rule.OperationNode;
 import com.yahoo.searchlib.rankingexpression.transform.TransformContext;
+import com.yahoo.tensor.TensorType;
+import com.yahoo.tensor.evaluation.TypeContext;
 import org.junit.jupiter.api.Test;
 
 import java.util.Map;
@@ -20,7 +23,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
 public class BooleanExpressionTransformerTestCase {
 
     @Test
-    public void testTransformer() throws Exception {
+    public void booleanTransformation() throws Exception {
         assertTransformed("if (a, b, false)", "a && b");
         assertTransformed("if (a, true, b)", "a || b");
         assertTransformed("if (a, true, b + c)", "a || b + c");
@@ -33,16 +36,17 @@ public class BooleanExpressionTransformerTestCase {
     }
 
     @Test
-    public void testIt() throws Exception {
-        assertTransformed("if(1 - 1, true, 1 - 1)", "1 - 1 || 1 - 1");
+    public void noTransformationOnTensorTypes() throws Exception {
+        var typeContext = new MapTypeContext();
+        typeContext.setType(Reference.fromIdentifier("tensorA"), TensorType.fromSpec("tensor(x{})"));
+        typeContext.setType(Reference.fromIdentifier("tensorB"), TensorType.fromSpec("tensor(x{})"));
+        assertUntransformed("tensorA && tensorB", typeContext);
+        assertTransformed("a && (tensorA * tensorB)","a && ( tensorA * tensorB)", typeContext);
     }
 
     @Test
     public void testNotSkewingNonBoolean() throws Exception {
-        assertTransformed("a + b + c * d + e + f", "a + b + c * d + e + f");
-        var expr = new BooleanExpressionTransformer()
-                .transform(new RankingExpression("a + b + c * d + e + f"),
-                        new TransformContext(Map.of(), new MapTypeContext()));
+        var expr = assertTransformed("a + b + c * d + e + f", "a + b + c * d + e + f");
         assertTrue(expr.getRoot() instanceof OperationNode);
         OperationNode root = (OperationNode) expr.getRoot();
         assertEquals(5, root.operators().size());
@@ -51,41 +55,53 @@ public class BooleanExpressionTransformerTestCase {
 
     @Test
     public void testTransformPreservesPrecedence() throws Exception {
-        assertUnTransformed("a");
-        assertUnTransformed("a + b");
-        assertUnTransformed("a + b + c");
-        assertUnTransformed("a * b");
-        assertUnTransformed("a + b * c + d");
-        assertUnTransformed("a + b + c * d + e + f");
-        assertUnTransformed("a * b + c + d + e * f");
-        assertUnTransformed("(a * b) + c + d + e * f");
-        assertUnTransformed("(a * b + c) + d + e * f");
-        assertUnTransformed("a * (b + c) + d + e * f");
-        assertUnTransformed("(a * b) + (c + (d + e)) * f");
+        assertUntransformed("a");
+        assertUntransformed("a + b");
+        assertUntransformed("a + b + c");
+        assertUntransformed("a * b");
+        assertUntransformed("a + b * c + d");
+        assertUntransformed("a + b + c * d + e + f");
+        assertUntransformed("a * b + c + d + e * f");
+        assertUntransformed("(a * b) + c + d + e * f");
+        assertUntransformed("(a * b + c) + d + e * f");
+        assertUntransformed("a * (b + c) + d + e * f");
+        assertUntransformed("(a * b) + (c + (d + e)) * f");
+    }
+
+    private void assertUntransformed(String input) throws Exception {
+        assertUntransformed(input, new MapTypeContext());
+    }
+
+    private void assertUntransformed(String input, MapTypeContext typeContext) throws Exception {
+        assertTransformed(input, input, typeContext);
     }
 
-    private void assertUnTransformed(String input) throws Exception {
-        assertTransformed(input, input);
+    private RankingExpression assertTransformed(String expected, String input) throws Exception {
+        return assertTransformed(expected, input, new MapTypeContext());
     }
 
-    private void assertTransformed(String expected, String input) throws Exception {
+    private RankingExpression assertTransformed(String expected, String input, MapTypeContext typeContext) throws Exception {
+        MapContext context = contextWithSingleLetterVariables(typeContext);
         var transformedExpression = new BooleanExpressionTransformer()
                                             .transform(new RankingExpression(input),
-                                                       new TransformContext(Map.of(), new MapTypeContext()));
+                                                       new TransformContext(Map.of(), typeContext));
 
         assertEquals(new RankingExpression(expected), transformedExpression, "Transformed as expected");
 
-        MapContext context = contextWithSingleLetterVariables();
         var inputExpression = new RankingExpression(input);
         assertEquals(inputExpression.evaluate(context).asBoolean(),
                      transformedExpression.evaluate(context).asBoolean(),
                      "Transform and original input are equivalent");
+        return transformedExpression;
     }
 
-    private MapContext contextWithSingleLetterVariables() {
+    private MapContext contextWithSingleLetterVariables(MapTypeContext typeContext) {
         var context = new MapContext();
-        for (int i = 0; i < 26; i++)
-            context.put(Character.toString(i + 97), Math.floorMod(i, 2));
+        for (int i = 0; i < 26; i++) {
+            String name = Character.toString(i + 97);
+            typeContext.setType(Reference.fromIdentifier(name), TensorType.empty);
+            context.put(name, Math.floorMod(i, 2));
+        }
         return context;
     }