User guided specialisation stage

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

@@ -50,6 +50,7 @@ class MLsCompiler(preludeFile: os.Path, mkOutput: ((Str => Unit) => Unit) => Uni
 
   // TODO adapt logic
   val etl = new TraceLogger{override def doTrace: Bool = false}
+  val stl = new TraceLogger{override def doTrace: Bool = false}
   val ltl = new TraceLogger{override def doTrace: Bool = false}
 
 
@@ -79,14 +80,16 @@ class MLsCompiler(preludeFile: os.Path, mkOutput: ((Str => Unit) => Unit) => Uni
     newCtx.nest(N).givenIn:
       val elab = Elaborator(etl, wd, newCtx)
       val parsed = mainParse.resultBlk
-      val (blk0, _) = elab.importFrom(parsed)
-      val blk = blk0.copy(stats = semantics.Import(State.runtimeSymbol, runtimeFile.toString) :: blk0.stats)
+      val (blk0, ctx) = elab.importFrom(parsed)
+      val blk: semantics.Term.Blk = blk0.copy(stats = semantics.Import(State.runtimeSymbol, runtimeFile.toString) :: blk0.stats)
+      val typ = new semantics.Specialiser(ctx, stl)
+      val spBlk = typ.topLevel(blk)
       val low = ltl.givenIn:
         new codegen.Lowering()
           with codegen.LoweringSelSanityChecks
       val jsb = ltl.givenIn:
         codegen.js.JSBuilder()
-      val le = low.program(blk)
+      val le = low.program(spBlk)
       val baseScp: utils.Scope =
         utils.Scope.empty
       val nestedScp = baseScp.nest

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

@@ -106,13 +106,14 @@ sealed abstract class Block extends Product with AutoLocated:
     case AssignDynField(_, _, _, rhs, rest) => rhs.subBlocks ::: rest :: Nil
     case Define(d, rest) => d.subBlocks ::: rest :: Nil
     case HandleBlock(_, _, par, args, _, handlers, body, rest) => par.subBlocks ++ args.flatMap(_.subBlocks) ++ handlers.map(_.body) :+ body :+ rest
+    case Label(_, body, rest) => body :: rest :: Nil
 
     // TODO rm Lam from values and thus the need for these cases
     case Return(r, _) => r.subBlocks
     case HandleBlockReturn(r) => r.subBlocks
     case Throw(r) => r.subBlocks
 
-    case _: Return | _: Throw | _: Label | _: Break | _: Continue | _: End | _: HandleBlockReturn => Nil
+    case _: Return | _: Throw | _: Break | _: Continue | _: End | _: HandleBlockReturn => Nil
 
   // Moves definitions in a block to the top. Only scans the top-level definitions of the block;
   // i.e, definitions inside other definitions are not moved out. Definitions inside `match`/`if`

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

@@ -85,13 +85,13 @@ class BlockTransformer(subst: SymbolSubst):
           (cls2 is cls) && (hdr2 is hdr) && (bod2 is bod) && (rst2 is rst)
         then b else HandleBlock(l2, res2, par2, args2, cls2, hdr2, bod2, rst2)
     case AssignDynField(lhs, fld, arrayIdx, rhs, rest) =>
-      val lhs2 = applyPath(lhs)
-      val fld2 = applyPath(fld)
-      val rhs2 = applyResult(rhs)
-      val rest2 = applyBlock(rest)
-      if (lhs2 is lhs) && (fld2 is fld) && (rhs2 is rhs) && (rest2 is rest)
-      then b
-      else AssignDynField(lhs2, fld2, arrayIdx, rhs2, rest2)
+      applyResult2(rhs): rhs2 =>
+        val lhs2 = applyPath(lhs)
+        val fld2 = applyPath(fld)
+        val rest2 = applyBlock(rest)
+        if (lhs2 is lhs) && (fld2 is fld) && (rhs2 is rhs) && (rest2 is rest)
+        then b
+        else AssignDynField(lhs2, fld2, arrayIdx, rhs2, rest2)
 
 
   def applyResult2(r: Result)(k: Result => Block): Block = k(applyResult(r))

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

@@ -59,6 +59,7 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
         .assignFieldN(state.res, tailIdent, state.res.tail.next)
         .ret(state.res))
   private val functionHandlerCtx = funcLikeHandlerCtx(N)
+  private val topLevelCtx = HandlerCtx(true, true, N, _ => rtThrowMsg("Unhandled effects"))
   private def ctorCtx(ctorThis: Path) = funcLikeHandlerCtx(S(ctorThis))
   private def handlerCtx(using HandlerCtx): HandlerCtx = summon
   private val predefPath: Path = State.globalThisSymbol.asPath.selN(Tree.Ident("Predef"))
@@ -267,6 +268,9 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
       case blk @ AssignField(lhs, nme, rhs, rest) =>
         val PartRet(head, parts) = go(rest)
         PartRet(AssignField(lhs, nme, rhs, head)(blk.symbol), parts)
+      case AssignDynField(lhs, fld, arrayIdx, rhs, rest) =>
+        val PartRet(head, parts) = go(rest)
+        PartRet(AssignDynField(lhs, fld, arrayIdx, rhs, head), parts)
       case Return(_, _) => PartRet(blk, Nil)
       // ignored cases
       case TryBlock(sub, finallyDo, rest) => ??? // ignore
@@ -336,7 +340,7 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
       override def applyLam(lam: Value.Lam): Value.Lam = Value.Lam(lam.params, translateBlock(lam.body, functionHandlerCtx))
       override def applyDefn(defn: Defn): Defn = defn match
         case f: FunDefn => translateFun(f)
-        case c: ClsLikeDefn => translateCls(c)
+        case c: ClsLikeDefn => translateCls(c, handlerCtx.isTopLevel)
         case _: ValDefn => super.applyDefn(defn)
     transformer.applyBlock(b)
 
@@ -356,9 +360,11 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
   private def translateFun(f: FunDefn): FunDefn =
     FunDefn(f.owner, f.sym, f.params, translateBlock(f.body, functionHandlerCtx))
 
-  private def translateCls(cls: ClsLikeDefn): ClsLikeDefn =
+  private def translateCls(cls: ClsLikeDefn, isTopLevel: Bool): ClsLikeDefn =
+    val curCtorCtx = if isTopLevel && (cls.k is syntax.Mod) then topLevelCtx else
+      ctorCtx(cls.sym.asClsLike.getOrElse(wat("asClsLike", cls.sym)).asPath)
     cls.copy(methods = cls.methods.map(translateFun),
-      ctor = translateBlock(cls.ctor, ctorCtx(cls.sym.asClsLike.getOrElse(wat("asClsLike", cls.sym)).asPath)))
+      ctor = translateBlock(cls.ctor, curCtorCtx))
 
   // Handle block becomes a FunDefn and CallPlaceholder
   private def translateHandleBlock(h: HandleBlock): Block =
