Add trace functionality (#244)

CrescentonC · web-flow · commit 76b129a2bb2d · 2024-12-04T11:34:33.000+08:00
diff --git a/hkmc2/jvm/src/test/scala/hkmc2/JSBackendDiffMaker.scala b/hkmc2/jvm/src/test/scala/hkmc2/JSBackendDiffMaker.scala
@@ -22,6 +22,7 @@ abstract class JSBackendDiffMaker extends MLsDiffMaker:
   val showRepl = NullaryCommand("showRepl")
   val silent = NullaryCommand("silent")
   val noSanityCheck = NullaryCommand("noSanityCheck")
+  val traceJS = NullaryCommand("traceJS")
   val expect = Command("expect"): ln =>
     ln.trim
   
@@ -50,7 +51,9 @@ abstract class JSBackendDiffMaker extends MLsDiffMaker:
     super.processTerm(blk, inImport)
     if js.isSet then
       val low = ltl.givenIn:
-        new codegen.Lowering with codegen.LoweringSelSanityChecks(noSanityCheck.isUnset)
+        new codegen.Lowering
+          with codegen.LoweringSelSanityChecks(noSanityCheck.isUnset)
+          with codegen.LoweringTraceLog(traceJS.isSet)
       given Elaborator.Ctx = curCtx
       val jsb = new JSBuilder
         with JSBuilderArgNumSanityChecks(noSanityCheck.isUnset)
@@ -95,7 +98,11 @@ abstract class JSBackendDiffMaker extends MLsDiffMaker:
                   case _ => output(s"$prefix= ${content}")
           case ReplHost.Empty =>
           case ReplHost.Unexecuted(message) => ???
-          case ReplHost.Error(isSyntaxError, message) =>
+          case ReplHost.Error(isSyntaxError, message, otherOutputs) =>
+            if otherOutputs.nonEmpty then
+              otherOutputs.splitSane('\n').foreach: line =>
+                output(s"> ${line}")
+            
             if (isSyntaxError) then
               // If there is a syntax error in the generated code,
               // it should be a code generation error.
@@ -107,8 +114,17 @@ abstract class JSBackendDiffMaker extends MLsDiffMaker:
                 source = Diagnostic.Source.Runtime))
         if stderr.nonEmpty then output(s"// Standard Error:\n${stderr}")
       
+      
+      if traceJS.isSet then
+        host.execute(
+          "globalThis.Predef.TraceLogger.enabled = true; " +
+          "globalThis.Predef.TraceLogger.resetIndent(0)")
+      
       mkQuery("", jsStr)
       
+      if traceJS.isSet then
+        host.execute("globalThis.Predef.TraceLogger.enabled = false")
+      
       import Elaborator.Ctx.*
       def definedValues = curCtx.env.iterator.flatMap:
         case (nme, e @ (_: RefElem | SelElem(RefElem(_: InnerSymbol), _, _))) =>
diff --git a/hkmc2/jvm/src/test/scala/hkmc2/ReplHost.scala b/hkmc2/jvm/src/test/scala/hkmc2/ReplHost.scala
@@ -48,12 +48,12 @@ class ReplHost(rootPath: Str)(using TL) {
         case None => reply
         case Some(uncaughtErrorLine) => {
           val message = uncaughtErrorLine.substring(ReplHost.uncaughtErrorHead.length)
-          ReplHost.Error(false, message)
+          ReplHost.Error(false, message, reply.take(reply.indexOf(uncaughtErrorLine)).trim())
         }
       }
       case Some(syntaxErrorLine) =>
         val message = syntaxErrorLine.substring(ReplHost.syntaxErrorHead.length)
-        ReplHost.Error(true, message)
+        ReplHost.Error(true, message, reply.take(reply.indexOf(syntaxErrorLine)).trim())
     }
     tl.log(s"REPL> Collected:\n${res}")
     res
