(*
 * tree_monadize.ml
 *
 * If you've got some block of code that uses `unit`s and `bind`s, and you
 * want to interpret it alternately using this monad, that monad, or another
 * monad, you can use OCaml's module system. You'd write your code like this:
 *) 

module Reader_monad = struct
    (* change this to suit your needs *)
    type env = int -> int;;

    type 'a monad = env -> 'a;;
    let unit a : 'a monad = fun e -> a;;
    let bind (u : 'a monad) (f : 'a -> 'b monad) : 'b monad =
      fun e -> f (u e) e;;
end

module State_monad = struct
    (* change this to suit your needs *)
    type store = int;;

    type 'a monad = store -> 'a * store;;
    let unit a : 'a monad  = fun s -> (a, s);;
    let bind (u : 'a monad) (f : 'a -> 'b monad) : 'b monad =
      fun s -> (let (a, s') = u s in (f a) s');;
end

module List_monad = struct
    type 'a monad = 'a list;;
    let unit a : 'a monad = [a];;
    let bind (u: 'a monad) (f : 'a -> 'b monad) : 'b monad =
      List.concat(List.map f u);;
end

(*
 * Then you can replace code that looks like this:
 *     ... reader_bind ...
 * with code that looks like this:
 *     ... Reader_monad.bind ...
 * and the latter can be reformulated like this:
 *     let open Reader_monad in ... bind ...
 * or equivalently, like this:
 *     Reader_monad.(... bind ...)
 * Then you can use literally the same `... bind ...` code when writing instead:
 *     State_monad.(... bind ...)
 *)

(* That's great, however it still requires us to repeat the
 * `... bind ...` code every time we want to change which monad we're working
 * with. Shouldn't there be a way to _parameterize_ the `... bind ...` code
 * on a monad, so that we only have to write the `... bind ...` code once,
 * but can invoke it alternately with the Reader_monad supplied as an
 * argument, or the State_monad, or another?
 *
 * There is a way to do this, but it requires putting the `... bind ...` code in
 * its own module, and making that module parameterized on some X_monad
 * module. Also we have to explicitly declare what commonality we're expecting
 * from X_monad modules we're going to use as parameters. We'll explain how to
 * do this in a moment.
 *
 * As preparation, a general observation:
 * 'a and so on are type variables in OCaml; they stand for arbitrary types.
 * What if you want a variable for a type constructor? For example, you want to
 * generalize this pattern:
 *      type ('a) t1 = 'a -> ('a) list
 *      type ('a) t2 = 'a -> ('a) option
 *      type ('a) t3 = 'a -> ('a) reader
 * and so on? OCaml won't let you do this:
 *      type ('a, 'b) t = 'a -> ('a) 'b
 * To generalize on the 'b position, we instead have to use OCaml's modules,
 * and in particular its ability to make modules parameterized on other modules
 * (OCaml calls these parameterized modules Functors, but that name is also
 * used in other ways in this literature, so I won't give in to it.)
 *
 * Here's how you'd have to define the t type from above:
 *      module T_maker(
 *      (* A sig...end block specifies the type of a module
 *       * What we're doing here is specifying the type of the 
 *       * module parameter that will choose
 *       * whether b = list or b = option or b = reader...
 *       * This module parameter may supply values as well as types *)
 *      X : sig
 *          type ('a) b
 *      end
 *      ) = 
 *      (* A struct...end block gives a module value
 *       * What we're doing here is building a new module that makes
 *       * use of the module that was supplied as X *)
 *      struct
 *          type ('a) t = 'a -> ('a) X.b
 *      end
 * And here's how you'd use it:
 *      module T_list = T_maker(struct type 'a b = 'a list end);;
 *      type 'a t1 = 'a T_list.t;;
 *      module T_option = T_maker(struct type 'a b = 'a option end);;
 *      type 'a t2 = 'a T_option.t;;
 *      (* and so on *)
 *
 * I know, it seems unnecessarily complicated. Nonetheless, that's how it
 * works. And that is also the technique we'll use to make our
 * `... bind ...` code parametric on some X_monad module.
 *)

type 'a tree = Leaf of 'a | Node of ('a tree) * ('a tree);;

let t1 = Node
           (Node
             (Leaf 2, Leaf 3),
            Node
             (Leaf 5,
              Node
                (Leaf 7, Leaf 11)));;


module Tree_monadizer(X : sig
  (* the module we're using as a parameter has to supply function values
   * for unit and bind, as well as a monadic type constructor *)
  type 'a monad
  val unit : 'a -> 'a monad
  val bind : 'a monad -> ('a -> 'b monad) -> 'b monad
end) = struct
  let rec monadize (t: 'a tree) (f: 'a -> 'b X.monad) : 'b tree X.monad =
    match t with
    | Leaf a -> X.bind (f a) (fun b -> X.unit (Leaf b))
    | Node(l, r) ->
        X.bind (monadize f l) (fun l' ->
          X.bind (monadize f r) (fun r' ->
            X.unit (Node (l', r'))))
end;;


(* Now we supply the Reader monad as a parameter to Tree_monadizer.
 * We'll get back a module TreeReader that contains a single value,
 * the monadize function specialized to the Reader monad *)
module TreeReader = Tree_monadizer(Reader_monad);;


(* Make a TreeState module containing monadize specialized to the State monad *)
module TreeState =  Tree_monadizer(State_monad);;


(* Make a TreeList module containing monadize specialized to the List monad *)
module TreeList =  Tree_monadizer(List_monad);;


(* The Continuation monad is a bit more complicated *)
module Continuation_monad = struct
    type ('a,'r) monad = ('a -> 'r) -> 'r;;
    let unit a : ('a,'r) monad = fun k -> k a;;
    let bind (u: ('a,'r) monad) (f: 'a -> ('b,'r) monad) : ('b,'r) monad =
      fun k -> u (fun a -> f a k);;
end

(* Since the Continuation monad is parameterized on two types---it's
 * ('a,'r) cont not ('a) cont---we can't match the type ('a) monad that
 * Tree_monadizer expects in its parameter. So we have to make a different
 * Tree_monadizer2 that takes a ('a,'z) monad type constructor in its
 * parameter instead *)
module Tree_monadizer2(X : sig
  type ('a,'z) monad
  val unit : 'a -> ('a,'z) monad
  val bind : ('a,'z) monad -> ('a -> ('b,'z) monad) -> ('b,'z) monad
end) = struct
  (* the body of the monadize function is the same; the only difference is in
   * the types *)
  let rec monadize (t: 'a tree) (f: 'a -> ('b,'x) X.monad) : ('b tree,'x) X.monad =
    match t with
    | Leaf a -> X.bind (f a) (fun b -> X.unit (Leaf b))
    | Node(l, r) ->
        X.bind (monadize f l) (fun l' ->
          X.bind (monadize f r) (fun r' ->
            X.unit (Node (l', r'))))
end;;

(* Make a TreeCont module containing monadize specialized to the Cont monad *)
module TreeCont =  Tree_monadizer2(Continuation_monad);;



(* 
 * Here are all the examples from
 * http://lambda.jimpryor.net/manipulating_trees_with_monads/
 *)


let int_readerize : int -> int Reader_monad.monad =
  fun (a : int) -> fun (env : int -> int) -> env a;;

(* int_readerize takes an int and returns a Reader monad that
 * "looks up" that int in the environment (i.e. modifies it)
 * this is structurally parallel to the function `lookup` we used
 * before to "look up" variables in the environment *)

(* double each leaf *)
let env = fun i -> i + i in
TreeReader.monadize t1 int_readerize env;;

(* You can also avoid declaring a separate toplevel TreeReader module
 * (or even a separate Reader_monad module) by using one of these forms:
 *     ...
 *     let module T = Tree_monadizer(Reader_monad) in
 *     T.monadize t1 int_readerize env;;
 * or:
 *     ...
 *     let env = fun i -> i + i in
 *     let module Monad = struct
 *       type env = int -> int;;
 *       type 'a monad = env -> 'a;;
 *       let unit a : 'a monad = fun e -> a;;
 *       let bind (u : 'a monad) (f : 'a -> 'b monad) : 'b monad =
 *         fun e -> f (u e) e;;
 *     end in
 *     let module T = Tree_monadizer(Monad) in
 *     T.monadize t1 int_readerize env;;
 * or:
 *     ...
 *     let module T = Tree_monadizer(struct
 *       type env = int -> int;;
 *       type 'a monad = env -> 'a;;
 *       let unit a : 'a monad = fun e -> a;;
 *       let bind (u : 'a monad) (f : 'a -> 'b monad) : 'b monad =
 *         fun e -> f (u e) e;;
 *     end) in
 *     T.monadize t1 int_readerize env;;
 *)


(* square each leaf *)
let env = fun i -> i * i in
TreeReader.monadize t1 int_readerize env;;



let incrementer : int -> int State_monad.monad =
  fun (a : int) -> fun s -> (a, s+1);;

(* incrementer takes an 'a and returns it wrapped in a
 * State monad that increments the store *)

(* count leaves *)
let initial_store = 0 in
TreeState.monadize t1 incrementer initial_store;;



(* replace leaves with list *)
TreeList.monadize t1 (fun i -> [ [i;i*i] ]);;



(* do nothing *)
let initial_continuation = fun t -> t in
TreeCont.monadize t1 Continuation_monad.unit initial_continuation;;

(* convert tree to list of leaves *)
let initial_continuation = fun t -> [] in
TreeCont.monadize t1 (fun a k -> a :: k a) initial_continuation;;

(* square each leaf using continuation *)
let initial_continuation = fun t -> t in
TreeCont.monadize t1 (fun a k -> k (a*a)) initial_continuation;;

(* replace leaves with list, using continuation *)
let initial_continuation = fun t -> t in
TreeCont.monadize t1 (fun a k -> k [a; a*a]) initial_continuation;;

(* count leaves, using continuation *)
let initial_continuation = fun t -> 0 in
TreeCont.monadize t1 (fun a k -> 1 + k a) initial_continuation;;

(*
(* Tree monad *)

(* type 'a tree defined above *)
let tree_unit (a: 'a) : 'a tree = Leaf a;;
let rec tree_bind (u : 'a tree) (f : 'a -> 'b tree) : 'b tree =
    match u with
    | Leaf a -> f a
    | Node (l, r) -> Node (tree_bind l f, tree_bind r f);;

type ('a) treeT_reader =
    'a tree reader;;

let unit (a: 'a) : 'a tree reader =
    reader_unit (Leaf a);;

let rec bind (u : 'a tree_reader) (f : 'a -> ('b, M) tree) : ('b, M) tree =
    match u with
    | Leaf a -> M.bind (f a) (fun b -> M.unit (Leaf b))
    | Node (l, r) -> M.bind (bind l f) (fun l' ->
                        M.bind (bind r f) (fun r' ->
                            M.unit (Node (l', r'));;

 *)