Skip to content

Instantly share code, notes, and snippets.

@laughinghan
Last active September 9, 2022 23:47
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save laughinghan/350ded589a2bd23a3b1f087f0960630b to your computer and use it in GitHub Desktop.
Save laughinghan/350ded589a2bd23a3b1f087f0960630b to your computer and use it in GitHub Desktop.

Simpler Stack Switching for Wasm

To me, #1360 "First-Class Stacks" seems like it started out great, in particular primitive stack-switching seems like a great low-level building block for everything you'd want to do including algebraic effects (#1359); but then it seems to me like it ended up both overly complex and yet also underspecified. In particular, stack extension and redirection seem like they should be extensions to stack-switching added later; whereas creating stacks seems like it would be immediately necessary for priority use cases like the Go and Erlang runtimes, and it seems like it would be a simple win to be able to send normal, non-exception values and thus remove the dependency on the exception-handling work.

Just like the TinyGo author suggested, it seems to me like it would be simplest to just have 3 instructions for creating, using, and destroying stack continuations, respectively:

  • (for brevity, (cont X … -> Y …) will be used as a shorthand for (cont (param X …) (result Y …)))
  • (cont.new (cont A … -> B …) (cont B … -> C …) $fn x …) will:
    • immediately, call $fn with a continuation of type A … -> B … (call it $cont) and the values x …
    • when that continuation $cont is resumed with some values a … of type A …, execution will return to the site of the (cont.new …) expression, proceeding as if the (cont.new …) expression evaluated to a continuation of type B … -> C … and the values a …
  • $fn can then do (cont.switch (cont B … -> C …) (cont C … -> D …) $cont a …), which will:
    • immediately, suspend execution of $fn and resume $cont with the values a …; execution returns to (cont.new …) as if it evaluated to a continuation B … -> C … (call it $cont2)and the valuesa …`
    • when someone resumes $cont2 with some values b … of type B …, by doing (cont.switch (cont C … -> D …) (cont D … -> E …) $cont2 b …), execution will return to $fn, proceeding as if the (cont.switch …) expression evaluated to a continuation of type C … -> D … and the values b …
  • (cont.switch (cont ...) (cont ...) $cont p1 p2 ...): given a continuation, a continuation type signature, and N parameters, resumes that continuation with N+1 parameters where the first parameter is a continuation with the specified type. There is no relationship between $cont and that (cont ...) type signature, rather, the current stack frame is suspended and a continuation for it is created and sent as that first parameter, and its type must match that type signature.
  • (cont.destroy $cont) will deallocate a continuation. Destroying the root continuation (the stack rooted in a call from JS to Wasm; there will always be exactly one at any given time) will trap immediately.

Continuations are one-shot, once used they can't be used again. Returning to a continuation uses it up same as switching to it. Continuations downward of a (cont.new …) are delimited. The only non-delimited continuation will be the root continuation. When a delimited continuation runs to completion instead of switching to another, it returns to the last continuation that switched to it (not the continuation that created it with (cont.new …), which will often have been used already).

That means that (cont.new …) can only be called on functions whose return types match the parameter types of the last continuation received in the function. This isn't as complicated as it sounds, see "Return Type Analysis" below for details, and anyway if we want to simplify things further we can require that only functions whose return type is unreachable can be passed to (cont.new …) (instead of returning, they'd have to cont.switch back).

I think this sidesteps the "grounding frame" problem that #1360 left up to the host, while keeping stack-switching very fast and primitive, just a matter of swapping the instruction pointer, stack pointer, and stack segment pointer, should be similar overhead to a function call?

Usage Examples

Generators

(func $generator (param $cont (cont (param i32))) (yields unreachable)
  (local $n i32)
  (local.set $n (i32.const 0))
  (loop $forever
    (local.set $cont (cont.switch (local.get $cont) (cont (result i32)) (local.get $n)))
    (local.set $n (i32.add (local.get $n) (i32.const 1)))
    (br $forever)
  )
)
(func $main
  (local $gen_cont (cont (result i32)))

  (cont.new $generator (cont (param i32)))
  (local.set $gen_cont) ;; store the generator continuation
  (call $console_log)   ;; logs '0'

  (cont.switch $gen_cont (cont (param i32)))
  (local.set $gen_cont) ;; store the new generator continuation
  (call $console_log)   ;; logs '1'

  (cont.switch $gen_cont (cont (param i32)))
  (local.set $gen_cont) ;; store the new generator continuation
  (call $console_log)   ;; logs '2'
)

setjmp/longjmp

(Based on the setjmp/longjmp example on Wikipedia)

(func $main
  (cont.new $first) ;; setjmp
  (cont.destroy)
  (call $console_log 0)
)
(func $first (param $cont (cont)) (yields)
  (call $second $cont)
  (call $console_log 1)
)
(func $second (param $cont (cont)) (yields)
  (call $console_log 2)
  (cont.switch (local.get $cont)) ;; longjmp
)

Coroutines

;; mutable global coroutine continuations, initially set to the null continuation
(global $producer (mut (cont)) (cont null))
(global $consumer (mut (cont)) (cont null))

(func $produce (param $cont (cont)) (yields unreachable)
  (loop $forever
    (loop $until_full
      (call $enqueue (call $create_item))
      (br_if $until_full (call $queue_full))
    )
    (local.set $cont (cont.switch (local.get $cont))
    (br $forever)
  )
)
(func $consume (param $cont (cont)) (yields unreachable)
  (loop $forever
    (loop $until_empty
      (call $use_item (call $dequeue))
      (br_if $until_empty (call $queue_empty))
    )
    (cont.switch (local.get $producer) (local.)
    (local.set $main)
    (local.set $producer)
    (br $forever)
  )
)
(func $main
  (global.set $producer (cont.new $produce))
  (global.set $consumer (cont.new $consume))
  ($cont.switch (local.get $producer) 
)

Effect Handlers (#1359)

By composing with

Return Type Analysis

Just 4 rules:

  • if a function has (cont.switch …) anywhere in its body, then it must have a (yields …) declaration after its parameter and result type declarations
  • if a function has a call to a function with a (yields …) declaration, then it must have a (yields …) declaration after its parameter and result type declarations
  • in a function's (yields T …) declaration, the types T … must match the last (cont.switch …) or function call with a (yields …) declaration (if there's more than one, they must all match each other)
    • matching a (cont.switch …) means that if it takes the form (cont.switch (cont A … -> B …) (cont B … -> C …) $cont $x …), then T … must match B … since that's what we promised to yield back
  • (cont.new …) can only be called on functions whose yields type matches its return type (or functions with no (yields …) type at all?)

Functions where neither apply should not have a (yields …) declaration.

Note that if this is too complicated we can avoid the whole issue by requiring that (cont.new …) can only be called on functions whose return type is unreachable (instead of ever running to completion, they have to cont.switch back).

Variants

  • Instead of (cont.new …) immediately calling the function with a continuation for the current stack, it could instead just initialize a delimited continuation for the function and do nothing else. It would only take one argument, the function, and all the arguments you currently pass to (cont.new …), you'd pass to the first (cont.switch …) call instead. I like the purity, it's truly orthogonal to cont.switch, but I couldn't think of any use-cases.

Open Questions

  • should (cont.new …) be allowed on functions that have no (yields …) type because they never call (cont.switch …)? I'm pretty sure they're totally useless, even if they stored the continuation they're given, they'll then run to completion and then use up that continuation
  • should destroying a continuation that's already been used-up be a no-op, or trap? I would normally err on the side of strictness, but I can't see much harm in no-opping
  • one of the V8 team presentations mentioned something about not wanting to have to save and restore stuff; if we require that the stack is empty when cont.new/switch are called (except for their arguments, of course), would that ensure that implementations don't need to save/restore registers?
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment