Merge pull request #20 from Julow/lwt_io_for_ocsipersist

Support more `Lwt_io` operations

Author: Julow

bin/lwt_to_direct_style/ast_rewrite.ml

@@ -280,6 +280,8 @@ let rewrite_apply_lwt ~backend ~state ident args =
         ("[<?>] can't be automatically translated."
        ^ backend#choose_comment_hint);
       return None
+  | "wrap" ->
+      take @@ fun f -> return (Some (Exp.apply f [ (Nolabel, mk_unit_val) ]))
   | _ -> return None
 
 let string_drop_suffix ~suffix s =
@@ -312,6 +314,17 @@ let rewrite_apply ~backend ~state full_ident args =
       (( "printf" | "eprintf" | "stdout" | "stderr" | "fprintf" | "kfprintf"
        | "ifprintf" | "ikfprintf" ) as ident) ) ->
       transparent [ "Format"; ident ]
+  | ( "Lwt_unix",
+      (( "socket" | "socketpair" | "listen" | "shutdown" | "getsockname"
+       | "getpeername" | "waitpid" | "wait4" | "wait" ) as ident) ) ->
+      transparent [ "Unix"; ident ]
+  | "Lwt_io", "read_value" -> transparent [ "Marshal"; "from_channel" ]
+  | "Lwt_io", "write_value" -> transparent [ "Marshal"; "to_channel" ]
+  | "Lwt_unix", (("connect" | "accept" | "bind") as ident) ->
+      Printf.ksprintf (add_comment state)
+        "This call to [Unix.%s] was [Lwt_unix.%s] before. It's now blocking."
+        ident ident;
+      transparent [ "Unix"; ident ]
   | "Lwt_list", ident -> (
       match string_drop_suffix ~suffix:"_s" ident with
       | Some ident -> transparent [ "List"; ident ]
@@ -356,12 +369,13 @@ let rewrite_apply ~backend ~state full_ident args =
   | "Lwt_mutex", "with_lock" ->
       take @@ fun t ->
       take @@ fun f -> return (Some (backend#mutex_with_lock t f))
-  | "Lwt_io", "read_into" ->
+  | "Lwt_io", (("read_into" | "read_into_exactly") as ident) ->
+      let exactly = ident = "read_into_exactly" in
       take @@ fun input ->
       take @@ fun buffer ->
       take @@ fun buf_off ->
       take @@ fun buf_len ->
-      return (Some (backend#io_read input buffer buf_off buf_len))
+      return (Some (backend#io_read ~exactly input buffer buf_off buf_len))
   | "Lwt_io", "of_fd" ->
       ignore_lblarg "buffer"
       @@ ignore_lblarg ~cmt:"Will behave as if it was [true]." "close"
@@ -384,8 +398,20 @@ let rewrite_apply ~backend ~state full_ident args =
       take @@ fun str -> return (Some (backend#io_write_str chan str))
   | "Lwt_io", "length" -> take @@ fun fd -> return (Some (backend#io_length fd))
   | "Lwt_io", "close" -> take @@ fun fd -> return (Some (backend#io_close fd))
+  | "Lwt_io", "flush" -> take @@ fun fd -> return (Some (backend#io_flush fd))
+  | "Lwt_io", "with_connection" ->
+      ignore_lblarg "fd" @@ ignore_lblarg "in_buffer"
+      @@ ignore_lblarg "out_buffer" @@ take
+      @@ fun sockaddr ->
+      take @@ fun f -> return (Some (backend#net_with_connection sockaddr f))
   | "Lwt_main", "run" ->
       take @@ fun promise -> return (Some (backend#main_run promise))
+  | "Lwt_preemptive", "detach" ->
+      take @@ fun f ->
+      take @@ fun arg ->
+      return
+        (Some
+           (backend#domain_detach (mk_thunk (Exp.apply f [ (Nolabel, arg) ]))))
   | _ -> return None
 
 (** Transform a [binding_op] into a [pattern] and an [expression] while
@@ -550,6 +576,7 @@ let rewrite_type ~backend ~state typ =
                  | "group_entry" ) as tname) ),
               params ) ->
               Some (mk_typ_constr ~params [ "Unix"; tname ])
+          | ("Lwt_io", "input_channel"), [] -> Some backend#type_in_channel
           | ("Lwt_io", "output_channel"), [] -> Some backend#type_out_channel
           | _ -> None)
   | _ -> None

bin/lwt_to_direct_style/concurrency_backend.ml

@@ -18,9 +18,10 @@ let eio ~eio_sw_as_fiber_var ~eio_env_as_fiber_var add_comment =
     [ i ]
   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
+  let add_comment_dropped_exp ~label ?(cmt = "") exp =
+    add_comment "Dropped expression (%s): [%s].%s" label
       (Ocamlformat_utils.format_expression exp)
+      cmt
   in
   (* If [--eio-sw-as-fiber-var] is passed on the command line, this will query
      the current switch. Otherwise, this will generate a comment.
@@ -206,7 +207,7 @@ let eio ~eio_sw_as_fiber_var ~eio_env_as_fiber_var add_comment =
       (*     ]) *)
       fd
 
-    method io_read input buffer buf_offset buf_len =
+    method io_read ~exactly input buffer buf_offset buf_len =
       add_comment "[%s] should be a [Cstruct.t]."
         (Ocamlformat_utils.format_expression buffer);
       add_comment
@@ -215,9 +216,12 @@ let eio ~eio_sw_as_fiber_var ~eio_env_as_fiber_var add_comment =
          contains an internal buffer) or change the call to \
          [Eio.Buf_read.of_flow] used to create the buffer."
         (Ocamlformat_utils.format_expression input);
-      add_comment_dropped_exp ~label:"buffer offset" buf_offset;
-      add_comment_dropped_exp ~label:"buffer length" buf_len;
-      mk_apply_simple [ "Eio"; "Flow"; "single_read" ] [ input; buffer ]
+      add_comment_dropped_exp ~label:"buffer offset"
+        ~cmt:" This will behave as if it was [0]." buf_offset;
+      add_comment_dropped_exp ~label:"buffer length"
+        ~cmt:" This will behave as if it was [Cstruct.length buffer]." buf_len;
+      let fun_ = if exactly then "read_exact" else "single_read" in
+      mk_apply_simple [ "Eio"; "Flow"; fun_ ] [ input; buffer ]
 
     method io_read_all input =
       mk_apply_simple [ "Eio"; "Buf_read"; "take_all" ] [ input ]
@@ -229,6 +233,9 @@ let eio ~eio_sw_as_fiber_var ~eio_env_as_fiber_var add_comment =
          unbuffered IO.";
       None
 
+    method io_flush output =
+      mk_apply_simple [ "Eio"; "Buf_write"; "flush" ] [ output ]
+
     method fd_close fd =
       (* TODO: See [of_unix_file_descr]. mk_apply_simple [ "Eio_unix"; "Fd" ] [ fd ] *)
       mk_apply_simple [ "Unix"; "close" ] [ fd ]
@@ -250,7 +257,8 @@ let eio ~eio_sw_as_fiber_var ~eio_env_as_fiber_var add_comment =
             in
             mk_apply_ident (switch_ident "run")
               [
-                (Labelled (mk_loc "name"), mk_const_string "main");
+                (* TODO: Add the [~name] argument. Currently commented-out because added in a too recent version of eio.
+                   (Labelled (mk_loc "name"), mk_const_string "main"); *)
                 (Nolabel, fun_sw);
               ]
         | None -> k
@@ -310,6 +318,7 @@ let eio ~eio_sw_as_fiber_var ~eio_env_as_fiber_var add_comment =
       mk_apply_simple [ "Eio"; "Buf_write"; "string" ] [ chan; str ]
 
     method io_close fd = mk_apply_simple [ "Eio"; "Resource"; "close" ] [ fd ]
+    method type_in_channel = mk_typ_constr [ "Eio"; "Buf_read"; "t" ]
     method type_out_channel = mk_typ_constr [ "Eio"; "Buf_write"; "t" ]
 
     method path_stat ~follow path =
@@ -318,4 +327,37 @@ let eio ~eio_sw_as_fiber_var ~eio_env_as_fiber_var add_comment =
           (Labelled (mk_loc "follow"), mk_constr_of_bool follow);
           (Nolabel, mk_apply_simple [ "Eio"; "Path"; "/" ] [ env "cwd"; path ]);
         ]
+
+    method domain_detach thunk =
+      mk_apply_ident
+        (fiber_ident "fork_promise")
+        [
+          get_current_switch_arg ();
+          ( Nolabel,
+            mk_thunk
+              (mk_apply_simple
+                 [ "Eio"; "Domain_manager"; "run" ]
+                 [ env "domain_mgr"; thunk ]) );
+        ]
+
+    method net_with_connection sockaddr f =
+      add_comment
+        "[%s] is of type [Unix.sockaddr] but it should be a \
+         [Eio.Net.Sockaddr.stream]."
+        (Ocamlformat_utils.format_expression sockaddr);
+      mk_apply_simple (switch_ident "run")
+        [
+          mk_fun ~arg_name:"sw" (fun sw ->
+              Exp.apply f
+                [
+                  ( Nolabel,
+                    mk_apply_ident
+                      [ "Eio"; "Net"; "connect" ]
+                      [
+                        (Labelled (mk_loc "sw"), sw);
+                        (Nolabel, env "net");
+                        (Nolabel, sockaddr);
+                      ] );
+                ]);
+        ]
   end

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

@@ -44,7 +44,7 @@ Make a writable directory tree:
   let () =
     Eio_main.run (fun env ->
         Fiber.with_binding Fiber_var.env env (fun () ->
-            Switch.run ~name:"main" (fun sw ->
+            Switch.run (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 ()))))