Arrow.ml

module StringMap = Map.Make (String)

module type MONAD = sig
  type 'a t

  val bind : 'a t -> f:('a -> 'b t) -> 'b t

  val return : 'a -> 'a t

  module Let_syntax : sig
    val bind : 'a t -> f:('a -> 'b t) -> 'b t
  end
end

module type BUILD_MONAD = sig
  include MONAD

  type location

  val unwrap : 'a t -> 'a

  val build : location -> 'a t -> 'a t

  val addRule : location -> (location -> 'a) -> 'a t
end

module StateMonad : BUILD_MONAD = struct
  type location = string

  type 'a t = {rules: ((location -> 'a) -> 'a) StringMap.t; result: 'a}

  let return v = {rules= StringMap.empty; result= v}

  let bind m ~f = f m.result

  let unwrap m = m.result

  let rec build location m =
    match StringMap.get location m.rules with
    | None -> failwith ("No rule for location: " ^ location)
    | Some rule -> rule (fun location -> build location m |> unwrap)
end