type (_, _) operation = ..

type 'a computation =
  | Return : 'a -> 'a computation
  | Call : ('arg, 'res) operation * 'arg * ('res -> 'a computation) -> 'a computation

type ('a, 'b) handler = {
  return : 'a -> 'b computation;
  operations : 'arg 'res. ('arg, 'res) operation ->
    'arg -> ('res -> 'b computation) -> 'b computation
}

let rec handle : ('a, 'b) handler -> 'a computation -> 'b computation
  = fun h -> function
    | Return x -> h.return x
    | Call (op, x, k) -> h.operations op x @@ fun y -> handle h (k y)

let rethrow_arm op' = fun x k -> Call (op', x, k)

let run (Return v) = v

let rec (>>=) c f =
  match c with
  | Return x -> f x
  | Call (op, x, k) -> Call (op, x, fun y -> k y >>= f)

let (let$) x k = Return x |> k
let return x = Return x
let perform op x = fun k -> Call(op, x, k)
let (let+) comp k = comp k
let (>>=) = (let+)

module Main = struct
  type (_, _) operation +=
    | Write : (string, unit) operation
    | Read : (unit, string) operation

  let operations : type arg res. (arg, res) operation -> arg -> (res -> 'b computation) -> 'b computation
    = function
      | Read -> fun () k -> read_line () |> k
      | Write -> fun x k -> print_endline x |> k
      | op -> rethrow_arm op

  let handler : ('a, 'a) handler = {
    return;
    operations
  }

  let comp () = run @@ handle handler @@
    let x = "hello, " in
    let+ y = perform Read () in
    perform Write @@ x ^ y >>= return
end