Final adjustments for dynamic staging

hkmc2/js/src/main/scala/hkmc2/Compiler.scala

@@ -64,7 +64,11 @@ class Compiler(paths: MLsCompiler.Paths)(using cctx: CompilerCtx):
     perFileDiagnostics
 
 @JSExportTopLevel("Paths")
-final class Paths(prelude: Str, runtime: Str, term: Str) extends MLsCompiler.Paths:
+final class Paths(prelude: Str, runtime: Str, term: Str, block: Str, spHelper: Str, option: Str, ss: Str) extends MLsCompiler.Paths:
   val preludeFile = Path(prelude)
   val runtimeFile = Path(runtime)
   val termFile = Path(term)
+  val blockFile = Path(block)
+  val specializeHelpersFile = Path(spHelper)
+  val optionFile = Path(option)
+  val shapeSetFile = Path(ss)

hkmc2/js/src/test/scala/hkmc2/CompilerTest.scala

@@ -21,7 +21,7 @@ class CompilerTest extends AnyFunSuite:
         None
     .toMap + ("/std/Prelude.mls" -> node.fs.readFileSync(preludePath, "utf-8"))
 
-  private val paths = new Paths("/std/Prelude.mls", "/std/Runtime.mjs", "/std/Term.mjs")
+  private val paths = new Paths("/std/Prelude.mls", "/std/Runtime.mjs", "/std/Term.mjs", "/std/Block.mls", "/std/SpecializeHelpers.mls", "/std/Option.mls", "/std/ShapeSet.mls")
 
   private def createCompiler(): (InMemoryFileSystem, Compiler) =
     val stdLib = loadStandardLibrary()

hkmc2/jvm/src/test/scala/hkmc2/CompileTestRunnerBase.scala

@@ -64,7 +64,11 @@ abstract class CompileTestRunnerBase(
           paths = new MLsCompiler.Paths:
             val preludeFile = mainTestDir / "mlscript" / "decls" / "Prelude.mls"
             val runtimeFile = mainTestDir / "mlscript-compile" / "Runtime.mjs"
-            val termFile = mainTestDir / "mlscript-compile" / "Term.mjs",
+            val termFile = mainTestDir / "mlscript-compile" / "Term.mjs"
+            val blockFile = mainTestDir / "mlscript-compile" / "Block.mjs"
+            val optionFile = mainTestDir / "mlscript-compile" / "Option.mjs"
+            val shapeSetFile = mainTestDir / "mlscript-compile" / "ShapeSet.mjs"
+            val specializeHelpersFile = mainTestDir / "mlscript-compile" / "SpecializeHelpers.mjs",
           mkRaise = report.mkRaise
         )
         compiler.compileModule(file)

hkmc2/shared/src/main/scala/hkmc2/MLsCompiler.scala

@@ -39,6 +39,10 @@ object MLsCompiler:
     def preludeFile: io.Path
     def runtimeFile: io.Path
     def termFile: io.Path
