Async JavaScript interop

hkmc2/shared/src/main/scala/hkmc2/AsyncLowering.scala

@@ -0,0 +1,82 @@
+package hkmc2
+package codegen
+
+import scala.annotation.tailrec
+import scala.collection.mutable
+import scala.util.boundary
+import sourcecode.{ Line, FileName, Name }
+
+import hkmc2.utils.*, shorthands.*
+import hkmc2.utils.*
+import hkmc2.utils.SymbolSubst
+import hkmc2.Message.MessageContext
+
+import syntax.{Literal, Tree}
+import semantics.*
+import semantics.Elaborator.ctx
+import semantics.Elaborator.State
+import hkmc2.Config.EffectHandlers
+import scala.collection.mutable.ArrayBuffer
+
+
+class AsyncLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx, Config):
+
+  object Rewriter extends BlockTransformer(SymbolSubst.Id):
+    val collectedFunDefn: ArrayBuffer[FunDefn] = new ArrayBuffer
+    var inAsyncFun: Bool = false
+
+    inline def wrapAwait[T](innerAllowAwait: Bool)(inline thunk: => T): T =
+      val saved = inAsyncFun
+      inAsyncFun = innerAllowAwait
+      val result = thunk
+      inAsyncFun = saved
+      result
+
+    override def applyFunDefn(fun: FunDefn): FunDefn =
+      if !fun.async then return wrapAwait(false)(super.applyFunDefn(fun))
+      val outerBms = BlockMemberSymbol(fun.sym.nme, Nil, fun.sym.nameIsMeaningful)
+      val outerDsym = TermSymbol(syntax.Fun, N, fun.dSym.id)
+      val outerParams = fun.params.flatMap: pl =>
+        pl.allParams.map: p =>
+          val v = p.sym
+          val nv = VarSymbol(v.id)
+          (p, p.copy(sym = nv))
+      val symMap = outerParams.iterator.map(p => p._1.sym -> p._2.sym).toMap[SimpleSymbol, SimpleSymbol]
+      val thisVar = VarSymbol(Tree.Ident("this"))
+      val thisParam = fun.owner.map(_ => Param.simple(thisVar))
+      val vars = fun.params.flatMap(_.paramSyms)
+      val noAsync = fun.annotations.filterNot(_ is Annot.Async)
+      val transformer = new BlockTransformer(SymbolSubst.Id):
+        override def applySimpleSymbol(sym: SimpleSymbol): SimpleSymbol =
+          symMap.getOrElse(sym, sym)
+        override def applyValue(v: Value)(k: Value => Block): Block = v match
+          case Value.This(sym) if fun.owner.contains(sym) =>
+            k(Value.SimpleRef(thisVar))
+          case _ => super.applyValue(v)(k)
+      val newBody = transformer.applyBlock(wrapAwait(true)(applyFunBodyLikeBlock(fun.body)))
+      collectedFunDefn += FunDefn(N, outerBms, outerDsym, PlainParamList((thisParam.iterator ++ outerParams.iterator.map(_._2)).toList) :: PlainParamList(Nil) :: Nil, newBody)(fun.configOverride, noAsync)
+      val callArgs = (fun.owner.iterator.map(s => Arg(N, Value.This(s))) ++ fun.params.iterator.flatMap(_.allParams.iterator.map(p => Arg(N, Value.SimpleRef(p.sym))))).toList
+      val tmp = TempSymbol(N, "tmp")
+      val wrapperBody = blockBuilder
+        .assignScoped(tmp, Call(Value.MemberRef(outerBms, outerDsym), callArgs ne_:: Nil)(CallMetadata.mlsFunWithEffect))
+        .ret(Call(Value.SimpleRef(State.runtimeSymbol).selSN("toJsAsync"), (tmp.asSimpleRef.asArg :: Nil) ne_:: Nil)(CallMetadata.defaultMlsFun))
+      FunDefn(fun.owner, fun.sym, fun.dSym, fun.params, wrapperBody)(fun.configOverride, noAsync)
+
+    override def applyMainBlock(main: Block): Block =
+      collectedFunDefn.foldRight(super.applyMainBlock(main)): (defn, acc) =>
+        Scoped(Set.single(defn.sym), Define(defn, acc))
+
+    override def applyPath(p: Path)(k: Path => Block): Block =
+      p match
+        case s: Select if s.symbol.contains(ctx.builtins.handlers.await) =>
+          if config.effectHandlers.isEmpty && !inAsyncFun then
+            raise(ErrorReport(
+              msg"Only await inside of async bodies are allowed if effect handlers are not enabled." ->
+              p.toLoc :: Nil,
+              source = Diagnostic.Source.Compilation))
+          k(State.runtimeSymbol.asSimpleRef.sel(new Tree.Ident("await"), ctx.builtins.handlers.await))
+        case _ =>
+          super.applyPath(p)(k)
+
+  def transform(prog: Program): Program =
+    Rewriter.applyProgram(prog)

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

