freer_ext2.ml

module Base = struct
  type _ operations = ..
end

module Freer : sig
  type _ t

  module Syntax : sig
    val ( let* ) : 'a t -> ('a -> 'b t) -> 'b t
  end

  module Handler : sig
    type nonrec 'a t = { handler : 'arg. 'arg Base.operations -> ('arg -> 'a t) -> 'a t }

    val tie : 'a t
    val ( <+> ) : ('a t -> 'a t) -> ('a t -> 'a t) -> 'a t -> 'a t
    val ( <+| ) : ('a t -> 'a t) -> ('a t -> 'a t) -> 'a t
  end

  module Operations : sig
    val bind : ('a -> 'b t) -> 'a t -> 'b t
    val app : 'a Base.operations -> 'a t
    val run : 'a Handler.t -> 'a t -> 'a
  end
end = struct
  open Base

  type _ t =
    | Pure : 'a -> 'a t
    | Impure : 'a operations * ('a -> 'b t) -> 'b t

  module Handler = struct
    type nonrec 'a t = { handler : 'arg. 'arg operations -> ('arg -> 'a t) -> 'a t }

    let tie = { handler = (fun _ _ -> failwith "bottom") }
    let ( <+> ) h1 h2 h3 = h1 (h2 h3)
    let ( <+| ) h1 h2 = h1 (h2 tie)
  end

  module Operations = struct
    let rec bind f m =
      match m with
      | Pure v -> f v
      | Impure (v, k) -> Impure (v, fun a -> bind f (k a))
    ;;

    let app a = Impure (a, fun x -> Pure x)

    let run Handler.{ handler } =
      let rec go = function
        | Pure a -> a
        | Impure (v, k) -> go @@ handler v k
      in
      go
    ;;
  end

  module Syntax = struct
    let ( let* ) m f = Operations.bind f m
  end
end

module type EFF = sig
  val handler' : 'a Freer.Handler.t -> 'a Freer.Handler.t
end

module State : sig
  include EFF

  module Operations : sig
    val ask : unit -> string Freer.t
    val put : string -> unit Freer.t
  end
end = struct
  open Base
  open Freer

  type _ operations += Get : string operations | Put : string -> unit operations

  module Operations = struct
    open Freer.Operations

    let ask () = app Get
    let put s = app (Put s)
  end

  let handler' =
    let open Handler in
    fun { handler = default } ->
      { handler =
          (let handler : type a. a operations -> (a -> 'b Freer.t) -> 'b Freer.t =
            fun v k ->
             match v with
             | Get ->
               let s = read_line () in
               k s
             | Put s ->
               print_endline s;
               k ()
             | e -> default e k
           in
           handler)
      }
  ;;
end

module Double : sig
  include EFF

  module Operations : sig
    val twice : unit -> unit Freer.t
  end
end = struct
  open Base
  open Freer

  type _ operations += Twice : unit operations

  module Operations = struct
    open Freer.Operations

    let twice () = app Twice
  end

  let handler' =
    let open Handler in
    let open Syntax in
    fun { handler = default } ->
      { handler =
          (let handler : type a. a operations -> (a -> 'b Freer.t) -> 'b Freer.t =
            fun v k ->
             match v with
             | Twice ->
               let* c = k () in
               k ()
             | e -> default e k
           in
           handler)
      }
  ;;
end

let exec () =
  Freer.(Operations.run Handler.(State.handler' <+| Double.handler'))
  @@
  let open State.Operations in
  let open Double.Operations in
  let open Freer.Syntax in
  let* a1 = ask () in
  let* a2 = ask () in
  let* () = twice () in
  put @@ a1 ^ a2
;;

let exec2 () =
  Freer.(Operations.run Handler.(Double.handler' <+| State.handler'))
  @@
  let open State.Operations in
  let open Double.Operations in
  let open Freer.Syntax in
  let* a1 = ask () in
  let* a2 = ask () in
  let* () = twice () in
  put @@ a1 ^ a2
;;