From b442996bb7aad753bd921e2a51f301854e583c67 Mon Sep 17 00:00:00 2001
From: Qexat <43090614+qexat@users.noreply.github.com>
Date: Wed, 22 Jan 2025 17:04:21 +0100
Subject: [PATCH] Add experimental stuff

---
 lib/Ansifmt.ml    |  9 +------
 lib/Formatting.ml | 60 +++++++++++++++++++++++++++++++++++++----------
 lib/IO.ml         |  9 +++----
 lib/Prelude.ml    |  8 +++++++
 lib/Styling.ml    | 12 ++++++++++
 lib/Styling.mli   |  4 ++++
 lib/Util.ml       | 23 ++++++++++++++++++
 7 files changed, 98 insertions(+), 27 deletions(-)
 create mode 100644 lib/Prelude.ml

diff --git a/lib/Ansifmt.ml b/lib/Ansifmt.ml
index 2396765..36a3664 100644
--- a/lib/Ansifmt.ml
+++ b/lib/Ansifmt.ml
@@ -1,12 +1,5 @@
 module Color = Color
 module IO = IO
 module Formatting = Formatting
+module Prelude = Prelude
 module Styling = Styling
-
-type color = Color.t
-type styling = Styling.t
-type stylizer = Formatting.Stylizer.t
-
-let make_styling = Styling.create
-let format = Formatting.Util.format
-let print_formatted = IO.print_formatted
diff --git a/lib/Formatting.ml b/lib/Formatting.ml
index 7f7146b..d06f270 100644
--- a/lib/Formatting.ml
+++ b/lib/Formatting.ml
@@ -1,3 +1,5 @@
+open Util
+
 module Token_type = struct
   type t =
     | Comment
@@ -86,11 +88,34 @@ module Token = struct
   let colon : t = Token_type.Punctuation_strong, ":"
 
   let format ?(stylizer : Stylizer.t = Stylizer.default) : t -> string =
-    fun (token_type, lexeme) ->
-    Printf.sprintf
-      "%s%s\x1b[22;23;24;25;39;49m"
-      (Styling.to_ansi (stylizer token_type))
-      lexeme
+    fun (token_type, lexeme) -> Styling.wrap ~contents:lexeme (stylizer token_type)
+  ;;
+end
+
+module Tree = struct
+  type t =
+    | Simple of Token.t list
+    | Parenthesized of t
+    | Block of t list
+
+  let simple : Token.t list -> t = fun tokens -> Simple tokens
+  let parenthesized : t -> t = fun tree -> Parenthesized tree
+  let block : t list -> t = fun trees -> Block trees
+
+  let rec format : ?parentheses:string * string -> ?stylizer:Stylizer.t -> t -> string =
+    fun ?(parentheses = "(", ")") ?(stylizer = Stylizer.default) tree ->
+    let opening_token, closing_token =
+      Pair.map_uniform (fun lexeme -> Token_type.Pair, lexeme) parentheses
+    in
+    tree
+    |> function
+    | Simple tokens -> tokens |> List.map (Token.format ~stylizer) |> String.concat ""
+    | Parenthesized subtree ->
+      Token.format ~stylizer opening_token
+      ^ format ~parentheses ~stylizer subtree
+      ^ Token.format ~stylizer closing_token
+    | Block subtrees ->
+      subtrees |> List.map (format ~parentheses ~stylizer) |> String.concat ""
   ;;
 end
 
@@ -103,28 +128,37 @@ module type TOKENIZABLE = sig
 
   (** [tokenize term] transforms [term] into a stream of
       formatter tokens. *)
-  val tokenize : t -> Token.t list
+  val tokenize : t -> Tree.t
 end
 
 module Util = struct
   (** [tokenize value ~using:(module M)] transforms [value] to a
     list of tokens. *)
-  let tokenize : type t. t -> using:(module TOKENIZABLE with type t = t) -> Token.t list =
+  let tokenize : type t. t -> using:(module TOKENIZABLE with type t = t) -> Tree.t =
     fun value ~using:(module M) -> M.tokenize value
   ;;
 
+  let parenthesize_if
+    : type t. (t -> bool) -> t -> using:(module TOKENIZABLE with type t = t) -> Tree.t
+    =
+    fun predicate value ~using:(module M) ->
+    let base = tokenize value ~using:(module M) in
+    if predicate value then Tree.parenthesized base else base
+  ;;
+
   (** [format ?stylizer value ~using:(module M)] transforms the
     [value] into a pretty-printable string using the [stylizer]
     if [M] provides tokenization for the [value] type. *)
   let format
     : type t.
-      ?stylizer:Stylizer.t -> t -> using:(module TOKENIZABLE with type t = t) -> string
+      ?stylizer:Stylizer.t
+      -> ?parentheses:string * string
+      -> t
+      -> using:(module TOKENIZABLE with type t = t)
+      -> string
     =
