From stdpp Require Import base relations.
From iris Require Import prelude.
From semantics.lib Require Import sets maps.
From semantics.ts.stlc Require Import lang notation types parallel_subst.
From Equations Require Import Equations.

Implicit Types
  (Γ : typing_context)
  (v : val)
  (e : expr)
  (A : type).

(** *** Big-Step Semantics for cbn *)
Inductive big_step : expr → val → Prop :=
  | bs_lit (n : Z) :
      big_step (LitInt n) (LitIntV n)
  | bs_lam (x : binder) (e : expr) :
      big_step (Lam x e) (LamV x e)
  | bs_add e1 e2 (z1 z2 : Z) :
      big_step e1 (LitIntV z1) →
      big_step e2 (LitIntV z2) →
      big_step (Plus e1 e2) (LitIntV (z1 + z2))%Z
  | bs_app e1 e2 x e v2 v :
      big_step e1 (@LamV x e) →
      big_step (subst' x e2 e) v →
      big_step (App e1 e2) v
    .
#[export] Hint Constructors big_step : core.


Lemma big_step_vals (v: val): big_step (of_val v) v.
Proof.
  induction v; econstructor.
Qed.

Lemma big_step_inv_vals (v w: val): big_step (of_val v) w → v = w.
Proof.
  destruct v; inversion 1; eauto.
Qed.



(* *** Definition of the logical relation. *)
(* We reuse most of these definitions. *)
Inductive val_or_expr : Type :=
| inj_val : val → val_or_expr
| inj_expr : expr → val_or_expr.

(* Note that we're using a slightly modified termination argument here. *)
Equations type_size (t : type) : nat :=
  type_size Int := 1;
  type_size (Fun A B) := type_size A + type_size B + 2.
Equations mut_measure (ve : val_or_expr) (t : type) : nat :=
  mut_measure (inj_val _) t := type_size t;
  mut_measure (inj_expr _) t := 1 + type_size t.

Equations type_interp (ve : val_or_expr) (t : type) : Prop by wf (mut_measure ve t) := {
  type_interp (inj_val v) Int =>
    ∃ z : Z, v = z ;
  type_interp (inj_val v) (A → B) =>
    ∃ x e, v = @LamV x e ∧ closed (x :b: nil) e ∧
      ∀ e',
        type_interp (inj_expr e') A →
        type_interp (inj_expr (subst' x e' e)) B;

  type_interp (inj_expr e) t =>
    (* we now need to explicitly require that expressions here are closed so
       that we can apply them to lambdas directly. *)
    ∃ v, big_step e v ∧ closed [] e ∧ type_interp (inj_val v) t
}.
Next Obligation.
  repeat simp mut_measure; simp type_size; lia.
Qed.
Next Obligation.
  simp mut_measure. simp type_size.
  destruct A; repeat simp mut_measure; repeat simp type_size; lia.
Qed.

(* We derive the expression/value relation. *)
Notation sem_val_rel t v := (type_interp (inj_val v) t).
Notation sem_expr_rel t e := (type_interp (inj_expr e) t).

Notation 𝒱 t v := (sem_val_rel t v).
Notation ℰ t v := (sem_expr_rel t v).


(* *** Semantic typing of contexts *)
Implicit Types
  (θ : gmap string expr).

Inductive sem_context_rel : typing_context → (gmap string expr) → Prop :=
  | sem_context_rel_empty : sem_context_rel ∅ ∅
  (* contexts may now contain arbitrary (semantically well-typed) expressions
     as opposed to just values. *)
  | sem_context_rel_insert Γ θ e x A :
    ℰ A e →
    sem_context_rel Γ θ →
    sem_context_rel (<[x := A]> Γ) (<[x := e]> θ).

Notation 𝒢 := sem_context_rel.

(* The semantic typing judgement. Note that we require e to be closed under Γ. *)
Definition sem_typed Γ e A :=
  closed (elements (dom Γ)) e ∧
  ∀ θ, 𝒢 Γ θ → ℰ A (subst_map θ e).
Notation "Γ ⊨ e : A" := (sem_typed Γ e A) (at level 74, e, A at next level).


(* We start by proving a couple of helper lemmas that will be useful later. *)

Lemma sem_expr_rel_of_val A v:
  ℰ A v → 𝒱 A v.
Proof.
  simp type_interp.
  intros (v' & ->%big_step_inv_vals & Hv').
  apply Hv'.
Qed.



Lemma val_rel_closed v A:
  𝒱 A v → closed [] v.
Proof.
  induction A; simp type_interp.
  - intros [z ->]. done.
  - intros (x & e & -> & Hcl & _). done.
Qed.
Lemma val_inclusion A v:
  𝒱 A v → ℰ A v.
Proof.
  intros H. simp type_interp. eauto using big_step_vals, val_rel_closed.
Qed.
Lemma expr_rel_closed e A :
  ℰ A e → closed [] e.
Proof.
  simp type_interp. intros (v & ? & ? & ?). done.
Qed.
Lemma sem_context_rel_closed Γ θ:
  𝒢 Γ θ → subst_closed [] θ.
Proof.
  induction 1; rewrite /subst_closed.
  - naive_solver.
  - intros y e'. rewrite lookup_insert_Some.
    intros [[-> <-]|[Hne Hlook]].
    + by eapply expr_rel_closed.
    + eapply IHsem_context_rel; last done.
Qed.


(* This is essentially an inversion lemma for 𝒢 *)
Lemma sem_context_rel_exprs Γ θ x A :
  sem_context_rel Γ θ →
  Γ !! x = Some A →
  ∃ e, θ !! x = Some e ∧ ℰ A e.
Proof.
  induction 1 as [|Γ θ e y B Hvals Hctx IH].
  - naive_solver.
  - rewrite lookup_insert_Some. intros [[-> ->]|[Hne Hlook]].
    + eexists; first by rewrite lookup_insert.
    + eapply IH in Hlook as (e' & Hlook & He).
      eexists; split; first by rewrite lookup_insert_ne.
      done.
Qed.

Lemma sem_context_rel_dom Γ θ :
  𝒢 Γ θ → dom Γ = dom θ.
Proof.
  induction 1.
  - by rewrite !dom_empty.
  - rewrite !dom_insert. congruence.
Qed.


(* *** Compatibility lemmas *)
Lemma compat_int Γ z : Γ ⊨ (LitInt z) : Int.
Proof.
  split; first done.
  intros θ _. simp type_interp.
  exists z. split; simpl.
  - constructor.
  - simp type_interp. eauto.
Qed.

Lemma compat_var Γ x A :
  Γ !! x = Some A →
  Γ ⊨ (Var x) : A.
Proof.
  intros Hx. split.
  { eapply bool_decide_pack, elem_of_elements, elem_of_dom_2, Hx. }
  intros θ Hctx; simpl.
  eapply sem_context_rel_exprs in Hx as (e & Heq & He); last done.
  rewrite Heq. done.
Qed.

Lemma compat_app Γ e1 e2 A B :
  Γ ⊨ e1 : (A → B) →
  Γ ⊨ e2 : A →
  Γ ⊨ (e1 e2) : B.
Proof.
  intros [Hfuncl Hfun] [Hargcl Harg]. split.
  { simpl. eauto. }
  intros θ Hctx; simpl.
  specialize (Hfun _ Hctx). simp type_interp in Hfun. destruct Hfun as (v1 & Hbs1 & Hcl1 & Hv1).
  simp type_interp in Hv1. destruct Hv1 as (x & e & -> & Hcl & Hv1).
  specialize (Harg _ Hctx). 
  (* compare this to the cbv logical relation: in that case we now extract a
     value from Harg. We don't need to do this here, since we don't need to
     evaluate [subst_map θ e2] before applying it to [λ: x, e]. However, we do
     need a closedness statement for e2, which we obtain from the definition of
     ℰ. *)
  specialize (expr_rel_closed _ _ Harg) as He2cl.

  apply Hv1 in Harg.
  simp type_interp in Harg. destruct Harg as (v & Hbsv & Hcl' & Hv).

  simp type_interp.
  exists v. repeat split.
  - eauto.
  - cbn. by apply andb_True. 
  - eauto.
Qed.


Lemma lam_closed Γ θ (x : string) A e :
  closed (elements (dom (<[x:=A]> Γ))) e → 
  𝒢 Γ θ → 
  closed [] (Lam x (subst_map (delete x θ) e)).
Proof.
  intros Hcl Hctxt.
  eapply subst_map_closed'_2.
  - eapply closed_weaken; first done.
    rewrite dom_delete dom_insert (sem_context_rel_dom Γ θ) //.
    intros y. destruct (decide (x = y)); set_solver.
  - eapply subst_closed_weaken, sem_context_rel_closed; last done.
    + set_solver.
    + apply map_delete_subseteq.
Qed.
Lemma compat_lam Γ x e A B :
  (<[ x := A ]> Γ) ⊨ e : B →
  Γ ⊨ (Lam (BNamed x) e) : (A → B).
Proof.
  intros [Hbodycl Hbody]. split.
  { simpl. eapply closed_weaken; first eassumption. set_solver. }
  intros θ Hctxt. simpl. simp type_interp.
  eexists.
  split; first by eauto.
  split; first by eapply lam_closed.
  simp type_interp.
  eexists _, _. split; first reflexivity.
  split; first by eapply lam_closed.
  intros e' He'.
  specialize (Hbody (<[ x := e']> θ)).
  simpl. rewrite subst_subst_map; last by eapply sem_context_rel_closed.
  apply Hbody.
  apply sem_context_rel_insert; done.
Qed.

Lemma compat_add Γ e1 e2 :
  Γ ⊨ e1 : Int →
  Γ ⊨ e2 : Int →
  Γ ⊨ (e1 + e2) : Int.
Proof.
  intros [Hcl1 He1] [Hcl2 He2]. split.
  { simpl. eauto. }
  intros θ Hctx.
  simp type_interp.
  specialize (He1 _ Hctx). specialize (He2 _ Hctx).
  simp type_interp in He1. simp type_interp in He2.

  destruct He1 as (v1 & Hb1 & Hv1).
  destruct He2 as (v2 & Hb2 & Hv2).
  simp type_interp in Hv1, Hv2.
  destruct Hv1 as (Hcl1' & z1 & ->). destruct Hv2 as (Hcl2' & z2 & ->).

  exists (z1 + z2)%Z. repeat split.
  - constructor; done.
  - simpl. by apply andb_True.
  - exists (z1 + z2)%Z. done.
Qed.

Lemma sem_soundness Γ e A :
  (Γ ⊢ e : A)%ty →
  Γ ⊨ e : A.
Proof.
  induction 1 as [ | Γ x e A B Hsyn IH | | | ].
  - by apply compat_var.
  - by apply compat_lam.
  - apply compat_int.
  - by eapply compat_app.
  - by apply compat_add.
Qed.

Lemma termination e A :
  (∅ ⊢ e : A)%ty →
  ∃ v, big_step e v.
Proof.
  (* You may want to add suitable intermediate lemmas, like we did for the cbv
     logical relation as seen in the lecture. *)
  intros [Hsemcl Hsem]%sem_soundness.
  specialize (Hsem ∅).
  simp type_interp in Hsem.
  rewrite subst_map_empty in Hsem.
  destruct Hsem as (v & Hbs & _); last eauto.
  constructor.
Qed.