treeprint.ml

gistfile1.fs

type tree = Tree of (string * tree list) | Leaf of string

(* Create a list of line repeated count times *)
let rec repeat_line line count =
  match count with 0 -> [] | _ -> line :: (repeat_line line (count - 1))

(* Appends blank lines which are filled with spaces *)
let append_blank_lines lines count =
  let width     = String.length (List.hd lines) in
  let blankline = String.make width ' '         in
  lines @ (repeat_line blankline count)

(* Given two lists of strings, make sure the lists are the same length.
   append_blank_lines if we have to. *)
let grow_lines leftl rightl =
  let lvert  = List.length leftl  in
  let rvert  = List.length rightl in
  if lvert < rvert then
    ((append_blank_lines leftl (rvert - lvert)), rightl)
  else if lvert > rvert then
    (leftl, (append_blank_lines rightl (lvert - rvert)))
  else
    (leftl, rightl)

(* Given two lists of strings, treat each as a rectangle of text
   and paste them together horizontally. *)
let paste_horz left right =
  (match grow_lines left right with
    (llines, rlines) ->
      List.map2 (fun l r -> l ^ " " ^ r)  llines rlines)

let string_center str col =
  let diff = col - (String.length str) in
  let pad = diff / 2 in
  if pad < 0 then failwith("Cannot center string it is too big")
  else (String.make pad ' ') ^ str ^
    (String.make (pad + (diff mod 2)) ' ')

(* Prepend a centered string to the top of a list of strings *)
let prepend_text str lines =
  let col = String.length (List.hd lines) in
  string_center str col :: lines

let prepend_bar lines =
  prepend_text "|" lines

(* The bar under a node with children looks like a swing... *)
let swing_bar topline =
  let left  = String.index topline '|' in
  let right = String.rindex topline '|' in
  let swing =
    if left == right then "!"
    else String.make (right - left + 1) '-' in
  (String.make left ' ') ^ swing ^
  (String.make ((String.length topline) - right - 1) ' ')

let prepend_swing childlines =
  let topline = List.hd childlines in
  swing_bar topline :: childlines

(* Renders a tree recursively into a list of strings. *)
let rec render_tree tree =
  prepend_bar begin
    match tree with
      Leaf( name ) -> [ name ]
    | Tree( name, children ) ->
        match List.map render_tree children with
          [] -> [ name ] (* You didn't use Leaf like you should have! *)
        | hd::tl ->
            let merged = List.fold_left paste_horz hd tl in
            (prepend_text name (prepend_swing merged))
  end

let print_tree tree =
  print_endline (String.concat "\n" (render_tree tree))

let _ =
  print_tree (Tree( "Hello",
                   [ Tree("How", [ Leaf("Are"); Leaf("You") ]) ;
                     Tree("Today?",
                          [ Tree("My", [ Leaf("BestestInTheWholeworld");
                                         Leaf("Friend") ])]
                            )]));