Output intermediate graph with type info on error

12 files changed, 299 insertions, 147 deletions

diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/ModelImporter.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/ModelImporter.java
index b88d7423a82..f22e89cc8bb 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/ModelImporter.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/ModelImporter.java
@@ -11,8 +11,7 @@ import com.yahoo.searchlib.rankingexpression.parser.ParseException;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.functions.Rename;
 import com.yahoo.tensor.functions.TensorFunction;
-import com.yahoo.text.ParenthesisExpressionPrettyPrinter;
-import com.yahoo.text.Text;
+import com.yahoo.text.ExpressionFormatter;
 import com.yahoo.yolean.Exceptions;
 
 import java.io.File;
@@ -54,7 +53,7 @@ public abstract class ModelImporter implements MlModelImporter {
     protected static ImportedModel convertIntermediateGraphToModel(IntermediateGraph graph, String modelSource) {
         ImportedModel model = new ImportedModel(graph.name(), modelSource);
         log.log(Level.FINE, () -> "Intermediate graph created from '" + modelSource + "':\n" +
-                                  ParenthesisExpressionPrettyPrinter.prettyPrint(graph.toFullString()));
+                                  ExpressionFormatter.inTwoColumnMode(20).format(graph.toFullString()));
         System.out.println("Intermediate graph created from '" + modelSource + "':\n" + graph.toFullString());
 
         graph.optimize();
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Argument.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Argument.java
index 7d1b6a61e2e..d5671889e01 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Argument.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Argument.java
@@ -56,6 +56,7 @@ public class Argument extends IntermediateOperation {
 
     @Override
     public String toString() {
-        return "Argument(" + standardNamingType + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + ":\tArgument(" + standardNamingType + ")";
     }
+
 }
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Const.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Const.java
index 6571e77a198..e7c4b355217 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Const.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Const.java
@@ -89,7 +89,7 @@ public class Const extends IntermediateOperation {
 
     @Override
     public String toString() {
-        return "Const(" + type + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + ":\tConst(" + type + ")";
     }
 
 }
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/ExpandDims.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/ExpandDims.java
index 8f7d3755005..18ff602148a 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/ExpandDims.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/ExpandDims.java
@@ -30,7 +30,7 @@ public class ExpandDims extends IntermediateOperation {
         if ( ! allInputTypesPresent(2)) return null;
 
         IntermediateOperation axisOperation = inputs().get(1);
-        if (!axisOperation.getConstantValue().isPresent()) {
+        if (  !axisOperation.getConstantValue().isPresent()) {
             throw new IllegalArgumentException("ExpandDims in " + name + ": Axis must be a constant.");
         }
         Tensor axis = axisOperation.getConstantValue().get().asTensor();
@@ -103,8 +103,8 @@ public class ExpandDims extends IntermediateOperation {
 
     @Override
     public String toFullString() {
-        return "ExpandDims(" + inputs().get(0).toFullString() + ", " +
-               inputs().get(1).toFullString() + ", " + expandDimensions + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + ":\tExpandDims(" + inputs().get(0).toFullString() + ", " +
+               inputs().get(1).toFullString() + ", " + expandDimensions + ")";
     }
 
 }
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Join.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Join.java
index c5e6ae49a25..5c7acc8a0ee 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Join.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Join.java
@@ -118,8 +118,8 @@ public class Join extends IntermediateOperation {
 
     @Override
     public String toFullString() {
-        return "Join(" + inputs().get(0).toFullString() + ", " +
-               inputs().get(1).toFullString() + ", " + operator + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + ":\tJoin(" + inputs().get(0).toFullString() + ", " +
+               inputs().get(1).toFullString() + ", " + operator + ")";
     }
 
 }
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Map.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Map.java
index 4f70c46e459..5fba7cc6da0 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Map.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Map.java
@@ -41,7 +41,7 @@ public class Map extends IntermediateOperation {
 
     @Override
     public String toFullString() {
-        return "Map(" + inputs().get(0).toFullString() + ", " + operator + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + "\tMap(" + inputs().get(0).toFullString() + ", " + operator + ")";
     }
 
 }
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/MatMul.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/MatMul.java
index 73aa40927be..3f7c32294ed 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/MatMul.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/MatMul.java
@@ -5,8 +5,7 @@ import ai.vespa.rankingexpression.importer.DimensionRenamer;
 import ai.vespa.rankingexpression.importer.OrderedTensorType;
 import com.yahoo.tensor.TensorType;
 import com.yahoo.tensor.functions.TensorFunction;
-import com.yahoo.text.ParenthesisExpressionPrettyPrinter;
-import com.yahoo.text.Text;
+import com.yahoo.text.ExpressionFormatter;
 
 import java.util.List;
 import java.util.Optional;
@@ -81,13 +80,12 @@ public class MatMul extends IntermediateOperation {
 
         throw new IllegalArgumentException("Expected 2 dimensions in the " + inputDescription + " to " + this +
                                            " but got just " + dimensions + " from\n" +
-                                           ParenthesisExpressionPrettyPrinter.prettyPrint(supplier.toFullString()));
+                                           ExpressionFormatter.inTwoColumnMode(70).format(supplier.toFullString()));
     }
 
     @Override
     public String toFullString() {
-        return "MatMul(" + inputs().get(0).toFullString() + ", " +
-                           inputs().get(1).toFullString() + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + ":\tMatMul(" + inputs().get(0).toFullString() + ", " + inputs().get(1).toFullString() + ")";
     }
 
     @Override
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Sum.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Sum.java
index df5c4e9cbfa..8b5d06e8d56 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Sum.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Sum.java
@@ -104,6 +104,7 @@ public class Sum extends IntermediateOperation {
                 builder.add(TensorType.Dimension.indexed(dimension.name(), 1L));
             }
         }
+        System.out.println("----------> Sum input type is " + inputType + ", keepDimensions: " + keepDimensions + ", result: " + builder.build());
         return builder.build();
     }
 
