From iris.proofmode Require Import tactics.
From iris.program_logic Require Export weakestpre.
From iris.base_logic Require Export big_op invariants.
From iris_logrel.F_mu_ref_conc Require Export rules_binary typing.
From iris.algebra Require Import list.
From stdpp Require Import tactics.
Import uPred.

(* HACK: move somewhere else *)
Ltac auto_equiv :=
  (* Deal with "pointwise_relation" *)
  repeat lazymatch goal with
  | |- pointwise_relation _ _ _ _ => intros ?
  end;
  (* Normalize away equalities. *)
  repeat match goal with
  | H : _ ≡{_}≡ _ |-  _ => apply (timeless_iff _ _) in H
  | _ => progress simplify_eq
  end;
  (* repeatedly apply congruence lemmas and use the equalities in the hypotheses. *)
  try (f_equiv; fast_done || auto_equiv).

(* TODO: modify this in coq-stdpp or something *)
(* this is /almost/ [solve_proper_core ltac:(fun _ => auto_equiv)] *)
(* but we do an additional simpl auto [simplify_eq] *)
Ltac solve_proper ::=
  intros;
  repeat lazymatch goal with
  | |- Proper _ _ => intros ???
  | |- (_ ==> _)%signature _ _ => intros ???
  | |- pointwise_relation _ _ _ _ => intros ?
  | |- ?R ?f _ => try let f' := constr:(λ x, f x) in intros ?
  end; simplify_eq/=;
  (solve_proper_unfold + idtac);
  solve [repeat first [eassumption | auto_equiv] ].

Definition logN : namespace := nroot .@ "logN".