@@ -548,5 +554,5 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
     transform.applyBlock(b)
 
   def translateTopLevel(b: Block): Block =
-    translateBlock(b, HandlerCtx(true, true, N, _ => rtThrowMsg("Unhandled effects")))
+    translateBlock(b, topLevelCtx)
 

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

@@ -98,7 +98,7 @@ class StackSafeTransform(depthLimit: Int)(using State):
       override def applyFunDefn(fun: FunDefn): FunDefn = rewriteFn(fun)
 
       override def applyDefn(defn: Defn): Defn = defn match
-        case defn: ClsLikeDefn => rewriteCls(defn)
+        case defn: ClsLikeDefn => rewriteCls(defn, isTopLevel)
         case _: FunDefn | _: ValDefn => super.applyDefn(defn)
 
       override def applyBlock(b: Block): Block = b match
@@ -143,12 +143,13 @@ class StackSafeTransform(depthLimit: Int)(using State):
     walker.applyBlock(b)
     trivial
 
-  def rewriteCls(defn: ClsLikeDefn): ClsLikeDefn = 
+  def rewriteCls(defn: ClsLikeDefn, isTopLevel: Bool): ClsLikeDefn = 
     val ClsLikeDefn(owner, isym, sym, k, paramsOpt, 
       parentPath, methods, privateFields, publicFields, preCtor, ctor) = defn
     ClsLikeDefn(
       owner, isym, sym, k, paramsOpt, parentPath, methods.map(rewriteFn), privateFields,
-      publicFields, rewriteBlk(preCtor), rewriteBlk(ctor)
+      publicFields, rewriteBlk(preCtor),
+      if isTopLevel && (defn.k is syntax.Mod) then transformTopLevel(ctor) else rewriteBlk(ctor)
     )
 
   def rewriteBlk(blk: Block) =

hkmc2/shared/src/main/scala/hkmc2/codegen/js/JSBuilder.scala

@@ -307,6 +307,7 @@ class JSBuilder(using TL, State, Ctx) extends CodeBuilder:
           case Elaborator.ctx.builtins.Num => doc"typeof $sd === 'number'"
           case Elaborator.ctx.builtins.Bool => doc"typeof $sd === 'boolean'"
           case Elaborator.ctx.builtins.Int => doc"globalThis.Number.isInteger($sd)"
+          case Elaborator.ctx.builtins.Function => doc"typeof $sd === 'function'"
           case _ => doc"$sd instanceof ${result(pth)}"
         case Case.Tup(len, inf) => doc"globalThis.Array.isArray($sd) && $sd.length ${if inf then ">=" else "==="} ${len}"
       val h = doc" # if (${ cond(hd._1) }) ${ braced(returningTerm(hd._2, endSemi = false)) }"

hkmc2/shared/src/main/scala/hkmc2/semantics/Elaborator.scala

@@ -1108,6 +1108,8 @@ extends Importer:
       ps
     case TypeDef(Pat, inner, N, N) =>
       param(inner, inUsing).map(_.mapSecond(p => p.copy(flags = p.flags.copy(pat = true))))
+    case Modified(Keyword.`spec`, _, inner) =>
+      param(inner, inUsing).map(_.mapSecond(p => p.copy(flags = p.flags.copy(spec = true))))
     case _ =>
       t.asParam(inUsing).map: (isSpd, p, t) =>
         isSpd -> Param(FldFlags.empty, fieldOrVarSym(ParamBind, p), t.map(term(_)))