@@ -114,8 +115,8 @@ public class Sum extends IntermediateOperation {
 
     @Override
     public String toFullString() {
-        return "Sum(" + inputs().get(0).toFullString() + ", " +
-               inputs().get(1).toFullString() + ", " + reduceDimensions + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + ":\tSum[keep_dims=" + shouldKeepDimensions() + "](" +
+               inputs().get(0).toFullString() + ", " + inputs().get(1).toFullString() + ", " + reduceDimensions + ")";
     }
 
 }
diff --git a/vespajlib/src/main/java/com/yahoo/text/ExpressionFormatter.java b/vespajlib/src/main/java/com/yahoo/text/ExpressionFormatter.java
new file mode 100644
index 00000000000..b7670ab70e6
--- /dev/null
+++ b/vespajlib/src/main/java/com/yahoo/text/ExpressionFormatter.java
@@ -0,0 +1,142 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.text;
+
+/**
+ * Formats any parenthesis expression.
+ * In addition to the obvious this can also operate in "two column mode",
+ * wherein each chunk that will be formatted on a separate line may optionally
+ * contain a prefix marked by a start and end tab sign which will be printed in a left column of the given fixed size.
+ * The prefix itself is not formatted but will be cut if too long.
+ *
+ * @author bratseth
+ */
+public class ExpressionFormatter {
+
+    private static final int indentUnit = 2;
+
+    /** The size of the first column, or 0 if none */
+    private int firstColumnSize = 0;
+
+    private ExpressionFormatter(int firstColumnSize) {
+        this.firstColumnSize = firstColumnSize;
+    }
+
+    public String format(String parenthesisExpression) {
+        StringBuilder b = new StringBuilder();
+        format(parenthesisExpression, 0, b);
+        return b.toString();
+    }
+
+    private void format(String expression, int indent, StringBuilder b) {
+        if (expression.isEmpty()) return;
+        expression = appendFirstColumn(expression, b);
+
+        Markup next = Markup.next(expression);
+
+        appendIndent( ! next.isClose() || next.position() > 0 ? indent : indent - 2, b);
+        if (next.isEmpty()) {
+            b.append(expression);
+        }
+        else if (next.isComma()) {
+            b.append(expression, 0, next.position() + 1).append("\n");
+            format(expression.substring(next.position() + 1), indent, b);
+        }
+        else {
+            if ( next.isClose() && next.position() > 0) { // content before end parenthesis: content, newline, then end parenthesis
+                b.append(expression, 0, next.position()).append("\n");
+                appendFirstColumn(")", b);
+                appendIndent(indent - 2, b);
+                b.append(")\n");
+            }
+            else {
+                b.append(expression, 0, next.position() + 1).append("\n");
+            }
+            format(expression.substring(next.position() + 1), indent + (next.isOpen() ? indentUnit : -indentUnit), b);
+        }
+    }
+
+    private String appendFirstColumn(String expression, StringBuilder b) {
+        if (firstColumnSize == 0) return expression;
+
+        while (expression.charAt(0) == ' ')
+            expression = expression.substring(1);
+
+        if (expression.charAt(0) == '\t') {
+            int tab2 = expression.indexOf('\t', 1);
+            if (tab2 >= 0) {
+                String firstColumn = expression.substring(1, tab2);
+                b.append(asSize(firstColumnSize, firstColumn)).append(" ");
+                return expression.substring(tab2 + 1);
+            }
+        }
+        appendIndent(firstColumnSize + 1, b);
+        return expression;
+    }
+
+    private void appendIndent(int indent, StringBuilder b) {
+        b.append(" ".repeat(Math.max(0, indent)));
+    }
+
+    private String asSize(int size, String s) {
+        if (s.length() > size)
+            return s.substring(0, size);
+        else
+            return s + " ".repeat(size - s.length());
+    }
+
+    /** Convenience method creating a formatter and using it to format the given expression */
+    public static String on(String parenthesisExpression) {
+        return new ExpressionFormatter(0).format(parenthesisExpression);
+    }
+
+    public static ExpressionFormatter inTwoColumnMode(int firstColumnSize) {
+        return new ExpressionFormatter(firstColumnSize);
+    }
+
+    /** Contains the next position of each kind of markup, or Integer.MAX_VALUE if not present */
+    private static class Markup {
+
+        final int open, close, comma;
+
+        private Markup(int open, int close, int comma) {
+            this.open = open;
+            this.close = close;
+            this.comma = comma;
+        }
+
+        int position() {
+            return Math.min(Math.min(open, close), comma);
+        }
+
+        boolean isOpen() {
+            return open < close && open < comma;
+        }
+
+        boolean isClose() {
+            return close < open && close < comma;
+        }
+
+        boolean isComma() {
+            return comma < open && comma < close;
+        }
+
+        boolean isEmpty() {
+            return open == Integer.MAX_VALUE && close == Integer.MAX_VALUE && comma == Integer.MAX_VALUE;
+        }
+
+        static Markup next(String expression) {
+            int nextOpen = expression.indexOf('(');
+            int nextClose = expression.indexOf(')');
+            int nextComma = expression.indexOf(',');
+            if (nextOpen < 0)
+                nextOpen = Integer.MAX_VALUE;
+            if (nextClose < 0)
+                nextClose = Integer.MAX_VALUE;
+            if (nextComma < 0)
+                nextComma = Integer.MAX_VALUE;
+            return new Markup(nextOpen, nextClose, nextComma);
+        }
+
+    }
+
+}
diff --git a/vespajlib/src/main/java/com/yahoo/text/ParenthesisExpressionPrettyPrinter.java b/vespajlib/src/main/java/com/yahoo/text/ParenthesisExpressionPrettyPrinter.java
deleted file mode 100644
index ad235d78679..00000000000
--- a/vespajlib/src/main/java/com/yahoo/text/ParenthesisExpressionPrettyPrinter.java
+++ /dev/null
@@ -1,47 +0,0 @@
-// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package com.yahoo.text;
-
-/**
- * Pretty prints any parenthesis expression
- *
- * @author bratseth
- */
-public class ParenthesisExpressionPrettyPrinter {
-
-    private static final int indentUnit = 2;
-
-    public static String prettyPrint(String parenthesisExpression) {
-        StringBuilder b = new StringBuilder();
-        prettyPrint(parenthesisExpression, 0, b);
-        return b.toString();
-    }
-
-    private static void prettyPrint(String expression, int indent, StringBuilder b) {
-        int nextStartParenthesis = expression.indexOf("(");
-        int nextEndParenthesis = expression.indexOf(")");
-        if (nextStartParenthesis < 0)
-            nextStartParenthesis = Integer.MAX_VALUE;
-        if (nextEndParenthesis < 0)
-            nextEndParenthesis = Integer.MAX_VALUE;
-
-        boolean start = nextStartParenthesis < nextEndParenthesis;
-        int nextParenthesis = Math.min(nextStartParenthesis, nextEndParenthesis);
-
-        int effectiveIndent = start || nextParenthesis > 0 ? indent : indent - 2;
-        b.append(" ".repeat(Math.max(0, effectiveIndent)));
-        if (nextParenthesis == Integer.MAX_VALUE) {
-            b.append(expression);
-        }
-        else {
-            if (! start && nextParenthesis > 0) {
-                b.append(expression, 0, nextParenthesis).append("\n");
-                b.append(" ".repeat(Math.max(0, indent - 2))).append(")\n");
-            }
-            else {
-                b.append(expression, 0, nextParenthesis + 1).append("\n");
-            }
-            prettyPrint(expression.substring(nextParenthesis + 1), indent + (start ? indentUnit : -indentUnit), b);
-        }
-    }
-
-}
diff --git a/vespajlib/src/test/java/com/yahoo/text/ExpressionFormatterTest.java b/vespajlib/src/test/java/com/yahoo/text/ExpressionFormatterTest.java
new file mode 100644
index 00000000000..6dfb2e6fc8a
--- /dev/null
+++ b/vespajlib/src/test/java/com/yahoo/text/ExpressionFormatterTest.java
@@ -0,0 +1,140 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.text;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author bratseth
+ */
+public class ExpressionFormatterTest {
+
+    @Test
+    public void testBasic() {
+        String expected =
+                "foo(\n" +
+                "  bar(\n" +
+                "    baz(\n" +
+                "    )\n" +
+                "  )\n" +
+                ")\n";
+        assertPrettyPrint(expected, "foo(bar(baz()))");
+    }
+
+    @Test
+    public void testArgument() {
+        String expected =
+                "foo(\n" +
+                "  bar(\n" +
+                "    baz(\n" +
+                "      hello world\n" +
+                "    )\n" +
+                "  )\n" +
+                ")\n";
+        assertPrettyPrint(expected, "foo(bar(baz(hello world)))");
+    }
+
+    @Test
+    public void testMultipleArguments() {
+        String expected =
+                "foo(\n" +
+                "  bar(\n" +
+                "    baz(\n" +
+                "      hello world,\n" +
+                "      37\n" +
+                "    )\n" +
+                "  )\n" +
+                ")\n";
+        assertPrettyPrint(expected, "foo(bar(baz(hello world,37)))");
+    }
+
+    @Test
+    public void testUnmatchedStart() {
+        String expected =
+                "foo(\n" +
+                "  (\n" +
+                "    bar(\n" +
+                "      baz(\n" +
+                "      )\n" +
+                "    )\n" +
+                "  )\n";
+        assertPrettyPrint(expected, "foo((bar(baz()))");
+    }
+
+    @Test
+    public void testUnmatchedEnd() {
+        String expected =
+                "foo(\n" +
+                "  bar(\n" +
+                "    baz(\n" +
+                "    )\n" +
+                "  )\n" +
+                ")\n" +
+                ")\n";
+        assertPrettyPrint(expected, "foo(bar(baz())))");
+    }
+
+    @Test
+    public void testNoParenthesis() {
+        String expected =
+                "foo bar baz";
+        assertPrettyPrint(expected, "foo bar baz");
+    }
+
+    @Test
+    public void testEmpty() {
+        String expected =
+                "";
+        assertPrettyPrint(expected, "");
+    }
+
+    @Test
+    public void test2ColumnMode() {
+        String expected =
+                "1:  foo(\n" +
+                "      bar(\n" +
+                "        baz(\n" +
+                "2:        hello world\n" +
+                "        )\n" +
+                "t(o   )\n" +
+                "    )\n";
+        ExpressionFormatter pp = ExpressionFormatter.inTwoColumnMode(3);
+        assertEquals(expected, pp.format("\t1:\tfoo(bar(baz(\t2:\thello world)\tt(o)@olong:\t))"));
+    }
+
+    @Test
+    public void test2ColumnModeMultipleArguments() {
+        String expected =
+                "1:  foo(\n" +
+                "      bar(\n" +
+                "        baz(\n" +
+                "2:        hello world,\n" +
+                "3:        37\n" +
+                "        )\n" +
+                "t(o   )\n" +
+                "    )\n";
+        ExpressionFormatter pp = ExpressionFormatter.inTwoColumnMode(3);
+        assertEquals(expected, pp.format("\t1:\tfoo(bar(baz(\t2:\thello world,\t3:\t37)\tt(o)@olong:\t))"));
+    }
+
+    @Test
+    public void test2ColumnModeMultipleArgumentsWithSpaces() {
+        String expected =
+                "1:  foo(\n" +
+                "      bar(\n" +
+                "        baz(\n" +
+                "2:        hello world,\n" +
+                "3:        37\n" +
+                "        )\n" +
+                "t(o   )\n" +
+                "    )\n";
+        ExpressionFormatter pp = ExpressionFormatter.inTwoColumnMode(3);
+        assertEquals(expected, pp.format("\t1:\tfoo(bar(baz(\t2:\thello world, \t3:\t37)\tt(o)@olong:\t))"));
+    }
+
+    private void assertPrettyPrint(String expected, String expression) {
+        assertEquals(expected, ExpressionFormatter.on(expression));
+    }
+
+}
diff --git a/vespajlib/src/test/java/com/yahoo/text/ParenthesisExpressionPrettyPrinterTest.java b/vespajlib/src/test/java/com/yahoo/text/ParenthesisExpressionPrettyPrinterTest.java
deleted file mode 100644
index 79bdc6a5318..00000000000
--- a/vespajlib/src/test/java/com/yahoo/text/ParenthesisExpressionPrettyPrinterTest.java
+++ /dev/null
@@ -1,82 +0,0 @@
-// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package com.yahoo.text;
-
-import org.junit.Test;
-
-import static org.junit.Assert.assertEquals;
-
-/**
- * @author bratseth
- */
-public class ParenthesisExpressionPrettyPrinterTest {
-
-    @Test
-    public void testBasic() {
-        String expected =
-                "foo(\n" +
-                "  bar(\n" +
-                "    baz(\n" +
-                "    )\n" +
-                "  )\n" +
-                ")\n";
-        assertPrettyPrint(expected, "foo(bar(baz()))");
-    }
-
-    @Test
-    public void testInnerContent() {
-        String expected =
-                "foo(\n" +
-                "  bar(\n" +
-                "    baz(\n" +
-                "      hello world\n" +
-                "    )\n" +
-                "  )\n" +
-                ")\n";
-        assertPrettyPrint(expected, "foo(bar(baz(hello world)))");
-    }
-    @Test
-    public void testUnmatchedStart() {
-        String expected =
-                "foo(\n" +
-                "  (\n" +
-                "    bar(\n" +
-                "      baz(\n" +
-                "      )\n" +
-                "    )\n" +
-                "  )\n" +
-                "  ";
-        assertPrettyPrint(expected, "foo((bar(baz()))");
-    }
-
-    @Test
-    public void testUnmatchedEnd() {
-        String expected =
-                "foo(\n" +
-                "  bar(\n" +
-                "    baz(\n" +
-                "    )\n" +
-                "  )\n" +
-                ")\n" +
-                ")\n";
-        assertPrettyPrint(expected, "foo(bar(baz())))");
-    }
-
-    @Test
-    public void testNoParenthesis() {
-        String expected =
-                "foo bar baz";
-        assertPrettyPrint(expected, "foo bar baz");
-    }
-
-    @Test
-    public void testEmpty() {
-        String expected =
-                "";
-        assertPrettyPrint(expected, "");
-    }
-
-    private void assertPrettyPrint(String expected, String expression) {
-        assertEquals(expected, ParenthesisExpressionPrettyPrinter.prettyPrint(expression));
-    }
-
-}