state_monad.ml

module type Monad = sig
  type 'a t
  val return : 'a -> 'a t
  val (>>=) : 'a t -> ('a -> 'b t) -> 'b t
end

module type StateS = sig
  include Monad
  type state
  val get : state t
  val put : state -> unit t
  val runS : 'a t -> init:state -> state * 'a
  val (let*) : 'a t -> ('a -> 'b t) -> 'b t
  val (and*) : 'a t -> 'b t -> ('a * 'b) t
end

module StateM (S : sig type t end) : StateS with type state = S.t
  = struct
  type state = S.t
  type 'a t = state -> state * 'a
  let return x = fun s -> s,x
  let (>>=) (type a) (type b) (x : a t) (f : a -> b t) (s : state) : state * b =
    let s', v = x s in
    f v s'
  let get s = s,s
  let put new_s _ = new_s, ()
  let runS f ~init = f init
  let (let*) = (>>=)
  let (and*) f1 f2 = fun state ->
    let s,va = f1 state in
    let s,vb = f2 state in
    s,(va,vb)
end

type 'a tree = Empty | Tree of 'a tree * 'a * 'a tree

module IState = StateM(struct type t = int end)

let fresh_name : string IState.t =
  let open IState in
  let* i = get in
  let* () = put (succ i) in
  return (string_of_int i)

let rec label_tree : 'a tree -> string tree IState.t =
  let open IState in function
  | Empty -> return Empty
  | Tree (l,_,r) ->
      let* l' = label_tree l in
      let* r' = label_tree r in
      let* name = fresh_name in
      return @@ Tree (l', name, r')

let unlabeled = label_tree (Tree (Tree (Empty, 3, Empty), 5, Tree (Empty, 4, Tree (Empty, 2, Empty))))

let make_labeled_tree () = IState.runS unlabeled ~init:0