-    fun ?(stylizer = Stylizer.default) value ~using:(module M) ->
+    fun ?stylizer ?parentheses value ~using:(module M) ->
     (* TODO: handle line breaks / width-aware formatting *)
-    value
-    |> tokenize ~using:(module M)
-    |> List.map (Token.format ~stylizer)
-    |> String.concat ""
+    value |> tokenize ~using:(module M) |> Tree.format ?parentheses ?stylizer
   ;;
 end
diff --git a/lib/IO.ml b/lib/IO.ml
index 7901419..b8e3946 100644
--- a/lib/IO.ml
+++ b/lib/IO.ml
@@ -11,20 +11,17 @@
 let print_formatted
   : type t.
     ?stylizer:Formatting.Stylizer.t
+    -> ?parentheses:string * string
     -> ?line_end:string
     -> ?out:out_channel
     -> t
     -> using:(module Formatting.TOKENIZABLE with type t = t)
     -> unit
   =
-  fun ?(stylizer = Formatting.Stylizer.default)
-    ?(line_end = "\n")
-    ?(out = stdout)
-    value
-    ~using:(module M) ->
+  fun ?stylizer ?parentheses ?(line_end = "\n") ?(out = stdout) value ~using:(module M) ->
   Printf.fprintf
     out
     "%s%s"
-    (Formatting.Util.format ~stylizer value ~using:(module M))
+    (Formatting.Util.format ?stylizer ?parentheses value ~using:(module M))
     line_end
 ;;
diff --git a/lib/Prelude.ml b/lib/Prelude.ml
new file mode 100644
index 0000000..2ef77dc
--- /dev/null
+++ b/lib/Prelude.ml
@@ -0,0 +1,8 @@
+type color = Color.t
+type styling = Styling.t
+type stylizer = Formatting.Stylizer.t
+
+let make_styling = Styling.create
+let format = Formatting.Util.format
+let parenthesize_if = Formatting.Util.parenthesize_if
+let print_formatted = IO.print_formatted
diff --git a/lib/Styling.ml b/lib/Styling.ml
index ea6c9d1..da06da8 100644
--- a/lib/Styling.ml
+++ b/lib/Styling.ml
@@ -57,3 +57,15 @@ let to_ansi : t -> string = function
     add_color_to_buffer buffer background ~ground:Color.background;
     Buffer.contents buffer
 ;;
+
+let wrap : contents:string -> t -> string =
+  fun ~contents -> function
+  | { foreground; background; bold; dim; italic; underlined } as styling ->
+    let buffer = Buffer.create 16 in
+    if bold || dim then Buffer.add_string buffer (make_sgr_sequence "22");
+    if italic then Buffer.add_string buffer (make_sgr_sequence "23");
+    if underlined then Buffer.add_string buffer (make_sgr_sequence "24");
+    if Option.is_some foreground then Buffer.add_string buffer (make_sgr_sequence "39");
+    if Option.is_some background then Buffer.add_string buffer (make_sgr_sequence "49");
+    to_ansi styling ^ contents ^ Buffer.contents buffer
+;;
diff --git a/lib/Styling.mli b/lib/Styling.mli
index a887f54..067cef7 100644
--- a/lib/Styling.mli
+++ b/lib/Styling.mli
@@ -25,3 +25,7 @@ val create
 (** [to_ansi styling] renders the [styling] to an ANSI escape
     sequence as a string. *)
 val to_ansi : t -> string
+
+(** [wrap ~contents styling] wraps [contents] in an ANSI escape
+    sequence using [styling]. *)
+val wrap : contents:string -> t -> string
diff --git a/lib/Util.ml b/lib/Util.ml
index b31bac3..84a5a3d 100644
--- a/lib/Util.ml
+++ b/lib/Util.ml
@@ -30,6 +30,29 @@ end = struct
   ;;
 end
 
+module Pair = struct
+  type ('a, 'b) t = 'a * 'b
+
+  let first : ('a, 'b) t -> 'a = fun (first, _) -> first
+  let second : ('a, 'b) t -> 'b = fun (_, second) -> second
+
+  let map : ('a -> 'c) -> ('b -> 'd) -> ('a, 'b) t -> ('c, 'd) t =
+    fun func_first func_second (first, second) -> func_first first, func_second second
+  ;;
+
+  let map_first : ('a -> 'c) -> ('a, 'b) t -> ('c, 'b) t =
+    fun func pair -> map func Fun.id pair
+  ;;
+
+  let map_second : ('b -> 'c) -> ('a, 'b) t -> ('a, 'c) t =
+    fun func pair -> map Fun.id func pair
+  ;;
+
+  let map_uniform : ('a -> 'b) -> ('a, 'a) t -> ('b, 'b) t =
+    fun func pair -> map func func pair
+  ;;
+end
+
 module Triplet = struct
   type ('a, 'b, 'c) t = 'a * 'b * 'c