fix(printer): Underscore at the end of the line (#77)

"\n" was missing when printing underscore on the last line

Author: potatomashed

cpp/printer.cc

@@ -476,10 +476,13 @@ inline void DocPrinter::Append(const Node &doc, const PrinterConfig &cfg) {
 }
 
 inline ::mlc::Str DocPrinter::GetString() const {
-  mlc::Str text = output_.str();
+  std::string text = output_.str();
   // Remove any trailing indentation
-  while (!text->empty() && std::isspace(text->back())) {
-    text->pop_back();
+  while (!text.empty() && text.back() == ' ') {
+    text.pop_back();
+  }
+  if (!text.empty() && text.back() != '\n') {
+    text.push_back('\n');
   }
   return DecorateText(text, line_starts_, options_, MergeAndExemptSpans(underlines_, underlines_exempted_));
 }

python/mlc/testing/toy_ir/ir.py

@@ -38,8 +38,8 @@ class Add(Expr):
     rhs: Expr
 
     def __ir_print__(self, printer: mlcp.IRPrinter, path: mlcp.ObjectPath) -> mlt.Node:
-        lhs: mlt.Expr = printer(self.lhs, path=path["a"])
-        rhs: mlt.Expr = printer(self.rhs, path=path["b"])
+        lhs: mlt.Expr = printer(self.lhs, path=path["lhs"])
+        rhs: mlt.Expr = printer(self.rhs, path=path["rhs"])
         return lhs + rhs
 
 
@@ -49,9 +49,9 @@ class Assign(Stmt):
     lhs: Var = mlcd.field(structure="bind")
 
     def __ir_print__(self, printer: mlcp.IRPrinter, path: mlcp.ObjectPath) -> mlt.Node:
-        rhs: mlt.Expr = printer(self.rhs, path=path["b"])
+        rhs: mlt.Expr = printer(self.rhs, path=path["rhs"])
         printer.var_def(self.lhs.name, obj=self.lhs)
-        lhs: mlt.Expr = printer(self.lhs, path=path["a"])
+        lhs: mlt.Expr = printer(self.lhs, path=path["lhs"])
         return mlt.Assign(lhs=lhs, rhs=rhs)
 
 

tests/python/test_printer_ir_printer.py

@@ -1,6 +1,8 @@
 import re
 
 import mlc.printer as mlcp
+import pytest
+from mlc.printer import ObjectPath
 from mlc.testing.toy_ir import Add, Assign, Func, Var
 
 
@@ -105,3 +107,108 @@ def f(a):
         r"  return a_0x[0-9A-Fa-f]+$",
         mlcp.to_python(f, mlcp.PrinterConfig(print_addr_on_dup_var=True)),
     )
+
+
+@pytest.mark.parametrize(
+    "path, expected",
+    [
+        (
+            ObjectPath.root()["args"][0],
+            """
+def f(a, b):
+      ^
+  c = a + b
+  return c
+""",
+        ),
+        (
+            ObjectPath.root()["args"][1],
+            """
+def f(a, b):
+         ^
+  c = a + b
+  return c
+""",
+        ),
+        (
+            ObjectPath.root()["stmts"][0],
+            """
+def f(a, b):
+  c = a + b
+  ^^^^^^^^^
+  return c
+""",
+        ),
+        (
+            ObjectPath.root()["stmts"][0],
+            """
+def f(a, b):
+  c = a + b
+  ^^^^^^^^^
+  return c
+""",
+        ),
+        (
+            ObjectPath.root()["stmts"][0]["lhs"],
+            """
+def f(a, b):
+  c = a + b
+  ^
+  return c
+""",
+        ),
+        (
+            ObjectPath.root()["stmts"][0]["rhs"],
+            """
+def f(a, b):
+  c = a + b
+      ^^^^^
+  return c
+""",
+        ),
+        (
+            ObjectPath.root()["stmts"][0]["rhs"]["lhs"],
+            """
+def f(a, b):
+  c = a + b
+      ^
+  return c
+""",
+        ),
+        (
+            ObjectPath.root()["stmts"][0]["rhs"]["rhs"],
+            """
+def f(a, b):
+  c = a + b
+          ^
+  return c
+""",
+        ),
+        (
+            ObjectPath.root()["ret"],
+            """
+def f(a, b):
+  c = a + b
+  return c
+         ^
+""",
+        ),
+    ],
+)
+def test_print_underscore(path: ObjectPath, expected: str) -> None:
+    a = Var(name="a")
+    b = Var(name="b")
+    c = Var(name="c")
+    f = Func(
+        name="f",
+        args=[a, b],
+        stmts=[
+            Assign(lhs=c, rhs=Add(a, b)),
+        ],
+        ret=c,
+    )
+    actual = mlcp.to_python(
+        f,
+        mlcp.PrinterConfig(path_to_underline=[path]),
+    )
+    assert actual.strip() == expected.strip()