Free monads in OCaml, for great referential transparency

Free.ml

module type Functor = sig
    type 'a t
    val map : ('a -> 'b) -> 'a t -> 'b t
end

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

module Free = functor (F : Functor) -> struct
    type 'a t = Return of 'a
              | Wrap of ('a t) F.t
    
    let return a = Return a

    let rec bind m f = match m with
        | Return x -> f x
        | Wrap x -> Wrap (F.map (fun m -> bind m f) x)

    let join m = bind m (fun x -> x)
    let map f m = bind m (fun x -> return (f x))
end

module MonadUtils = functor (M : Monad) -> struct
    type 'a t = 'a M.t
    let map = M.map
    let return = M.return
    let join = M.join

    let bind m f = join (map f m)
    let seq m n = bind m (fun _ -> n)
end

module IOOp = struct
    type 'a t = Print_string of (string * 'a)
              | Read_string of (string -> 'a)
    let map f x = match x with
        | Print_string (str, cont) -> Print_string (str, f cont)
        | Read_string cont -> Read_string (fun str -> f (cont str))
end

module FreeIO = Free(IOOp)
module IO = MonadUtils(FreeIO)

(* Run a computation in the IO monad *)
module IOInterp = struct
    let rec unsafePerform m = match m with
        | FreeIO.Return x -> x
        | FreeIO.Wrap f -> match f with
            | IOOp.Print_string (msg, cont) -> print_string msg; unsafePerform cont
            | IOOp.Read_string cont -> unsafePerform (cont (read_line ()))
end

let main =
    let print_string msg = FreeIO.Wrap (IOOp.Print_string (msg, FreeIO.Return ())) in
    let read_string = FreeIO.Wrap (IOOp.Read_string (fun msg -> FreeIO.Return msg)) in
    let (>>=) m n = IO.bind m n in
    let (>>) m n = IO.seq m n in

    print_string "What's your name? " >>
    read_string >>= fun name ->
    print_string "Hello, " >>
    print_string name >>
    print_string "\n"

let _ = IOInterp.unsafePerform main