From iris.proofmode Require Import tactics.
From iris_logrel Require Export logrel examples.lock.
Import uPred.

(* Stack τ = μ x. Unit + (τ * x), essentially a type of lists *)
(* writing nil and cons for "constructors" *)
Definition CG_StackType τ :=
  TRec (TSum TUnit (TProd τ.[ren (+1)] (TVar 0))).

Notation Conse h t := (Fold (InjR (Pair h t))).
Notation Nile := (Fold (InjL #())).

(* Coarse-grained push *)
Program Definition CG_push : val := λ: "st" "x",
  "st" <- Conse "x" (!"st").

Definition CG_locked_push : val := λ: "st" "l" "x",
  acquire "l";; CG_push "st" "x";; release "l".

(* pop s = λ x. match (load s) with
                | nil => InjL ()
                | cons y ys => s <- ys ;; InjR y 
                end *)
Definition CG_pop : val := λ: "st" <>,
  match: Unfold !"st" with
    InjL <> => InjL #()
  | InjR "y" => "st" <- (Snd "y");; InjR (Fst "y")
  end.


Definition CG_locked_pop : val := λ: "st" "l" <>,
  acquire "l";; (let: "v" := CG_pop "st" #() in (release "l";; "v")).

(* snap st l = with_lock (λ _, load st) l *)
Definition CG_snap : val := λ: "st" "l" <>,
  acquire "l";; let: "v" := !"st" in (release "l";; "v").

(* iter f = λ s. match s with
                 | nil => Unit
                 | cons x xs => (f x) ;; iter f xs
                 end *)
Definition CG_iter : val := rec: "iter" "f" := λ: "s",
  match: (Unfold "s") with
    InjL <> => #()
  | InjR "x" => "f" (Fst "x");; "iter" "f" (Snd "x")
  end.

(* snap_iter st l := λ f. iter f (snap st l #()) *)
Definition CG_snap_iter : val := λ: "st" "l" "f",
  CG_iter "f" (CG_snap "st" "l" #()).

(* stack_body st l :=
   ⟨locked_push st l, locked_pop st l, snap_iter st l⟩ *)
Definition CG_stack_body : val := λ: "st" "l",
  (CG_locked_push "st" "l", CG_locked_pop "st" "l", CG_snap_iter "st" "l").

Definition CG_stack : val :=
  Λ: let: "l" := ref #false in 
     let: "st" := ref Nile in
     CG_stack_body "st" "l".

(** Coarse-grained stack presented as a module *)
(* type s α := (ref (list α), lockτ) *)
Definition sτ α := TProd (Tref (CG_StackType α)) LockType.
(* newStack : ∀ α, s α *)
Definition newStack : val := Λ: λ: <>,
  (ref Nile, ref #false)%E.
(* popStack : ∀ α, s α → MAYBE α *)
Definition popStack : val := Λ: λ: "x",
  CG_locked_pop (Fst "x") (Snd "x") #().
(* pushStack : ∀ α, s α → α → () *)
Definition pushStack : val := Λ: λ: "x" "a",
  CG_locked_push (Fst "x") (Snd "x") "a".
Definition stackmod : val := Λ:
  Pack (TApp newStack, TApp popStack, TApp pushStack).

Section typing.
  Hint Unfold sτ : typeable.
  Lemma newStack_typed Γ :
    Γ ⊢ₜ newStack : TForall (TArrow TUnit (sτ (TVar 0))).
  Proof.
    unlock sτ newStack. (* TODO need to explicitly unlock newStack here *)
    solve_typed.
  Qed.
  Hint Resolve newStack_typed : typeable.

  Lemma popStack_typed Γ :
    Γ ⊢ₜ popStack : TForall $ TArrow (sτ (TVar 0)) (TSum TUnit (TVar 0)).
  Proof.
    unlock sτ popStack. (* TODO need to explicitly unlock newStack here *)
    unlock CG_locked_pop CG_pop.
    repeat (econstructor; solve_typed).
  Qed.
  Hint Resolve popStack_typed : typeable.

  Lemma pushStack_typed Γ :
    Γ ⊢ₜ pushStack : TForall $ TArrow (sτ (TVar 0)) (TArrow (TVar 0) TUnit).
  Proof.
    unlock sτ pushStack. (* TODO need to explicitly unlock newStack here *)
    unlock CG_locked_push CG_push.
    repeat (econstructor; solve_typed).
  Qed.
  Hint Resolve pushStack_typed : typeable.

  Lemma stackmod_typed Γ :
    Γ ⊢ₜ stackmod : TForall $ TExists $ TProd (TProd
                                                 (TArrow TUnit (TVar 0))
                                                 (TArrow (TVar 0) (TSum TUnit (TVar 1))))
                                                 (TArrow (TVar 0) (TArrow (TVar 1) TUnit)).
  Proof.
    unlock stackmod.
    econstructor.
    eapply TPack with (sτ (TVar 0)).
    econstructor; [econstructor | ].
    - simpl.
      replace (TArrow TUnit (sτ (TVar 0))) with (TArrow TUnit (sτ (TVar 0))).[TVar 0/]; last first.
      { autosubst. }
      solve_typed.
    - simpl.
      replace (TArrow (sτ (TVar 0)) (TSum TUnit (ids 0)))
        with (TArrow (sτ (TVar 0)) (TSum TUnit (TVar 0))).[TVar 0/]; last first.
      { autosubst. }
      solve_typed.
    - simpl.
      replace (TArrow (sτ (TVar 0)) (TArrow (ids 0) TUnit))
        with (TArrow (sτ (TVar 0)) (TArrow (ids 0) TUnit)).[TVar 0/] by autosubst.
      solve_typed.
  Qed.
  Hint Resolve stackmod_typed.
End typing.

Section CG_Stack.
  Context `{logrelG Σ}.

  Lemma CG_push_type Γ τ :
    typed Γ CG_push (TArrow (Tref (CG_StackType τ)) (TArrow τ TUnit)).
  Proof.
    unfold CG_push. unlock.
    repeat econstructor. eauto 10 with typeable.
    (* TODO: make eauto call asimpl? *)
    asimpl. eauto 10 with typeable.
  Qed.

  Hint Resolve CG_push_type : typeable.

  Lemma CG_locked_push_type Γ τ :
    typed Γ CG_locked_push (TArrow (Tref (CG_StackType τ)) (TArrow LockType (TArrow τ TUnit))).
  Proof.
    unfold CG_locked_push. unlock.
    eauto 20 with typeable.
  Qed.
  
  Hint Resolve CG_locked_push_type : typeable.

  Lemma CG_push_r st' (v w : val) l E1 E2 Δ Γ t K τ :
    nclose logrelN ⊆ E1 →
    st' ↦ₛ v -∗ l ↦ₛ #false -∗
    (st' ↦ₛ FoldV (InjRV (w, v)) -∗ l ↦ₛ #false
     -∗ {E1,E2;Δ;Γ} ⊨ t ≤log≤ fill K #() : τ) -∗
    {E1,E2;Δ;Γ} ⊨ t ≤log≤ fill K ((CG_locked_push $/ (LitV (Loc st')) $/ (LitV (Loc l))) w) : τ.
  Proof.
    iIntros (?)"Hst' Hl Hlog".
    unlock CG_locked_push CG_push. simpl_subst/=.
    rel_let_r.
    rel_apply_r (bin_log_related_acquire_r with "Hl").
    { solve_ndisj. }
    iIntros "Hl /=".
    rel_seq_r.
    do 2 rel_let_r.
    rel_load_r.
    rel_store_r.
    rel_seq_r.
    rel_apply_r (bin_log_related_release_r with "Hl").
    { solve_ndisj. }
    iIntros "Hl /=".
    iApply ("Hlog" with "Hst' Hl").
  Qed.

  (* Coarse-grained pop *)
  Lemma CG_pop_type Γ τ :
    typed Γ CG_pop (TArrow (Tref (CG_StackType τ)) (TArrow TUnit (TSum TUnit τ))).
  Proof.
    unfold CG_pop. unlock.
    repeat econstructor; eauto 20 with typeable.
    asimpl. eauto 20 with typeable.
  Qed.

  Hint Resolve CG_pop_type : typeable.
  Global Opaque CG_pop.

  Lemma CG_locked_pop_type Γ τ :
    typed Γ CG_locked_pop (TArrow (Tref (CG_StackType τ)) (TArrow LockType (TArrow TUnit (TSum TUnit τ)))).
  Proof.
    unfold CG_locked_pop. unlock.
    eauto 20 with typeable.
  Qed.

  Hint Resolve CG_locked_pop_type : typeable.

  Lemma CG_pop_suc_r st' l (w v : val) E1 E2 Δ Γ t K τ :
    nclose logrelN ⊆ E1 →
    st' ↦ₛ FoldV (InjRV (w, v)) -∗
    l ↦ₛ #false -∗
    (st' ↦ₛ v -∗ l ↦ₛ #false
    -∗ {E1,E2;Δ;Γ} ⊨ t ≤log≤ fill K (InjR w) : τ) -∗
    {E1,E2;Δ;Γ} ⊨ t ≤log≤ fill K ((CG_locked_pop $/ LitV (Loc st') $/ LitV (Loc l)) #()) : τ.
  Proof.
    iIntros (?) "Hst' Hl Hlog".
    unlock CG_locked_pop CG_pop. simpl_subst/=.
    rel_seq_r.
    rel_apply_r (bin_log_related_acquire_r with "Hl").
    { solve_ndisj. }
    iIntros "Hl /=".
    repeat rel_rec_r.
    rel_load_r.
    rel_fold_r.
    rel_case_r.
    rel_let_r.
    rel_proj_r.
    rel_store_r.
    rel_seq_r.
    rel_proj_r.
    rel_rec_r.
    rel_apply_r (bin_log_related_release_r with "Hl").
    { solve_ndisj. }
    iIntros "Hl /=".
    rel_rec_r.
    iApply ("Hlog" with "Hst' Hl").
  Qed.

  Lemma CG_pop_fail_r st' l E1 E2 Δ Γ t K τ :
    nclose logrelN ⊆ E1 →
    st' ↦ₛ FoldV (InjLV #()) -∗
    l ↦ₛ #false -∗
    (st' ↦ₛ FoldV (InjLV #()) -∗ l ↦ₛ #false
    -∗ {E1,E2;Δ;Γ} ⊨ t ≤log≤ fill K (InjL #()) : τ) -∗
    {E1,E2;Δ;Γ} ⊨ t ≤log≤ fill K ((CG_locked_pop $/ LitV (Loc st') $/ LitV (Loc l)) #()) : τ.
  Proof.
    iIntros (?) "Hst' Hl Hlog".
    unlock CG_locked_pop CG_pop. simpl_subst/=.
    rel_seq_r.
    rel_apply_r (bin_log_related_acquire_r with "Hl").
    { solve_ndisj. }
    iIntros "Hl /=".
    repeat rel_rec_r.
    rel_load_r.
    rel_fold_r.
    rel_case_r.
    repeat rel_let_r.
    rel_apply_r (bin_log_related_release_r with "Hl").
    { solve_ndisj. }
    iIntros "Hl /=".
    rel_rec_r.
    iApply ("Hlog" with "Hst' Hl").
  Qed.

  Global Opaque CG_locked_pop.

  Lemma CG_snap_type Γ τ :
    typed Γ CG_snap (TArrow (Tref (CG_StackType τ)) (TArrow LockType (TArrow TUnit (CG_StackType τ)))).
  Proof.
    unfold CG_snap. unlock.
    eauto 20 with typeable.
  Qed.

  Hint Resolve CG_snap_type : typeable.
  Global Opaque CG_snap.

  Lemma CG_iter_type Γ τ :
    typed Γ CG_iter (TArrow (TArrow τ TUnit) (TArrow (CG_StackType τ) TUnit)).
  Proof.
    unfold CG_iter. unlock.
    repeat econstructor; eauto 50 with typeable.
    asimpl. eauto with typeable.
  Qed.

  Hint Resolve CG_iter_type : typeable.
  Global Opaque CG_iter.

  Lemma CG_snap_iter_type Γ τ :
    typed Γ CG_snap_iter (TArrow (Tref (CG_StackType τ)) (TArrow LockType (TArrow (TArrow τ TUnit) TUnit))).
  Proof.
    unfold CG_snap_iter. unlock.
    eauto 50 with typeable.
  Qed.

  Hint Resolve CG_snap_iter_type : typeable.
  
  Lemma CG_stack_body_type Γ τ :
    typed Γ CG_stack_body (TArrow (Tref (CG_StackType τ)) (TArrow LockType
          (TProd
             (TProd (TArrow τ TUnit) (TArrow TUnit (TSum TUnit τ)))
             (TArrow (TArrow τ TUnit) TUnit)
          ))).
  Proof.
    unfold CG_stack_body. unlock.
    eauto 50 with typeable.
  Qed.

  Hint Resolve CG_stack_body_type : typeable.

  Opaque CG_snap_iter.

  (* CG_stack 
    : ∀ α. ((α → Unit) * (Unit → Unit + α) * ((α → Unit) → Unit))  *)
  Lemma CG_stack_type Γ :
    typed Γ CG_stack
          (TForall
             (TProd
                (TProd
                   (TArrow (TVar 0) TUnit)
                   (TArrow TUnit (TSum TUnit (TVar 0)))
                )
                (TArrow (TArrow (TVar 0) TUnit) TUnit)
          )).
  Proof.
    unfold CG_stack. unlock.
    eauto 50 with typeable.
  Qed.

End CG_Stack.