feat(printer): Print address for duplicated variable names (#66)

Should fix #62

Author: potatomashed

include/mlc/printer/ast.h

@@ -14,13 +14,16 @@ struct PrinterConfigObj : public Object {
   int32_t indent_spaces = 2;
   int8_t print_line_numbers = 0;
   int32_t num_context_lines = -1;
+  bool print_addr_on_dup_var = false;
   mlc::List<ObjectPath> path_to_underline;
 
   PrinterConfigObj() = default;
   explicit PrinterConfigObj(bool def_free_var, int32_t indent_spaces, int8_t print_line_numbers,
-                            int32_t num_context_lines, mlc::List<ObjectPath> path_to_underline)
+                            int32_t num_context_lines, bool print_addr_on_dup_var,
+                            mlc::List<ObjectPath> path_to_underline)
       : def_free_var(def_free_var), indent_spaces(indent_spaces), print_line_numbers(print_line_numbers),
-        num_context_lines(num_context_lines), path_to_underline(path_to_underline) {}
+        num_context_lines(num_context_lines), print_addr_on_dup_var(print_addr_on_dup_var),
+        path_to_underline(path_to_underline) {}
   MLC_DEF_DYN_TYPE(MLC_EXPORTS, PrinterConfigObj, Object, "mlc.printer.PrinterConfig");
 };
 
@@ -31,11 +34,13 @@ struct PrinterConfig : public ObjectRef {
       .Field("print_line_numbers", &PrinterConfigObj::print_line_numbers)
       .Field("num_context_lines", &PrinterConfigObj::num_context_lines)
       .Field("path_to_underline", &PrinterConfigObj::path_to_underline)
-      .StaticFn("__init__", InitOf<PrinterConfigObj, bool, int32_t, int8_t, int32_t, mlc::List<ObjectPath>>);
+      .Field("print_addr_on_dup_var", &PrinterConfigObj::print_addr_on_dup_var)
+      .StaticFn("__init__", InitOf<PrinterConfigObj, bool, int32_t, int8_t, int32_t, bool, mlc::List<ObjectPath>>);
   explicit PrinterConfig(bool def_free_var = true, int32_t indent_spaces = 2, int8_t print_line_numbers = 0,
-                         int32_t num_context_lines = -1, mlc::List<ObjectPath> path_to_underline = {})
+                         int32_t num_context_lines = -1, bool print_addr_on_dup_var = false,
+                         mlc::List<ObjectPath> path_to_underline = {})
       : PrinterConfig(PrinterConfig::New(def_free_var, indent_spaces, print_line_numbers, num_context_lines,
-                                         path_to_underline)) {}
+                                         print_addr_on_dup_var, path_to_underline)) {}
 };
 
 } // namespace printer

include/mlc/printer/ir_printer.h

@@ -69,9 +69,18 @@ struct IRPrinterObj : public Object {
       }
     }
     // Find a unique name
-    Str name = name_hint;
-    for (int i = 1; defined_names.count(name) > 0; ++i) {
-      name = name_hint.ToStdString() + '_' + std::to_string(i);
+    std::string name_hint_str = name_hint.ToStdString();
+    Str name(name_hint_str);
+    if (defined_names->count(name)) {
+      if (this->cfg->print_addr_on_dup_var) {
+        std::ostringstream os;
+        os << name_hint_str << "_0x" << std::setfill('0') << std::setw(12) << std::hex << (uintptr_t)(obj.get());
+        name = Str(os.str());
+      } else {
+        for (int i = 1; defined_names.count(name) > 0; ++i) {
+          name = name_hint_str + '_' + std::to_string(i);
+        }
+      }
     }
     defined_names->Set(name, 1);
     this->_VarDef(VarInfo(name, Func([name]() { return Id(name); })), obj, frame);

python/mlc/printer/ast.py

@@ -13,6 +13,7 @@ class PrinterConfig(Object):
     indent_spaces: int = 2
     print_line_numbers: int = 0
     num_context_lines: int = -1
+    print_addr_on_dup_var: bool = False
     path_to_underline: list[ObjectPath] = mlcd.field(default_factory=list)
 
 

tests/python/test_printer_ir_printer.py

@@ -1,3 +1,5 @@
+import re
+
 import mlc.printer as mlcp
 from mlc.testing.toy_ir import Add, Assign, Func, Var
 
@@ -75,3 +77,31 @@ def test_print_bool() -> None:
     path = mlcp.ObjectPath.root()
     node = printer(True, path)
     assert node.to_python() == "True"
+
+
+def test_duplicated_vars() -> None:
+    a = Var(name="a")
+    b = Var(name="a")
+    stmts = [
+        Assign(lhs=b, rhs=Add(a, a)),
+    ]
+    f = Func(
+        name="f",
+        args=[a],
+        stmts=stmts,
+        ret=b,
+    )
+    assert (
+        mlcp.to_python(f)
+        == """
+def f(a):
+  a_1 = a + a
+  return a_1
+""".strip()
+    )
+    assert re.fullmatch(
+        r"^def f\(a\):\n"
+        r"  a_0x[0-9A-Fa-f]+ = a \+ a\n"
+        r"  return a_0x[0-9A-Fa-f]+$",
+        mlcp.to_python(f, mlcp.PrinterConfig(print_addr_on_dup_var=True)),
+    )