Priliminary support for staged keyword (#348)

ChingLongTin · yeungsinchun · LPTK · web-flow · commit 481546fffb5b · 2025-11-06T21:26:57.000+08:00
Co-authored-by: ChingLongTin &lt;26105652+Sunny-Ching@users.noreply.github.com&gt;
Co-authored-by: TYeung &lt;scyeungaf@connect.ust.hk&gt;
Co-authored-by: Lionel Parreaux &lt;lionel.parreaux@gmail.com&gt;
diff --git a/hkmc2/shared/src/main/scala/hkmc2/Config.scala b/hkmc2/shared/src/main/scala/hkmc2/Config.scala
@@ -20,6 +20,7 @@ case class Config(
   sanityChecks: Opt[SanityChecks],
   effectHandlers: Opt[EffectHandlers],
   liftDefns: Opt[LiftDefns],
+  stageCode: Bool,
   target: CompilationTarget,
 ):
   
@@ -35,6 +36,7 @@ object Config:
     // sanityChecks = S(SanityChecks(light = true)),
     effectHandlers = N,
     liftDefns = N,
+    stageCode = false,
     target = CompilationTarget.JS
   )
   
diff --git a/hkmc2/shared/src/main/scala/hkmc2/codegen/Instrumentation.scala b/hkmc2/shared/src/main/scala/hkmc2/codegen/Instrumentation.scala
@@ -0,0 +1,15 @@
+package hkmc2
+package codegen
+
+import utils.*
+import hkmc2.Message.MessageContext
+
+class Instrumentation(using Raise) extends BlockTransformer(new SymbolSubst()):
+  def transform(prgm: Program) = Program(prgm.imports, applyBlock(prgm.main))
+
+  override def applyDefn(d: Defn)(k: Defn => Block): Block = d match 
+    case defn: ClsLikeDefn =>
+      if defn.sym.defn.exists(_.hasStagedModifier.isDefined) && defn.companion.isDefined
+      then raise(WarningReport(msg"`staged` keyword doesn't do anything currently." -> defn.sym.toLoc :: Nil))
+      super.applyDefn(defn)(k)
+    case b => super.applyDefn(b)(k)
diff --git a/hkmc2/shared/src/main/scala/hkmc2/codegen/Lowering.scala b/hkmc2/shared/src/main/scala/hkmc2/codegen/Lowering.scala
@@ -11,6 +11,8 @@ import utils.*
 
 import hkmc2.Message.MessageContext
 
+import codegen.Instrumentation
+
 import semantics.*, ucs.FlatPattern
 import hkmc2.{semantics => sem}
 import semantics.{Term => st}
@@ -1068,7 +1070,11 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
     
     val bufferable = BufferableTransform().transform(lifted)
     
