open Effect
open Effect.Deep

type _ Effect.t += Xchg: int -> int t

let comp1 () = let a= Xchg 0 in let x= perform a in let b = Xchg 1 in let y = perform b in x+ y
let comp2 () = let _ = perform (Xchg 0) in raise Not_found

let comp3 () = let _ = perform (Xchg 0) in int_of_string "fdjsl"

let handle comp =
(*  try*)
Format.printf "%d@."  @@
match_with comp ()
{ retc = (fun x -> x - 30);
  exnc = (fun _ -> 42);
  effc = fun (type a) (eff: a t) ->
    match eff with
    | Xchg n -> Some (fun (k: (a, _) continuation) ->
        continue k (n+17))
    | _ -> None }
(*with Not_found -> assert false*)

let () = handle comp1; handle comp2; handle comp3

type 'a status =
  Complete of 'a
| Suspended of {msg: int; cont: (int, 'a status) continuation}


let step (f : unit -> 'a) () : 'a status =
  match_with f ()
  { retc = (fun v -> Complete v);
    exnc = raise;
    effc = fun (type a) (eff: a t) ->
      match eff with
      | Xchg msg -> Some (fun (cont: (a, _) continuation) ->
          Suspended {msg; cont})
      | _ -> None }


let rec run_both a b =
  match a (), b () with
  | Complete va, Complete vb -> (va, vb)
  | Suspended {msg = m1; cont = k1},
    Suspended {msg = m2; cont = k2} ->
      run_both (fun () -> continue k1 m2)
               (fun () -> continue k2 m1)
  | _ -> failwith "Improper synchronization"


let comp2 () = perform (Xchg 21) * perform (Xchg 21)

let () = let x, y = run_both (step comp1) (step comp2) in Format.printf ">> %d %d@." x y


type _ Effect.t += Fork : (unit -> unit) -> unit t
                 | Yield : unit t

let fork f = perform (Fork f)
let yield () = perform Yield
let xchg v = perform (Xchg v)


(* A concurrent round-robin scheduler *)
let run (main : unit -> unit) : unit =
  let exchanger = ref None in (* waiting exchanger *)
  let run_q = Queue.create () in (* scheduler queue *)
  let enqueue k v =
    let task () = continue k v in
    Queue.push task run_q
  in
  let dequeue () =
    if Queue.is_empty run_q then () (* done *)
    else begin
      let task = Queue.pop run_q in
      task ()
    end
  in
  let rec spawn (f : unit -> unit) : unit =
    match_with f () {
      retc = dequeue;
      exnc = (fun e ->
        print_endline (Printexc.to_string e);
        dequeue ());
      effc = fun (type a) (eff : a t) ->
        match eff with
        | Yield -> Some (fun (k : (a, unit) continuation) ->
            enqueue k (); dequeue ())
        | Fork f -> Some (fun (k : (a, unit) continuation) ->
            enqueue k (); spawn f)
        | Xchg n -> Some (fun (k : (int, unit) continuation) ->
            begin match !exchanger with
            | Some (n', k') ->
                exchanger := None; enqueue k' n; continue k n'
            | None -> exchanger := Some (n, k); dequeue ()
            end)
        | _ -> None
    }
  in
  spawn main

let _ = run (fun _ ->
  fork (fun _ ->
    Format.printf "[t1] Sending 0@.";
    let v = xchg 0 in
    Format.printf "[t1] received %d@." v);
  fork (fun _ ->
    Format.printf "[t2] Sending 1@.";
    let v = xchg 1 in
    Format.printf "[t2] received %d@." v))

(*****)

type _ Effect.t += E : string t
                 | F : string t

let foo () = perform F ^ " " ^ perform E ^ " " ^ perform F

let bar () =
  try_with foo ()
  { effc = fun (type a) (eff: a t) ->
      match eff with
      | E -> Some (fun (k: (a,_) continuation) ->
          continue k "Coucou!")
      | _ -> None }

let baz () =
  try_with bar ()
  { effc = fun (type a) (eff: a t) ->
      match eff with
      | F -> Some (fun (k: (a,_) continuation) ->
          continue k "Hello, world!")
      | _ -> None }

let () = Format.printf "%s@." (baz())

(****)

let () =
  Format.printf "%s@."
    (try_with (fun () -> try perform F with Not_found -> "Discontinued") ()
         { effc = fun (type a) (eff: a t) ->
                    Some (fun k -> discontinue k Not_found) })
let () =
  Format.printf "%s@."
    (try_with (fun () -> try perform F with Unhandled _ -> "Unhandled") ()
         { effc = fun (type a) (eff: a t) -> None })

let () =
  Format.printf "%s@." (try bar () with Unhandled _ -> "Saw unhandled exception")

let () =
  try
    Format.printf "%d@." @@
    try_with perform (Xchg 0)
  { effc = fun (type a) (eff : a t) ->
      match eff with
      | Xchg n -> Some (fun (k: (a, _) continuation) ->
          continue k 21 + continue k 21)
      | _ -> None }
  with Continuation_already_resumed ->
    Format.printf "One-shot@."

(****)

let invert (type a) ~(iter : (a -> unit) -> unit) : a Seq.t =
  let module M = struct
    type _ Effect.t += Yield : a -> unit t
  end in
  let yield v = perform (M.Yield v) in
  fun () -> match_with iter yield
  { retc = (fun _ -> Seq.Nil);
    exnc = raise;
    effc = fun (type b) (eff : b Effect.t) ->
      match eff with
      | M.Yield v -> Some (fun (k: (b,_) continuation) ->
          Seq.Cons (v, continue k))
      | _ -> None }

let s = invert ~iter:(Fun.flip String.iter "OCaml")
let next = Seq.to_dispenser s;;

let rec loop () =
  match next() with Some c -> Format.printf "%c" c; loop() | None -> Format.printf "@."
let () = loop()

(****)

type _ Effect.t += Send : int -> unit Effect.t
                 | Recv : int Effect.t

open! Effect.Shallow

let run (comp: unit -> unit) : unit =
  let rec loop_send : type a. (a,unit) continuation -> a -> unit = fun k v ->
    continue_with k v
      { retc = Fun.id;
        exnc = raise;
        effc = fun (type b) (eff : b Effect.t) ->
          match eff with
          | Send n -> Some (fun (k: (b,_) continuation) ->
              loop_recv n k ())
          | Recv -> failwith "protocol violation"
          | _ -> None }
  and loop_recv : type a. int -> (a,unit) continuation -> a -> unit = fun n k v ->
    continue_with k v
      { retc = Fun.id;
        exnc = raise;
        effc = fun (type b) (eff : b Effect.t) ->
          match eff with
          | Recv -> Some (fun (k: (b,_) continuation) ->
              loop_send k n)
          | Send v -> failwith "protocol violation"
          | _ -> None }
  in
  loop_send (fiber comp) ()

let () = run (fun () ->
  Format.printf "Send 42@.";
  perform (Send 42);
  Format.printf "Recv: %d@." (perform Recv);
  Format.printf "Send 43@.";
  perform (Send 43);
  Format.printf "Recv: %d@." (perform Recv))