(* state-passing style -> monadic-style programming
   great when you're in a no-state world, like Haskell or
     client-server computing

   The interim steps below are admittedly tricky, but where we start
   is straightforward and where we end up is elegant and easy-to-use.
 *)

 module List =
  let assoc key l = List.pick (fun (k,v) -> if k = key then Some v else None) l

let print_string (s : string) = System.Console.Write s
let print_newline () = System.Console.WriteLine ""

let string_of_int i = sprintf "%i" i

(* an interface for a functional heap that returns an answer and a new heap *)
(* written by a guru who understands this will be a convenient interface *)
let empty_heap = []
let lookup str heap = ((try List.assoc str heap with _ -> 0), heap)
let update str v heap = ((), (str,v)::heap)
(* ... could have more operations ... *)

(* increment z, if original z is positive set x to y, else set x to 37 *)
(* written in a stylized way to make sure things are sequenced just right,
   manually shadowing each heap with the next one *)
let example1 heap =
  let x1,heap = lookup "z" heap in
  let x2,heap = update "z" (x1+1) heap in
  let x3,heap = (if x1 > 0 then lookup "y" heap else (37,heap)) in
  update "x" x3 heap

(* two lines below written by guru *)
(* f1: function from heap to result and heap
   f2: function from arg and heap to result and heap
   result: function from heap to sequenced result and heap
*)
let bind f1 f2 = (fun heap -> let x,heap' = f1 heap in f2 x heap')

(* in monad-land called "return", but this is confusing *)
let ret e = (fun heap -> (e,heap))

(* naively using the helper functions looks like a mess, but bear with me *)
let example2 heap =
  (bind (fun heap -> lookup "z" heap)
        (fun x1 ->
   bind (fun heap -> update "z" (x1+1) heap)
        (fun x2 ->
   bind (fun heap -> if x1 > 0
                       then lookup "y" heap
                       else ret 37 heap)
         (fun x3 ->
           (fun heap -> update "x" x3 heap)))))
    heap

(* thanks to fun x -> e1 ... en x being e1 ... en, we can now "hide"
  _every_ explicit use of heap-passing (just like in imperative programming)
 *)
let example3 =
   bind (lookup "z")
        (fun x1 ->
   bind (update "z" (x1+1))
              (fun x2 ->
     bind (if x1 > 0
           then lookup "y"
           else ret 37)
       (fun x3 ->
         (update "x" x3))))

(* and now let's not change anything except spacing and indentation.
   The result looks like imperative programming in "Hebrew":
     * bind starts each line
        (vs. semicolon at end of line)
     * variable holding result of operation at end
        (vs. declaration at beginning)
     * the ret part is mandatory though -- "lifting into the heap-passing"
 *)
let example4 =
  bind (lookup "z")        (fun x1 ->
    bind (update "z" (x1+1)) (fun x2 ->
      bind (if x1 > 0 then lookup "y"
            else ret 37) (fun x3 ->
               (update "x" x3))))

(* syntactic sugar (not in Caml; see Haskell)
    x <- e1 ; e2 ==> bind e1 (fun x -> e2)
   now programmers think they're in an "imperative" setting

let example5 =
   x1 <- lookup "z" ;
   x2 <- update "z" (x1+1) ;
   x3 <- if x1 > 0 then lookup "y" else ret 37 ;
   update "x" x3

   actually also have e1; e2 ==> bind e1 (fun _ -> e2), letting us write

let example6 =
   x1 <- lookup "z" ;
   update "z" (x1+1) ;
   x3 <- if x1 > 0 then lookup "y" else ret 37 ;
   update "x" x3

*)

let pi i = print_string (string_of_int i); print_newline ()
let test f = pi (fst (lookup "x" (snd (f empty_heap))))

let _ = test example1
let _ = test example2
let _ = test example3
let _ = test example4

(* F# only! *)

(* This is the general "pattern" for creating your own way to "reuse"
   workflow syntax *)

type HeapBuilder() =
  (* Here we are using F#'s ability to define a new type of computation
     expression. Here our monad type M<'U> is (heap -> ('U * heap)). It isn't hard
     to see that the types of the bind and ret functions previously are pretty close
     to that expected by F#.*)

  (*
    The Bind function is expected to have the type (M<'U> * ('U -> M<'T>)) -> M<'U>.
    If you expand out our definition of M<'U> you get:

    (heap -> ('U * heap)) * ('U -> (heap -> ('T * heap))) -> (heap -> ('T * heap))

    This is really close to our bind function's type, which is

    (heap -> ('U * heap)) -> ('U -> (heap -> ('T * heap))) -> (heap -> ('T * heap))

    Our implementation of Bind(susp, func) just smooths over this different calling convention. *)
  member this.Bind(susp, func) = bind susp func

  (* Our implementation of ret has exactly the expected type, no uncurrying required! *)
  member this.Return(x) = ret x

  (* We need this to tell F# the value x is the final result of the computation *)
  member this.ReturnFrom(x) = x

let heap_monad = new HeapBuilder()

(* Now let's use our new computation expression: *)
let example5 =
  heap_monad
    {
      (* There's nothing magic happening here. Behind the scenes, F# transforms this: *)
      let! x1 = lookup "z"
      (* Into:
          heap_monad.Bind(lookup "z",
            fun x1 -> heap_monad.Bind(update "z" (x1+1),
              fun x2 -> ...))

         Notice that the translation occurs recursively.

         Now, if we expand out the definition of heap_monad.Bind we get:

         bind (lookup "z") (fun x1 -> bind (update "z" (x1 + 1) (fun x2 -> ...)

         Which is exactly what we saw in example4!
       *)
      let! x2 = update "z" (x1+1)

      (* F# does not automatically perform the desugaring transformation on
         the righthand sides of let! expressions, which is why we have the nested
         heap_monad { } block here. *)
      let! x3 = heap_monad {
          if x1 > 0 then
            (* return! is transformed into ReturnFrom which is the identity function.
               But lookup "y" has the expected type (heap -> (int * heap)) so all is well *)
            return! lookup "y"

           (* The return keword is transformed into heap_monad.Return(37), which itself is just
             ret 37, which we've seen earlier. This lifts the int value 37 into the
             expected monadic type*)
          else return 37
        }

      (* As before, this return! indicates the computation stops here *)
      return! update "x" x3
    }

(* If you're still confused about how the above works, consider applying
   the transformation rules on the "Computation Expressions" page linked from the
   class homepage on the above code until you've convinced yourself it is nearly
   equivalent to example4 *)

(* Same thing without the comments: *)
let example6 =
  heap_monad
    {
      let! x1 = lookup "z"
      let! x2 = update "z" (x1+1)
      let! x3 = heap_monad {
          if x1 > 0 then lookup "y"
          else return 37
        }
      return! update "x" x3
    }

let _ = test example5
let _ = test example6