feat(AST): Introduce mlc.printer

Author: potatomashed

.github/workflows/ci.yml

@@ -4,7 +4,7 @@ on: [push, pull_request]
 env:
   CIBW_BUILD_VERBOSITY: 3
   CIBW_TEST_REQUIRES: "pytest"
-  CIBW_TEST_COMMAND: "pytest -svv {project}/tests/python/"
+  CIBW_TEST_COMMAND: "pytest -svv --durations=20 {project}/tests/python/"
   MLC_CIBW_VERSION: "2.20.0"
   MLC_PYTHON_VERSION: "3.9"
   MLC_CIBW_WIN_BUILD: "cp39-win_amd64"

CMakeLists.txt

@@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.15)
 
 project(
   mlc
-  VERSION 0.0.12
+  VERSION 0.0.13
   DESCRIPTION "MLC-Python"
   LANGUAGES C CXX
 )
@@ -41,6 +41,9 @@ if (MLC_BUILD_REGISTRY)
   add_library(mlc_registry_objs OBJECT
     "${CMAKE_CURRENT_SOURCE_DIR}/cpp/c_api.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/cpp/c_api_tests.cc"
+    "${CMAKE_CURRENT_SOURCE_DIR}/cpp/printer.cc"
+    "${CMAKE_CURRENT_SOURCE_DIR}/cpp/json.cc"
+    "${CMAKE_CURRENT_SOURCE_DIR}/cpp/structure.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/cpp/traceback.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/cpp/traceback_win.cc"
   )

README.md

