Basic Rose Trees in OCaml

rose.ml

open StdLabels

type 'a tree = Node of ('a * 'a tree list)

let rec make ~f init =
  let (label, forest) = f init in
  Node (label, (List.map ~f:(make ~f) forest))

and draw ~f (Node (label, forest)) =
  let rec inner = function
    | [] -> []
    | t :: ts ->
      let (start, other) = match ts with
        | [] -> ("`- ", "   ")
        | _  -> ("+- ", "|  ")
      in
      let lines = draw ~f t in
      "|" :: shift start other lines @ inner ts
  in f label :: inner forest

and shift start other =
  let rec inner acc = function
    | [] -> List.rev acc
    | l :: lines ->
      let markup = if acc = [] then start else other in
      inner (String.concat ~sep:"" [markup; l] :: acc) lines
  in inner []

and flatten t =
  let rec inner (Node (label, forest)) acc =
    label :: List.fold_right ~f:inner ~init:acc forest
  in inner t []

and map ~f (Node (label, forest)) =
  Node (f label, List.map ~f:(map ~f) forest)

rose.mli

(** Simple implementation of multi-way trees aka Rose Trees. *)

type 'a tree = Node of ('a * 'a tree list)

val make : f:('a -> ('b * 'a list)) -> 'a -> 'b tree
(** [make ~f seeds] builds a tree from the projection function and a
    list of seed values. *)

val draw : f:('a -> string) -> 'a tree -> string list
(** [draw ~f tree] returns a neat two dimentional drawing of tree;
    shamelessly ripped of from Haskell 'containers' package. *)

val map : f:('a -> 'b) -> 'a tree -> 'b tree
(** [map ~f tree] map a function over each node of the [tree] and return
    a new tree with the results, returned by [f]. *)

val flatten : 'a tree -> 'a list
(** [flatten tree] returns a list of elements of a tree in a pre-order. *)