Add Buf_write.printf

Co-authored-by: Thomas Leonard <talex5@gmail.com>

lib_eio/buf_write.ml

@@ -163,6 +163,7 @@ type t =
   ; mutable bytes_written  : int        (* Total written bytes. Wraps. *)
   ; mutable state          : state
   ; mutable wake_writer    : unit -> unit
+  ; mutable printf         : (Format.formatter * bool ref) option
   }
 (* Invariant: [write_pos >= scheduled_pos] *)
 
@@ -377,6 +378,7 @@ let of_buffer ?sw buffer =
           ; bytes_written   = 0
           ; state           = Active
           ; wake_writer     = ignore
+          ; printf          = None
           }
   in
   begin match sw with
@@ -416,6 +418,39 @@ let flush t =
     Promise.await_exn p
   )
 
+let make_formatter t =
+  Format.make_formatter
+    (fun buf off len -> write_gen t buf ~off ~len ~blit:Bigstringaf.blit_from_string)
+    (fun () -> flush t)
+
+let printf t =
+  let ppf, is_formatting =
+    match t.printf with
+    | Some (_, is_formatting as x) ->
+      is_formatting := true;
+      x
+    | None ->
+      let is_formatting = ref true in
+      let ppf =
+        Format.make_formatter
+          (fun buf off len -> write_gen t buf ~off ~len ~blit:Bigstringaf.blit_from_string)
+          (fun () ->
+             (* As per the Format module manual, an explicit flush writes to the
+                output channel and ensures that "all pending text is displayed"
+                and "these explicit flush calls [...] could dramatically impact efficiency".
+                Therefore it is clear that we need to call `flush t` instead of `flush_buffer t`. *)
+             if !is_formatting then flush t)
+      in
+      t.printf <- Some (ppf, is_formatting);
+      ppf, is_formatting
+  in
+  Format.kfprintf (fun ppf ->
+      if not !is_formatting then raise (Sys_error "Buf_write.printf: invalid concurrent access");
+      (* Ensure that [ppf]'s internal buffer is flushed to [t], but without flushing [t] itself: *)
+      is_formatting := false;
+      Format.pp_print_flush ppf ()
+    ) ppf
+
 let rec shift_buffers t written =
   match Buffers.dequeue_exn t.scheduled with
   | { Cstruct.len; _ } as iovec ->

lib_eio/buf_write.mli

@@ -46,7 +46,7 @@
     or application-specific output APIs.
 
     A Buf_write serializer manages an internal buffer and a queue of output
-    buffers. The output bufferes may be a sub range of the serializer's
+    buffers. The output buffers may be a sub range of the serializer's
     internal buffer or one that is user-provided. Buffered writes such as
     {!string}, {!char}, {!cstruct}, etc., copy the source bytes into the
     serializer's internal buffer. Unbuffered writes are done with
@@ -123,6 +123,18 @@ val cstruct : t -> Cstruct.t -> unit
     It is safe to modify [cs] after this call returns.
     For large cstructs, it may be more efficient to use {!schedule_cstruct}. *)
 
+val printf : t -> ('a, Format.formatter, unit) format -> 'a
+(** [printf t fmt x y z] formats the arguments according to the format string [fmt].
+    It supports all formatting and pretty-printing features of the Format module.
+    Accordingly, explicit flushes using [@.] or [%!] perform a full (blocking) flush. *)
+
+val make_formatter : t -> Format.formatter
+(** [make_formatter t] creates a new formatter that writes to [t].
+
+    Flushing the formatter also flushes [t] itself.
+    This function is useful to mutate formatter settings by calling, for example,
+    [Format.pp_set_margin] or [Format.pp_set_geometry]. *)
+
 val write_gen
   :  t
   -> blit:('a -> src_off:int -> Cstruct.buffer -> dst_off:int -> len:int -> unit)

tests/buf_write.md

@@ -146,6 +146,44 @@ With pausing
 - : unit = ()
 ```
 
+## Formatting
+
+```ocaml
+# Eio_mock.Backend.run @@ fun () ->
+  Write.with_flow flow @@ fun t ->
+  Write.printf t "Write.printf can force a full flush@.@[<v2>It also@,flushes to [t] automatically";
+  Write.string t " at the end, but without flushing [t] itself.\n";
+  (* Create a formatter for full control: *)
+  let f = Write.make_formatter t in
+  Format.pp_set_geometry f ~max_indent:4 ~margin:10;
+  (*
+    "@ "      breakable space
+    "@[<v 6>" open vertical box, indentation: 6 (overriden by our geometry settings)
+    "%s"      print string
+    "@ "      breakable space
+    "%i"      print int
+    "@."      print newline + explicit flush
+    "%a"      print arbitrary type
+    "@]"      close box
+    "@ "      breakable space
+  *)
+  Fmt.pf f "Space@ @[<v 6>%s@ %i@.%a@]@ "
+    "This is a test" 123
+    Eio.Net.Sockaddr.pp (`Tcp (Eio.Net.Ipaddr.V6.loopback, 8080));
+  Write.printf t "This is a %s call to printf" "second";
+  Fmt.pf f "@.Flushed. %s@." "Goodbye"
++flow: wrote "Write.printf can force a full flush\n"
++flow: wrote "It also\n"
++            "  flushes to [t] automatically at the end, but without flushing [t] itself.\n"
++            "Space\n"
++            "This is a test\n"
++            "    123\n"
++flow: wrote "tcp:[::1]:8080This is a second call to printf\n"
++            "\n"
++flow: wrote "Flushed. Goodbye\n"
+- : unit = ()
+```
+
 ## Flushing
 
 ```ocaml