open Core
open Core.O

type 'a outcome =
  | Cancelled of 'a
  | Not_cancelled

type handlers =
  | End_of_handlers
  | Handler of
      { ivar : unit outcome Ivar.t
      ; mutable next : handlers
      ; mutable prev : handlers
      }

module State = struct
  type t =
    | Cancelled
    | Not_cancelled of { mutable handlers : handlers }
end

type t = { mutable state : State.t }

let create () = { state = Not_cancelled { handlers = End_of_handlers } }

let rec invoke_handlers = function
  | Handler { ivar; next; prev = _ } ->
    let* () = Ivar.fill ivar (Cancelled ()) in
    invoke_handlers next
  | End_of_handlers -> return ()
;;

let fire t =
  of_thunk (fun () ->
    match t.state with
    | Cancelled -> return ()
    | Not_cancelled { handlers } ->
      t.state <- Cancelled;
      invoke_handlers handlers)
;;

let rec fills_of_handlers acc = function
  | Handler { ivar; next; prev = _ } ->
    fills_of_handlers (Scheduler.Fill (ivar, Cancelled ()) :: acc) next
  | End_of_handlers -> List.rev acc
;;

let fire' t =
  match t.state with
  | Cancelled -> []
  | Not_cancelled { handlers } ->
    t.state <- Cancelled;
    fills_of_handlers [] handlers
;;

let fired t =
  match t.state with
  | Cancelled -> true
  | Not_cancelled _ -> false
;;

let with_handler t f ~on_cancel =
  match t.state with
  | Cancelled ->
    let+ x, y = fork_and_join f on_cancel in
    x, Cancelled y
  | Not_cancelled h ->
    let ivar = Ivar.create () in
    let node = Handler { ivar; next = h.handlers; prev = End_of_handlers } in
    (match h.handlers with
     | End_of_handlers -> ()
     | Handler first -> first.prev <- node);
    h.handlers <- node;
    fork_and_join
      (fun () ->
        let* y = f () in
        match t.state with
        | Cancelled -> return y
        | Not_cancelled h ->
          (match node with
           | End_of_handlers ->
             (* We could avoid this [assert false] with GADT sorcery given that
                we created [node] just above and we know for sure it is the
                [Handler _] case, but it's not worth the code complexity. *)
             assert false
           | Handler node ->
             (match node.prev with
              | End_of_handlers -> h.handlers <- node.next
              | Handler prev -> prev.next <- node.next);
             (match node.next with
              | End_of_handlers -> ()
              | Handler next -> next.prev <- node.prev);
             let+ () = Ivar.fill ivar Not_cancelled in
             y))
      (fun () ->
        Ivar.read ivar
        >>= function
        | Cancelled () ->
          let+ x = on_cancel () in
          Cancelled x
        | Not_cancelled -> return Not_cancelled)
;;