+    def blockFile: io.Path
+    def specializeHelpersFile: io.Path
+    def optionFile: io.Path
+    def shapeSetFile: io.Path
 
 /**
   * The compiler that compiles MLscript code into JavaScript modules.
@@ -103,14 +107,27 @@ class MLsCompiler
         case Term.Quoted(_) | Term.Unquoted(_) => true
         case Term.Ref(sym) => sym === State.termSymbol
         case _ => t.subTerms.exists(findQuote)
+      def findStage(t: semantics.Statement): Bool = t match
+        case d: semantics.Definition => d.hasStagedModifier.isDefined
+        case _ => t.subStatements.exists(findStage)
+
       val hasQuote = findQuote(blk0)
+      val isStaged = findStage(blk0)
+      // println(s"yydz: ${blk0.subTerms}")
       val blk = new Term.Blk(
         Import(State.runtimeSymbol, runtimeFile.toString, runtimeFile) ::
-          // Only import `Term.mls` when necessary.
+          // Only import files when necessary.
           (if hasQuote then
-            Import(State.termSymbol, termFile.toString, termFile) :: blk0.stats
+            Import(State.termSymbol, termFile.toString, termFile) :: Nil
           else
-            blk0.stats),
+            Nil) :::
+          (if isStaged then
+            Import(State.optionSymbol, optionFile.toString, optionFile) ::
+            Import(State.shapeSetSymbol, shapeSetFile.toString, shapeSetFile) ::
+            Import(State.blockSymbol, blockFile.toString, blockFile) ::
+            Import(State.specializeHelpersSymbol, specializeHelpersFile.toString, specializeHelpersFile) :: Nil
+          else
+            Nil) ::: blk0.stats,
         blk0.res
       )
       val low = ltl.givenIn:

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

@@ -761,17 +761,15 @@ final case class ClsLikeBody(
     ctor.freeVars ++ methods.flatMap(_.freeVars)
   lazy val size = 1 + methods.map(_.size).sum + ctor.size
 
-/*
 object ClsLikeBody:
-  // TODO rm `empty`? it's currently unused
   def empty(id: Tree.Ident)(using State) = ClsLikeBody(
     isym = ModuleOrObjectSymbol(Tree.DummyTypeDef(syntax.Mod), id),
     methods = Nil,
     privateFields = Nil,
     publicFields = Nil,
     ctor = End(),
+    annotations = Nil,
   )
-*/
 
 final case class Handler(
     sym: BlockMemberSymbol,

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

@@ -36,7 +36,8 @@ class FirstClassFunctionTransformer
         End()), End(), None, None)(N, annotations = Nil)
 
   private def getParamList(l: BlockMemberSymbol): Option[ParamList] = funDefns.get(l) match
-    case Some(fd) => fd.params.headOption
+    case Some(fd) => fd.params.headOption.map(pl =>
+      ParamList(pl.flags, pl.params.map(p => Param(p.flags, VarSymbol(p.sym.id), p.sign, p.modulefulness)), pl.restParam))
     case _ => l.tsym.flatMap(getParamList)
 
   private def getParamList(ts: TermSymbol): Option[ParamList] =
@@ -117,4 +118,19 @@ class FirstClassFunctionTransformer
   def transform(b: Block): Block =
     val desugared = new DesugarMultipleParamList().applyBlock(b)
     new CollectFunDefns().applyBlock(desugared)
-    applyBlock(desugared)
+    new LabelTransformer().applyBlock(applyBlock(desugared))
+
+
+class LabelTransformer(using State, Raise) extends BlockTransformer(new SymbolSubst()):
+  private val contMap = HashMap.empty[Symbol, Symbol]
+
+  override def applyBlock(b: Block): Block = b match
+    case Label(label, false, body, rest) =>
+      val contSym = BlockMemberSymbol("cont$", Nil, false)
+      val contFun = FunDefn.withFreshSymbol(N, contSym, PlainParamList(Nil) :: Nil, rest)(N, Nil)
+      contMap.addOne(label -> contSym)
+      super.applyBlock(Scoped(Set(contSym), Define(contFun, body)))
+    case Break(label) => contMap.get(label) match
+      case Some(sym: Symbol) => Return(Call(Value.Ref(sym, N), Nil ne_:: Nil)(true, false, false))
+      case _ => super.applyBlock(b)
+    case _ => super.applyBlock(b)

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

@@ -281,16 +281,13 @@ class Lowering()(using Config, TL, Raise, State, Ctx, SymbolPrinter):
               mod.classCompanion match
               case S(comp) => comp.defn.getOrElse(wat("Module companion without definition", mod.companion))
               case N =>
-                val stagedAnnots = mod.annotations.collect { 
-                  case Annot.Modifier(Keyword.`staged`) => Annot.Modifier(Keyword.`staged`) 
-                }
                 ClassDef.Plain(mod.owner, syntax.Cls, new ClassSymbol(Tree.DummyTypeDef(syntax.Cls), mod.sym.id),
                   mod.bsym,
                   Nil,
                   N,
                   ObjBody(Blk(Nil, UnitVal())),
                   S(mod.sym),
-                  stagedAnnots,
+                  Nil,
                   Nil,
                   N,
                 )