Rose trees implemented in OCaml via binary trees

rose.ml

[@@@warning "-8"]

(** Abstract interface for monoids.
  *)
module type Monoid_S = sig
  type t
  val zero : t
  val ( + ) : t -> t -> t

(* Required laws:
 *   x + zero = zero + x = x
 *   x + (y + z) = (x + y) + z   *)
end

(** Example operation on monoids.
  *)
let sum_by (type a m) (module M : Monoid_S with type t = m)
      (f : a -> m)
      (xs : a list) : m =
  List.fold_left (fun m x -> M.(m + f x)) M.zero xs

(* * *)

(** Abstract interface for rose trees.
  *)
module type Rose_S = sig
  type 'a t

  (** constructor *)
  val node : 'a            (* data *)
             -> 'a t list  (* children *)
             -> 'a t

  (** catamorphism. *)
  val fold : (module Monoid_S with type t = 'm)
             -> ('a -> 'm -> 'm)
             -> 'a t
             -> 'm

(* Required law:
 *   fold f (make x l)
 *   =
 *   f x (sum_by (cata f) l)   *)
end

(** Rose tree implementation via binary tree encoding.
  *)
module Rose : Rose_S = struct
  type 'a t =
    { data    : 'a
    ; sibling : 'a t option
    ; child   : 'a t option }

  let node data ts =
    { data
    ; sibling = None
    ; child   = List.fold_right
                  (fun l r -> Some { l with sibling = r })
                  ts
                  None }

  let fold (type m) (module M : Monoid_S with type t = m) f t =
    let open M in
    let rec accum = function
      | None -> zero
      | Some { data; sibling; child } ->
         f data (accum child) + accum sibling
    in
    accum (Some t)
end

(* * *)

(** Example monoid.  *)
module CommaSep = struct
  type t = string
  let zero = ""
  let ( + ) x y = match x, y with
    | _,  "" -> x
    | "", _  -> y
    | _,  _  -> x ^ ", " ^ y
end

let "1, 2, 3" =
  sum_by (module CommaSep) string_of_int
    [1; 2; 3]

(* Use-case: represent simple function-call expressions using rose trees.  *)

open Rose

let render =
  fold (module CommaSep)
    (fun hd args ->
      if args = "" then hd
      else Printf.sprintf "%s(%s)" hd args)

let _ =
  (* prints: "f(A, g(h(B, C), D))" *)
  ( node "f" [node "A" [];
              node "g" [node "h" [node "B" [];
                                  node "C" []];
                        node "D" []]] )
  |> render
  |> print_endline