Monads in OCaml (chapter 11)

(** The OCaml signature for a monad *)

module type MONAD = sig
  type 'a mon  (* the type of monadic computations producing a value of type 'a *)
  val ret: 'a -> 'a mon
  val bind: 'a mon -> ('a -> 'b mon) -> 'b mon
end

(** The nondeterminism monad (list monad) *)

module Nondet = struct
  type 'a mon = 'a list
  let ret x = [x]
  let rec bind m f =
    match m with
    | [] -> []
    | x :: m -> f x @ bind m f
  let choose (m: 'a mon) (n: 'a mon) : 'a mon = m @ n
  let fail : 'a mon = []
end

(** The exception monad *)

module Exception = struct
  type 'a mon = V of 'a | E of exn
  let ret x = V x
  let bind m f =
    match m with V x -> f x | E _ -> m
  let raise (e: exn) : 'a mon = E e
  let handle (m: 'a mon) (f: exn -> 'a mon) : 'a mon =
    match m with V _ -> m | E e -> f e
end

(** The continuation monad *)

module Continuation = struct
  type res = unit    (* the result type *)
  type 'a mon = ('a -> res) -> res
  let ret x = fun k -> k x
  let bind m f = fun k -> m (fun x -> f x k)
  type 'a cont = 'a -> res
  let callcc (f: 'a cont -> 'a mon) : 'a mon = fun k -> f k k
  let throw (k: 'a cont) (v: 'a) : 'b mon = fun k' -> k v
end

(** An implementation of a polymorphic store, for use with the state monad *)

module Store = struct
  module IMap = Map.Make(struct type t = int let compare = compare end)
  type value = ..
  type t = { nextkey: int; contents: value IMap.t }
  type 'a ref = {
    key: int;
    inj: 'a -> value;
    proj: value -> 'a
  }
  let empty = { nextkey = 0; contents = IMap.empty }
  let newref (type a) (v: a) (s: t) : a ref * t =
    let module M = struct type value += V of a end in
    let r =
      { key = s.nextkey;
        inj = (fun x -> M.V x);
        proj = (function M.V x -> x | _ -> assert false) } in
    let s' =
      { nextkey = s.nextkey + 1;
        contents = IMap.add r.key (r.inj v) s.contents } in
    (r, s')
  let get (r: 'a ref) (s: t) : 'a =
    r.proj (IMap.find r.key s.contents)
  let set (r: 'a ref) (v: 'a) (s: t) : t =
    { s with contents = IMap.add r.key (r.inj v) s.contents }
end

(** The state monad *)

module State = struct
  type 'a ref = 'a Store.ref
  type 'a mon = Store.t -> 'a * Store.t
  let ret x = fun s -> (x, s)
  let bind m f = fun s -> let (x, s1) = m s in f x s1
  let newref (v: 'a) : 'a ref mon = fun s -> Store.newref v s
  let get (r: 'a ref) : 'a mon = fun s -> (Store.get r s, s)
  let set (r: 'a ref) (v: 'a) : unit mon = fun s -> ((), Store.set r v s)
end

(** The reader monad (environment monad) *)

module Reader = struct
  type key = string   (* for example *)
  type value = int    (* for example *)
  type env = (key * value) list
  type 'a mon = env -> 'a
  let ret x = fun e -> x
  let bind m f = fun e -> f (m e) e
  let lookup (k: key) : value option mon = fun e -> List.assoc_opt k e
  let assign (k: key) (v: value) (m: 'a mon) : 'a mon = fun e -> m ((k, v) :: e)
end

(** A type of ropes, for use with the writer monad. *)

module Rope = struct
  type t = Zero | One of string | Two of t * t
  let empty = Zero
  let of_string s = One s
  let concat r1 r2 = Two(r1, r2)
  let extract (r: t) : string =
    let b = Buffer.create 100 in
    let rec extract r =
      match r with
      | Zero -> ()
      | One s -> Buffer.add_string b s
      | Two(r1, r2) -> extract r1; extract r2 in
    extract r; Buffer.contents b
end

(** The writer monad (output monad) *)

module Writer = struct
  type 'a mon = 'a * Rope.t
  let ret x = (x, Rope.empty)
  let bind m f =
    let (x, s1) = m in let (y, s2) = f x in (y, Rope.concat s1 s2)
  let print (s: string) : unit mon = ((), Rope.of_string s)
end

(** Probability monads *)

module Distribution = struct
  type 'a mon = ('a * float) list
  let ret x = [(x, 1.0)]
  let rec bind m f =
    match m with
    | [] -> []
    | (x, p) :: m -> List.map (fun (y, p') -> (y, p *. p')) (f x) @ bind m f
  let flip (p: float) : bool mon = [(true, p); (false, 1. -. p)]
end

module Expectation = struct
  type 'a mon = ('a -> float) -> float
  let ret x = fun k -> k x
  let bind m f = fun k -> m (fun x -> f x k)
  let flip (p: float) : bool mon = fun k -> p *. k true +. (1. -. p) *. k false
end

(** The undoable monad *)

module Undoable = struct
  type 'a mon = 'a * (unit -> unit)
  let ret x = (x, fun () -> ())
  let bind m f =
    let (x, undo1) = m in
    let (y, undo2) = f x in
    (y, (fun () -> undo2(); undo1()))
  let undo (m: 'a mon) : unit =
    let (_x, undo) = m in undo ()
  let ref_set (r: 'a ref) (new_v: 'a) : unit mon =
    let old_v = !r in r := new_v; ((), fun () -> r := old_v)
  let array_set (a: 'a array) (idx: int) (new_v: 'a) : unit mon =
    let old_v = a.(idx) in a.(idx) <- new_v; ((), fun () -> a.(idx) <- old_v)
end

(** The freer monad *)

type _ eff = ..

module Freer = struct
  type 'a mon =
    | Pure: 'a -> 'a mon
    | Op: 'b eff * ('b -> 'a mon) -> 'a mon
  let ret x = Pure x
  let rec bind m f =
    match m with
    | Pure x -> f x
    | Op(phi, k) -> Op(phi, fun x -> bind (k x) f)
  let perform (phi: 'b eff) : 'b mon =
    Op(phi, ret)
  type ('a, 'b) deep_handler =
    { retc: 'a -> 'b;
      effc: 'c. 'c eff -> (('c -> 'b) -> 'b) option }
  let rec fold (h: ('a, 'b) deep_handler) (m: 'a mon) : 'b =
    match m with
    | Pure x -> h.retc x
    | Op(phi, k) ->
        match h.effc phi with
        | Some g -> g (fun x -> fold h (k x))
        | None -> Op(phi, fun x -> fold h (k x))
  type ('a, 'b) shallow_handler =
    { retc: 'a -> 'b;
      effc: 'c. 'c eff -> (('c -> 'a mon) -> 'b) option }
  let case (h: ('a, 'b) shallow_handler) (m: 'a mon) : 'b =
    match m with
    | Pure x -> h.retc x
    | Op(phi, k) ->
        match h.effc phi with
        | Some g -> g k
        | None -> Op(phi, k)
end