-    val res = MergeMatchArmTransformer.applyBlock(bufferable)
+    val merged = MergeMatchArmTransformer.applyBlock(bufferable)
+
+    val res = 
+      if config.stageCode then Instrumentation(using summon).applyBlock(merged)
+      else merged
     
     Program(
       imps.map(imp => imp.sym -> imp.str),
@@ -1095,6 +1101,7 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
   def reportAnnotations(target: Statement, annotations: Ls[Annot]): Unit =
     annotations.foreach:
       case Annot.Untyped => ()
+      case Annot.Modifier(syntax.Keyword("staged")) => ()
       case annot => raise:
         WarningReport(msg"This annotation has no effect." -> annot.toLoc :: Nil)
 
diff --git a/hkmc2/shared/src/main/scala/hkmc2/codegen/Printer.scala b/hkmc2/shared/src/main/scala/hkmc2/codegen/Printer.scala
@@ -82,7 +82,8 @@ object Printer:
       val docPubFlds = if publicFields.isEmpty then doc"" else doc" # ${pubFields}"
       val docBody = if publicFields.isEmpty && privateFields.isEmpty then doc"" else doc" { #{ ${docPrivFlds}${docPubFlds} #}  # }"
       val docCtorParams = if clsParams.isEmpty then doc"" else doc"(${ctorParams.mkDocument(", ")})"
-      doc"class ${own.fold("")(_.toString+"::")}${sym.nme}${docCtorParams}${docBody}"
+      val docStaged = if sym.defn.forall(_.hasStagedModifier.isEmpty) then doc"" else doc"staged "
+      doc"${docStaged}class ${own.fold("")(_.toString+"::")}${sym.nme}${docCtorParams}${docBody}"
   
   def mkDocument(arg: Arg)(using Raise, Scope): Document =
     val doc = mkDocument(arg.value)
diff --git a/hkmc2/shared/src/main/scala/hkmc2/semantics/Elaborator.scala b/hkmc2/shared/src/main/scala/hkmc2/semantics/Elaborator.scala
@@ -233,6 +233,8 @@ object Elaborator:
     val definitionMetadataSymbol = TempSymbol(N, "definitionMetadata")
     val prettyPrintSymbol = TempSymbol(N, "prettyPrint")
     val termSymbol = TempSymbol(N, "Term")
+    val blockSymbol = TempSymbol(N, "Block")
+    val shapeSymbol = TempSymbol(N, "Shape")
     val wasmSymbol = TempSymbol(N, "wasm")
     val effectSigSymbol = ClassSymbol(DummyTypeDef(syntax.Cls), Ident("EffectSig"))
     val nonLocalRetHandlerTrm =
@@ -319,7 +321,7 @@ extends Importer:
     case _ => N
   
   def annot(tree: Tree): Ctxl[Opt[Annot]] = tree match
-    case Keywrd(kw @ (Keyword.`abstract` | Keyword.`declare` | Keyword.`data`)) => S(Annot.Modifier(kw))
+    case Keywrd(kw @ (Keyword.`abstract` | Keyword.`declare` | Keyword.`data` | Keyword.`staged`)) => S(Annot.Modifier(kw))
     case _ => term(tree) match
       case Term.Error => N
       case trm =>
diff --git a/hkmc2/shared/src/main/scala/hkmc2/semantics/Term.scala b/hkmc2/shared/src/main/scala/hkmc2/semantics/Term.scala
@@ -678,6 +678,8 @@ sealed abstract class Definition extends Declaration, Statement:
   val annotations: Ls[Annot]
   def hasDeclareModifier: Opt[Annot.Modifier] = annotations.collectFirst:
     case mod @ Annot.Modifier(Keyword.`declare`) => mod
+  def hasStagedModifier: Opt[Annot.Modifier] = annotations.collectFirst:
+    case mod @ Annot.Modifier(Keyword.`staged`) => mod
 
 sealed trait CompanionValue extends Definition
 
diff --git a/hkmc2/shared/src/main/scala/hkmc2/syntax/Keyword.scala b/hkmc2/shared/src/main/scala/hkmc2/syntax/Keyword.scala
@@ -113,6 +113,7 @@ object Keyword:
   val `abstract` = Keyword("abstract", N, N)
   val `constructor` = Keyword("constructor", N, N)
   val `virtual` = Keyword("virtual", N, N)
+  val `staged` = Keyword("staged", N, N)
   val `true` = Keyword("true", N, N)
   val `false` = Keyword("false", N, N)
   val `public` = Keyword("public", N, N)
@@ -168,5 +169,5 @@ object Keyword:
   type LetLike = `let`.type | `set`.type
   
   type Modifier = `in`.type | `out`.type | `mut`.type | `abstract`.type | `declare`.type | `data`.type | `virtual`.type | `override`.type |
-    `public`.type | `private`.type
+    `public`.type | `private`.type | `staged`.type
 
diff --git a/hkmc2/shared/src/main/scala/hkmc2/syntax/Lexer.scala b/hkmc2/shared/src/main/scala/hkmc2/syntax/Lexer.scala
@@ -576,7 +576,8 @@ object Lexer:
     "undefined",
     "abstract",
     "constructor",
-    "virtual"
+    "virtual",
+    "staged"
   )
   
   private val SEP = "┊"
diff --git a/hkmc2/shared/src/main/scala/hkmc2/syntax/ParseRule.scala b/hkmc2/shared/src/main/scala/hkmc2/syntax/ParseRule.scala
@@ -377,6 +377,7 @@ class ParseRules(using State):
     modified(`public`),
     modified(`private`),
     modified(`out`),
+    modified(`staged`),
     singleKw(`true`)(BoolLit(true)),
     singleKw(`false`)(BoolLit(false)),
     singleKw(`undefined`)(UnitLit(false)),
diff --git a/hkmc2/shared/src/main/scala/hkmc2/syntax/Tree.scala b/hkmc2/shared/src/main/scala/hkmc2/syntax/Tree.scala
@@ -243,6 +243,8 @@ enum Tree extends AutoLocated:
           Annotated(kw, s.desugared)
         case Modified(kw @ Keywrd(Keyword.`abstract`), s) =>
           Annotated(kw, s.desugared)
+        case Modified(kw @ Keywrd(Keyword.`staged`), s) =>
+          Annotated(kw, s.desugared)
         case Modified(kw @ Keywrd(Keyword.`mut`), TermDef(ImmutVal, anme, rhs)) =>
           TermDef(MutVal, anme, rhs).withLocOf(this).desugared
         case _ => m