@@ -8,14 +8,16 @@
 * [:key: Key Features](#key-key-features)
   + [:building_construction: MLC Dataclass](#building_construction-mlc-dataclass)
   + [:dart: Structure-Aware Tooling](#dart-structure-aware-tooling)
-  + [:snake: Text Formats in Python AST](#snake-text-formats-in-python-ast)
+  + [:snake: Text Formats in Python](#snake-text-formats-in-python)
   + [:zap: Zero-Copy Interoperability with C++ Plugins](#zap-zero-copy-interoperability-with-c-plugins)
 * [:fuelpump: Development](#fuelpump-development)
   + [:gear: Editable Build](#gear-editable-build)
   + [:ferris_wheel: Create Wheels](#ferris_wheel-create-wheels)
 
 
-MLC is a Python-first toolkit that makes it more ergonomic to build AI compilers, runtimes, and compound AI systems with Pythonic dataclass, rich tooling infra and zero-copy interoperability with C++ plugins.
+MLC is a Python-first toolkit that streamlines the development of AI compilers, runtimes, and compound AI systems with its Pythonic dataclasses, structure-aware tooling, and Python-based text formats.
+
+Beyond pure Python, MLC natively supports zero-copy interoperation with C++ plugins, and enables a smooth engineering practice transitioning from Python to hybrid or Python-free development.
 
 ## :inbox_tray: Installation
 
@@ -41,7 +43,7 @@ class MyClass(mlcd.PyClass):
 instance = MyClass(12, "test", c=None)
 ```
 
-**Type safety**. MLC dataclass checks type strictly in Cython and C++.
+**Type safety**. MLC dataclass enforces strict type checking using Cython and C++.
 
 ```python
 >>> instance.c = 10; print(instance)
@@ -68,6 +70,8 @@ demo.MyClass(a=12, b='test', c=None)
 
 An extra `structure` field are used to specify a dataclass's structure, indicating def site and scoping in an IR.
 
+<details><summary> Define a toy IR with `structure`. </summary>
+
 ```python
 import mlc.dataclasses as mlcd
 
@@ -92,18 +96,15 @@ class Let(Expr):
   body: Expr
 ```
 
+</details>
+
 **Structural equality**. Member method `eq_s` compares the structural equality (alpha equivalence) of two IRs represented by MLC's structured dataclass.
 
 ```python
-"""
-L1: let z = x + y; z
-L2: let x = y + z; x
-L3: let z = x + x; z
-"""
 >>> x, y, z = Var("x"), Var("y"), Var("z")
->>> L1 = Let(rhs=x + y, lhs=z, body=z)
->>> L2 = Let(rhs=y + z, lhs=x, body=x)
->>> L3 = Let(rhs=x + x, lhs=z, body=z)
+>>> L1 = Let(rhs=x + y, lhs=z, body=z)  # let z = x + y; z
+>>> L2 = Let(rhs=y + z, lhs=x, body=x)  # let x = y + z; x
+>>> L3 = Let(rhs=x + x, lhs=z, body=z)  # let z = x + x; z
 >>> L1.eq_s(L2)
 True
 >>> L1.eq_s(L3, assert_mode=True)
@@ -118,9 +119,98 @@ ValueError: Structural equality check failed at {root}.rhs.b: Inconsistent bindi
 >>> assert L1_hash != L3_hash
 ```
 
-### :snake: Text Formats in Python AST
+### :snake: Text Formats in Python
 
-TBD
+**IR Printer.** By defining an `__ir_print__` method, which converts an IR node to MLC's Python-style AST, MLC's `IRPrinter` handles variable scoping, renaming and syntax highlighting automatically for a text format based on Python syntax.
+
+<details><summary>Defining Python-based text format on a toy IR using `__ir_print__`.</summary>
+
+```python
+import mlc.dataclasses as mlcd
+import mlc.printer as mlcp
+from mlc.printer import ast as mlt
+
+@mlcd.py_class
+class Expr(mlcd.PyClass): ...
+
+@mlcd.py_class
+class Stmt(mlcd.PyClass): ...
+
+@mlcd.py_class
+class Var(Expr):
+  name: str
+  def __ir_print__(self, printer: mlcp.IRPrinter, path: mlcp.ObjectPath) -> mlt.Node:
+    if not printer.var_is_defined(obj=self):
+      printer.var_def(obj=self, frame=printer.frames[-1], name=self.name)
+    return printer.var_get(obj=self)
+
+@mlcd.py_class
+class Add(Expr):
+  lhs: Expr
+  rhs: Expr
+  def __ir_print__(self, printer: mlcp.IRPrinter, path: mlcp.ObjectPath) -> mlt.Node:
+    lhs: mlt.Expr = printer(obj=self.lhs, path=path["a"])
+    rhs: mlt.Expr = printer(obj=self.rhs, path=path["b"])
+    return lhs + rhs
+
+@mlcd.py_class
+class Assign(Stmt):
+  lhs: Var
+  rhs: Expr
+  def __ir_print__(self, printer: mlcp.IRPrinter, path: mlcp.ObjectPath) -> mlt.Node:
+    rhs: mlt.Expr = printer(obj=self.rhs, path=path["b"])
+    printer.var_def(obj=self.lhs, frame=printer.frames[-1], name=self.lhs.name)
+    lhs: mlt.Expr = printer(obj=self.lhs, path=path["a"])
+    return mlt.Assign(lhs=lhs, rhs=rhs)
+
+@mlcd.py_class
+class Func(mlcd.PyClass):
+  name: str
+  args: list[Var]
+  stmts: list[Stmt]
+  ret: Var
+  def __ir_print__(self, printer: mlcp.IRPrinter, path: mlcp.ObjectPath) -> mlt.Node:
+    with printer.with_frame(mlcp.DefaultFrame()):
+      for arg in self.args:
+        printer.var_def(obj=arg, frame=printer.frames[-1], name=arg.name)
+      args: list[mlt.Expr] = [printer(obj=arg, path=path["args"][i]) for i, arg in enumerate(self.args)]
+      stmts: list[mlt.Expr] = [printer(obj=stmt, path=path["stmts"][i]) for i, stmt in enumerate(self.stmts)]
+      ret_stmt = mlt.Return(printer(obj=self.ret, path=path["ret"]))
+      return mlt.Function(
+        name=mlt.Id(self.name),
+        args=[mlt.Assign(lhs=arg, rhs=None) for arg in args],
+        decorators=[],
+        return_type=None,
+        body=[*stmts, ret_stmt],
+      )
+
+# An example IR:
+a, b, c, d, e = Var("a"), Var("b"), Var("c"), Var("d"), Var("e")
+f = Func(
+  name="f",
+  args=[a, b, c],
+  stmts=[
+    Assign(lhs=d, rhs=Add(a, b)),  # d = a + b
+    Assign(lhs=e, rhs=Add(d, c)),  # e = d + c
+  ],
+  ret=e,
+)
+```
+
+</details>
+
+Two printer APIs are provided for Python-based text format:
+- `mlc.printer.to_python` that converts an IR fragment to Python text, and
+- `mlc.printer.print_python` that further renders the text with proper syntax highlighting.
+
+```python
+>>> print(mlcp.to_python(f))  # Stringify to Python
+def f(a, b, c):
+  d = a + b
+  e = d + c
+  return e
+>>> mlcp.print_python(f)  # Syntax highlighting
+```
 
 ### :zap: Zero-Copy Interoperability with C++ Plugins
 

cpp/c_api.cc

@@ -1,5 +1,5 @@
 #include "./registry.h"
-#include <mlc/all.h>
+#include <mlc/core/all.h>
 
 namespace mlc {
 namespace registry {
@@ -20,27 +20,6 @@ using ::mlc::registry::TypeTable;
 namespace {
 thread_local Any last_error;
 MLC_REGISTER_FUNC("mlc.ffi.LoadDSO").set_body([](std::string name) { TypeTable::Get(nullptr)->LoadDSO(name); });
-MLC_REGISTER_FUNC("mlc.core.JSONLoads").set_body([](AnyView json_str) {
-  if (json_str.type_index == kMLCRawStr) {
-    return ::mlc::core::JSONLoads(json_str.operator const char *());
-  } else {
-    ::mlc::Str str = json_str;
-    return ::mlc::core::JSONLoads(str);
-  }
-});
-MLC_REGISTER_FUNC("mlc.core.JSONSerialize").set_body(::mlc::core::Serialize); // TODO: `AnyView` as function argument
-MLC_REGISTER_FUNC("mlc.core.JSONDeserialize").set_body([](AnyView json_str) {
-  if (json_str.type_index == kMLCRawStr) {
-    return ::mlc::core::Deserialize(json_str.operator const char *());
-  } else {
-    return ::mlc::core::Deserialize(json_str.operator ::mlc::Str());
-  }
-});
-MLC_REGISTER_FUNC("mlc.core.StructuralEqual").set_body(::mlc::core::StructuralEqual);
-MLC_REGISTER_FUNC("mlc.core.StructuralHash").set_body([](::mlc::Object *obj) -> int64_t {
-  uint64_t ret = ::mlc::core::StructuralHash(obj);
-  return static_cast<int64_t>(ret);
-});
 } // namespace
 
 MLC_API MLCAny MLCGetLastError() {

cpp/c_api_tests.cc

@@ -1,4 +1,4 @@
-#include <mlc/all.h>
+#include <mlc/core/all.h>
 
 namespace mlc {
 namespace {

include/mlc/core/json.h renamed to cpp/json.cc

@@ -1,31 +1,19 @@
-#ifndef MLC_CORE_JSON_H_
-#define MLC_CORE_JSON_H_
-#include "./field_visitor.h"
-#include "./func.h"
-#include "./str.h"
-#include "./udict.h"
-#include "./ulist.h"
 #include <cstdint>
 #include <iomanip>
-#include <mlc/base/all.h>
+#include <mlc/core/all.h>
 #include <sstream>
 
 namespace mlc {
 namespace core {
+namespace {
 
 mlc::Str Serialize(Any any);
 Any Deserialize(const char *json_str, int64_t json_str_len);
 Any JSONLoads(const char *json_str, int64_t json_str_len);
 MLC_INLINE Any Deserialize(const char *json_str) { return Deserialize(json_str, -1); }
 MLC_INLINE Any Deserialize(const Str &json_str) { return Deserialize(json_str->data(), json_str->size()); }
-MLC_INLINE Any Deserialize(const std::string &json_str) {
-  return Deserialize(json_str.data(), static_cast<int64_t>(json_str.size()));
-}
 MLC_INLINE Any JSONLoads(const char *json_str) { return JSONLoads(json_str, -1); }
 MLC_INLINE Any JSONLoads(const Str &json_str) { return JSONLoads(json_str->data(), json_str->size()); }
-MLC_INLINE Any JSONLoads(const std::string &json_str) {
-  return JSONLoads(json_str.data(), static_cast<int64_t>(json_str.size()));
-}
 
 inline mlc::Str Serialize(Any any) {
   using mlc::base::TypeTraits;
@@ -46,7 +34,7 @@ inline mlc::Str Serialize(Any any) {
     // clang-format off
       MLC_INLINE void operator()(MLCTypeField *, const Any *any) { EmitAny(any); }
       MLC_INLINE void operator()(MLCTypeField *, ObjectRef *obj) { if (Object *v = obj->get()) EmitObject(v); else EmitNil(); }
-      MLC_INLINE void operator()(MLCTypeField *, Optional<Object> *opt) { if (Object *v = opt->get()) EmitObject(v); else EmitNil(); }
+      MLC_INLINE void operator()(MLCTypeField *, Optional<ObjectRef> *opt) { if (Object *v = opt->get()) EmitObject(v); else EmitNil(); }
       MLC_INLINE void operator()(MLCTypeField *, Optional<int64_t> *opt) { if (const int64_t *v = opt->get()) EmitInt(*v); else EmitNil(); }
       MLC_INLINE void operator()(MLCTypeField *, Optional<double> *opt) { if (const double *v = opt->get())  EmitFloat(*v); else EmitNil(); }
       MLC_INLINE void operator()(MLCTypeField *, Optional<DLDevice> *opt) { if (const DLDevice *v = opt->get()) EmitDevice(*v); else EmitNil(); }
@@ -119,25 +107,21 @@ inline mlc::Str Serialize(Any any) {
     } else {
       os->put(',');
     }
-    int32_t type_index = type_info->type_index;
-    if (type_index == kMLCStr) {
-      StrObj *str = reinterpret_cast<StrObj *>(object);
+    if (StrObj *str = object->TryCast<StrObj>()) {
       str->PrintEscape(*os);
       return;
     }
     (*os) << '[' << (*get_json_type_index)(type_info->type_key);
-    if (type_index == kMLCList) {
-      UListObj *list = reinterpret_cast<UListObj *>(object); // TODO: support Downcast
+    if (UListObj *list = object->TryCast<UListObj>()) {
       for (Any &any : *list) {
         emitter(nullptr, &any);
       }
-    } else if (type_index == kMLCDict) {
-      UDictObj *dict = reinterpret_cast<UDictObj *>(object); // TODO: support Downcast
+    } else if (UDictObj *dict = object->TryCast<UDictObj>()) {
       for (auto &kv : *dict) {
         emitter(nullptr, &kv.first);
         emitter(nullptr, &kv.second);
       }
-    } else if (type_index == kMLCFunc || type_index == kMLCError) {
+    } else if (object->IsInstance<FuncObj>() || object->IsInstance<ErrorObj>()) {
       MLC_THROW(TypeError) << "Unserializable type: " << object->GetTypeKey();
     } else {
       VisitFields(object, type_info, emitter);
@@ -182,9 +166,9 @@ inline Any Deserialize(const char *json_str, int64_t json_str_len) {
   MLCVTableHandle init_vtable;
   MLCVTableGetGlobal(nullptr, "__init__", &init_vtable);
   // Step 0. Parse JSON string
-  UDict json_obj = JSONLoads(json_str, json_str_len).operator UDict(); // TODO: impl "Any -> UDict"
+  UDict json_obj = JSONLoads(json_str, json_str_len);
   // Step 1. type_key => constructors
-  UList type_keys = json_obj->at(Str("type_keys")).operator UList(); // TODO: impl `UDict::at(Str)`
+  UList type_keys = json_obj->at("type_keys");
   std::vector<Func> constructors;
   constructors.reserve(type_keys.size());
   for (Str type_key : type_keys) {
@@ -205,7 +189,7 @@ inline Any Deserialize(const char *json_str, int64_t json_str_len) {
     return ret;
   };
   // Step 2. Translate JSON object to objects
-  UList values = json_obj->at(Str("values")).operator UList(); // TODO: impl `UDict::at(Str)`
+  UList values = json_obj->at("values");
   for (int64_t i = 0; i < values->size(); ++i) {
     Any obj = values[i];
     if (obj.type_index == kMLCList) {
@@ -492,7 +476,22 @@ inline Any JSONLoads(const char *json_str, int64_t json_str_len) {
   return JSONParser{0, json_str_len, json_str}.Parse();
 }
 
+MLC_REGISTER_FUNC("mlc.core.JSONLoads").set_body([](AnyView json_str) {
+  if (json_str.type_index == kMLCRawStr) {
+    return ::mlc::core::JSONLoads(json_str.operator const char *());
+  } else {
+    ::mlc::Str str = json_str;
+    return ::mlc::core::JSONLoads(str);
+  }
+});
+MLC_REGISTER_FUNC("mlc.core.JSONSerialize").set_body(::mlc::core::Serialize);
+MLC_REGISTER_FUNC("mlc.core.JSONDeserialize").set_body([](AnyView json_str) {
+  if (json_str.type_index == kMLCRawStr) {
+    return ::mlc::core::Deserialize(json_str.operator const char *());
+  } else {
+    return ::mlc::core::Deserialize(json_str.operator ::mlc::Str());
+  }
+});
+} // namespace
 } // namespace core
 } // namespace mlc
-
-#endif // MLC_CORE_JSON_H_