Instrumentation into Block IR

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

@@ -109,26 +109,26 @@ class InstrumentationImpl(using State):
     blockCtor("End", Ls(), symName)(k)
 
   def ruleBranches(x: Path, p: Path, arms: Ls[Case -> Block], dflt: Opt[Block], symName: String = "branches")(using Context)(k: (Path, Context) => Block): Block =
-    def applyRuleBranch(cse: Case, block: Block)(f: Path => (Context, Path) => Block)(ctx: Context, x: Path): Block =
+    def applyRuleBranch(cse: Case, block: Block)(f: Path => Context => Block)(ctx: Context): Block =
       transformCase(cse): cse =>
         transformBlock(block)(using ctx.clone() += p -> x): (y, ctx) =>
           // TODO: use Arm type instead of Tup
           tuple(Ls(cse, y), "branch"): cde =>
-            f(cde)(ctx.clone() -= p, x)
-
-    val a = arms.map(applyRuleBranch).collectApply
-    ((f: (Ls[Path], Context) => Block) => a(ys => (ctx, _) => f(ys, ctx))(summon, x)): (arms, ctx) =>
-      tuple(arms): arms =>
-        ruleEnd(): e =>
-          // TODO: use transformOption here
-          def dfltStaged(k: (Path, Context) => Block) =
-            dflt match
-            case S(dflt) =>
-              transformBlock(dflt)(using ctx.clone() += p -> x): (dflt, ctx) =>
-                optionSome(dflt)(k(_, ctx.clone() -= p))
-            case N => optionNone()(k(_, ctx))
-          dfltStaged: (dflt, ctx) =>
-            blockCtor("Match", Ls(x, arms, dflt, e), symName)(k(_, ctx))
+            f(cde)(ctx.clone() -= p)
+
+    (arms.map(applyRuleBranch).collectApply(_: Ls[Path] => Context => Block)(summon)): arms =>
+      ctx =>
+        tuple(arms): arms =>
+          ruleEnd(): e =>
+            // TODO: use transformOption here
+            def dfltStaged(k: (Path, Context) => Block) =
+              dflt match
+              case S(dflt) =>
+                transformBlock(dflt)(using ctx.clone() += p -> x): (dflt, ctx) =>
+                  optionSome(dflt)(k(_, ctx.clone() -= p))
+              case N => optionNone()(k(_, ctx))
+            dfltStaged: (dflt, ctx) =>
+              blockCtor("Match", Ls(x, arms, dflt, e), symName)(k(_, ctx))
 
   // transformations of Block
 
@@ -137,15 +137,18 @@ class InstrumentationImpl(using State):
     ctx.get(p).map(k).getOrElse:
       p match
       case Value.Ref(l, disamb) =>
-        // not in formalization
         transformSymbol(disamb.getOrElse(l)): sym =>
           blockCtor("ValueRef", Ls(sym), "var")(k)
       case l: Value.Lit =>
         blockCtor("ValueLit", Ls(l), "lit")(k)
-      case Select(p, i @ Tree.Ident(name)) =>
+      case s @ Select(p, i @ Tree.Ident(name)) =>
         transformPath(p): x =>
-          blockCtor("Symbol", Ls(toValue(name))): name =>
-            blockCtor("Select", Ls(x, name), "sel")(k)
+          val sym =
+            if s.symbol.isDefined
+            then transformSymbol(s.symbol.get)
+            else blockCtor("Symbol", Ls(toValue(name)))
+          sym: sym =>
+            blockCtor("Select", Ls(x, sym), "sel")(k)
       case DynSelect(qual, fld, arrayIdx) =>
         transformPath(qual): x =>
           transformPath(fld): y =>
@@ -212,8 +215,7 @@ class InstrumentationImpl(using State):
       transformResult(r): y =>
         transformSymbol(x): xSym =>
           blockCtor("ValueRef", Ls(xSym)): xStaged =>
-            // x should always be defined, either as an argument to the function or in a Scope Block
-            assert(ctx.get(x.asPath).isDefined)
+            assert(ctx.get(x.asPath).isDefined, "x should always be defined, either as an argument to the function or in a Scope Block")
             (Assign(x, xStaged, _)):
               given Context = ctx.clone() += x.asPath -> xStaged
               transformBlock(b): (z, ctx) =>
@@ -228,14 +230,13 @@ class InstrumentationImpl(using State):
                 blockCtor("ClsLikeDefn", Ls(c, none)): cls =>
                   blockCtor("Define", Ls(cls, p)): p =>
                     ruleEnd(): end =>
-                      fnPrintCode(p)(k(end, summon))
-    case End(_) => ruleEnd()(k(_, summon))
+                      fnPrintCode(p)(k(end, ctx))
+    case End(_) => ruleEnd()(k(_, ctx))
     case Match(p, ks, dflt, rest) =>
       transformPath(p): x =>
-        ruleBranches(x, p, ks, dflt): (stagedMatch, ctx1) =>
-          transformBlock(rest)(using ctx1): (z, ctx2) =>
-            fnConcat(stagedMatch, z, "match"): cde =>
-              k(cde, ctx2)
+        ruleBranches(x, p, ks, dflt): (stagedMatch, ctx) =>
+          transformBlock(rest)(using ctx): (z, ctx) =>
+            fnConcat(stagedMatch, z, "match")(k(_, ctx))
     case Begin(sub, rest) =>
       // TODO: This is untested as there is no test case that generates the Begin block yet
       transformBlock(sub): (sub, ctx) =>
