Merge pull request #18 from Julow/eio-env-as-var

eio: Add '--eio-env-as-fiber-var'

Eio_main.run now also sets the variable if passed. It also creates a switch and initializes the sw variable if specified with --eio-sw-as-fiber-var.

Author: Julow

bin/lwt_to_direct_style/concurrency_backend.ml

@@ -4,16 +4,15 @@ open Parsetree
 open Ast_helper
 open Ocamlformat_utils.Ast_utils
 
-let eio ~eio_sw_as_fiber_var add_comment =
+let eio ~eio_sw_as_fiber_var ~eio_env_as_fiber_var add_comment =
   let used_eio_std = ref false in
-  let fiber_ident i =
+  let eio_std_ident mod_ i =
     used_eio_std := true;
-    [ "Fiber"; i ]
-  in
-  let promise_ident i =
-    used_eio_std := true;
-    [ "Promise"; i ]
+    [ mod_; i ]
   in
+  let fiber_ident = eio_std_ident "Fiber"
+  and promise_ident = eio_std_ident "Promise"
+  and switch_ident = eio_std_ident "Switch" in
   let add_comment fmt = Format.kasprintf add_comment fmt in
   let add_comment_dropped_exp ~label exp =
     add_comment "Dropped expression (%s): [%s]." label
@@ -25,12 +24,29 @@ let eio ~eio_sw_as_fiber_var add_comment =
   let get_current_switch () =
     match eio_sw_as_fiber_var with
     | Some ident ->
-        mk_apply_simple [ "Option"; "get" ]
-          [ mk_apply_simple [ "Fiber"; "get" ] [ Exp.ident (mk_loc ident) ] ]
+        mk_apply_simple
+          [ "Stdlib"; "Option"; "get" ]
+          [ mk_apply_simple (fiber_ident "get") [ Exp.ident (mk_loc ident) ] ]
     | None ->
         add_comment "[sw] (of type Switch.t) must be propagated here.";
         mk_exp_ident [ "sw" ]
   in
+  let get_current_switch_arg () =
+    (Labelled (mk_loc "sw"), get_current_switch ())
+  in
+  let env field =
+    let env_exp =
+      match eio_env_as_fiber_var with
+      | Some ident ->
+          mk_apply_simple
+            [ "Stdlib"; "Option"; "get" ]
+            [ mk_apply_simple (fiber_ident "get") [ Exp.ident (mk_loc ident) ] ]
+      | None ->
+          add_comment "[env] must be propagated from the main loop";
+          mk_exp_ident [ "env" ]
+    in
+    Exp.send env_exp (mk_loc field)
+  in
   object
     method both ~left ~right =
       mk_apply_simple (fiber_ident "pair") [ left; right ]
@@ -40,9 +56,7 @@ let eio ~eio_sw_as_fiber_var add_comment =
     method async process_f =
       Exp.apply
         (mk_exp_ident (fiber_ident "fork"))
-        [
-          (Labelled (mk_loc "sw"), get_current_switch ()); (Nolabel, process_f);
-        ]
+        [ get_current_switch_arg (); (Nolabel, process_f) ]
 
     method wait () =
       add_comment
@@ -73,9 +87,9 @@ let eio ~eio_sw_as_fiber_var add_comment =
     method sleep d = mk_apply_simple [ "Eio_unix"; "sleep" ] [ d ]
 
     method with_timeout d f =
-      add_comment "[env] must be propagated from the main loop";
-      let clock = Exp.send (mk_exp_ident [ "env" ]) (mk_loc "mono_clock") in
-      mk_apply_simple [ "Eio"; "Time"; "with_timeout_exn" ] [ clock; d; f ]
+      mk_apply_simple
+        [ "Eio"; "Time"; "with_timeout_exn" ]
+        [ env "mono_clock"; d; f ]
 
     method timeout_exn = mk_longident [ "Eio"; "Time"; "Timeout" ]
 
