Support transpiling more test features: iterate, IRPrinter, assertions, GraphEvaluator

4 files changed, 158 insertions, 57 deletions

diff --git a/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/Arg.java b/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/Arg.java
deleted file mode 100644
index 4334641f..00000000
--- a/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/Arg.java
+++ /dev/null
@@ -1,23 +0,0 @@
-package com.seaofnodes.simple.transpile;
-
-public sealed interface Arg {
-    record IntBot() implements Arg {
-    }
-
-    record IntTop() implements Arg {
-    }
-
-    record IntConstant(long value) implements Arg {
-    }
-
-    static Arg parse(String s) {
-        if (s.equals("TypeInteger.BOT"))
-            return new IntBot();
-        if (s.equals("TypeInteger.TOP"))
-            return new IntTop();
-        var parts = s.split("\\(|\\)");
-        if (parts.length != 2 | !parts[0].equals("TypeInteger.constant"))
-            throw new RuntimeException("Unexpected Argument " + s);
-        return new IntConstant(Long.parseLong(parts[1]));
-    }
-}
diff --git a/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/AstVisitor.java b/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/AstVisitor.java
index 581969c0..36b3d368 100644
--- a/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/AstVisitor.java
+++ b/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/AstVisitor.java
@@ -45,33 +45,58 @@ class AstVisitor extends TreeScanner<Void, Void> {
     @Override
     public Void visitMethodInvocation(MethodInvocationTree node, Void unused) {
         if (current != null) {
-            if (node.getMethodSelect().toString().endsWith(".parse")) {
-                current.showAfterParse = switch (node.getArguments().size()) {
+            var methodSelect = node.getMethodSelect().toString();
+            var args = node.getArguments();
+
+            if (methodSelect.endsWith(".parse")) {
+                current.showAfterParse = switch (args.size()) {
                     case 0 -> false;
-                    case 1 -> (Boolean) ((LiteralTree) node.getArguments().getFirst()).getValue();
+                    case 1 -> (Boolean) literal(args.getFirst());
                     default -> throw new RuntimeException("unexpected parse arguments" + node);
                 };
             }
-            if (node.getMethodSelect().toString().endsWith("assertEquals")) {
-                var args = node.getArguments();
+            if (methodSelect.endsWith(".iterate")) {
+                current.iterate = true;
+                current.showAfterIterate = switch (args.size()) {
+                    case 0 -> false;
+                    case 1 -> (Boolean) literal(args.getFirst());
+                    default -> throw new RuntimeException("unexpected parse arguments" + node);
+                };
+            }
+            if (methodSelect.endsWith("assertEquals")) {
                 if (inCatch) {
                     if (args.size() != 2)
                         throw new RuntimeException("Unexpected number of arguments " + node);
-                    if (args.get(0).toString().contains("e.getMessage()")) {
-                        current.parseErrorMessage = args.get(1).toString();
-                    } else if (args.get(1).toString().contains("e.getMessage()")) {
-                        current.parseErrorMessage = args.get(0).toString();
+                    if (args.get(1).toString().contains("e.getMessage()")) {
+                        current.parseErrorMessage = (String) literal(args.get(0));
                     }
                 } else {
                     if (args.size() != 2)
                         throw new RuntimeException("Unexpected number of arguments " + node);
-                    if (args.get(0).toString().contains("stop.toString()")) {
-                        current.assertStopEquals = args.get(1).toString();
-                    } else if (args.get(1).toString().contains("stop.toString()")) {
-                        current.assertStopEquals = args.get(0).toString();
+                    if (args.get(1).toString().endsWith(".toString()") || args.get(1).toString().endsWith(".print()")) {
+                        current.assertStopEquals = (String) literal(args.get(0));
+                    } else if (args.get(1).toString().contains("evaluate(")) {
+                        Number result = (Number) literal(args.get(0));
+                        var evalArgs = ((MethodInvocationTree) args.get(1)).getArguments();
+                        if (evalArgs.size() != 2 || !evalArgs.get(0).toString().equals("stop"))
+                            throw new RuntimeException("Unexpected eval arguments " + node);
+                        long parameter = Long.parseLong(literal(evalArgs.get(1)).toString());
+                        current.evaluations.add(new TestMethod.Evaluation(result, parameter));
                     }
                 }
             }
+            if (methodSelect.endsWith("assertTrue")) {
+                var arg = args.getFirst().toString();
+                current.assertStopRetCtrlIsProj = arg.contains("stop.ret().ctrl() instanceof ProjNode");
+                current.assertStopRetCtrlIsCProj = arg.contains("stop.ret().ctrl() instanceof CProjNode");
+                current.assertStopRetCtrlIsRegion = arg.contains("stop.ret().ctrl() instanceof RegionNode");
+            }
+            if (methodSelect.endsWith("prettyPrint")) {
+                if (args.size() < 2 || args.size() > 3 || !args.get(0).toString().equals("stop") || !args.get(1).toString().equals("99"))
+                    throw new RuntimeException("unexpected prettyPrint args " + node);
+                current.irPrinter = true;
+                current.irPrinterLLVM = args.size() == 3 && (Boolean) literal(args.get(2));
+            }
         }
         return super.visitMethodInvocation(node, unused); // otherwise we would miss the first part of new Parser("").parse()
     }
@@ -84,10 +109,22 @@ class AstVisitor extends TreeScanner<Void, Void> {
                 switch (args.size()) {
                     case 1 -> {
                     }
-                    case 2 -> current.parserArg = Arg.parse(args.get(1).toString());
+                    case 2 -> current.parserArg = switch (args.get(1).toString()) {
+                        case "TypeInteger.BOT" -> new TestMethod.Arg.IntBot();
+                        case "TypeInteger.TOP" -> new TestMethod.Arg.IntTop();
+                        case String s -> {
+                            var parts = s.split("\\(|\\)");
+                            if (parts.length != 2 | !parts[0].equals("TypeInteger.constant"))
+                                throw new RuntimeException("Unexpected Argument " + s);
+                            yield new TestMethod.Arg.IntConstant(Long.parseLong(parts[1]));
+                        }
+                    };
                     default -> throw new RuntimeException("Parser constructor with unexpected arguments: " + node);
                 }
-                current.parserInput = args.get(0).toString();
+                current.parserInput = (String) literal(args.get(0));
+            } else if (node.getIdentifier().toString().equals("GraphVisualizer")) {
+                // only happens in Chapter02Tests
+                current.showAfterParse = true;
             }
         }
         return super.visitNewClass(node, unused);
@@ -100,4 +137,18 @@ class AstVisitor extends TreeScanner<Void, Void> {
         inCatch = false;
         return null;
     }
+
+    @Override
+    public Void visitAssignment(AssignmentTree node, Void unused) {
+        if (current != null && node.getVariable().toString().equals("Node._disablePeephole")) {
+            if ((Boolean) literal(node.getExpression())) {
+                current.disablePeephole = true;
+            }
+        }
+        return super.visitAssignment(node, unused);
+    }
+
+    private Object literal(ExpressionTree t) {
+        return ((LiteralTree) t).getValue();
+    }
 }
diff --git a/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/TestMethod.java b/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/TestMethod.java
index 56357c4a..107b5f56 100644
--- a/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/TestMethod.java
+++ b/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/TestMethod.java
@@ -1,10 +1,36 @@
 package com.seaofnodes.simple.transpile;
 
+import java.util.ArrayList;
+
 public class TestMethod {
     public String name;
     public String parserInput;
     public Arg parserArg;
+    // if this is set you can ignore the rest
     public String parseErrorMessage;
+    public boolean disablePeephole;
     public boolean showAfterParse;
+    public boolean iterate;
+    public boolean showAfterIterate;
+    public boolean irPrinter;
+    public boolean irPrinterLLVM;
     public String assertStopEquals;
+    public boolean assertStopRetCtrlIsProj;
+    public boolean assertStopRetCtrlIsCProj;
+    public boolean assertStopRetCtrlIsRegion;
+    public final ArrayList<Evaluation> evaluations = new ArrayList<>();
+
+    public sealed interface Arg {
+        record IntBot() implements Arg {
+        }
+
+        record IntTop() implements Arg {
+        }
+
+        record IntConstant(long value) implements Arg {
+        }
+    }
+
+    public record Evaluation(Number result, long parameter) {
+    }
 }
diff --git a/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/printers/RustPrinter.java b/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/printers/RustPrinter.java
index 7362ef60..5397eb5e 100644
--- a/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/printers/RustPrinter.java
+++ b/transpile-tests/src/main/java/com/seaofnodes/simple/transpile/printers/RustPrinter.java
@@ -1,6 +1,5 @@
 package com.seaofnodes.simple.transpile.printers;
 
-import com.seaofnodes.simple.transpile.Arg;
 import com.seaofnodes.simple.transpile.JUnitParser;
 import com.seaofnodes.simple.transpile.TestClass;
 import com.seaofnodes.simple.transpile.TestMethod;
@@ -61,11 +60,17 @@ public class RustPrinter {
                 use crate::sea_of_nodes::types::Types;
                 """);
 
-        for (TestMethod test : tests) {
-            if (test.parseErrorMessage != null) {
-                out.println("use crate::sea_of_nodes::tests::test_error;");
-                break;
-            }
+        if (tests.stream().anyMatch(t -> t.parseErrorMessage != null)) {
+            out.println("use crate::sea_of_nodes::tests::test_error;");
+        }
+        if (tests.stream().anyMatch(t -> t.irPrinter && !t.irPrinterLLVM)) {
+            out.println("use crate::sea_of_nodes::ir_printer::pretty_print;");
+        }
+        if (tests.stream().anyMatch(t -> t.irPrinterLLVM)) {
+            out.println("use crate::sea_of_nodes::ir_printer::pretty_print_llvm;");
+        }
+        if (tests.stream().anyMatch(t -> !t.evaluations.isEmpty())) {
+            out.println("use crate::sea_of_nodes::graph_evaluator::evaluate;");
         }
 
         for (TestMethod test : tests) {
@@ -76,33 +81,75 @@ public class RustPrinter {
     private void printTest(TestMethod test, PrintStream out) {
         out.println();
         out.println("#[test]");
-        out.println("fn " + test.name + "() {");
+        out.println("fn " + snakeCase(test.name) + "() {");
         if (test.parseErrorMessage != null) {
-            out.println("test_error(" + test.parserInput + ", " + test.parseErrorMessage + ")");
+            out.println("test_error(" + string(test.parserInput) + ", " + string(test.parseErrorMessage) + ");");
         } else if (test.parserInput != null) {
             out.println("let arena = DroplessArena::new();");
             out.println("let types = Types::new(&arena);");
-            out.printf("let mut parser = Parser::new(%s%s, &types);\n", multiline(test.parserInput), switch (test.parserArg) {
-                case null -> "";
-                case Arg.IntBot ignored -> ", types.ty_int_bot";
-                case Arg.IntTop ignored -> ", types.ty_int_top";
-                case Arg.IntConstant c -> ", types.get_int(" + c.value() + ")";
-            });
+            if (test.parserArg == null) {
+                out.println("let mut parser = Parser::new(" + string(test.parserInput, true) + ", &types);");
+            } else {
+                out.printf("let mut parser = Parser::new_with_arg(%s, &types, %s);\n", string(test.parserInput, true), switch (test.parserArg) {
+                    case TestMethod.Arg.IntBot ignored -> "types.ty_int_bot";
+                    case TestMethod.Arg.IntTop ignored -> "types.ty_int_top";
+                    case TestMethod.Arg.IntConstant c -> "types.get_int(" + c.value() + ")";
+                });
+            }
+            if (test.disablePeephole) {
+                out.println("parser.nodes.disable_peephole = true;");
+            }
             out.println("let stop = parser.parse().unwrap();");
             if (test.showAfterParse) {
                 out.println("parser.show_graph();");
             }
+            if (test.iterate) {
+                out.println("parser.iterate(stop);");
+            }
+            if (test.showAfterIterate) {
+                out.println("parser.show_graph();");
+            }
+            if (test.irPrinter) {
+                out.println("println!(\"{}\", pretty_print" + (test.irPrinterLLVM ? "_llvm" : "") + "(&parser.nodes, stop, 99));");
+            }
             if (test.assertStopEquals != null) {
-                out.printf("assert_eq!(parser.print(stop), %s);\n", test.assertStopEquals);
+                out.println("assert_eq!(parser.print(stop), " + string(test.assertStopEquals) + ");");
+            }
+            if (test.assertStopRetCtrlIsProj) {
+                out.println("assert!(matches!(parser.nodes.ret_ctrl(stop), Node::Proj(_)));");
+            }
+            if (test.assertStopRetCtrlIsCProj) {
+                out.println("assert!(matches!(parser.nodes.ret_ctrl(stop), Node::CProj(_)));");
+            }
+            if (test.assertStopRetCtrlIsRegion) {
+                out.println("assert!(matches!(parser.nodes.ret_ctrl(stop), Node::Region{..}));");
+            }
+            for (var evaluation : test.evaluations) {
+                out.println("assert_eq!(evaluate(&parser.nodes, stop, Some(" + evaluation.parameter() + "), None), " + evaluation.result() + ");");
             }
+        } else {
+            out.println("todo!();");
         }
         out.println("}");
     }
 
-    private String multiline(String string) {
-        if (!string.contains("\\n")) {
-            return string;
+    private String string(String value) {
+        return string(value, false);
+    }
+
+    private String string(String value, boolean multiline) {
+        var escaped = value.replace("\\", "\\\\")
+                .replace("\r", "\\r")
+                .replace("\t", "\\t")
+                .replace("\"", "\\\"");
+        if (multiline && escaped.contains("\n")) {
+            return "\"\\\n" + escaped + "\"";
+        } else {
+            return "\"" + escaped.replace("\n", "\\n") + "\"";
         }
-        return string.replaceFirst("^\"", "\"\\\\\\\\n").replace("\\n", "\n");
+    }
+
+    private String snakeCase(String identifier) {
+        return identifier.replaceAll("[A-Z]+", "_$0").toLowerCase();
     }
 }