@@ -255,18 +256,15 @@ class InstrumentationImpl(using State):
     // NOTE: this debug printing only works for top-level modules, nested modules don't work
     // TODO: remove it. only for test
     val debug = blockCtor("ValueLit", Ls(Value.Lit(Tree.UnitLit(false)))): undef =>
-      // TODO: put correct parameters instead of End
-      val argsList = f.params.map(ps => List.fill(ps.params.length)(undef))
-      def makeCalls(k: Path => Block) =
-        argsList.foldRight(k)((args, cont) => res => call(res, args)(cont))(sym)
-      makeCalls(fnPrintCode(_)(End()))
+      // TODO: put correct parameters instead of undefined
+      f.params.map(ps => List.fill(ps.params.length)(undef))
+        .foldRight((p: Path) => fnPrintCode(p)(End()))((args, cont) => call(_, args)(cont))(sym)
 
     val dSym = TermSymbol(f.dSym.k, f.dSym.owner, Tree.Ident(f.sym.nme + "_gen"))
     val args = f.params.flatMap(_.params).map(_.sym)
     val newBody =
-      ruleEnd(): end =>
-        given Context = HashMap(args.map(s => Value.Ref(s, N) -> Value.Ref(s, N))*)
-        transformBlock(f.body)(p => Return(p, false))
+      given Context = HashMap(args.map(s => Value.Ref(s, N) -> Value.Ref(s, N))*)
+      transformBlock(f.body)(Return(_, false))
     val newFun = f.copy(sym = genSym, dSym = dSym, body = newBody)(false)
     (newFun, debug)
 

hkmc2/shared/src/test/mlscript/staging/Functions.mls

@@ -33,17 +33,6 @@ staged module Expressions with
 //│ > Assign(Symbol("scrut"), Tuple([...Ref(Symbol("x"))]), Match(Ref(Symbol("scrut")), [Tup(0, false) -> Return(Lit(1), false), Tup(1, true) -> Assign(Symbol("element0$"), Call(Select(Select(Ref(Symbol("runtime")), Symbol("Tuple")), Symbol("get")), [Ref(Symbol("scrut")), Lit(0)]), Assign(Symbol("middleElements"), Call(Select(Select(Ref(Symbol("runtime")), Symbol("Tuple")), Symbol("slice")), [Ref(Symbol("scrut")), Lit(1), Lit(0)]), Assign(Symbol("a"), Ref(Symbol("element0$")), Return(Lit(3), false))))], Return(Lit(0), false), End))
 //│ x = [1, 2, 3]
 
-:js
-:staging
-:fixme
-staged module OtherBlocks with
-  fun breakAndLabel() =
-    if 1 is
-      2 then 0
-      3 then 0
-      else 0
-//│ /!!!\ Uncaught error: scala.NotImplementedError: an implementation is missing
-
 :js
 :staging
 class Outside(a)
@@ -57,9 +46,9 @@ staged module ClassInstrumentation with
 //│ > Define(ClsLikeDefn(ClassSymbol("NoArg"), TODO), End)
 //│ > Define(ClsLikeDefn(ClassSymbol("Inside":[Symbol("a"), Symbol("b")]), TODO), End)
 //│ > Return(Instantiate(Ref(ClassSymbol("Outside":[Symbol("a")])), [Lit(1)]), false)
-//│ > Return(Instantiate(Select(Ref(Symbol("ClassInstrumentation")), Symbol("NoArg")), []), false)
+//│ > Return(Instantiate(Select(Ref(Symbol("ClassInstrumentation")), ClassSymbol("NoArg")), []), false)
 //│ > Return(Call(Ref(ClassSymbol("Outside":[Symbol("a")])), [Lit(1)]), false)
-//│ > Return(Call(Select(Ref(Symbol("ClassInstrumentation")), Symbol("Inside")), [Lit(1), Lit(2)]), false)
+//│ > Return(Call(Select(Ref(Symbol("ClassInstrumentation")), ClassSymbol("Inside":[Symbol("a"), Symbol("b")])), [Lit(1), Lit(2)]), false)
 
 :js
 :staging
@@ -71,6 +60,17 @@ staged module Arguments with
 //│ > Assign(Symbol("x"), Lit(1), Return(Ref(Symbol("x")), false))
 //│ > Return(Lit(undefined), false)
 
+:js
+:staging
+:fixme
+staged module OtherBlocks with
+  fun breakAndLabel() =
+    if 1 is
+      2 then 0
+      3 then 0
+      else 0
+//│ /!!!\ Uncaught error: scala.NotImplementedError: an implementation is missing
+
 // debug printing fails, collision with class name?
 :js
 :staging

hkmc2DiffTests/src/test/scala/hkmc2/MLsDiffMaker.scala

@@ -167,7 +167,9 @@ abstract class MLsDiffMaker extends DiffMaker:
     if stageCode.isSet then
       given Config = mkConfig
       processTrees(
-        PrefixApp(Keywrd(`import`), StrLit(blockFile.toString)) :: Nil)
+        PrefixApp(Keywrd(`import`), StrLit(blockFile.toString))
+        :: PrefixApp(Keywrd(`import`), StrLit(optionFile.toString))
+        :: Nil)
     super.init()