module type Basic = sig
  type 'a t
  val bind : 'a t -> ('a -> 'b t) -> 'b t
  val return : 'a -> 'a t

  (* The [map] argument to [Monad.Make] says how to implement the monad's [map] function.
     [`Define_using_bind] means to define [map t ~f = bind t (fun a -> return (f a))].
     [`Custom] overrides the default implementation, presumably with something more
     efficient.

     Some other functions returned by [Monad.Make] are defined in terms of [map], so
     passing in a more efficient [map] will improve their efficiency as well. *)
  val map : [ `Define_using_bind
            | `Custom of ('a t -> f:('a -> 'b) -> 'b t)
            ]
end

module type Infix = sig
  type 'a t

  (** [t >>= f] returns a computation that sequences the computations represented by two
      monad elements.  The resulting computation first does [t] to yield a value [v], and
      then runs the computation returned by [f v]. *)
  val (>>=) : 'a t -> ('a -> 'b t) -> 'b t

  (** [t >>| f] is [t >>= (fun a -> return (f a))]. *)
  val (>>|) : 'a t -> ('a -> 'b) -> 'b t

end

module type S = sig
  (** A monad is an abstraction of the concept of sequencing of computations.  A value of
      type 'a monad represents a computation that returns a value of type 'a. *)
  include Infix

  module Monad_infix : Infix with type 'a t := 'a t

  (** [bind t f] = [t >>= f] *)
  val bind : 'a t -> ('a -> 'b t) -> 'b t

  (** [return v] returns the (trivial) computation that returns v. *)
  val return : 'a -> 'a t

  (** [map t ~f] is t >>| f. *)
  val map : 'a t -> f:('a -> 'b) -> 'b t

  (** [join t] is [t >>= (fun t' -> t')]. *)
  val join : 'a t t -> 'a t

  (** [ignore t] = map t ~f:(fun _ -> ()). *)
  val ignore : 'a t -> unit t

  val all : 'a t list -> 'a list t
  val all_ignore : unit t list -> unit t
end

module Make (M : Basic) : S with type 'a t := 'a M.t = struct

  let bind   = M.bind
  let return = M.return

  let map_via_bind ma ~f = M.bind ma (fun a -> M.return (f a))

  let map =
    match M.map with
    | `Define_using_bind -> map_via_bind
    | `Custom x -> x

  module Monad_infix = struct

    let (>>=) = bind

    let (>>|) t f = map t ~f
  end

  include Monad_infix

  let join t = t >>= fun t' -> t'

  let ignore t = map t ~f:(fun _ -> ())

  let all =
    let rec loop vs = function
      | [] -> return (List.rev vs)
      | t :: ts -> t >>= fun v -> loop (v :: vs) ts
    in
    fun ts -> loop [] ts

  let rec all_ignore = function
    | [] -> return ()
    | t :: ts -> t >>= fun () -> all_ignore ts

end

(**
   Multi parameter monad.
   The second parameter get unified across all the computation. This is used
   to encode monads working on a multi parameter data structure like
   ([('a,'b result)]).
*)
module type Basic2 = sig
  type ('a, 'd) t
  val bind : ('a, 'd) t -> ('a -> ('b, 'd) t) -> ('b, 'd) t
  val map : [ `Define_using_bind
            | `Custom of (('a, 'd) t -> f:('a -> 'b) -> ('b, 'd) t)
            ]
  val return : 'a -> ('a, _) t
end

(** Same as Infix, except the monad type has two arguments. The second is always just
    passed through. *)
module type Infix2 = sig
  type ('a, 'd) t
  val (>>=) : ('a, 'd) t -> ('a -> ('b, 'd) t) -> ('b, 'd) t
  val (>>|) : ('a, 'd) t -> ('a -> 'b) -> ('b, 'd) t
end

(** The same as S except the monad type has two arguments. The second is always just
    passed through. *)
module type S2 = sig
  include Infix2

  module Monad_infix : Infix2 with type ('a, 'd) t := ('a, 'd) t

  val bind : ('a, 'd) t -> ('a -> ('b, 'd) t) -> ('b, 'd) t

  val return : 'a -> ('a, _) t

  val map : ('a, 'd) t -> f:('a -> 'b) -> ('b, 'd) t

  val join : (('a, 'd) t, 'd) t -> ('a, 'd) t

  val ignore : (_, 'd) t -> (unit, 'd) t

  val all : ('a, 'd) t list -> ('a list, 'd) t

  val all_ignore : (unit, 'd) t list -> (unit, 'd) t
end

module Check_S2_refines_S (X : S) : (S2 with type ('a, 'd) t = 'a X.t) =
struct
  type ('a, 'd) t = 'a X.t
  include struct
    open X
    let (>>=)      = (>>=)
    let (>>|)      = (>>|)
    let bind       = bind
    let return     = return
    let map        = map
    let join       = join
    let ignore     = ignore
    let all        = all
    let all_ignore = all_ignore
  end
  module Monad_infix = struct
    open X.Monad_infix
    let (>>=) = (>>=)
    let (>>|) = (>>|)
  end
end

module Make2 (M : Basic2) : S2 with type ('a, 'd) t := ('a, 'd) M.t = struct

  let bind   = M.bind
  let return = M.return

  let map_via_bind ma ~f = M.bind ma (fun a -> M.return (f a))

  let map =
    match M.map with
    | `Define_using_bind -> map_via_bind
    | `Custom x -> x

  module Monad_infix = struct

    let (>>=) = bind

    let (>>|) t f = map t ~f
  end

  include Monad_infix

  let join t = t >>= fun t' -> t'

  let ignore t = map t ~f:(fun _ -> ())

  let all =
    let rec loop vs = function
      | [] -> return (List.rev vs)
      | t :: ts -> t >>= fun v -> loop (v :: vs) ts
    in
    fun ts -> loop [] ts

  let rec all_ignore = function
    | [] -> return ()
    | t :: ts -> t >>= fun () -> all_ignore ts

end
