github-tui/lib/pretty.ml

(* A simple pretty printing combinator library *)

type styles = ANSITerminal.style list

type doc =
  | Empty
  | Str of styles * string
  | Vertical of doc * doc
  | Horizontal of doc * doc

let str string = Str ([], string)
let fmt styles string = Str (styles, string)

let row = function 
  | [] -> Empty
  | hd :: tl -> List.fold_left (fun l r -> Horizontal (l, r)) hd tl

let col = function 
  | [] -> Empty
  | hd :: tl -> List.fold_left (fun l r -> Vertical (l, r)) hd tl

type chunk =
  {
    styles: styles;
    string: string;
  }

let fmt_chunk {styles; string} =
  ANSITerminal.sprintf styles "%s" string

let mk_padding_chunk n =
  let padding = String.make n ' ' in
  {styles = []; string = padding}

type line = 
  { 
    chunks: chunk list
  }

let fmt_line line =
  line.chunks
  |> List.map fmt_chunk
  |> String.concat "" 

let line_len line = 
  List.fold_left
    (fun acc {string; _} -> acc + String_extra.graphemes_len string)
    0
    line.chunks

let zip_lines (l: line list) (r: line list) =
  let max_len_l = List.map line_len l |> List.fold_left max 0 in

  let rec zip l r =
    match (l, r) with
    | (l, []) -> 
      l

    | ([], r) -> 
      (* Optimisation: Add extra chunk only if padding is needed *)
      if max_len_l > 0 then
        let padding_chunk = mk_padding_chunk max_len_l in
        List.map (fun line -> { chunks = padding_chunk :: line.chunks }) r
      else
        r

    | (hd_l :: tl_l, hd_r :: tl_r) ->
      let left_len = line_len hd_l in

      (* Optimisation: Combine chunks when left is already max len *)
      if left_len >= max_len_l then
        let new_line = { chunks = hd_l.chunks @ hd_r.chunks } in
        new_line :: zip tl_l tl_r
      else
        let padding_chunk = mk_padding_chunk (max_len_l - left_len) in
        let new_line = { chunks = hd_l.chunks @ [padding_chunk] @ hd_r.chunks } in
        new_line :: zip tl_l tl_r
  in

  zip l r

let rec render_to_lines = function
  | Empty ->
    []
  | Str (styles, string) ->
    [{ chunks = [{styles; string}] }]
  | Vertical (top, bottom) ->
    render_to_lines top @ render_to_lines bottom
  | Horizontal (left, right) ->
    zip_lines (render_to_lines left) (render_to_lines right)

let render doc = 
  doc
  |> render_to_lines
  |> List.map fmt_line
  |> String_extra.unlines