(** interp : is a unary logical relation. *)
Section logrel.
  Context `{heapIG Σ, cfgSG Σ}.
  Notation D := (prodC valC valC -n> iProp Σ).
  Implicit Types τi : D.
  Implicit Types Δ : listC D.
  Implicit Types interp : listC D → D.

  Definition interp_expr (E1 E2 : coPset) (τi : listC D -n> D) (Δ : listC D)
      (ee : expr * expr) : iProp Σ := (∀ j K,
    j ⤇ fill K (ee.2) -∗
    |={E1,E2}=> WP ee.1 {{ v, ∃ v', j ⤇ fill K (of_val v') ∗ τi Δ (v, v') }})%I.
  Global Instance interp_expr_ne n E1 E2:
    Proper (dist n ==> dist n ==> (=) ==> dist n) (interp_expr E1 E2).
  Proof. solve_proper. Qed.       

  Program Definition ctx_lookup (x : var) : listC D -n> D := λne Δ ww,
    (□ from_option id (cconst True)%I (Δ !! x) ww)%I.
  Solve Obligations with solve_proper.

  Program Definition interp_unit : listC D -n> D := λne Δ ww,
    (⌜ww.1 = UnitV⌝ ∧ ⌜ww.2 = UnitV⌝)%I.
  Solve Obligations with solve_proper_alt.

  Program Definition interp_nat : listC D -n> D := λne Δ ww,
    (∃ n : nat, ⌜ww.1 = #nv n⌝ ∧ ⌜ww.2 = #nv n⌝)%I.
  Solve Obligations with solve_proper. 

  Program Definition interp_bool : listC D -n> D := λne Δ ww,
    (∃ b : bool, ⌜ww.1 = #♭v b⌝ ∧ ⌜ww.2 = #♭v b⌝)%I.
  Solve Obligations with solve_proper. 

  Program Definition interp_prod
      (interp1 interp2 : listC D -n> D) : listC D -n> D := λne Δ ww,
    (∃ vv1 vv2, ⌜ww = (PairV (vv1.1) (vv2.1), PairV (vv1.2) (vv2.2))⌝ ∧
                interp1 Δ vv1 ∧ interp2 Δ vv2)%I.
  Solve Obligations with solve_proper. 

  Program Definition interp_sum
      (interp1 interp2 : listC D -n> D) : listC D -n> D := λne Δ ww,
    ((∃ vv, ⌜ww = (InjLV (vv.1), InjLV (vv.2))⌝ ∧ interp1 Δ vv) ∨
     (∃ vv, ⌜ww = (InjRV (vv.1), InjRV (vv.2))⌝ ∧ interp2 Δ vv))%I.
  Solve Obligations with solve_proper. 

  Program Definition interp_arrow (E1 E2 : coPset)
          (interp1 interp2 : listC D -n> D) : listC D -n> D :=
    λne Δ ww,
    (□ ∀ vv, interp1 Δ vv →
             interp_expr E1 E2 
               interp2 Δ (App (of_val (ww.1)) (of_val (vv.1)),
                          App (of_val (ww.2)) (of_val (vv.2))))%I.
  Solve Obligations with solve_proper. 

  Program Definition interp_forall (E1 E2 : coPset)
      (interp : listC D -n> D) : listC D -n> D := λne Δ ww,
    (□ ∀ τi,
          ⌜∀ ww, PersistentP (τi ww)⌝ →
          interp_expr E1 E2
            interp (τi :: Δ) (TApp (of_val (ww.1)), TApp (of_val (ww.2))))%I.
  Solve Obligations with solve_proper.

  Program Definition interp_exists
      (interp : listC D -n> D) : listC D -n> D := λne Δ ww,
    (∃ vv, ⌜ww = (PackV (vv.1), PackV (vv.2))⌝ 
     ∧ □ ∃ τi : D, ⌜∀ ww, PersistentP (τi ww)⌝ ∧ interp (τi :: Δ) vv)%I.
  Solve Obligations with solve_proper.

  Program Definition interp_rec1
      (interp : listC D -n> D) (Δ : listC D) (τi : D) : D := λne ww,
    (□ ∃ vv, ⌜ww = (FoldV (vv.1), FoldV (vv.2))⌝ ∧ ▷ interp (τi :: Δ) vv)%I.
  Solve Obligations with solve_proper.

  Global Instance interp_rec1_contractive
    (interp : listC D -n> D) (Δ : listC D) : Contractive (interp_rec1 interp Δ).
  Proof. solve_contractive. Qed.

  Program Definition interp_rec (interp : listC D -n> D) : listC D -n> D := λne Δ,
    fixpoint (interp_rec1 interp Δ).
  Next Obligation.
    intros interp n Δ1 Δ2 HΔ; apply fixpoint_ne => τi ww. solve_proper.
  Qed.

  Program Definition interp_ref_inv (ll : loc * loc) : D -n> iProp Σ := λne τi,
    (∃ vv, ll.1 ↦ᵢ vv.1 ∗ ll.2 ↦ₛ vv.2 ∗ τi vv)%I.
  Solve Obligations with solve_proper.

  Program Definition interp_ref
      (interp : listC D -n> D) : listC D -n> D := λne Δ ww,
    (∃ ll, ⌜ww = (LocV (ll.1), LocV (ll.2))⌝ ∧
           inv (logN .@ ll) (interp_ref_inv ll (interp Δ)))%I.
  Solve Obligations with solve_proper.

  Program Definition interp_singleton (v v' : val) : listC D -n> D :=
    λne Δ ww, ⌜ww = (v,v')⌝%I.
  Solve Obligations with solve_proper.       

  Program Definition interp_iref
      (interp : listC D -n> D) : listC D -n> D := λne Δ ww,
    (∃ (ll : loc * loc), 
       interp_ref (interp_singleton (LocV (ll.1)) (LocV (ll.2))) Δ ww ∧
       inv (logN .@ ll) (interp_ref_inv ll (interp Δ)))%I.
  Solve Obligations with solve_proper.

  Fixpoint interp (E1 E2 : coPset) (τ : type) : listC D -n> D :=
    match τ return _ with
    | TUnit => interp_unit
    | TNat => interp_nat
    | TBool => interp_bool
    | TProd τ1 τ2 => interp_prod (interp E1 E2 τ1) (interp E1 E2 τ2)
    | TSum τ1 τ2 => interp_sum (interp E1 E2 τ1) (interp E1 E2 τ2)
    | TArrow τ1 τ2 =>
      interp_arrow ⊤ ⊤ (interp E1 E2 τ1) (interp ⊤ ⊤ τ2)
    | TVar x => ctx_lookup x
    | TForall τ' => interp_forall ⊤ ⊤ (interp E1 E2 τ')
    | TExists τ' => interp_exists (interp E1 E2 τ')
    | TRec τ' => interp_rec (interp E1 E2 τ')
    | Tref τ' => interp_ref (interp E1 E2 τ')
    end.
  Notation "⟦ τ ⟧" := (interp ⊤ ⊤ τ).

  Definition interp_env (Γ : stringmap type) (E1 E2 : coPset)
      (Δ : listC D) (vvs : stringmap (val * val)) : iProp Σ :=
    (⌜dom _ Γ = dom _ vvs⌝ 
    ∗ [∗ map] x ↦ τvv ∈ (map_zip Γ vvs), interp E1 E2 (fst τvv) Δ (snd τvv))%I.
  Notation "⟦ Γ ⟧*" := (interp_env Γ ⊤ ⊤).

  Global Instance interp_persistent τ E1 E2 Δ vv :
    PersistentP (interp E1 E2 τ Δ vv).
  Proof.
    revert vv Δ; induction τ=> vv Δ; simpl; try apply _.
    rewrite /PersistentP /interp_rec fixpoint_unfold /interp_rec1 /=.
    by apply always_intro'.
  Qed.
  Global Instance interp_env_persistent Γ E1 E2 Δ vvs :
    PersistentP (interp_env Γ E1 E2 Δ vvs) := _.

  (* DF: automate this proof some more *)
  Lemma interp_weaken Δ1 Π Δ2 E1 E2 τ :
    interp E1 E2 (τ.[upn (length Δ1) (ren (+ length Π))]) (Δ1 ++ Π ++ Δ2)
    ≡ interp E1 E2 τ (Δ1 ++ Δ2).
  Proof.
    revert Δ1 Π Δ2 E1 E2. induction τ=> Δ1 Π Δ2 E1 E2; simpl; auto.
    - intros ww; simpl; properness; auto. by apply IHτ1. by apply IHτ2.
    - intros ww; simpl; properness; auto. by apply IHτ1. by apply IHτ2.
    - unfold interp_expr.
      intros ww; simpl; properness; auto. by apply IHτ1. by apply IHτ2.
    - apply fixpoint_proper=> τi ww /=.
      properness; auto. apply (IHτ (_ :: _)).
    - intros ww; simpl; properness; auto.
      rewrite iter_up; destruct lt_dec as [Hl | Hl]; simpl.
      { by rewrite !lookup_app_l. }
      rewrite !lookup_app_r; [|lia ..]. do 4 f_equiv. lia. 
    - unfold interp_expr.
      intros ww; simpl; properness; auto. by apply (IHτ (_ :: _)).
    - intros ww; simpl; properness; auto. by apply (IHτ (_ :: _)).
    - intros ww; simpl; properness; auto. by apply IHτ.
  Qed.

  Lemma interp_subst_up Δ1 Δ2 τ τ' :
    interp ⊤ ⊤ τ (Δ1 ++ interp ⊤ ⊤ τ' Δ2 :: Δ2)
    ≡ interp ⊤ ⊤ (τ.[upn (length Δ1) (τ' .: ids)]) (Δ1 ++ Δ2).
  Proof.
    revert Δ1 Δ2; induction τ=> Δ1 Δ2; simpl; auto.
    - intros ww; simpl; properness; auto. by apply IHτ1. by apply IHτ2.
    - intros ww; simpl; properness; auto. by apply IHτ1. by apply IHτ2.
    - unfold interp_expr.
      intros ww; simpl; properness; auto. by apply IHτ1. by apply IHτ2.
    - apply fixpoint_proper=> τi ww /=.
      properness; auto. apply (IHτ (_ :: _)).
    - intros ww; simpl.
      rewrite iter_up; destruct lt_dec as [Hl | Hl]; simpl.
      { by rewrite !lookup_app_l. }
      rewrite !lookup_app_r; [|lia ..].
      destruct (x - length Δ1) as [|n] eqn:?; simpl.
      { symmetry. rewrite always_always. asimpl. apply (interp_weaken [] Δ1 Δ2 _ _ τ'). }
      rewrite !lookup_app_r; [|lia ..]. do 4 f_equiv. lia. 
    - unfold interp_expr.
      intros ww; simpl; properness; auto. apply (IHτ (_ :: _)).
    - intros ww; simpl; properness; auto. apply (IHτ (_ :: _)).
    - intros ww; simpl; properness; auto. by apply IHτ.
  Qed.

  Lemma interp_subst Δ2 τ τ' :
    interp ⊤ ⊤ τ ((interp ⊤ ⊤ τ' Δ2) :: Δ2) ≡ interp ⊤ ⊤ (τ.[τ'/]) Δ2.
  Proof. apply (interp_subst_up []). Qed.

  Lemma interp_env_dom Δ Γ E1 E2 vvs : interp_env Γ E1 E2 Δ vvs ⊢ ⌜dom _ Γ = dom _ vvs⌝.
  Proof. by iIntros "[% ?]". Qed.

  Lemma interp_env_Some_l Δ Γ E1 E2 vvs x τ :
    Γ !! x = Some τ → interp_env Γ E1 E2 Δ vvs ⊢ ∃ vv, ⌜vvs !! x = Some vv⌝ ∧ interp E1 E2 τ Δ vv.
  Proof.
    iIntros (Hτ) "[Hdom HΓ]"; iDestruct "Hdom" as %Hdom.
    destruct (elem_of_dom vvs x) as [[v Hv] ].
    { rewrite -Hdom. apply elem_of_dom. by exists τ. }
    assert (map_zip Γ vvs !! x = Some (τ, v)) as Hτv.
    { rewrite map_lookup_zip_with.
      by rewrite Hτ /= Hv /=. }    
    iExists v; iSplit. done.    
    iApply (big_sepM_lookup _ _ _ _ Hτv with "HΓ").
  Qed.

  Lemma interp_env_nil Δ E1 E2 : True ⊢ interp_env ∅ E1 E2 Δ ∅.
  Proof.
    iIntros ""; iSplit. 
    - iPureIntro. unfold_leibniz. by rewrite ?dom_empty.
    - rewrite map_zip_with_empty. auto.
  Qed.

  Lemma interp_env_cons (Δ : list D) (Γ : stringmap type)
     (vvs : stringmap (val * val)) E1 E2 (τ : type) (vv : val * val) (x : string) :
    interp E1 E2 τ Δ vv ∗ interp_env Γ E1 E2 Δ vvs 
    ⊢ interp_env (<[x:=τ]> Γ) E1 E2 Δ (<[x:=vv]> vvs).
  Proof.
    iIntros "[Hτ [Hdom HΓ]]". iDestruct "Hdom" as %Hdom. iSplit.
    - iPureIntro. by rewrite !dom_insert_L Hdom. (* TODO: RK: look at why this is so slow *)
    - rewrite -(map_insert_zip_with pair Γ vvs x (τ,vv)); auto.
      rewrite -insert_delete.
      rewrite big_sepM_insert; last by rewrite lookup_delete.
      iFrame "Hτ".
      iApply (big_sepM_mono _ _ (map_zip Γ vvs) with "HΓ").
      { apply delete_subseteq. }
      done.
  Qed.

  Lemma interp_env_cons_bi (Δ : list D) (Γ : stringmap type) (vvs : stringmap (val * val)) E1 E2 (τ : type) (vv : val * val) (x : string) :
    x ∉ dom (gset string) Γ →
    x ∉ dom (gset string) vvs →
    interp_env (<[x:=τ]> Γ) E1 E2 Δ (<[x:=vv]> vvs) ⊣⊢ interp E1 E2 τ Δ vv ∗ interp_env Γ E1 E2 Δ vvs.
  Proof.
    intros Hx1 Hx2.
    rewrite /interp_env /= (assoc _ (⟦ _ ⟧ _ _)) -(comm _ ⌜(_ = _)⌝%I) -assoc.
    apply sep_proper; [apply pure_proper|].
    - unfold_leibniz. rewrite ?dom_insert.
      split. 
      + intros Hd y. destruct (Hd y) as [Hd1 Hd2].
        split; intro Hy.
        * destruct (elem_of_union {[x]} (dom (gset string) Γ) y) as [Hy1 Hy2].
          pose (Hy2':=Hd1 (Hy2 (or_intror Hy))).
          destruct (decide (x = y)) as [?|neq]; subst.
          { exfalso. by apply Hx1. }
          apply elem_of_union in Hy2'.
          destruct Hy2' as [F | ?]; auto.
          exfalso. apply neq. symmetry.
          (* what the hell am i even doing here TODO *)
          apply (and_rec (fun x1 _ => x1 F) (elem_of_singleton y x)).
        * destruct (elem_of_union {[x]} (dom (gset string) vvs) y) as [Hy1 Hy2 ].
          pose (Hy2':=Hd2 (Hy2 (or_intror Hy))).
          destruct (decide (x = y)) as [?|neq]; subst.
          { exfalso. by apply Hx2. }
          apply elem_of_union in Hy2'.
          destruct Hy2' as [F | ?]; auto.
          exfalso. apply neq. symmetry.
          apply (and_rec (fun x1 _ => x1 F) (elem_of_singleton y x)).
      + intros Heq. fold_leibniz. by f_equal.
    - rewrite -(map_insert_zip_with _ _ _ _ (τ, vv)); auto.
      rewrite big_sepM_insert /=; auto.
      apply not_elem_of_dom in Hx1.
      apply not_elem_of_dom in Hx2.
      unfold map_zip. erewrite lookup_merge.
      by rewrite Hx1 /=. 
      by compute.
  Qed.

  Lemma interp_env_ren Δ (Γ : stringmap type) E1 E2 (vvs : stringmap (val * val)) (τi : D) :
    interp_env (subst (ren (+1)) <$> Γ) E1 E2 (τi :: Δ) vvs
               ⊣⊢
    interp_env Γ E1 E2 Δ vvs.
  Proof.
    apply sep_proper; [apply pure_proper|].
    - unfold_leibniz. by rewrite dom_fmap. 
    - rewrite map_zip_with_fmap_1. 
      rewrite map_zip_with_map_zip.
      generalize (map_zip Γ vvs).
      apply map_ind. 
      + rewrite fmap_empty. done.
      + intros i [σ ww] m Hm IH.
        rewrite fmap_insert.
        rewrite ?big_sepM_insert /=; auto; last first.
        { rewrite lookup_fmap. rewrite Hm. done. }
        rewrite IH. apply sep_proper; auto.
        apply (interp_weaken [] [τi] Δ).
  Qed.

  Lemma interp_EqType_agree τ v v' E1 E2 Δ :
    EqType τ → interp E1 E2 τ Δ (v, v') ⊢ ⌜v = v'⌝.
  Proof.
    intros Hτ; revert v v'; induction Hτ; iIntros (v v') "#H1 /=".
    - by iDestruct "H1" as "[% %]"; subst.
    - by iDestruct "H1" as (n) "[% %]"; subst.
    - by iDestruct "H1" as (b) "[% %]"; subst.
    - iDestruct "H1" as ([??] [??]) "[% [H1 H2]]"; simplify_eq/=.
      rewrite IHHτ1 IHHτ2.
      by iDestruct "H1" as "%"; iDestruct "H2" as "%"; subst.
    - iDestruct "H1" as "[H1|H1]".
      + iDestruct "H1" as ([??]) "[% H1]"; simplify_eq/=.
        rewrite IHHτ1. by iDestruct "H1" as "%"; subst.
      + iDestruct "H1" as ([??]) "[% H1]"; simplify_eq/=.
        rewrite IHHτ2. by iDestruct "H1" as "%"; subst.
  Qed.
End logrel.

Typeclasses Opaque interp_env.
Notation "⟦ τ ⟧" := (interp ⊤ ⊤ τ).
Notation "⟦ τ ⟧ₑ" := (interp_expr ⊤ ⊤ ⟦ τ ⟧).
Notation "⟦ Γ ⟧*" := (interp_env Γ).

Section bin_log_def.
  Context `{heapIG Σ, cfgSG Σ}.
  Notation D := (prodC valC valC -n> iProp Σ).

  Definition bin_log_related (E1 E2 : coPset) (Γ : stringmap type) (e e' : expr) (τ : type) : iProp Σ := (∀ Δ (vvs : stringmap (val * val)) ρ,  
    spec_ctx ρ -∗
    □ interp_env Γ ⊤ ⊤ Δ vvs 
    -∗ interp_expr E1 E2 (interp ⊤ ⊤ τ) Δ
             (env_subst (fst <$> vvs) e, env_subst (snd <$> vvs) e'))%I.
End bin_log_def.

Notation "'{' E1 ',' E2 ';' Γ '}' ⊨ e '≤log≤' e' : τ" :=
  (bin_log_related E1 E2 Γ e e' τ) (at level 74, E1,E2, e, e', τ at next level).
Notation "Γ ⊨ e '≤log≤' e' : τ" :=
  (bin_log_related ⊤ ⊤ Γ e e' τ) (at level 74, e, e', τ at next level).