Require Import Coq.Lists.List.
Require Import Coq.PArith.PArith.
Require Import Computation.
Require Import Events.

Import ListNotations.
Local Open Scope type.

(** Heap of values indexed by positives. *)
Module Heap.
  (** A heap is a binary tree. *)
  Inductive t (A : Type) : Type :=
  | Empty : t A
  | Node : A -> t A -> t A -> t A.
  Arguments Empty {A}.
  Arguments Node {A} _ _ _.

  (** Add an element at the position `n` if this is `Empty`. Do nothing else. *)
  Fixpoint add {A : Type} (heap : t A) (x : A) (n : positive) : t A :=
    match (heap, n) with
    | (Empty, xH) => Node x Empty Empty
    | (Node y heap1 heap2, xO n) => Node y (add heap1 x n) heap2
    | (Node y heap1 heap2, xI n) => Node y heap1 (add heap2 x n)
    | (Empty, _) | (Node _ _ _, xH) => heap
    end.

  (** Return the element at the position `n`, or `None`. *)
  Fixpoint find {A : Type} (heap : t A) (n : positive) : option A :=
    match (heap, n) with
    | (Empty, _) => None
    | (Node x _ _, xH) => Some x
    | (Node x heap _, xO n) => find heap n
    | (Node x _ heap, xI n) => find heap n
    end.

  (** Update the element at the position `n`, if it exists. *)
  Fixpoint update {A : Type} (heap : t A) (n : positive) (x : A) : t A :=
    match (heap, n) with
    | (Empty, _) => heap
    | (Node _ heap1 heap2, xH) => Node x heap1 heap2
    | (Node y heap1 heap2, xO n) => Node y (update heap1 n x) heap2
    | (Node y heap1 heap2, xI n) => Node y heap1 (update heap2 n x)
    end.

  Module Tests.
    Definition h :=
      let h := add Empty 1 1 in
      let h := add h 2 2 in
      let h := add h 3 3 in
      let h := add h 4 4 in
      add h 5 5.

    Definition test1 : h = Node 1 (Node 2 (Node 4 Empty Empty) Empty)
      (Node 3 (Node 5 Empty Empty) Empty) :=
      eq_refl.

    Definition test2 : find Empty 3 = None (A := nat) :=
      eq_refl.

    Definition test3 : find h 3 = Some 3 :=
      eq_refl.

    Definition test4 : update h 3 7 =
      Node 1 (Node 2 (Node 4 Empty Empty) Empty)
        (Node 7 (Node 5 Empty Empty) Empty) :=
      eq_refl.
  End Tests.
End Heap.

(** Describe one callback. *)
Module CallBack.
  (** A callback is an initial state value and a function to react to an
      answer. If the answer is `None`, the callback is terminated. *)
  Record t (sig : Signature.t) (command : Command.t) : Type := New {
    A : Type;
    a : A;
    handler : A -> Command.answer command -> C.t sig (option A) }.
  Arguments New {sig command A} _ _.
  Arguments A {sig command} _.
  Arguments a {sig command} _.
  Arguments handler {sig command} _ _ _.
End CallBack.