@@ -146,7 +160,7 @@ let eio ~eio_sw_as_fiber_var add_comment =
       in
       mk_apply_ident
         [ "Eio_unix"; "Fd"; "of_unix" ]
-        ([ (Labelled (mk_loc "sw"), get_current_switch ()) ]
+        ([ get_current_switch_arg () ]
         @ blocking_arg
         @ [
             (Labelled (mk_loc "close_unix"), mk_constr_exp [ "true" ]);
@@ -163,19 +177,37 @@ let eio ~eio_sw_as_fiber_var add_comment =
     method fd_close fd = mk_apply_simple [ "Eio_unix"; "Fd" ] [ fd ]
 
     method main_run promise =
+      let with_binding var_ident x body =
+        let var = Exp.ident (mk_loc var_ident) in
+        mk_apply_simple (fiber_ident "with_binding") [ var; x; mk_thunk body ]
+      in
       add_comment
         "[Eio_main.run] argument used to be a [Lwt] promise and is now a \
          [fun]. Make sure no asynchronous or IO calls are done outside of this \
          [fun].";
-      (match eio_sw_as_fiber_var with
-      | Some ident ->
-          add_comment
-            "Make sure to create a [Switch.t] and store it in fiber variable \
-             [%a]."
-            Ocamlformat_utils.Parsing.Printast.fmt_longident ident
-      | None -> ());
+      let wrap_sw_fiber_var k =
+        match eio_sw_as_fiber_var with
+        | Some var_ident ->
+            let fun_sw =
+              mk_fun ~arg_name:"sw" (fun sw -> with_binding var_ident sw k)
+            in
+            mk_apply_ident (switch_ident "run")
+              [
+                (Labelled (mk_loc "name"), mk_const_string "main");
+                (Nolabel, fun_sw);
+              ]
+        | None -> k
+      in
+      let wrap_env_fiber_var env k =
+        match eio_env_as_fiber_var with
+        | Some var_ident -> with_binding var_ident env k
+        | None -> k
+      in
       mk_apply_simple [ "Eio_main"; "run" ]
-        [ mk_fun ~arg_name:"env" (fun _env -> promise) ]
+        [
+          mk_fun ~arg_name:"env" (fun env ->
+              wrap_env_fiber_var env (wrap_sw_fiber_var promise));
+        ]
 
     method input_io_of_fd fd =
       Exp.constraint_

bin/lwt_to_direct_style/main.ml

@@ -1,5 +1,7 @@
-let main migrate eio_sw_as_fiber_var =
-  let backend = Concurrency_backend.eio ~eio_sw_as_fiber_var in
+let main migrate eio_sw_as_fiber_var eio_env_as_fiber_var =
+  let backend =
+    Concurrency_backend.eio ~eio_sw_as_fiber_var ~eio_env_as_fiber_var
+  in
   let modify_ast ~fname = Ast_rewrite.rewrite_lwt_uses ~fname ~backend in
   let units = function
     | "Lwt" -> true
@@ -31,6 +33,17 @@ let opt_eio_sw_as_fiber_var =
     & opt (some ident_conv) None
     & info ~doc ~docv:"Fiber.key" [ "eio-sw-as-fiber-var" ])
 
+let opt_eio_env_as_fiber_var =
+  let doc =
+    "Eio only: Pass the environment as a Fiber variable. It will be queried \
+     everytime the environment is needed. Argument must be a fully-qualified \
+     OCaml identifier pointing to a value of type 'Switch.t Fiber.key'."
+  in
+  Arg.(
+    value
+    & opt (some ident_conv) None
+    & info ~doc ~docv:"Fiber.key" [ "eio-env-as-fiber-var" ])
+
 let opt_migrate =
   let doc = "Modify the source code instead of printing occurrences of Lwt." in
   Arg.(value & flag & info ~doc [ "migrate" ])
@@ -41,6 +54,9 @@ let cmd =
   in
   let info = Cmd.info "lwt-to-direct-style" ~version:"%%VERSION%%" ~doc in
   Cmd.v info
-    Term.(term_result (const main $ opt_migrate $ opt_eio_sw_as_fiber_var))
+    Term.(
+      term_result
+        (const main $ opt_migrate $ opt_eio_sw_as_fiber_var
+       $ opt_eio_env_as_fiber_var))
 
 let () = exit (Cmd.eval cmd)

test/lwt_to_direct_style/eio-switch.t/run.t

@@ -3,25 +3,36 @@ Make a writable directory tree:
   $ cd out
 
   $ dune build @ocaml-index
-  $ lwt-to-direct-style --migrate --eio-sw-as-fiber-var Fiber_var.sw
+  $ lwt-to-direct-style --migrate --eio-sw-as-fiber-var Fiber_var.sw --eio-env-as-fiber-var Fiber_var.env
   Formatted 1 files
-  Warning: main.ml: 2 occurrences have not been rewritten.
-    Lwt_io.read (line 9 column 12)
-    Lwt_io.printf (line 10 column 3)
+  Warning: main.ml: 5 occurrences have not been rewritten.
+    Lwt_io.open_file (line 8 column 13)
+    Lwt_io.input (line 8 column 36)
+    Lwt_io.close (line 9 column 3)
+    Lwt_io.read (line 15 column 12)
+    Lwt_io.printf (line 16 column 3)
 
   $ cat main.ml
   open Eio.Std
 
   let async_process _ = ()
 
+  let _f _ =
+    Eio.Time.with_timeout_exn
+      (Stdlib.Option.get (Fiber.get Fiber_var.env))#mono_clock 1.0 (fun () -> 42)
+  
+  let _f fname =
+    let fd = Lwt_io.open_file ~mode:Lwt_io.input fname in
+    Lwt_io.close fd
+  
   let main () =
     Fiber.fork
-      ~sw:(Option.get (Fiber.get Fiber_var.sw))
+      ~sw:(Stdlib.Option.get (Fiber.get Fiber_var.sw))
       (fun () -> async_process 1);
     let fd =
      fun ?blocking:x1 ?set_flags:x2 ->
       Eio_unix.Fd.of_unix
-        ~sw:(Option.get (Fiber.get Fiber_var.sw))
+        ~sw:(Stdlib.Option.get (Fiber.get Fiber_var.sw))
         ?blocking:x1 ~close_unix:true
         (* TODO: lwt-to-direct-style: Labelled argument ?set_flags was dropped. *)
         Unix.stdin
@@ -34,6 +45,8 @@ Make a writable directory tree:
   
   let () =
     Eio_main.run (fun env ->
-        (* TODO: lwt-to-direct-style: [Eio_main.run] argument used to be a [Lwt] promise and is now a [fun]. Make sure no asynchronous or IO calls are done outside of this [fun]. *)
-        (* TODO: lwt-to-direct-style: Make sure to create a [Switch.t] and store it in fiber variable ["Fiber_var.sw"]. *)
-        main ())
+        Fiber.with_binding Fiber_var.env env (fun () ->
+            Switch.run ~name:"main" (fun sw ->
+                Fiber.with_binding Fiber_var.sw sw (fun () ->
+                    (* TODO: lwt-to-direct-style: [Eio_main.run] argument used to be a [Lwt] promise and is now a [fun]. Make sure no asynchronous or IO calls are done outside of this [fun]. *)
+                    main ()))))

test/lwt_to_direct_style/eio-switch.t/src/main.ml

@@ -2,6 +2,12 @@ open Lwt.Syntax
 
 let async_process _ = Lwt.return ()
 
+let _f _ = Lwt_unix.with_timeout 1.0 (fun () -> Lwt.return 42)
+
+let _f fname =
+  let* fd = Lwt_io.open_file ~mode:Lwt_io.input fname in
+  Lwt_io.close fd
+
 let main () =
   Lwt.async (fun () -> async_process 1);
   let fd = Lwt_unix.of_unix_file_descr Unix.stdin in