@@ -673,8 +673,9 @@ final case class FunDefn(
   val asPath = sym.asMemberRef(dSym)
   lazy val tailRec: Bool = annotations.contains(Annot.TailRec)
   lazy val inline: Bool = annotations.contains(Annot.Inline)
-  lazy val noInline: Bool = annotations.contains(Annot.NoInline) || generator
+  lazy val noInline: Bool = annotations.contains(Annot.NoInline) || generator || async
   lazy val generator: Bool = annotations.contains(Annot.Generator)
+  lazy val async: Bool = annotations.contains(Annot.Async)
   lazy val visibility: Visibility = annotations.collectFirst:
     case Annot.Modifier(Keyword.`private`) => Visibility.Private
     case Annot.Modifier(Keyword.`public`) => Visibility.Public

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

@@ -47,8 +47,8 @@ class CompilationPipeline(using Config, Raise, State, Ctx, SymbolPrinter):
         blockPass(Lifter(_).transform)(prog)
       else prog
     runPass("HandlerLowering"): prog =>
-      config.effectHandlers.fold(prog): opt =>
-        HandlerLowering(new HandlerPaths, opt).translateProgram(prog)
+      HandlerLowering(new HandlerPaths, config.effectHandlers).translateProgram(prog)
+    runPass("AsyncLowering")(AsyncLowering().transform)
     runPass("Flattening")(blockPass(_.flattened))
     runPass("BufferableTransform")(BufferableTransform().transform)
     runPass("MergeMatchArmTransformer")(MergeMatchArmTransformer.applyProgram)

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

@@ -53,11 +53,14 @@ object HandlerLowering:
       case Ctor => true
       case ModCtor(trulyNested) => trulyNested
       case TopLevel => false
+    def inAsync = this match
+      case FunctionLike(ctx) => ctx.inAsync
+      case _ => false
 
 
   // currentFun: path to the current function for resumption
   // thisPath: path to `this` binding if the function is a method, `this` will be rebinded on resumption
-  private case class FunctionCtx(currentFun: Path, thisPath: Option[Path], resumeInfo: ResumeInfo, debugInfo: DebugInfo, inGetter: Bool):
+  private case class FunctionCtx(currentFun: Path, thisPath: Option[Path], resumeInfo: ResumeInfo, debugInfo: DebugInfo, inGetter: Bool, inAsync: Bool):
     def doUnwind(loc: Value, state: Path, restoreList: List[LocalVarSymbol])(using paths: HandlerPaths) =
       Return(Call(paths.unwindPath, (
         currentFun ::
@@ -117,7 +120,10 @@ class HandlerPaths(using Elaborator.State):
   val resumeValueIdent = new Tree.Ident("resumeValue")
   val resumeValue: Path = runtimePath.selN(resumeValueIdent)
 
-class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise, Elaborator.State, Elaborator.Ctx, Config):
+class HandlerLowering(paths: HandlerPaths, opt: Opt[EffectHandlers])(using TL, Raise, Elaborator.State, Elaborator.Ctx, Config):
+
+  val debugEnabled = opt.exists(_.debug)
+  val stackSafety = opt.flatMap(_.stackSafety)
 
   private def freshTmp(dbgNme: Str = "tmp") = new TempSymbol(N, dbgNme)
   private def freshLabel(nme: Str) = new LabelSymbol(N, nme)
@@ -320,7 +326,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
 
     val initId = allocId()
     // Note: initial part will only be resumed if stack safety is on.
-    val initPart = BlockPartition(go(blk)(using N, false), opt.stackSafety.isDefined)
+    val initPart = BlockPartition(go(blk)(using N, false), stackSafety.isDefined)
     result(initId) = initPart
 
     val replaceStaleLabels = new BlockTransformerShallow(SymbolSubst.Id):
@@ -511,7 +517,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
     ret.sortBy(x => x.headOption.getOrElse(BigInt(-1))).toList
 
   private def lifterReport(using Line, FileName)(msgs: Ls[Message -> Opt[Loc]])(using Name) =
-    if opt.softLifterError then
+    if opt.fold(false)(_.softLifterError) then
       WarningReport(msgs, source = Diagnostic.Source.Compilation)
     else
       InternalError(msgs, source = Diagnostic.Source.Compilation)
@@ -543,7 +549,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
       val rtArgLists = intLit(fun.params.length) :: fun.params.flatMap: pl =>
         intLit(pl.params.length) :: pl.params.map(p => p.sym.asSimpleRef)
       val newCtx = HandlerCtx.FunctionLike(FunctionCtx(funcPath, thisPath, ResumeInfo(rtArgLists, sortedVars, L(fun.sym)),
-        DebugInfo(debugNme, if opt.debug then debugInfoSym.asSimpleRef else unit), thisPath.isDefined && fun.params.isEmpty))
+        DebugInfo(debugNme, if debugEnabled then debugInfoSym.asSimpleRef else unit), thisPath.isDefined && fun.params.isEmpty, fun.async))
       val bod2 = translateBlock(fun.body, newCtx, scopedVars)
       val fun2 = if fun.body is bod2 then fun else
         FunDefn(fun.owner, fun.sym, fun.dSym, fun.params, bod2)(fun.configOverride, fun.annotations)
@@ -559,21 +565,21 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
         case _ => super.applyResult(r)(k)
       override def applyDefn(defn: Defn)(k: Defn => Block): Block = defn match
         case fun: FunDefn =>
-          if h.currentBlockIsTrulyNested then
+          if h.currentBlockIsTrulyNested && opt.isDefined then
             raise(lifterReport(msg"Unexpected nested function: lambdas may not function correctly." -> fun.sym.toLoc :: Nil))
           val (debugInfoSym, debugInfo, fun2) = translateFunLike(fun, fun.sym.asMemberRef(fun.dSym), N, fun.sym.nme)
-          if opt.debug then Scoped(Set.single(debugInfoSym), Assign(debugInfoSym, Tuple(false, debugInfo), k(fun2))) else k(fun2)
+          if debugEnabled then Scoped(Set.single(debugInfoSym), Assign(debugInfoSym, Tuple(false, debugInfo), k(fun2))) else k(fun2)
         case defn @ ClsLikeDefn(owner, isym, sym, ctorSym, kind, paramsOpt, auxParams, parentPath, methods, privateFields, publicFields, preCtor, ctor, companion, bufferable) =>
-          if h.currentBlockIsTrulyNested then
+          if h.currentBlockIsTrulyNested && opt.isDefined then
             raise(lifterReport(msg"Unexpected nested class: lambdas may not function correctly." -> isym.toLoc :: Nil))
           val debugInfos = mutable.ArrayBuffer.empty[(TempSymbol, List[Arg])]
-          val newMtds = methods.map: f =>
+          val newMtds = methods.mapConserve: f =>
             val (debugInfoSym, debugInfo, fun2) = translateFunLike(f, isym.asThis.sel(new Tree.Ident(f.sym.nme), f.dSym),
               S(isym.asThis), s"${sym.nme}#${f.sym.nme}")
             debugInfos += debugInfoSym -> debugInfo
             fun2
-          val companion2 = companion.map: bod =>
-            val newMtds = bod.methods.map: f =>
+          val companion2 = companion.mapConserve: bod =>
+            val newMtds = bod.methods.mapConserve: f =>
               val (debugInfoSym, debugInfo, fun2) = translateFunLike(f, bod.isym.asThis.sel(new Tree.Ident(f.sym.nme), f.dSym),
                 S(bod.isym.asThis), s"${sym.nme}.${f.sym.nme}")
               debugInfos += debugInfoSym -> debugInfo
@@ -583,18 +589,28 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
             // TODO: Companion's ctor is more well behaved so it is possible to handle it
             // However, JSBuilder inserts extra statements between preCtor and ctor and it's not possible to replicate the exact behavior
             // without many special handling.
-            val newCtor = if opt.doNotInstrumentTopLevelModCtor && !h.currentBlockIsTrulyNested then bod.ctor else
+            val newCtor = if opt.fold(true)(_.doNotInstrumentTopLevelModCtor) && !h.currentBlockIsTrulyNested then bod.ctor else
               translateCtorLike(bod.ctor, bod.isym.asThis, true)
             tl.log(s"companion name: ${bod.isym.nme}")
-            ClsLikeBody(bod.isym, newMtds, bod.privateFields, bod.publicFields, newCtor, bod.annotations)
-          val c2 = ClsLikeDefn(owner, isym, sym, ctorSym, kind, paramsOpt, auxParams, parentPath, newMtds, privateFields, publicFields,
-            translateCtorLike(preCtor, isym.asThis, false), translateCtorLike(ctor, isym.asThis, false), companion2, bufferable)(defn.configOverride, defn.annotations)
-          if opt.debug then
+            if (bod.methods is newMtds) && (bod.ctor is newCtor) then
+              bod
+            else
+              ClsLikeBody(bod.isym, newMtds, bod.privateFields, bod.publicFields, newCtor, bod.annotations)
+          val newPreCtor = translateCtorLike(preCtor, isym.asThis, false)
+          val newCtor = translateCtorLike(ctor, isym.asThis, false)
+          val c2 =
+            if (methods is newMtds) && (preCtor is newPreCtor) && (ctor is newCtor) && (companion is companion2) then
+              defn
+            else
+              defn.copy(methods = newMtds, preCtor = newPreCtor, ctor = newCtor, companion = companion2)(defn.configOverride, defn.annotations)
+          if debugEnabled then
             Scoped(debugInfos.map(_._1).toSet, debugInfos.foldRight(k(c2)): (elem, blk) =>
               Assign(elem._1, Tuple(false, elem._2), blk))
           else k(c2)
         case _ => super.applyDefn(defn)(k)
     val b = preTransform.applyBlock(blk)
+    if !opt.isDefined && !h.inAsync then
+      return b
     if !h.currentBlockIsTrulyNested then
       return postTranslateTopLevelCtx(b)
     if h.inCtor then
@@ -604,11 +620,11 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
       return postTranslateIllegalEffectCtx(b, "in a getter")
     given FunctionCtx = ctx
     val parts = partitionBlock(b)
-    val needsStackSafety = parts.needsStackSafety && opt.stackSafety.isDefined
+    val needsStackSafety = parts.needsStackSafety && stackSafety.isDefined
     val oneState = parts.states.size <= 1
     if oneState && !parts.containsError && !needsStackSafety then
       return b
-    val vars = if opt.debug then ctx.resumeInfo.currentLocals else computeRestoreList(parts)
+    val vars = if debugEnabled then ctx.resumeInfo.currentLocals else computeRestoreList(parts)
 
     val pcVar = freshTmp("pc")
     val curDepth = freshTmp("curDepth")
@@ -734,7 +750,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
    */
 
   private def postTranslateTopLevelCtx(b: Block)(using HandlerCtx): Block =
-    postTranslateIllegalEffectCtx(b, Call.raw(paths.topLevelEffectPath, (Value.Lit(Tree.BoolLit(opt.debug)).asArg :: Nil) ne_:: Nil)(CallMetadata.defaultMlsFun), opt.stackSafety.map(_.stackLimit))
+    postTranslateIllegalEffectCtx(b, Call.raw(paths.topLevelEffectPath, (Value.Lit(Tree.BoolLit(debugEnabled)).asArg :: Nil) ne_:: Nil)(CallMetadata.defaultMlsFun), stackSafety.map(_.stackLimit))
 
   private def postTranslateIllegalEffectCtx(b: Block, reason: Str)(using HandlerCtx): Block =
     postTranslateIllegalEffectCtx(b, Call.raw(paths.illegalEffectPath, (Value.Lit(Tree.StrLit(reason)).asArg :: Nil) ne_:: Nil)(CallMetadata.defaultMlsFun), N)
@@ -781,7 +797,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
     val ctx = HandlerCtx.TopLevel
     val transformed = blockBuilder
         .staticif(
-          !opt.doNotInstrumentTopLevelModCtor,
+          opt.fold(false)(!_.doNotInstrumentTopLevelModCtor),
           _.assign(NoSymbol, Call(paths.resetEffects, Nil ne_:: Nil)(CallMetadata.defaultMlsFun))
         )
         .rest(translateBlock(prog.main, ctx, Set.empty))

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

@@ -707,6 +707,28 @@ class Lowering()(using Config, TL, Raise, State, Ctx, SymbolPrinter):
 
     @tailrec
     def extractAnnots(t: st, acc: List[Annot]): (List[Annot], st) = t match
+      case st.Annotated(Annot.Async, trm) =>
+        val ident = new Tree.Ident("asyncBody").withLocOf(trm)
+        val bms = BlockMemberSymbol(ident.name, Nil, false)
+        val dsym = TermSymbol(syntax.Fun, N, ident)
+        val td: TermDefinition = TermDefinition(
+          syntax.Fun,
+          bms,
+          dsym,
+          sem.ParamList(sem.ParamListFlags.empty, Nil, N) :: Nil,
+          N,
+          N,
+          S(trm),
+          TermDefFlags.empty,
+          Modulefulness.none,
+          Annot.RaiseEffects :: Nil,
+          N,
+        )
+        val rewritten = st.App(
+          st.SynthSel(State.runtimeSymbol.ref(), Tree.Ident("toJsAsync"))(N, FlowSymbol.sel("toJsAsync"), N, N),
+          st.Tup(PlainFld(st.Blk(td :: Nil, bms.ref(ident).resolved(dsym))) :: Nil)(Tree.DummyTup)
+        )(Tree.DummyApp, N, FlowSymbol.app())
+        extractAnnots(rewritten, acc)
       case st.Annotated(annot, trm) => extractAnnots(trm.instantiated, annot :: acc)
       case _ => (acc, t)
 
@@ -913,7 +935,7 @@ class Lowering()(using Config, TL, Raise, State, Ctx, SymbolPrinter):
           conclude(Select(p, definitionIdent(nme, sym))(S(sym))(false).withLocOf(sel))
       case _ => subTerm(baseF)(conclude)
     case h @ Handle(lhs, rhs, as, cls, defs, bod) =>
-      if !lowerHandlers then
+      if config.effectHandlers.isEmpty then
         return fail:
           ErrorReport(
             msg"Effect handlers are not enabled" ->
@@ -1416,12 +1438,14 @@ class Lowering()(using Config, TL, Raise, State, Ctx, SymbolPrinter):
         case N => WarningReport(msg"This annotation has no effect." -> annot.toLoc :: Nil)
     annotations.foreach:
       case Annot.Untyped => ()
-      case a @ (Annot.TailRec | Annot.Inline | Annot.NoInline | Annot.Generator) =>
+      case a @ (Annot.TailRec | Annot.Inline | Annot.NoInline | Annot.Generator | Annot.Async | Annot.RaiseEffects) =>
         val annot = a match
           case Annot.TailRec => "@tailrec"
           case Annot.Inline => "@inline"
           case Annot.NoInline => "@noInline"
           case Annot.Generator => "@generator"
+          case Annot.Async => "@async"
+          case Annot.RaiseEffects => "@raiseEffects"
         target match
           case TermDefinition(body = S(bod), k = syntax.Fun) => ()
           case TermDefinition(k = syntax.Fun) => warn(a, S(msg"Only functions with a body may be marked as $annot."))