Reverse State Monad in OCaml.

RevState.ml

(* technically this is the lazy state monad.  In particular, note the return type of get *)
module type STATE_MONAD = 
  sig
    type ('a,'s) st

    val return : 'a -> ('a,'s) st
    val bind : ('a,'s) st -> ('a -> ('b, 's) st) -> ('b, 's) st
    val get : ('s Lazy.t, 's) st
    val modify : ('s -> 's) -> (unit, 's) st
    val put : 's -> (unit, 's) st
    val lazyPut : 's Lazy.t -> (unit, 's) st

    val run : ('a, 's) st -> 's -> 'a * 's
  end

(* A few convenience functions *)
module LazyUtils =
  struct
    type 'a l = 'a Lazy.t

    let fst (p : ('a * 'b) l) : 'a l =
      lazy (match p
            with lazy (x,y) -> x)

    let snd (p : ('a * 'b) l) : 'b l =
      lazy (match p
            with lazy (x,y) -> y)

    (* collapse a lazy lazy value to a lazy value.
       The important bit here is to only force the outer computation
       within a lazy expression.  The trick to laziness is to never ever eta-contract.
     *)
    let join (xll : 'a l l) : 'a l = lazy (Lazy.force (Lazy.force xll))

  end

module RevState : STATE_MONAD =
  struct
    type 's l = 's Lazy.t

    type ('a, 's) st = 's l -> 'a * 's l

    let return x s = (x, s)

    let bind (mx : ('a, 's) st) (f : ('a -> ('b, 's) st)) (s : 's l) : ('b * 's l) =
      (* conceptually we want

             let rec (lazy (y, s'')) = lazy (mx s')
                 and (lazy (z, s')) = lazy (f y s)
             in (force z, s'')

         but that's not legal Caml.
         
         So instead we get back lazy pairs of type ('a * 's l) l and we
         lazily project out the pieces that we need.
       *)
      let rec ys'' = lazy (mx (LazyUtils.join (LazyUtils.snd zs')))
        and (zs' : ('b * 's l) l) = lazy (f (Lazy.force (LazyUtils.fst ys'')) s)
      in (Lazy.force (LazyUtils.fst zs'), LazyUtils.join (LazyUtils.snd ys''))

    let get st = (st, st)

    let modify f st =
      let st' = lazy (f (Lazy.force st))
      in ((), st')

    let put st = modify (fun _ -> st)

    let lazyPut st _ = ((), st)

    let run comp st = 
      let (ans, lazy st) = comp (Lazy.from_val st)
      in (ans, st)
  end

module SimpleExample =
  struct
    open RevState

    let example = bind get (fun ans ->
                            bind (put "a") (fun () -> return ans))
    (* try:  RevState.run SimpleExample.example "xyz" *)
  end

(* a module of lazy infinite lists, just for the upcoming example *)
module LazyList =
  struct
    type 'a lazy_list = | Cons of 'a * 'a lazy_list Lazy.t

    let scanl (f : 'a -> 'b -> 'b) : ('b -> 'a lazy_list -> 'b lazy_list) =
      let rec go y0 (Cons (x,xs)) =
        Cons (y0, lazy (go (f x y0) (Lazy.force xs)))
      in go
           
    let rec ones = Cons (1, lazy ones)
                                            
    let rec take n (Cons (x, xs) : 'a lazy_list) : 'a list =
      if n = 0
      then []
      else x :: (take (n - 1) (Lazy.force xs))
  end

(* transcibed from https://lukepalmer.wordpress.com/2008/08/10/mindfuck-the-reverse-state-monad/ *)
module Fibs =
  struct
    open RevState
    let cumulativeSums = LazyList.scanl (+) 0

    let fibs_comp = 
          bind get (fun fibs ->
          bind (modify cumulativeSums) (fun () ->
          bind (lazyPut (lazy (LazyList.Cons (1, fibs)))) (fun () ->
          return fibs)))

    let computeFibs = let (_, s) = run fibs_comp LazyList.ones
                         in s

    (* try:   LazyList.take 10 computeFibs *)
  end