diff --git a/hkmc2/shared/src/test/mlscript-compile/Block.mls b/hkmc2/shared/src/test/mlscript-compile/Block.mls
diff --git a/hkmc2/shared/src/test/mlscript-compile/Shape.mls b/hkmc2/shared/src/test/mlscript-compile/Shape.mls
diff --git a/hkmc2/shared/src/test/mlscript/staging/Syntax.mls b/hkmc2/shared/src/test/mlscript/staging/Syntax.mls
@@ -0,0 +1,28 @@
+
+:pt
+staged module A
+//│ Parsed tree:
+//│ Modified:
+//│   modifier = Keywrd of keyword 'staged'
+//│   body = TypeDef:
+//│     k = Mod
+//│     head = Ident of "A"
+//│     rhs = N
+
+// TODO: reject these annotations?
+staged object Foo
+
+staged class Foo
+
+staged fun f() = 0
+
+:js
+:slot
+:staging
+:w
+staged module A
+//│ ╔══[WARNING] `staged` keyword doesn't do anything currently.
+//│ ║  l.23: 	staged module A
+//│ ╙──      	       ^^^^^^^^
+//│ Pretty Lowered:
+//│   define staged class A in set block$res = undefined in end
diff --git a/hkmc2DiffTests/src/test/scala/hkmc2/JSBackendDiffMaker.scala b/hkmc2DiffTests/src/test/scala/hkmc2/JSBackendDiffMaker.scala
@@ -32,6 +32,8 @@ abstract class JSBackendDiffMaker extends MLsDiffMaker:
   
   val runtimeNme = baseScp.allocateName(Elaborator.State.runtimeSymbol)
   val termNme = baseScp.allocateName(Elaborator.State.termSymbol)
+  val blockNme = baseScp.allocateName(Elaborator.State.blockSymbol)
+  val shapeNme = baseScp.allocateName(Elaborator.State.shapeSymbol)
   val definitionMetadataNme = baseScp.allocateName(Elaborator.State.definitionMetadataSymbol)
   val prettyPrintNme = baseScp.allocateName(Elaborator.State.prettyPrintSymbol)
   
@@ -57,6 +59,9 @@ abstract class JSBackendDiffMaker extends MLsDiffMaker:
     h.execute(s"const $definitionMetadataNme = Symbol.for(\"mlscript.definitionMetadata\");")
     h.execute(s"const $prettyPrintNme = Symbol.for(\"mlscript.prettyPrint\");")
     if importQQ.isSet then importRuntimeModule(termNme, termFile)
+    if stageCode.isSet then
+      importRuntimeModule(blockNme, blockFile)
+      importRuntimeModule(shapeNme, shapeFile)
     h
   
   private var hostCreated = false
diff --git a/hkmc2DiffTests/src/test/scala/hkmc2/MLsDiffMaker.scala b/hkmc2DiffTests/src/test/scala/hkmc2/MLsDiffMaker.scala
@@ -20,6 +20,8 @@ abstract class MLsDiffMaker extends DiffMaker:
   val predefFile: os.Path // * Contains MLscript standard library definitions
   val runtimeFile: os.Path = predefFile/os.up/"Runtime.mjs" // * Contains MLscript runtime definitions
   val termFile: os.Path = predefFile/os.up/"Term.mjs" // * Contains MLscript runtime term definitions
+  val blockFile: os.Path = predefFile/os.up/"Block.mjs" // * Contains MLscript runtime block definitions
+  val shapeFile: os.Path = predefFile/os.up/"Shape.mjs" // * Contains MLscript runtime shape definitions
   
   val wd = file / os.up
   
@@ -68,6 +70,7 @@ abstract class MLsDiffMaker extends DiffMaker:
   val stackSafe = Command("stackSafe")(_.trim)
   val liftDefns = NullaryCommand("lift")
   val importQQ = NullaryCommand("qq")
+  val stageCode = NullaryCommand("staging")
   
   def mkConfig: Config =
     import Config.*
@@ -93,6 +96,7 @@ abstract class MLsDiffMaker extends DiffMaker:
         ,
       )),
       liftDefns = Opt.when(liftDefns.isSet)(LiftDefns()),
+      stageCode = stageCode.isSet,
       target = if wasm.isSet then CompilationTarget.Wasm else CompilationTarget.JS,
     )
   
@@ -156,6 +160,12 @@ abstract class MLsDiffMaker extends DiffMaker:
       given Config = mkConfig
       processTrees(
         PrefixApp(Keywrd(`import`), StrLit(termFile.toString)) :: Nil)
+    if stageCode.isSet then
+      given Config = mkConfig
+      processTrees(
+        PrefixApp(Keywrd(`import`), StrLit(blockFile.toString))
+        :: PrefixApp(Keywrd(`import`), StrLit(shapeFile.toString))
+        :: Nil)
     super.init()
   
   

Original file line number	Diff line number	Diff line change
`@@ -576,7 +576,8 @@ object Lexer:`
`576`	`576`	`"undefined",`
`577`	`577`	`"abstract",`
`578`	`578`	`"constructor",`
`579`		`- "virtual"`
	`579`	`+ "virtual",`
	`580`	`+ "staged"`
`580`	`581`	`)`
`581`	`582`
`582`	`583`	`private val SEP = "┊"`