Make resuming a continuation more efficient (#1765)

* Compiler: Optimize stack resuming in the js runtime, fix 1658

Author: hhugo

.github/workflows/build.yml

@@ -57,6 +57,11 @@ jobs:
             skip-effects: false
             skip-test: false
             skip-doc: false
+          - os: ubuntu-latest
+            ocaml-compiler: "5.1"
+            skip-effects: false
+            skip-test: false
+            skip-doc: true
             # Note this OCaml compiler is bytecode only
           - os: ubuntu-latest
             ocaml-compiler: "ocaml-variants.5.2.0+options,ocaml-option-32bit"

CHANGES.md

@@ -6,7 +6,10 @@
 * Merged Wasm_of_ocaml (#1724)
 * Lib: removed no longer relevant Js.optdef type annotations (#1769)
 * Misc: drop support for IE
-* Effects: add an optional feature of "dynamic switching" between CPS and direct style, resulting in better performance when no effect handler is installed
+* Effects: add an optional feature of "dynamic switching" between CPS
+  and direct style, resulting in better performance when
+  no effect handler is installed
+* Compiler/Runtime: Make resuming a continuation more efficient in js (#1765)
 
 ## Bug fixes
 * Fix small bug in global data flow analysis (#1768)

compiler/lib/effects.ml

@@ -613,11 +613,34 @@ let cps_block ~st ~k ~orig_pc block =
   in
 
   let rewrite_last_instr (x : Var.t) (e : expr) : (k:Var.t -> instr list * last) option =
-    let perform_effect ~effect_ ~continuation =
+    let perform_effect ~effect_ continuation_and_tail =
       Some
         (fun ~k ->
           let e =
-            Prim (Extern "caml_perform_effect", [ Pv effect_; continuation; Pv k ])
+            match Config.target () with
+            | `JavaScript -> (
+                match continuation_and_tail with
+                | None -> Prim (Extern "caml_perform_effect", [ Pv effect_; Pv k ])
+                | Some (continuation, tail) ->
+                    Prim
+                      ( Extern "caml_reperform_effect"
+                      , [ Pv effect_; continuation; tail; Pv k ] ))
+            | `Wasm -> (
+                (* temporary until we finish the change to the wasmoo
+                   runtime *)
+                match continuation_and_tail with
+                | None ->
+                    Prim
+                      ( Extern "caml_perform_effect"
+                      , [ Pv effect_
+                        ; Pc (Int Targetint.zero)
+                        ; Pc (Int Targetint.zero)
+                        ; Pv k
+                        ] )
+                | Some (continuation, tail) ->
+                    Prim
+                      ( Extern "caml_perform_effect"
+                      , [ Pv effect_; continuation; tail; Pv k ] ))
           in
           let x = Var.fresh () in
           [ Let (x, e) ], Return x)
@@ -628,22 +651,22 @@ let cps_block ~st ~k ~orig_pc block =
           (fun ~k ->
             let exact = exact || call_exact st.flow_info f (List.length args) in
             tail_call ~st ~exact ~in_cps:true ~check:true ~f (args @ [ k ]))
-    | Prim (Extern "%resume", [ Pv stack; Pv f; Pv arg ]) ->
+    | Prim (Extern "%resume", [ Pv stack; Pv f; Pv arg; tail ]) ->
         Some
           (fun ~k ->
             let k' = Var.fresh_n "cont" in
             tail_call
               ~st
-              ~instrs:[ Let (k', Prim (Extern "caml_resume_stack", [ Pv stack; Pv k ])) ]
+              ~instrs:
+                [ Let (k', Prim (Extern "caml_resume_stack", [ Pv stack; tail; Pv k ])) ]
               ~exact:(call_exact st.flow_info f 1)
               ~in_cps:true
               ~check:true
               ~f
               [ arg; k' ])
-    | Prim (Extern "%perform", [ Pv effect_ ]) ->
-        perform_effect ~effect_ ~continuation:(Pc (Int Targetint.zero))
-    | Prim (Extern "%reperform", [ Pv effect_; continuation ]) ->
-        perform_effect ~effect_ ~continuation
+    | Prim (Extern "%perform", [ Pv effect_ ]) -> perform_effect ~effect_ None
+    | Prim (Extern "%reperform", [ Pv effect_; continuation; tail ]) ->
+        perform_effect ~effect_ (Some (continuation, tail))
     | _ -> None
   in
 
@@ -712,14 +735,14 @@ let rewrite_direct_block ~st ~cps_needed ~closure_info ~pc block =
           ; Let (cps_c, Closure (cps_params, cps_cont))
           ; Let (x, Prim (Extern "caml_cps_closure", [ Pv direct_c; Pv cps_c ]))
           ]
-      | Let (x, Prim (Extern "%resume", [ Pv stack; Pv f; Pv arg ])) ->
-          [ Let (x, Prim (Extern "caml_resume", [ Pv f; Pv arg; Pv stack ])) ]
-      | Let (x, Prim (Extern "%perform", [ Pv effect_ ])) ->
+      | Let (x, Prim (Extern "%resume", [ stack; f; arg; tail ])) ->
+          [ Let (x, Prim (Extern "caml_resume", [ f; arg; stack; tail ])) ]
+      | Let (x, Prim (Extern "%perform", [ effect_ ])) ->
           (* In direct-style code, we just raise [Effect.Unhandled]. *)
-          [ Let (x, Prim (Extern "caml_raise_unhandled", [ Pv effect_ ])) ]
-      | Let (x, Prim (Extern "%reperform", [ Pv effect_; Pv _continuation ])) ->
+          [ Let (x, Prim (Extern "caml_raise_unhandled", [ effect_ ])) ]
+      | Let (x, Prim (Extern "%reperform", [ effect_; _continuation; _tail ])) ->
           (* Similar to previous case *)
-          [ Let (x, Prim (Extern "caml_raise_unhandled", [ Pv effect_ ])) ]
+          [ Let (x, Prim (Extern "caml_raise_unhandled", [ effect_ ])) ]
       | Let (x, Prim (Extern "caml_assume_no_perform", [ Pv f ])) ->
           (* We just need to call [f] in direct style. *)
           let unit = Var.fresh_n "unit" in

compiler/lib/parse_bytecode.ml

@@ -2368,7 +2368,6 @@ and compile infos pc state (instrs : instr list) =
         let func = State.peek 0 state in
         let arg = State.peek 1 state in
         let x, state = State.fresh_var state in
-
         if debug_parser ()
         then
           Format.printf
@@ -2381,23 +2380,30 @@ and compile infos pc state (instrs : instr list) =
             func
             Var.print
             arg;
-        let state =
+        let state, tail =
           match Ocaml_version.compare Ocaml_version.current [ 5; 2 ] < 0 with
-          | true -> State.pop 2 state
-          | false -> State.pop 3 state
+          | true -> State.pop 2 state, Pc (Int (Targetint.of_int_exn 0))
+          | false ->
+              let tail = State.peek 2 state in
+              State.pop 3 state, Pv tail
         in
-
         compile
           infos
           (pc + 1)
           state
-          (Let (x, Prim (Extern "%resume", [ Pv stack; Pv func; Pv arg ])) :: instrs)
+          (Let (x, Prim (Extern "%resume", [ Pv stack; Pv func; Pv arg; tail ])) :: instrs)
     | RESUMETERM ->
         let stack = State.accu state in
         let func = State.peek 0 state in
         let arg = State.peek 1 state in
         let x, state = State.fresh_var state in
-
+        let tail =
+          match Ocaml_version.compare Ocaml_version.current [ 5; 2 ] < 0 with
+          | true -> Pc (Int (Targetint.of_int_exn 0))
+          | false ->
+              let tail = State.peek 2 state in
+              Pv tail
+        in
         if debug_parser ()
         then
           Format.printf
@@ -2408,7 +2414,7 @@ and compile infos pc state (instrs : instr list) =
             func
             Var.print
             arg;
-        ( Let (x, Prim (Extern "%resume", [ Pv stack; Pv func; Pv arg ])) :: instrs
+        ( Let (x, Prim (Extern "%resume", [ Pv stack; Pv func; Pv arg; tail ])) :: instrs
         , Return x
         , state )
     | PERFORM ->
@@ -2425,13 +2431,13 @@ and compile infos pc state (instrs : instr list) =
     | REPERFORMTERM ->
         let eff = State.accu state in
         let stack = State.peek 0 state in
-        (* We don't need [State.peek 1 state] *)
+        let tail = State.peek 1 state in
         let state = State.pop 2 state in
         let x, state = State.fresh_var state in
 
         if debug_parser ()
         then Format.printf "return reperform(%a, %a)@." Var.print eff Var.print stack;
-        ( Let (x, Prim (Extern "%reperform", [ Pv eff; Pv stack ])) :: instrs
+        ( Let (x, Prim (Extern "%reperform", [ Pv eff; Pv stack; Pv tail ])) :: instrs
         , Return x
         , state )
     | EVENT | BREAK | FIRST_UNIMPLEMENTED_OP -> assert false)

compiler/tests-check-prim/main.4.14.output

@@ -35,9 +35,6 @@ caml_bigstring_blit_string_to_ba
 caml_bigstring_memcmp
 caml_hash_mix_bigstring
 
-From +effect.js:
-jsoo_effect_not_supported
-
 From +fs.js:
 caml_ba_map_file
 caml_ba_map_file_bytecode

compiler/tests-check-prim/unix-Unix.4.14.output

@@ -144,9 +144,6 @@ caml_bigstring_blit_string_to_ba
 caml_bigstring_memcmp
 caml_hash_mix_bigstring
 
-From +effect.js:
-jsoo_effect_not_supported
-
 From +fs.js:
 caml_ba_map_file
 caml_ba_map_file_bytecode

compiler/tests-check-prim/unix-Win32.4.14.output

@@ -109,9 +109,6 @@ caml_bigstring_blit_string_to_ba
 caml_bigstring_memcmp
 caml_hash_mix_bigstring
 
-From +effect.js:
-jsoo_effect_not_supported
-
 From +fs.js:
 caml_ba_map_file
 caml_ba_map_file_bytecode

compiler/tests-full/stdlib.cma.expected.js

@@ -36136,8 +36136,8 @@
       /*<<effect.ml:28:11>>*/ return 0;
     var
      x =  /*<<effect.ml:22:16>>*/ param[2],
-     _i_ =  /*<<effect.ml:25:12>>*/ caml_call1(Stdlib_Printexc[26], x),
-     msg =  /*<<effect.ml:24:18>>*/ caml_call2(Stdlib_Printf[4], _a_, _i_);
+     _o_ =  /*<<effect.ml:25:12>>*/ caml_call1(Stdlib_Printexc[26], x),
+     msg =  /*<<effect.ml:24:18>>*/ caml_call2(Stdlib_Printf[4], _a_, _o_);
      /*<<effect.ml:27:8>>*/ return [0, msg];
     /*<<effect.ml:28:15>>*/ }
     /*<<effect.ml:30:2>>*/ caml_call1(Stdlib_Printexc[9], printer);
@@ -36155,22 +36155,28 @@
      "Effect.Continuation_already_resumed",
      Continuation_already_resumed);
    function continue$0(k, v){
-    var _h_ =  /*<<effect.ml:62:11>>*/ caml_continuation_use_noexc(k);
-    function _g_(x){
+    var
+     _l_ =  /*<<effect.ml:62:11>>*/ k[2],
+     _n_ = caml_continuation_use_noexc(k);
+    function _m_(x){
       /*<<effect.ml:62:41>>*/ return x;
      /*<<effect.ml:62:42>>*/ }
      /*<<effect.ml:62:30>>*/ return jsoo_effect_not_supported() /*<<effect.ml:62:65>>*/ ;
    }
    function discontinue(k, e){
-    var _f_ =  /*<<effect.ml:65:11>>*/ caml_continuation_use_noexc(k);
-    function _e_(e){
+    var
+     _i_ =  /*<<effect.ml:65:11>>*/ k[2],
+     _k_ = caml_continuation_use_noexc(k);
+    function _j_(e){
       /*<<effect.ml:65:41>>*/ throw caml_maybe_attach_backtrace(e, 1);
      /*<<effect.ml:65:48>>*/ }
      /*<<effect.ml:65:30>>*/ return jsoo_effect_not_supported() /*<<effect.ml:65:71>>*/ ;
    }
    function discontinue_with_backtrace(k, e, bt){
-    var _d_ =  /*<<effect.ml:68:11>>*/ caml_continuation_use_noexc(k);
-    function _c_(e){
+    var
+     _f_ =  /*<<effect.ml:68:11>>*/ k[2],
+     _h_ = caml_continuation_use_noexc(k);
+    function _g_(e){
       /*<<effect.ml:68:41>>*/ caml_restore_raw_backtrace(e, bt);
      throw caml_maybe_attach_backtrace(e, 0);
      /*<<effect.ml:68:75>>*/ }
@@ -36187,8 +36193,9 @@
     }
     var
      s =
-        /*<<effect.ml:87:12>>*/ caml_alloc_stack(handler[1], handler[2], effc);
-     /*<<effect.ml:88:4>>*/ return jsoo_effect_not_supported() /*<<effect.ml:88:23>>*/ ;
+        /*<<effect.ml:87:12>>*/ caml_alloc_stack(handler[1], handler[2], effc),
+     _e_ =  /*<<effect.ml:88:4>>*/ 0;
+    return jsoo_effect_not_supported() /*<<effect.ml:88:23>>*/ ;
    }
    function try_with(comp, arg, handler){
     function effc(eff, k, last_fiber){
@@ -36208,8 +36215,9 @@
          function(e){
            /*<<effect.ml:101:47>>*/ throw caml_maybe_attach_backtrace(e, 1);
           /*<<effect.ml:101:54>>*/ },
-         effc);
-     /*<<effect.ml:102:4>>*/ return jsoo_effect_not_supported() /*<<effect.ml:102:23>>*/ ;
+         effc),
+     _d_ =  /*<<effect.ml:102:4>>*/ 0;
+    return jsoo_effect_not_supported() /*<<effect.ml:102:23>>*/ ;
    }
    var
     Deep =
@@ -36243,7 +36251,7 @@
     var s =  /*<<effect.ml:135:12>>*/ caml_alloc_stack(error, error, effc);
      /*<<effect.ml:136:4>>*/ try{
       /*<<effect.ml:136:10>>*/ jsoo_effect_not_supported();
-     var _b_ =  /*<<effect.ml:136:26>>*/ 0;
+     var _b_ =  /*<<effect.ml:136:26>>*/ 0, _c_ = 0;
     }
     catch(exn$0){
      var exn =  /*<<?>>*/ caml_wrap_exception(exn$0);
@@ -36263,6 +36271,7 @@
       /*<<effect.ml:160:10>>*/ return caml_call1(f, k) /*<<effect.ml:161:42>>*/ ;
     }
     var
+     last_fiber =  /*<<effect.ml:163:4>>*/ k[2],
      stack =
         /*<<effect.ml:164:16>>*/ runtime.caml_continuation_use_and_update_handler_noexc
         (k, handler[1], handler[2], effc);

compiler/tests-jsoo/lib-effects/deep_state.ml

@@ -0,0 +1,73 @@
+(* deep_state.ml *)
+
+open Effect
+open Effect.Shallow
+
+module type State = sig
+  type a
+
+  type _ Effect.t += Get : a Effect.t
+
+  type _ Effect.t += Set : a -> unit Effect.t
+end
+
+module Make (S : State) = struct
+  let rec loop : type x y. S.a -> (x, y) continuation -> x -> y =
+   fun s k x ->
+    continue_with
+      k
+      x
+      { retc = (fun y -> y)
+      ; exnc = raise
+      ; effc =
+          (fun (type b) (e : b Effect.t) ->
+            match e with
+            | S.Get -> Some (fun (k : (b, _) continuation) -> loop s k s)
+            | S.Set s -> Some (fun (k : (b, _) continuation) -> loop s k ())
+            | _ -> None)
+      }
+
+  let handle (s : S.a) (f : unit -> 'a) : 'a = loop s (fiber f) ()
+
+  let get () = perform S.Get
+
+  let set v = perform (S.Set v)
+end
+
+module IntState = struct
+  type a = int
+
+  type _ Effect.t += Get : int Effect.t
+
+  type _ Effect.t += Set : int -> unit Effect.t
+end
+
+module StringState = struct
+  type a = string
+
+  type _ Effect.t += Get : string Effect.t
+
+  type _ Effect.t += Set : string -> unit Effect.t
+end
+
+let main () =
+  let depth = int_of_string Sys.argv.(1) in
+  let ops = int_of_string Sys.argv.(2) in
+  Printf.printf "Running deepstate: depth=%d ops=%d\n" depth ops;
+  let module SS = Make (StringState) in
+  let rec setup_deep_state n () =
+    if n = 0
+    then
+      for _ = 1 to ops do
+        (*         SS.set (SS.get () ^ "_" ^ (string_of_int i)) *)
+        SS.set (SS.get ())
+      done
+    (*       print_endline @@ SS.get() *)
+    else
+      let module IS = Make (IntState) in
+      IS.handle 0 @@ setup_deep_state (n - 1)
+  in
+
+  SS.handle "Hello, world!" @@ setup_deep_state depth
+
+let _ = main ()

compiler/tests-jsoo/lib-effects/dune

@@ -15,7 +15,8 @@
    \
    assume_no_perform
    assume_no_perform_unhandled
-   assume_no_perform_nested_handler))
+   assume_no_perform_nested_handler
+   deep_state))
  (preprocess
   (pps ppx_expect)))
 
@@ -37,3 +38,10 @@
     0
     (run node %{test}))))
  (modes js wasm))
+
+(executable
+ (name deep_state)
+ (enabled_if
+  (>= %{ocaml_version} 5))
+ (modules deep_state)
+ (modes js))