From d36c57459e331ba5d386bad92e9077c0388fb621 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Thu, 30 Mar 2023 23:20:41 -0400
Subject: [PATCH] feat(Pp): add `wrap` for extension nodes

---
 src/pp/containers_pp.ml  | 52 ++++++++++++++++++++++++++++++++++++++--
 src/pp/containers_pp.mli | 43 +++++++++++++++++++++++++++------
 tests/core/t_pp.ml       | 19 +++++++++++++++
 3 files changed, 105 insertions(+), 9 deletions(-)

diff --git a/src/pp/containers_pp.ml b/src/pp/containers_pp.ml
index 57f73135..a5f06df2 100644
--- a/src/pp/containers_pp.ml
+++ b/src/pp/containers_pp.ml
@@ -1,3 +1,33 @@
+module B = Buffer
+
+module Ext = struct
+  module type OUT = sig
+    val char : char -> unit
+    val sub_string : string -> int -> int -> unit
+    val string : string -> unit
+    val newline : unit -> unit
+  end
+
+  type out = (module OUT)
+
+  let out_of_buf (buf : Buffer.t) : out =
+    (module struct
+      let char = B.add_char buf
+      let sub_string = B.add_substring buf
+      let string = B.add_string buf
+      let newline () = B.add_char buf '\n'
+    end)
+
+  module type S = sig
+    type t
+
+    val pre : out -> t -> unit
+    val post : out -> t -> unit
+  end
+
+  type 'a t = (module S with type t = 'a)
+end
+
 type t = {
   view: view;  (** Document view *)
   wfl: int;  (** Width if flattened *)
@@ -13,6 +43,7 @@ and view =
   | Text_sub of string * int * int
   | Group of t
   | Fill of { sep: t; l: t list }
+  | Wrap : 'a Ext.t * 'a * t -> view
 
 let nil : t = { view = Nil; wfl = 0 }
 let newline : t = { view = Newline; wfl = 1 }
@@ -42,6 +73,7 @@ let group d : t =
   | Group _ -> d
   | _ -> { view = Group d; wfl = d.wfl }
 
+let wrap ext v d : t = { view = Wrap (ext, v, d); wfl = d.wfl }
 let ( ^ ) = append
 let text_sub_ s i len : t = { view = Text_sub (s, i, len); wfl = len }
 
@@ -76,8 +108,6 @@ let text (str : string) : t =
 let textpf fmt = Printf.ksprintf text fmt
 let textf fmt = Format.kasprintf text fmt
 
-module B = Buffer
-
 module Flatten = struct
   let to_buffer buf (self : t) : unit =
     let rec loop (d : t) =
@@ -98,7 +128,13 @@ module Flatten = struct
             if i > 0 then loop sep;
             loop x)
           l
+      | Wrap ((module E), v, d) ->
+        let out = Ext.out_of_buf buf in
+        E.pre out v;
+        loop d;
+        E.post out v
     in
+
     loop self
 
   let to_string self : string =
@@ -145,6 +181,12 @@ module Pretty = struct
           n := !n + pp_flatten st x)
         l;
       !n
+    | Wrap ((module E), v, d) ->
+      let out = Ext.out_of_buf st.buf in
+      E.pre out v;
+      let n = pp_flatten st d in
+      E.post out v;
+      n
 
   (** Does [x] fit in the current line when flattened, given that [k] chars
           are already on the line? *)
@@ -194,6 +236,12 @@ module Pretty = struct
       ) else
         pp_rec_top st ~k ~i x kont
     | Fill { sep; l } -> pp_fill st ~k ~i sep l kont
+    | Wrap ((module E), v, d) ->
+      let out = Ext.out_of_buf st.buf in
+      E.pre out v;
+      pp_rec_top st ~k ~i d (fun k ->
+          E.post out v;
+          kont k)
 
   and pp_fill st ~k ~i sep l (kont : int -> unit) : unit =
     (* [k] is the current offset in the line *)
diff --git a/src/pp/containers_pp.mli b/src/pp/containers_pp.mli
index 30d86aee..5564fa6e 100644
--- a/src/pp/containers_pp.mli
+++ b/src/pp/containers_pp.mli
@@ -55,13 +55,42 @@ val newline : t
 val nl : t
 (** Alias for {!newline} *)
 
-(* TODO:
-val ext : pre:(unit -> string) -> post:(unit -> string) -> t -> t
-(** Extension. This is a custom-rendered document.
-    TODO: customize how long it is?
-    TODO: generic output, not [unit -> string]
- *)
- *)
+(** Extension node.
+
+    In here, we can stuff custom printer nodes. *)
+module Ext : sig
+  module type OUT = sig
+    val char : char -> unit
+    val sub_string : string -> int -> int -> unit
+    val string : string -> unit
+    val newline : unit -> unit
+  end
+
+  type out = (module OUT)
+
+  module type S = sig
+    type t
+
+    val pre : out -> t -> unit
+    val post : out -> t -> unit
+  end
+
+  type 'a t = (module S with type t = 'a)
+end
+
+val wrap : 'a Ext.t -> 'a -> t -> t
+(** [wrap ext v d] wraps [d] with value [v].
+
+    It is a document that has the same
+    shape (and size) as [d], except that additional data will
+    be output when it is rendered.
+
+    Let [(module Ext)] be [ext], and [out]
+    be the output buffer/stream for rendering.
+
+    When this is rendered, first [Ext.pre out v] is called;
+    then [d] is printed; then [Exp.post out v] is called.
+*)
 
 (** {2 Pretty print and rendering} *)
 
diff --git a/tests/core/t_pp.ml b/tests/core/t_pp.ml
index b782a942..c5ce3309 100644
--- a/tests/core/t_pp.ml
+++ b/tests/core/t_pp.ml
@@ -39,3 +39,22 @@ let () =
          [ text "bar"; text "baaz"; sexp_apply "g" [ int 42; int 10 ] ]
      in
      Pretty.to_string ~width:10 d)
+
+module Ext_coucou : Ext.S with type t = unit = struct
+  type t = unit
+
+  let pre (module O : Ext.OUT) () = O.string "<coucou>"
+  let post (module O : Ext.OUT) () = O.string "</coucou>"
+end
+
+let () =
+  eq ~name:"wrap1" ~printer:(spf "%S")
+    "(foo\n bar\n <coucou>(g 42 10)</coucou>)"
+    (let d =
+       sexp_apply "foo"
+         [
+           text "bar";
+           wrap (module Ext_coucou) () (sexp_apply "g" [ int 42; int 10 ]);
+         ]
+     in
+     Pretty.to_string ~width:10 d)