Add pretty printer for lowered tree(Block) (#254)

hkmc2/jvm/src/test/scala/hkmc2/JSBackendDiffMaker.scala

@@ -85,6 +85,9 @@ abstract class JSBackendDiffMaker extends MLsDiffMaker:
       if showLoweredTree.isSet then
         output(s"Lowered:")
         output(le.showAsTree)
+      if ppLoweredTree.isSet then
+        output(s"Pretty Lowered:")
+        output(Printer.mkDocument(le)(using summon[Raise], baseScp.nest).toString)
 
       // * Note that the codegen scope is not in sync with curCtx in terms of its `this` symbol.
       // * We do not nest TopLevelSymbol in codegen `Scope`s

hkmc2/jvm/src/test/scala/hkmc2/MLsDiffMaker.scala

@@ -39,6 +39,7 @@ abstract class MLsDiffMaker extends DiffMaker:
   val showElab = NullaryCommand("el")
   val showElaboratedTree = DebugTreeCommand("elt")
   val showLoweredTree = NullaryCommand("lot")
+  val ppLoweredTree = NullaryCommand("slot")
   val showContext = NullaryCommand("ctx")
   val parseOnly = NullaryCommand("parseOnly")
 

hkmc2/shared/src/main/scala/hkmc2/codegen/Printer.scala

@@ -0,0 +1,103 @@
+package hkmc2.codegen
+
+import scala.collection.mutable.{Map => MutMap}
+
+import mlscript.utils._, shorthands._
+
+import hkmc2._
+import hkmc2.Message.MessageContext
+import hkmc2.document._
+import hkmc2.semantics.Elaborator.State
+import hkmc2.codegen.js.Scope
+
+object Printer:
+  def getVar(l: Local)(using Raise, Scope): String = l match
+    case ts: semantics.TermSymbol =>
+      ts.id.name
+    case ts: semantics.BlockMemberSymbol => // this means it's a locally-defined member
+      ts.nme
+      // ts.trmTree
+    case ts: semantics.InnerSymbol => ts.nme
+    case ts: semantics.BuiltinSymbol => ts.nme
+    case _ => summon[Scope].lookup_!(l)
+
+  def mkDocument(blk: Block)(using Raise, Scope): Document = blk match
+    case Match(scrut, arms, dflt, rest) =>
+      def case_doc(c: Case) = c match
+        case Case.Lit(lit) => doc"${lit.idStr}"
+        case Case.Cls(cls, path) => doc"${cls.nme}"
+        case Case.Tup(len, inf) => doc"tuple$len"
+      val docCases = arms
+        .map{ case (c, b) => doc"${case_doc(c)} => #{  # ${mkDocument(b)} #} " }
+        .mkDocument(sep = doc" # ")
+      val docDefault = dflt.map(mkDocument).getOrElse(doc"")
+      doc"match ${mkDocument(scrut)} #{  # ${docCases} # else #{  # ${docDefault} #}  #} "
+    case Return(res, implct) => doc"return ${mkDocument(res)}"
+    case Throw(exc) => doc"throw ${mkDocument(exc)}"
+    case Label(label, body, rest) =>
+      val l2 = summon[Scope].allocateName(label)
+      doc"label $l2 = ${mkDocument(body)} in # ${mkDocument(rest)}"
+    case Break(label) =>
+      doc"break ${getVar(label)}"
+    case Continue(label) =>
+      doc"continue ${getVar(label)}"
+    case Begin(sub, rest) =>
+      doc"begin #{  # ${mkDocument(sub)}; # ${mkDocument(rest)} #} "
+    case TryBlock(sub, finallyDo, rest) =>
+      doc"try #{  # ${mkDocument(sub)} #  #} finally #  #{ ${mkDocument(finallyDo)} in #  #} ${mkDocument(rest)}"
+    case Assign(lhs, rhs, rest) =>
+      val docLhs = summon[Scope].lookup(lhs).getOrElse(summon[Scope].allocateName(lhs))
+      doc"set $docLhs = ${mkDocument(rhs)} in # ${mkDocument(rest)}"
+    case AssignField(lhs, nme, rhs, rest) =>
+      doc"set ${mkDocument(lhs)}.${nme.name} = ${mkDocument(rhs)} in # ${mkDocument(rest)}"
+    case Define(defn, rest) => {
+      doc"define ${mkDocument(defn)} in # ${mkDocument(rest)}"
+    }
+    case End("") => doc"end"
+    case End(msg) => doc"end ${msg}"
+
+  def mkDocument(defn: Defn)(using Raise, Scope): Document = defn match
+    case FunDefn(sym, params, body) =>
+      val docParams = doc"${params.map(_.params.map(x => summon[Scope].allocateName(x.sym)).mkString("(", ", ", ")")).mkString}"
+      val docBody = mkDocument(body)
+      doc"fun ${sym.nme}${docParams} { #{  # ${docBody} #}  # }"
+    case ValDefn(owner, k, sym, rhs) =>
+      doc"val ${sym.nme} = ${mkDocument(rhs)}"
+    case ClsLikeDefn(sym, k, methods, privateFields, publicFields, ctor) =>
+      doc"class ${sym.nme} #{  #} "
+
+  def mkDocument(arg: Arg)(using Raise, Scope): Document =
+    val doc = mkDocument(arg.value)
+    if arg.spread
+      then doc"...${doc}"
+      else doc
+
+  def mkDocument(value: Value)(using Raise, Scope): Document = value match
+    case Value.Ref(l) => getVar(l)
+    case Value.This(sym) => doc"this"
+    case Value.Lit(lit) => doc"${lit.idStr}"
+    case Value.Lam(params, body) =>
+      val docParams = params.params.map(x => summon[Scope].allocateName(x.sym)).mkString(", ")
+      doc"(${docParams}) => ${mkDocument(body)}"
+    case Value.Arr(elems) =>
+      val docElems = elems.map(x => mkDocument(x)).mkString(", ")
+      doc"[${docElems}]"
+
+  def mkDocument(path: Path)(using Raise, Scope): Document = path match
+    case Select(qual, name) =>
+      val docQual = mkDocument(qual)
+      doc"${docQual}.${name.name}"
+    case x: Value => mkDocument(x)
+
+  def mkDocument(result: Result)(using Raise, Scope): Document = result match
+    case Call(fun, args) => doc"${mkDocument(fun)}(${args.map(mkDocument).mkString(", ")})"
+    case Instantiate(cls, args) => doc"new ${mkDocument(cls)}(${args.map(mkDocument).mkString(", ")})"
+    case x: Path => mkDocument(x)
+
+  def mkDocument(prog: Program)(using Raise, Scope): Document = {
+    val docImports = prog.imports.map:
+      case (local, path) =>
+        val docLocal = summon[Scope].allocateName(local)
+        doc"import ${docLocal}"
+    doc" ${docImports.mkDocument(sep = doc" # ")} # ${mkDocument(prog.main)}"
+  }

hkmc2/shared/src/test/mlscript/codegen/BlockPrinter.mls

@@ -0,0 +1,37 @@
+:js
+
+:slot
+let x = 1
+x + 1
+//│ Pretty Lowered:
+//│   set x = 1 in return +(x, 1)
+//│ = 2
+//│ x = 1
+
+:slot
+fun incr(n) = n + 1
+fun (|>) pipe(x, f) = f(x)
+//│ Pretty Lowered:
+//│   define fun incr(n) { return +(n, 1) } in define fun pipe(x, f) { return f(x) } in return null
+
+:slot
+let x = 1
+let x = if x == 0 then 1 else 0
+let x = x + 1
+//│ Pretty Lowered:
+//│  
+//│ set x = 1 in
+//│ begin
+//│   set scrut = ==(x, 0) in
+//│   match scrut
+//│     true =>
+//│       set tmp = 1 in
+//│       end
+//│     else
+//│       set tmp = 0 in
+//│       end;
+//│   set x1 = tmp in
+//│   set tmp1 = +(x, 1) in
+//│   set x2 = tmp1 in
+//│   return null
+//│ x = 1