Effects based on Stack Switching proposal

vouillon · vouillon · commit 0e9cbd00a2f7 · 2025-02-13T17:43:06.000+01:00
diff --git a/compiler/lib-wasm/binaryen.ml b/compiler/lib-wasm/binaryen.ml
@@ -37,6 +37,7 @@ let common_options () =
     ; "--enable-nontrapping-float-to-int"
     ; "--enable-strings"
     ; "--enable-multimemory" (* To keep wasm-merge happy *)
+    ; "--enable-stack-switching"
     ]
   in
   if Config.Flag.pretty () then "-g" :: l else l
diff --git a/compiler/tests-jsoo/lib-effects/dune b/compiler/tests-jsoo/lib-effects/dune
@@ -1,10 +1,10 @@
 (env
  (with-effects-double-translation)
  (with-effects)
- (wasi
-  (wasm_of_ocaml
-   (flags
-    (:standard --enable effects))))
+; (wasi
+;  (wasm_of_ocaml
+;   (flags
+;    (:standard --enable effects))))
  (_
   (js_of_ocaml
    (flags
diff --git a/compiler/tests-ocaml/effect-syntax/dune b/compiler/tests-ocaml/effect-syntax/dune
@@ -1,12 +1,12 @@
 (env
  (with-effects-double-translation)
  (with-effects)
- (wasi
-  (flags
-   (:standard -w -38))
-  (wasm_of_ocaml
-   (flags
-    (:standard --enable=effects))))
+; (wasi
+;  (flags
+;   (:standard -w -38))
+;  (wasm_of_ocaml
+;   (flags
+;    (:standard --enable=effects))))
  (_
   (js_of_ocaml
    (flags
diff --git a/compiler/tests-ocaml/effects/dune b/compiler/tests-ocaml/effects/dune
@@ -1,12 +1,12 @@
 (env
  (with-effects-double-translation)
  (with-effects)
- (wasi
-  (flags
-   (:standard -w -38))
-  (wasm_of_ocaml
-   (flags
-    (:standard --enable=effects))))
+; (wasi
+;  (flags
+;   (:standard -w -38))
+;  (wasm_of_ocaml
+;   (flags
+;    (:standard --enable=effects))))
  (_
   (js_of_ocaml
    (flags
diff --git a/runtime/wasm/effect-native.wat b/runtime/wasm/effect-native.wat
@@ -0,0 +1,239 @@
+(module
+(@if wasi
+(@then
+   (import "fail" "caml_raise_constant"
+      (func $caml_raise_constant (param (ref eq))))
+   (import "fail" "caml_raise_with_arg"
+      (func $caml_raise_with_arg (param $tag (ref eq)) (param $arg (ref eq))))
+   (import "obj" "caml_fresh_oo_id"
+     (func $caml_fresh_oo_id (param (ref eq)) (result (ref eq))))
+   (import "obj" "cont_tag" (global $cont_tag i32))
+   (import "stdlib" "caml_named_value"
+      (func $caml_named_value (param (ref eq)) (result (ref null eq))))
+   (import "fail" "ocaml_exception" (tag $ocaml_exception (param (ref eq))))
+   (import "fail" "javascript_exception"
+      (tag $javascript_exception (param externref)))
+   (import "jslib" "caml_wrap_exception"
+      (func $caml_wrap_exception (param externref) (result (ref eq))))
+   (import "bindings" "start_fiber" (func $start_fiber (param (ref eq))))
+   (import "bindings" "suspend_fiber"
+      (func $suspend_fiber
+         (param externref) (param $f funcref) (param $env eqref)
+         (result eqref)))
+   (import "bindings" "resume_fiber"
+      (func $resume_fiber (param externref) (param (ref eq))))
+   (import "stdlib" "caml_main_wrapper"
+      (global $caml_main_wrapper (mut (ref null $wrapper_func))))
+
+   (type $block (array (mut (ref eq))))
+   (type $string (array (mut i8)))
+   (type $function_1 (func (param (ref eq) (ref eq)) (result (ref eq))))
+   (type $closure (sub (struct (;(field i32);) (field (ref $function_1)))))
+   (type $function_3
+      (func (param (ref eq) (ref eq) (ref eq) (ref eq)) (result (ref eq))))
+   (type $closure_3
+      (sub $closure
+         (struct (field (ref $function_1)) (field (ref $function_3)))))
+
+   ;; Effect types
+
+   (tag $effect (param (ref eq)) (result (ref eq) (ref eq)))
+
+   (type $cont_function (func (param (ref eq) (ref eq)) (result (ref eq))))
+
+   (type $cont (cont $cont_function))
+
+   (type $handlers
+      (struct
+         (field $value (ref eq))
+         (field $exn (ref eq))
+         (field $effect (ref eq))))
+
+   (type $fiber
+      (struct
+         (field $handlers (mut (ref $handlers)))
+         (field $cont (ref $cont))))
+
+   (@string $effect_unhandled "Effect.Unhandled")
+
+   (func $raise_unhandled
+      (param $eff (ref eq)) (param (ref eq)) (result (ref eq))
+      (block $null
+         (call $caml_raise_with_arg
+            (br_on_null $null
+               (call $caml_named_value (global.get $effect_unhandled)))
+            (local.get $eff)))
+      (call $caml_raise_constant
+         (array.new_fixed $block 3 (ref.i31 (i32.const 248))
+            (global.get $effect_unhandled)
+            (call $caml_fresh_oo_id (ref.i31 (i32.const 0)))))
+      (ref.i31 (i32.const 0)))
+
+   (global $raise_unhandled (ref $closure)
+      (struct.new $closure (ref.func $raise_unhandled)))
+
+   (type $func (func (result (ref eq))))
+   (type $wrapper_func (func (param (ref $func))))
+   (type $func_closure (struct (field (ref $func))))
+
+   (func $wrapper_cont
+      (param $f (ref eq)) (param (ref eq)) (result (ref eq))
+      (return_call_ref $func
+         (local.get $f)
+         (struct.get $func_closure 0
+            (ref.cast (ref $func_closure) (local.get $f)))))
+
+   (func $unhandled_effect_wrapper (param $start (ref $func))
+      (local $cont (ref $cont))
+      (local $f (ref eq)) (local $v (ref eq))
+      (local $resume_res (tuple (ref eq) (ref $cont)))
+      (local.set $cont (cont.new $cont (ref.func $wrapper_cont)))
+      (local.set $f (struct.new $func_closure (local.get $start)))
+      (local.set $v (ref.i31 (i32.const 0)))
+      (loop $loop
+         (local.set $resume_res
+            (block $handle_effect (result (ref eq) (ref $cont))
+               (resume $cont (on $effect $handle_effect)
+                  (local.get $f) (local.get $v) (local.get $cont))
+               (return)))
+         (local.set $cont (tuple.extract 2 1 (local.get $resume_res)))
+         (local.set $v (tuple.extract 2 0 (local.get $resume_res)))
+         (local.set $f (global.get $raise_unhandled))
+         (br $loop)))
+
+   (func $init
+      (global.set $caml_main_wrapper (ref.func $unhandled_effect_wrapper)))
+
+   (start $init)
+
+   ;; Resume
+
+   (@string $already_resumed "Effect.Continuation_already_resumed")
+
+   (func $resume (export "%resume")
+      (param $vfiber (ref eq)) (param $f (ref eq)) (param $v (ref eq))
+      (param $tail (ref eq)) (result (ref eq))
+      (local $fiber (ref $fiber))
+      (local $res (ref eq))
+      (local $exn (ref eq))
+      (local $resume_res (tuple (ref eq) (ref $cont)))
+      (if (ref.eq (local.get $vfiber) (ref.i31 (i32.const 0)))
+         (then
+            (call $caml_raise_constant
+               (ref.as_non_null
+                  (call $caml_named_value (global.get $already_resumed))))))
+      (local.set $fiber (ref.cast (ref $fiber) (local.get $vfiber)))
+      (local.set $exn
+         (block $handle_exception (result (ref eq))
+            (local.set $resume_res
+               (block $handle_effect (result (ref eq) (ref $cont))
+                  (local.set $res
+                     (try (result (ref eq))
+                        (do
+                           (resume $cont
+                               (on $effect $handle_effect)
+                               (local.get $f) (local.get $v)
+                               (struct.get $fiber $cont (local.get $fiber))))
+                        (catch $javascript_exception
+                           (br $handle_exception
+                              (call $caml_wrap_exception (pop externref))))
+                        (catch $ocaml_exception
+                           (br $handle_exception (pop (ref eq))))))
+                  ;; handle return
+                  (return_call_ref $function_1 (local.get $res)
+                     (local.tee $f
+                        (struct.get $handlers $value
+                           (struct.get $fiber $handlers (local.get $fiber))))
+                     (struct.get $closure 0
+                        (ref.cast (ref $closure) (local.get $f))))))
+            ;; handle effect
+            (return_call_ref $function_3
+               (tuple.extract 2 0 (local.get $resume_res))
+               (array.new_fixed $block 3 (ref.i31 (global.get $cont_tag))
+                  (struct.new $fiber
+                     (struct.get $fiber $handlers (local.get $fiber))
+                     (tuple.extract 2 1 (local.get $resume_res)))
+                  (ref.i31 (i32.const 0)))
+               (ref.i31 (i32.const 0)) ;; unused
+               (local.tee $f
+                  (struct.get $handlers $effect
+                     (struct.get $fiber $handlers (local.get $fiber))))
+               (struct.get $closure_3 1
+                  (ref.cast (ref $closure_3) (local.get $f))))))
+      ;; handle exception
+      (return_call_ref $function_1 (local.get $exn)
+         (local.tee $f
+            (struct.get $handlers $exn
+               (struct.get $fiber $handlers (local.get $fiber))))
+         (struct.get $closure 0 (ref.cast (ref $closure) (local.get $f)))))
+
+   ;; Perform
+
+   (func (export "%reperform")
+      (param $eff (ref eq)) (param $cont (ref eq)) (param $tail (ref eq))
+      (result (ref eq))
+      (local $res (tuple (ref eq) (ref eq)))
+      (local.set $res (suspend $effect (local.get $eff)))
+      (return_call $resume
+         (ref.as_non_null
+            (array.get $block
+               (ref.cast (ref $block) (local.get $cont))
+               (i32.const 1)))
+         (tuple.extract 2 0 (local.get $res))
+         (tuple.extract 2 1 (local.get $res))
+         (local.get $tail)))
+
+   (func (export "%perform") (param $eff (ref eq)) (result (ref eq))
+      (local $res (tuple (ref eq) (ref eq)))
+      (local.set $res (suspend $effect (local.get $eff)))
+      (return_call_ref $function_1 (tuple.extract 2 1 (local.get $res))
+         (tuple.extract 2 0 (local.get $res))
+         (struct.get $closure 0
+            (ref.cast (ref $closure) (tuple.extract 2 0 (local.get $res))))))
+
+   ;; Allocate a stack
+
+   (func $initial_cont
+      (param $f (ref eq)) (param $x (ref eq)) (result (ref eq))
+      (return_call_ref $function_1 (local.get $x)
+         (local.get $f)
+         (struct.get $closure 0 (ref.cast (ref $closure) (local.get $f)))))
+
+   (func (export "caml_alloc_stack")
+      (param $hv (ref eq)) (param $hx (ref eq)) (param $hf (ref eq))
+      (result (ref eq))
+      (struct.new $fiber
+         (struct.new $handlers (local.get $hv) (local.get $hx) (local.get $hf))
+         (cont.new $cont (ref.func $initial_cont))))
+
+   ;; Other functions
+
+   (func $caml_continuation_use_noexc (export "caml_continuation_use_noexc")
+      (param (ref eq)) (result (ref eq))
+      (local $cont (ref $block))
+      (local $stack (ref eq))
+      (drop (block $used (result (ref eq))
+         (local.set $cont (ref.cast (ref $block) (local.get 0)))
+         (local.set $stack
+            (br_on_cast_fail $used (ref eq) (ref $fiber)
+               (array.get $block (local.get $cont) (i32.const 1))))
+         (array.set $block (local.get $cont) (i32.const 1)
+            (ref.i31 (i32.const 0)))
+         (return (local.get $stack))))
+      (ref.i31 (i32.const 0)))
+
+   (func (export "caml_continuation_use_and_update_handler_noexc")
+      (param $cont (ref eq)) (param $hval (ref eq)) (param $hexn (ref eq))
+      (param $heff (ref eq)) (result (ref eq))
+      (local $stack (ref eq))
+      (local.set $stack (call $caml_continuation_use_noexc (local.get $cont)))
+      (drop (block $used (result (ref eq))
+         (struct.set $fiber $handlers
+            (br_on_cast_fail $used (ref eq) (ref $fiber)
+               (local.get $stack))
+            (struct.new $handlers
+               (local.get $hval) (local.get $hexn) (local.get $heff)))
+         (ref.i31 (i32.const 0))))
+      (local.get $stack))
+))
+)
diff --git a/runtime/wasm/effect.wat b/runtime/wasm/effect.wat
@@ -211,6 +211,12 @@
             (br $loop)))
       (local.get $k))
 
+   (@string $already_resumed "Effect.Continuation_already_resumed")
+
+   (global $effect_allowed (mut i32) (i32.const 1))
+
+(@if (not wasi)
+(@then
    ;; Resume
 
    (func $do_resume (param $k (ref $cont)) (param $vp (ref eq))
@@ -393,6 +399,7 @@
                (local.get $hval) (local.get $hexn) (local.get $heff)))
          (ref.i31 (i32.const 0))))
       (local.get $stack))
+))
 
    (func (export "caml_get_continuation_callstack")
       (param (ref eq) (ref eq)) (result (ref eq))
diff --git a/runtime/wasm/stdlib.wat b/runtime/wasm/stdlib.wat
@@ -215,6 +215,11 @@
       (call $caml_main (ref.func $reraise_exception)))
 ))
 
+   (type $wrapper_func (func (param (ref $func))))
+   (global $caml_main_wrapper (export "caml_main_wrapper")
+      (mut (ref null $wrapper_func))
+      (ref.null $wrapper_func))
+
    (func $caml_main (export "caml_main") (param $start (ref func))
       (local $exn (ref eq))
       (local $msg (ref eq))
@@ -227,6 +232,11 @@
 ))
       (try
          (do
+            (block $fallback
+               (call_ref $wrapper_func
+                  (ref.cast (ref $func) (local.get $start))
+                  (br_on_null $fallback (global.get $caml_main_wrapper)))
+               (return))
             (drop (call_ref $func (ref.cast (ref $func) (local.get $start)))))
          (catch $ocaml_exit
             (call $exit (pop i32)))

Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,7 @@ let common_options () =`
`37`	`37`	`; "--enable-nontrapping-float-to-int"`
`38`	`38`	`; "--enable-strings"`
`39`	`39`	`; "--enable-multimemory" (* To keep wasm-merge happy *)`
	`40`	`+ ; "--enable-stack-switching"`
`40`	`41`	`]`
`41`	`42`	`in`
`42`	`43`	`if Config.Flag.pretty () then "-g" :: l else l`