From cf5a813b89815c96d749c6fb1e36daa7277a838a Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Mon, 13 Nov 2023 13:27:20 -0500
Subject: [PATCH] add `Containers_pp.newline_or_spaces`

---
 src/pp/containers_pp.ml  | 36 ++++++++++++++++++++++++------------
 src/pp/containers_pp.mli |  6 ++++++
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/src/pp/containers_pp.ml b/src/pp/containers_pp.ml
index f10a8f1e..0b9ea16e 100644
--- a/src/pp/containers_pp.ml
+++ b/src/pp/containers_pp.ml
@@ -37,7 +37,7 @@ type t = {
 
 and view =
   | Nil
-  | Newline
+  | Newline of int
   | Nest of int * t
   | Append of t * t
   | Char of char
@@ -52,7 +52,8 @@ and view =
 let rec debug out (self : t) : unit =
   match self.view with
   | Nil -> Format.fprintf out "nil"
-  | Newline -> Format.fprintf out "nl"
+  | Newline 1 -> Format.fprintf out "nl"
+  | Newline i -> Format.fprintf out "nl(%d)" i
   | Nest (i, x) -> Format.fprintf out "(@[nest %d@ %a@])" i debug x
   | Append (a, b) -> Format.fprintf out "@[%a ^@ %a@]" debug a debug b
   | Char c -> Format.fprintf out "%C" c
@@ -65,7 +66,12 @@ let rec debug out (self : t) : unit =
   | Wrap (_, _, d) -> Format.fprintf out "(@[ext@ %a@])" debug d
 
 let nil : t = { view = Nil; wfl = 0 }
-let newline : t = { view = Newline; wfl = 1 }
+let newline : t = { view = Newline 1; wfl = 1 }
+
+let newline_or_spaces n : t =
+  if n < 0 then invalid_arg "Containers_pp.newline_or_spaces";
+  { view = Newline n; wfl = n }
+
 let nl = newline
 
 let char c =
@@ -131,9 +137,13 @@ module Flatten = struct
   let to_out (out : Out.t) (self : t) : unit =
     let rec loop (d : t) =
       match d.view with
-      | Nil -> ()
+      | Nil | Newline 0 -> ()
       | Char c -> out.char c
-      | Newline -> out.char ' '
+      | Newline 1 -> out.char ' '
+      | Newline n ->
+        for _i = 1 to n do
+          out.char ' '
+        done
       | Nest (_, x) -> loop x
       | Append (x, y) ->
         loop x;
@@ -175,13 +185,15 @@ module Pretty = struct
 
   let rec pp_flatten (st : st) (self : t) : int =
     match self.view with
-    | Nil -> 0
+    | Nil | Newline 0 -> 0
     | Char c ->
       st.out.char c;
       1
-    | Newline ->
-      st.out.char ' ';
-      1
+    | Newline n ->
+      for _i = 1 to n do
+        st.out.char ' '
+      done;
+      n
     | Nest (_i, x) -> pp_flatten st x
     | Append (x, y) ->
       let n = pp_flatten st x in
@@ -229,15 +241,15 @@ module Pretty = struct
       pp_rec_top st ~k ~i d (fun k -> pp_rec st k stack_tl)
 
   (** Print [d] at indentation [i], with [k] chars already printed
-        on the current line, then calls [kont] with the
-        new [k]. *)
+      on the current line, then calls [kont] with the
+      new [k]. *)
   and pp_rec_top st ~k ~i d (kont : int -> unit) : unit =
     match d.view with
     | Nil -> kont k
     | Char c ->
       st.out.char c;
       kont (k + 1)
-    | Newline ->
+    | Newline _ ->
       pp_newline st i;
       kont i
     | Nest (j, x) -> pp_rec_top st ~k ~i:(i + j) x kont
diff --git a/src/pp/containers_pp.mli b/src/pp/containers_pp.mli
index 9c669fad..fe502348 100644
--- a/src/pp/containers_pp.mli
+++ b/src/pp/containers_pp.mli
@@ -74,6 +74,12 @@ val append : t -> t -> t
 val newline : t
 (** A line break. *)
 
+val newline_or_spaces : int -> t
+(** [newline_or_spaces n] either prints a newline (respecting indentation),
+    or prints [n] spaces. {!newline} is basically [newline_or_spaces 1].
+    @raise Invalid_argument if [n < 0].
+    @since NEXT_RELEASE *)
+
 val nl : t
 (** Alias for {!newline} *)