JS runtime: Extract RESUME function (#708)

Co-authored-by: Jonathan Brachthäuser <[email protected]>

Author: dvdvgt

effekt/shared/src/main/scala/effekt/cps/Transformer.scala

@@ -138,8 +138,8 @@ object Transformer {
     case core.Stmt.Resume(cont, body) =>
       val ks2 = Id("ks")
       val k2 = Id("k")
-      val thunk: BlockLit = Block.BlockLit(Nil, Nil, ks2, k2, transform(body, ks2, Continuation.Dynamic(k2)))
-      App(Block.BlockVar(cont.id), Nil, List(thunk), MetaCont(ks), k.reify)
+      Resume(cont.id, Block.BlockLit(Nil, Nil, ks2, k2, transform(body, ks2, Continuation.Dynamic(k2))),
+        MetaCont(ks), k.reify)
 
     case core.Stmt.Hole() => Hole()
 

effekt/shared/src/main/scala/effekt/cps/Tree.scala

@@ -139,6 +139,9 @@ enum Stmt extends Tree {
   // shift(p, { (resume, ks, k) => STMT }, ks, k)
   case Shift(prompt: Id, body: BlockLit, ks: MetaCont, k: Cont)
 
+  // resume(r, (ks, k) => STMT, ks, k)
+  case Resume(resumption: Id, body: BlockLit, ks: MetaCont, k: Cont)
+
   // Others
   case Hole()
 }
@@ -219,6 +222,7 @@ object Variables {
 
     case Stmt.Reset(prog, ks, k) => free(prog) ++ free(ks) ++ free(k)
     case Stmt.Shift(prompt, body, ks, k) => block(prompt) ++ free(body) ++ free(ks) ++ free(k)
+    case Stmt.Resume(r, body, ks, k) => block(r) ++ free(body) ++ free(ks) ++ free(k)
     case Stmt.Hole() => empty
   }
 

effekt/shared/src/main/scala/effekt/generator/js/TransformerCps.scala

@@ -14,6 +14,7 @@ object TransformerCps extends Transformer {
   val RUN_TOPLEVEL = Variable(JSName("RUN_TOPLEVEL"))
   val RESET = Variable(JSName("RESET"))
   val SHIFT = Variable(JSName("SHIFT"))
+  val RESUME = Variable(JSName("RESUME"))
   val THUNK = Variable(JSName("THUNK"))
   val DEALLOC = Variable(JSName("DEALLOC"))
   val TRAMPOLINE = Variable(JSName("TRAMPOLINE"))
@@ -257,6 +258,9 @@ object TransformerCps extends Transformer {
     case cps.Stmt.Shift(prompt, body, ks, k) =>
       pure(js.Return(Call(SHIFT, nameRef(prompt), noThunking { toJS(body) }, toJS(ks), toJS(k))))
 
+    case cps.Stmt.Resume(r, b, ks2, k2) =>
+      pure(js.Return(js.Call(RESUME, nameRef(r), toJS(b), toJS(ks2), toJS(k2))))
+
     case cps.Stmt.Hole() =>
       pure(js.Return($effekt.call("hole")))
   }

libraries/js/effekt_runtime.js

@@ -113,22 +113,23 @@ function SHIFT(p, body, ks, k) {
   cont = { stack: meta.stack, prompt: meta.prompt, backup: meta.arena.backup(), rest: cont }
   meta = meta.rest
 
-  function resumeComp(c, ks, k) {
-    let meta = { stack: k, prompt: ks.prompt, arena: ks.arena, rest: ks.rest }
-    let toRewind = cont
-    while (!!toRewind) {
-      meta = { stack: toRewind.stack, prompt: toRewind.prompt, arena: toRewind.backup(), rest: meta }
-      toRewind = toRewind.rest
-    }
+  const k1 = meta.stack
+  meta.stack = null
+  return body(cont, meta, k1)
+}
 
-    const k2 = meta.stack // TODO instead copy meta here, like elsewhere?
-    meta.stack = null
-    return () => c(meta, k2)
+// Rewind stack `cont` back onto `k` :: `ks` and resume with c
+function RESUME(cont, c, ks, k) {
+  let meta = { stack: k, prompt: ks.prompt, arena: ks.arena, rest: ks.rest }
+  let toRewind = cont
+  while (!!toRewind) {
+    meta = { stack: toRewind.stack, prompt: toRewind.prompt, arena: toRewind.backup(), rest: meta }
+    toRewind = toRewind.rest
   }
 
-  let k1 = meta.stack
+  const k1 = meta.stack // TODO instead copy meta here, like elsewhere?
   meta.stack = null
-  return body(resumeComp, meta, k1)
+  return () => c(meta, k1)
 }
 
 function RUN_TOPLEVEL(comp) {
@@ -210,4 +211,4 @@ $effekt.run = RUN
  *
  * If a runtime is available, use `$effekt.run`, instead.
  */
-$effekt.runToplevel = RUN_TOPLEVEL
+$effekt.runToplevel = RUN_TOPLEVEL