Experiment: Lazy arena allocation on JS

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

@@ -23,6 +23,8 @@ object TransformerCps extends Transformer {
   val THUNK = Variable(JSName("THUNK"))
   val DEALLOC = Variable(JSName("DEALLOC"))
   val TRAMPOLINE = Variable(JSName("TRAMPOLINE"))
+  val VAR = Variable(JSName("VAR"))
+  val REGION = Variable(JSName("REGION"))
 
   class RecursiveUsage(var jumped: Boolean)
   case class RecursiveDefInfo(id: Id, label: Id, vparams: List[Id], bparams: List[Id], ks: Id, k: Id, used: RecursiveUsage)
@@ -397,24 +399,25 @@ object TransformerCps extends Transformer {
       val args = vargs.map(toJS) ++ bargs.map(argumentToJS) ++ List(toJS(ks), toJS(k))
       pure(js.Return(MethodCall(toJS(callee), memberNameRef(method), args:_*)) :: Nil)
 
-    // const r = ks.arena.newRegion(); body
+    // const r = REGION(ks); body
     case cps.Stmt.Region(id, ks, body) =>
       Binding { k =>
-        js.Const(nameDef(id), js.MethodCall(js.Member(toJS(ks), JSName("arena")), JSName("newRegion"))) ::
+        js.Const(nameDef(id), js.Call(REGION, List(toJS(ks)))) ::
           toJS(body).run(k)
       }
 
     // const x = r.alloc(init); body
     case cps.Stmt.Alloc(id, init, region, body) =>
       Binding { k =>
+        // region should be non-null here as it's constructed via `REGION`
         js.Const(nameDef(id), js.MethodCall(nameRef(region), JSName("fresh"), toJS(init))) ::
           toJS(body).run(k)
       }
 
-    // const x = ks.arena.fresh(1); body
+    // const x = VAR(1, ks); body
     case cps.Stmt.Var(id, init, ks, body) =>
       Binding { k =>
-        js.Const(nameDef(id), js.MethodCall(js.Member(toJS(ks), JSName("arena")), JSName("fresh"), toJS(init))) ::
+        js.Const(nameDef(id), js.Call(VAR, List(toJS(init), toJS(ks)))) ::
           toJS(body).run(k)
       }
 

libraries/js/effekt_runtime.js

@@ -82,6 +82,9 @@ function restore(store, snap) {
   store.generation = snap.generation + 1
 }
 
+// Sentinel snapshot used for null arenas — restoring it is always a no-op.
+const NULL_SNAP = { root: { value: Mem }, generation: -1 }
+
 // Common Runtime
 // --------------
 let _prompt = 1;
@@ -94,6 +97,16 @@ function THUNK(f) {
   return f
 }
 
+function VAR(init, ks) {
+  if (ks.arena === null) ks.arena = Arena()
+  return ks.arena.fresh(init)
+}
+
+function REGION(ks) {
+  if (ks.arena === null) ks.arena = Arena()
+  return ks.arena.newRegion()
+}
+
 function CAPTURE(body) {
   return (ks, k) => {
     const res = body(x => TRAMPOLINE(() => k(x, ks)))
@@ -108,7 +121,8 @@ const RETURN = (x, ks) => ks.rest.stack(x, ks.rest)
 function RESET(prog, ks, k) {
   const prompt = _prompt++;
   const rest = { stack: k, prompt: ks.prompt, arena: ks.arena, rest: ks.rest }
-  return prog(prompt, { prompt, arena: Arena([]), rest }, RETURN)
+  // arena: null; only materialised on first VAR in this handler scope
+  return prog(prompt, { stack: null, prompt, arena: null, rest }, RETURN)
 }
 
 function SHIFT(p, body, ks, k) {
@@ -119,14 +133,20 @@ function SHIFT(p, body, ks, k) {
 
   while (!!meta && meta.prompt !== p) {
     let store = meta.arena
-    cont = { stack: meta.stack, prompt: meta.prompt, arena: store, backup: snapshot(store), rest: cont }
+    // If this handler scope never called VAR, its arena is null.
+    // We still capture the frame, but use a no-op sentinel snapshot.
+    cont = { stack: meta.stack, prompt: meta.prompt, arena: store,
+             backup: store !== null ? snapshot(store) : NULL_SNAP,
+             rest: cont }
     meta = meta.rest
   }
   if (!meta) { throw `Prompt not found ${p}` }
 
   // package the prompt itself
   let store = meta.arena
-  cont = { stack: meta.stack, prompt: meta.prompt, arena: store, backup: snapshot(store), rest: cont }
+  cont = { stack: meta.stack, prompt: meta.prompt, arena: store,
+           backup: store !== null ? snapshot(store) : NULL_SNAP,
+           rest: cont }
   meta = meta.rest
 
   const k1 = meta.stack
@@ -139,7 +159,10 @@ function RESUME(cont, c, ks, k) {
   let meta = { stack: k, prompt: ks.prompt, arena: ks.arena, rest: ks.rest }
   let toRewind = cont
   while (!!toRewind) {
-    restore(toRewind.arena, toRewind.backup)
+    // arena is null when this handler scope never called VAR ~> skip restore entirely
+    if (toRewind.arena !== null) {
+      restore(toRewind.arena, toRewind.backup)
+    }
     meta = { stack: toRewind.stack, prompt: toRewind.prompt, arena: toRewind.arena, rest: meta }
     toRewind = toRewind.rest
   }