(** Describe the store of callbacks. *)
Module CallBacks.
  (** The store of callbacks is a heap of callbacks plus the `next` free
      position where a callback should be added. *)
  Record t (sig : Signature.t) : Type := New {
    heap : Heap.t (option {command : Command.t & CallBack.t sig command});
    next : positive }.
  Arguments New {sig} _ _.
  Arguments heap {sig} _.
  Arguments next {sig} _.

  (** The empty store of callbacks. *)
  Definition empty (sig : Signature.t) : t sig := {|
    heap := Heap.Empty;
    next := xH |}.

  (** Add a new callback. *)
  Definition add {sig : Signature.t} (call_backs : t sig) {command : Command.t}
    (call_back : CallBack.t sig command) : positive * t sig :=
    let id := next call_backs in
    (id, {|
      heap := Heap.add (heap call_backs) (Some (existT _ command call_back)) id;
      next := id + 1 |}).

  (** Find a callback at position `id`, else return `None`. *)
  Definition find {sig : Signature.t} (call_backs : t sig) {command : Command.t}
    (id : positive) : option (CallBack.t sig command) :=
    match Heap.find (heap call_backs) id with
    | None | Some None => None
    | Some (Some (
        existT <%= "_" unless version[0..2] == "8.4" %> command' call_back)) =>
      match Command.eq_dec command command' with
      | left Heq =>
          eq_rect_r (fun c => option (CallBack.t sig c)) (Some call_back) Heq
      | right _ => None
      end
    end.

  (** Update a callback. *)
  Definition update {sig : Signature.t} (call_backs : t sig) (id : positive)
    {command : Command.t} (call_back : CallBack.t sig command) : t sig := {|
    heap := Heap.update (heap call_backs) id
      (Some (existT _ command call_back));
    next := next call_backs |}.

  (** Remove a callback. *)
  Definition remove {sig : Signature.t} (call_backs : t sig) (id : positive)
    : t sig := {|
    heap := Heap.update (heap call_backs) id None;
    next := next call_backs |}.
End CallBacks.

(** Run a computation on an initialized shared memory. *)
Fixpoint run {sig : Signature.t} {A : Type} (call_backs : CallBacks.t sig)
  (mem : Memory.t sig) (outputs : list Output.t) (x : C.t sig A)
  : option A * CallBacks.t sig * Memory.t sig * list Output.t :=
  match x in C.t _ A return option A  * _ * _ * _ with
  | C.Ret <%= "_" if version[0..2] == "8.4" %> x => (Some x, call_backs, mem, outputs)
  | C.Bind <%= "_ _" if version[0..2] == "8.4" %> x f =>
    match run call_backs mem outputs x with
    | (Some x, call_backs, mem, outputs) => run call_backs mem outputs (f x)
    | (None, call_backs, mem, outputs) => (None, call_backs, mem, outputs)
    end
  | C.Read <%= "_" if version[0..2] == "8.4" %> _ => (Some (Ref.read mem), call_backs, mem, outputs)
  | C.Write <%= "_" if version[0..2] == "8.4" %> _ v => (Some tt, call_backs, Ref.write mem v, outputs)
  | C.Send <%= "_" if version[0..2] == "8.4" %> _ request a handler =>
    let call_back := CallBack.New a handler in
    let (id, call_backs) := CallBacks.add call_backs call_back in
    let output := Output.New _ id request in
    (Some tt, call_backs, mem, output :: outputs)
  | C.Exit <%= "_" if version[0..2] == "8.4" %> => (None, call_backs, mem, outputs)
  end.

(** Run a set of callbacks on a list of inputs. Return if the program exited and
    the list of outputs. *)
Fixpoint run_call_backs {sig : Signature.t} (call_backs : CallBacks.t sig)
  (mem : Memory.t sig) (outputs : list Output.t) (inputs : list Input.t)
  : bool * list Output.t :=
  match inputs with
  | [] => (false, outputs)
  | input :: inputs =>
    let (command, id, argument) := input in
    match CallBacks.find call_backs id with
    | None => run_call_backs call_backs mem outputs inputs
    | Some (CallBack.New <%= "_" if version[0..2] == "8.4" %> a handler) =>
      match run call_backs mem outputs (handler a argument) with
      | (None, _, _, outputs) => (true, outputs)
      | (Some a, call_backs, mem, outputs) =>
        let call_backs :=
          match a with
          | None => call_backs
          | Some a =>
            let call_back := CallBack.New a handler in
            CallBacks.update call_backs id call_back
          end in
        run_call_backs call_backs mem outputs inputs
      end
    end
  end.

(** Run a program on a list of inputs. Return if the program exited and the list
    of outputs. *)
Definition run_on_inputs {sig : Signature.t} (x : C.t sig unit)
  (mem : Memory.t sig) (inputs : list Input.t) : bool * list Output.t :=
  match run (CallBacks.empty _) mem [] x with
  | (None, _, _, outputs) => (true, outputs)
  | (Some _, call_backs, mem, outputs) =>
    run_call_backs call_backs mem outputs inputs
  end.