@@ -85,7 +85,7 @@ class ReplHost(rootPath: Str)(using TL) {
         if begin >= 0 && end >= 0 then
           // `console.log` inserts a space between every two arguments,
           // so + 1 and - 1 is necessary to get correct length.
-          ReplHost.Error(false, reply.substring(begin + 1, end))
+          ReplHost.Error(false, reply.substring(begin + 1, end), reply.takeWhile(_ != 0x200b).trim())
         else reply
       case error: ReplHost.Error => error
     tl.log(s"REPL> Parsed:\n${parsed}")
@@ -194,7 +194,7 @@ object ReplHost {
     *               runtime error
     * @param message the error message
     */
-  final case class Error(syntax: Bool, message: Str) extends Reply {
+  final case class Error(syntax: Bool, message: Str, otherOutputs: Str) extends Reply {
     override def map(f: Str => Reply): Reply = this
     override def toString(): Str =
       if syntax then
diff --git a/hkmc2/shared/src/main/scala/hkmc2/codegen/Lowering.scala b/hkmc2/shared/src/main/scala/hkmc2/codegen/Lowering.scala
@@ -110,7 +110,8 @@ class Lowering(using TL, Raise, Elaborator.State):
               Define(ValDefn(td.owner, knd, td.sym, r),
                 term(st.Blk(stats, res))(k)))
           case syntax.Fun =>
-            Define(FunDefn(td.sym, td.params, returnedTerm(bod)),
+            val (paramLists, bodyBlock) = setupFunctionDef(td.params, bod, S(td.sym.nme))
+            Define(FunDefn(td.sym, paramLists, bodyBlock),
               term(st.Blk(stats, res))(k))
       // case cls: ClassDef =>
       case cls: ClassLikeDef =>
@@ -126,7 +127,8 @@ class Lowering(using TL, Raise, Elaborator.State):
         Define(ClsLikeDefn(cls.sym, syntax.Cls,
             mtds.flatMap: td =>
               td.body.map: bod =>
-                FunDefn(td.sym, td.params, term(bod)(Ret))
+                val (paramLists, bodyBlock) = setupFunctionDef(td.params, bod, S(td.sym.nme))
+                FunDefn(td.sym, paramLists, bodyBlock)
             ,
             privateFlds,
             publicFlds,
@@ -165,7 +167,8 @@ class Lowering(using TL, Raise, Elaborator.State):
       term(st.Blk(stats, res))(k)
       
     case st.Lam(params, body) =>
-      k(Value.Lam(params, returnedTerm(body)))
+      val (paramLists, bodyBlock) = setupFunctionDef(params :: Nil, body, N)
+      k(Value.Lam(paramLists.head, bodyBlock))
     
     /* 
     case t @ st.If(Split.Let(sym, trm, tail)) =>
@@ -330,6 +333,8 @@ class Lowering(using TL, Raise, Elaborator.State):
     subTerm(prefix): p =>
       k(Select(p, nme)(sym))
   
+  def setupFunctionDef(paramLists: List[ParamList], bodyTerm: Term, name: Option[Str])(using Subst): (List[ParamList], Block) =
+    (paramLists, returnedTerm(bodyTerm))
 
 
 trait LoweringSelSanityChecks
@@ -357,3 +362,79 @@ trait LoweringSelSanityChecks
       super.setupSelection(prefix, nme, sym)(k)
 
 
+
+trait LoweringTraceLog
+    (instrument: Bool)(using TL, Raise, Elaborator.State)
+    extends Lowering:
+      
+  private def selFromGlobalThis(path: Str*): Path =
+      path.foldLeft[Path](Value.Ref(State.globalThisSymbol)):
+        (qual, name) => Select(qual, Tree.Ident(name))(N)
+    
+  private def assignStmts(stmts: (Local, Result)*)(rest: Block) =
+    stmts.foldRight(rest):
+      case ((sym, res), acc) => Assign(sym, res, acc)
+  
+  extension (k: Block => Block)
+    def |>: (b: Block): Block = k(b)
+
+  private val traceLogFn = selFromGlobalThis("Predef", "TraceLogger", "log")
+  private val traceLogIndentFn = selFromGlobalThis("Predef", "TraceLogger", "indent")
+  private val traceLogResetFn = selFromGlobalThis("Predef", "TraceLogger", "resetIndent")
+  private val strConcatFn = selFromGlobalThis("String", "prototype", "concat", "call")
+  private val inspectFn = selFromGlobalThis("util", "inspect")
+  
+
+  override def setupFunctionDef(paramLists: List[ParamList], bodyTerm: st, name: Option[Str])(using Subst): (List[ParamList], Block) = 
+    if instrument then
+      val (ps, bod) = handleMultipleParamLists(paramLists, bodyTerm)
+      val instrumentedBody = setupFunctionBody(ps, bod, name)
+      (ps :: Nil, instrumentedBody)
+    else
+      super.setupFunctionDef(paramLists, bodyTerm, name)
+
+  def handleMultipleParamLists(paramLists: List[ParamList], bod: Term) =
+    def go(paramLists: List[ParamList], bod: Term): (ParamList, Term) =
+      paramLists match
+        case Nil => ???
+        case h :: Nil => (h, bod)
+        case h :: t => go(t, Term.Lam(h, bod))
+    go(paramLists.reverse, bod)
+  
+  def setupFunctionBody(params: ParamList, bod: Term, name: Option[Str])(using Subst): Block =
+    val enterMsgSym = TempSymbol(N, dbgNme = "traceLogEnterMsg")
+    val prevIndentLvlSym = TempSymbol(N, dbgNme = "traceLogPrevIndent")
+    val resSym = TempSymbol(N, dbgNme = "traceLogRes")
+    val retMsgSym = TempSymbol(N, dbgNme = "traceLogRetMsg")
+    val psInspectedSyms = params.params.map(p => TempSymbol(N, dbgNme = s"traceLogParam_${p.sym.nme}") -> p.sym)
+    val resInspectedSym = TempSymbol(N, dbgNme = "traceLogResInspected")
+    
+    val psSymArgs = psInspectedSyms.zipWithIndex.foldRight[Ls[Arg]](Arg(false, Value.Lit(Tree.StrLit(")"))) :: Nil):
+      case (((s, p), i), acc) => if i == psInspectedSyms.length - 1
+        then Arg(false, Value.Ref(s)) :: acc
+        else Arg(false, Value.Ref(s)) :: Arg(false, Value.Lit(Tree.StrLit(", "))) :: acc
+    
+    assignStmts(psInspectedSyms.map: (pInspectedSym, pSym) =>
+      pInspectedSym -> Call(inspectFn, Arg(false, Value.Ref(pSym)) :: Nil)
+    *) |>:
+    assignStmts(
+      enterMsgSym -> Call(
+        strConcatFn,
+        Arg(false, Value.Lit(Tree.StrLit(s"CALL ${name.getOrElse("[arrow function]")}("))) :: psSymArgs
+      ),
+      TempSymbol(N) -> Call(traceLogFn, Arg(false, Value.Ref(enterMsgSym)) :: Nil),
+      prevIndentLvlSym -> Call(traceLogIndentFn, Nil)
+    ) |>: 
+    term(bod)(r =>
+    assignStmts(
+      resSym -> r,
+      resInspectedSym -> Call(inspectFn, Arg(false, Value.Ref(resSym)) :: Nil),
+      retMsgSym -> Call(
+        strConcatFn,
+        Arg(false, Value.Lit(Tree.StrLit("=> "))) :: Arg(false, Value.Ref(resInspectedSym)) :: Nil
+      ),
+      TempSymbol(N) -> Call(traceLogResetFn, Arg(false, Value.Ref(prevIndentLvlSym)) :: Nil),
+      TempSymbol(N) -> Call(traceLogFn, Arg(false, Value.Ref(retMsgSym)) :: Nil)
+    ) |>:
+      Ret(Value.Ref(resSym))
+    )
diff --git a/hkmc2/shared/src/test/mlscript-compile/Predef.mjs b/hkmc2/shared/src/test/mlscript-compile/Predef.mjs
@@ -16,6 +16,51 @@ const Predef$class = class Predef {
       }
       toString() { return "MatchFailure(" + this.errors + ")"; }
     };
+    const TraceLogger$class = class TraceLogger {
+      constructor() {
+        this.enabled = false;
+        this.indentLvl = 0;
+      }
+      indent() {
+        let scrut, prev, tmp;
+        scrut = this.enabled;
+        if (scrut) {
+          prev = this.indentLvl;
+          tmp = prev + 1;
+          this.indentLvl = tmp;
+          return prev;
+        } else {
+          return undefined;
+        }
+      } 
+      resetIndent(n) {
+        let scrut;
+        scrut = this.enabled;
+        if (scrut) {
+          this.indentLvl = n;
+          return undefined;
+        } else {
+          return undefined;
+        }
+      } 
+      log(msg) {
+        let scrut, tmp, tmp1, tmp2, tmp3, tmp4;
+        scrut = this.enabled;
+        if (scrut) {
+          tmp = "| ".repeat(this.indentLvl);
+          tmp1 = "  ".repeat(this.indentLvl);
+          tmp2 = "\n" + tmp1;
+          tmp3 = msg.replaceAll("\n", tmp2);
+          tmp4 = tmp + tmp3;
+          return console.log(tmp4);
+        } else {
+          return undefined;
+        }
+      }
+      toString() { return "TraceLogger"; }
+    };
+    this.TraceLogger = new TraceLogger$class;
+    this.TraceLogger.class = TraceLogger$class;
     this.Test = class Test {
       constructor() {
         this.y = 1;
diff --git a/hkmc2/shared/src/test/mlscript-compile/Predef.mls b/hkmc2/shared/src/test/mlscript-compile/Predef.mls
@@ -40,6 +40,24 @@ fun checkArgs(functionName, expected, isUB, got) =
     //   + " arguments but got " + got)
   else ()
 
+module TraceLogger with
+  mut val enabled = false
+  mut val indentLvl = 0
+  fun indent() =
+    if enabled then
+      let prev = indentLvl
+      set indentLvl = prev + 1
+      prev
+    else ()
+  fun resetIndent(n) =
+    if enabled then
+      set indentLvl = n
+    else ()
+  fun log(msg) =
+    if enabled then
+      console.log("| ".repeat(indentLvl) + msg.replaceAll("\n", "\n" + "  ".repeat(indentLvl)))
+    else ()
+
 class Test with
   val y = 1
 
diff --git a/hkmc2/shared/src/test/mlscript/basics/Classes.mls b/hkmc2/shared/src/test/mlscript/basics/Classes.mls
@@ -35,6 +35,8 @@ class Foo(x: Int) { log("Hello!") }
 //│ ╙──      	                    ^^^^^^^^^^^^^
 //│ JS:
 //│ this.<error> = class <error> { constructor() {  } toString() { return "<error>"; } }; undefined
+//│ > try { this.<error> = class <error> { constructor() {  } toString() { return "<error>"; } }; undefined } catch (e) { console.log('\u200B' + e + '\u200B'); }
+//│ >            ^
 //│ ═══[COMPILATION ERROR] [Uncaught SyntaxError] Unexpected token '<'
 
 
diff --git a/hkmc2/shared/src/test/mlscript/codegen/Hygiene.mls b/hkmc2/shared/src/test/mlscript/codegen/Hygiene.mls
@@ -88,6 +88,8 @@ module Test with
   let f = () => x
   let x = 2
   log(f())
+//│ > try { const Test$class = class Test {   #x;   #f;   #x;   constructor() {     let tmp;     this.#x = 1;     this.#f = (...args) => {       globalThis.Predef.checkArgs("", 0, true, args.length);       return this.#x;     };     this.#x = 2;     tmp = this.#f() ?? null;     globalThis.log(tmp) ?? null   }   toString() { return "Test"; } }; this.Test = new Test$class; this.Test.class = Test$class; undefined } catch (e) { console.log('\u200B' + e + '\u200B'); }
+//│ >                                                       ^
 //│ ═══[COMPILATION ERROR] [Uncaught SyntaxError] Identifier '#x' has already been declared
 
 
diff --git a/hkmc2/shared/src/test/mlscript/codegen/Primes.mls b/hkmc2/shared/src/test/mlscript/codegen/Primes.mls
@@ -9,7 +9,11 @@ let x = 1
 let x' = 1
 //│ JS:
 //│ this.x' = 1; undefined
+//│ > try { this.x' = 1; undefined } catch (e) { console.log('\u200B' + e + '\u200B'); }
+//│ >             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 //│ ═══[COMPILATION ERROR] [Uncaught SyntaxError] Unexpected string
+//│ > try { this.x' } catch (e) { console.log('\u200B' + e + '\u200B'); }
+//│ >             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 //│ ═══[COMPILATION ERROR] [Uncaught SyntaxError] Unexpected string
 
 
diff --git a/hkmc2/shared/src/test/mlscript/codegen/ReboundLet.mls b/hkmc2/shared/src/test/mlscript/codegen/ReboundLet.mls
@@ -44,6 +44,8 @@ class Foo with
   print(x)
   print(f())
   print(g())
+//│ > try { this.Foo = class Foo {   #x;   #f;   #x;   constructor() {     let tmp, tmp1, tmp2, tmp3, tmp4, tmp5;     this.#x = 1;     this.#f = (...args) => {       globalThis.Predef.checkArgs("", 0, true, args.length);       return this.#x;     };     tmp = this.#f() ?? null;     tmp1 = Predef.print(tmp) ?? null;     this.#x = 2;     tmp2 = Predef.print(this.#x) ?? null;     tmp3 = this.#f() ?? null;     tmp4 = Predef.print(tmp3) ?? null;     tmp5 = this.g() ?? null;     Predef.print(tmp5) ?? null   }   g(...args) {     globalThis.Predef.checkArgs("g", 0, true, args.length);     return this.#x;   }   toString() { return "Foo"; } }; undefined } catch (e) { console.log('\u200B' + e + '\u200B'); }
+//│ >                                              ^
 //│ ═══[COMPILATION ERROR] [Uncaught SyntaxError] Identifier '#x' has already been declared
 
 
diff --git a/hkmc2/shared/src/test/mlscript/codegen/TraceLog.mls b/hkmc2/shared/src/test/mlscript/codegen/TraceLog.mls
diff --git a/hkmc2/shared/src/test/mlscript/codegen/TraceLogIndent.mls b/hkmc2/shared/src/test/mlscript/codegen/TraceLogIndent.mls
diff --git a/hkmc2/shared/src/test/mlscript/codegen/While.mls b/hkmc2/shared/src/test/mlscript/codegen/While.mls
diff --git a/hkmc2/shared/src/test/mlscript/syntax/KeywordStutters.mls b/hkmc2/shared/src/test/mlscript/syntax/KeywordStutters.mls
