From stdpp Require Import gmap base relations.
From iris Require Import prelude.
From semantics.ts.systemf Require Import lang notation parallel_subst types logrel tactics.

(** Exercise 1 (LN Exercise 19): De Bruijn Terms *)
Module dbterm.
  (** Your type of expressions only needs to encompass the operations of our base lambda calculus. *)
  Inductive expr :=
    | Lit (l : base_lit)
    | Var (n : nat)
    | Lam (e : expr)
    | Plus (e1 e2 : expr)
    | App (e1 e2 : expr)
  .

  (** Formalize substitutions and renamings as functions. *)
  Definition subt := nat → expr.
  Definition rent := nat → nat.

  Implicit Types
    (σ : subt)
    (δ : rent)
    (n x : nat)
    (e : expr).

  

  (* Operations on renamings *)
  Definition RCons n δ x :=
    match x with
    | 0 => n
    | S x => δ x
    end.
  Definition RComp δ1 δ2 := δ1 ∘ δ2.
  Definition RUp δ := RCons 0 (RComp S δ).

  Fixpoint ren_expr δ e :=
    match e with
    | Lit l => Lit l
    | Var x => Var (δ x)
    | Lam e => Lam (ren_expr (RUp δ) e)
    | App e1 e2 => App (ren_expr δ e1) (ren_expr δ e2)
    | Plus e1 e2 => Plus (ren_expr δ e1) (ren_expr δ e2)
    end.

  (* Operations on substitutions *)
  Definition Cons e σ x :=
    match x with
    | 0 => e
    | S x => σ x
    end.
  Definition ren_subst δ σ n := ren_expr δ (σ n).
  Definition Up σ := Cons (Var 0) (ren_subst S σ).

  Fixpoint subst σ e :=
    match e with
    | Lit l => Lit l
    | Var x => σ x
    | Lam e => Lam (subst (Up σ) e)
    | App e1 e2 => App (subst σ e1) (subst σ e2)
    | Plus e1 e2 => Plus (subst σ e1) (subst σ e2)
    end.

  Goal (subst
    (λ n, match n with
          | 0 => Lit (LitInt 42)
          | 1 => Var 0
          | _ => Var n
          end)
    (Lam (Plus (Plus (Var 2) (Var 1)) (Var 0)))) = 
    Lam (Plus (Plus (Var 1) (Lit 42%Z)) (Var 0)).
  Proof.
    cbn. 
    (* Should be by reflexivity. *)
    reflexivity.
  Qed.

End dbterm.

Section church_encodings.
  (** Exercise 2 (LN Exercise 24): Church encoding, sum types *)
  (* a) Define your encoding *)
  Definition sum_type (A B : type) : type := 
    ∀: (A.[ren (+1)] → #0) → (B.[ren (+1)] → #0) → #0.

  (* b) Implement inj1, inj2, case *)
  Definition injl_val (v : val) : val := Λ, λ: "f" "g", "f" v.
  Definition injl_expr (e : expr) : expr := let: "x" := e in Λ, λ: "f" "g", "f" "x".
  Definition injr_val (v : val) : val := Λ, λ: "f" "g", "g" v.
  Definition injr_expr (e : expr) : expr := let: "x" := e in Λ, λ: "f" "g", "g" "x".

  (* You may want to use the variables x1, x2 for the match arms to fit the typing statements below. *)
  Definition match_expr (e : expr) (e1 e2 : expr) : expr := 
    (e <> (λ: "x1", e1) (λ: "x2", e2))%E.

  (* c) Reduction behavior *)
  (* Some lemmas about substitutions might be useful. Look near the end of the lang.v file! *)
  Lemma match_expr_red_injl e e1 e2 (vl v' : val) :
    is_closed [] vl →
    is_closed ["x1"] e1 →
    big_step e (injl_val vl) →
    big_step (subst' "x1" vl e1) v' →
    big_step (match_expr e e1 e2) v'.
  Proof.
    intros. bs_step_det.
    erewrite (lang.subst_is_closed ["x1"] _ "g"); [ done | done | rewrite elem_of_list_singleton; done].
  Qed.

  Lemma match_expr_red_injr e e1 e2 (vl v' : val) :
    is_closed [] vl →
    big_step e (injr_val vl) →
    big_step (subst' "x2" vl e2) v' →
    big_step (match_expr e e1 e2) v'.
  Proof.
    intros. bs_step_det.
    Qed.

  Lemma injr_expr_red e v :
    big_step e v →
    big_step (injr_expr e) (injr_val v).
  Proof.
    intros. bs_step_det.
    Qed.

  Lemma injl_expr_red e v :
    big_step e v →
    big_step (injl_expr e) (injl_val v).
  Proof.
    intros. bs_step_det.
    Qed.


  (* d) Typing rules *)
  Lemma sum_injl_typed n Γ (e : expr) A B :
    type_wf n B →
    type_wf n A →
    TY n; Γ ⊢ e : A →
    TY n; Γ ⊢ injl_expr e : sum_type A B.
  Proof.
    intros. solve_typing.
    Qed.

  Lemma sum_injr_typed n Γ e A B :
    type_wf n B →
    type_wf n A →
    TY n; Γ ⊢ e : B →
    TY n; Γ ⊢ injr_expr e : sum_type A B.
  Proof.
    intros. solve_typing.
    Qed.

  Lemma sum_match_typed n Γ A B C e e1 e2 :
    type_wf n A →
    type_wf n B →
    type_wf n C →
    TY n; Γ ⊢ e : sum_type A B →
    TY n; <["x1" := A]> Γ ⊢ e1 : C →
    TY n; <["x2" := B]> Γ ⊢ e2 : C →
    TY n; Γ ⊢ match_expr e e1 e2 : C.
  Proof.
    intros. solve_typing.
    Qed.


  (** Exercise 3 (LN Exercise 25): church encoding, list types *)

  (* a) translate the type of lists into De Bruijn. *)
  Definition list_type (A : type) : type := 
    ∀: #0 → (A.[ren (+1)] → #0 → #0) → #0.

  (* b) Implement nil and cons. *)
  Definition nil_val : val := Λ, λ: "e" "c", "e".
  Definition cons_val (v1 v2 : val) : val := Λ, λ: "e" "c", "c" v1 (v2 <> "e" "c").
  Definition cons_expr (e1 e2 : expr) : expr :=
    let: "p" := (e1, e2) in Λ, λ: "e" "c", "c" (Fst "p") ((Snd "p") <> "e" "c").

  (* c) Define typing rules and prove them *)
  Lemma nil_typed n Γ A :
    type_wf n A →
    TY n; Γ ⊢ nil_val : list_type A.
  Proof.
    intros. solve_typing.
    Qed.

  Lemma cons_typed n Γ (e1 e2 : expr) A :
    type_wf n A →
    TY n; Γ ⊢ e2 : list_type A →
    TY n; Γ ⊢ e1 : A →
    TY n; Γ ⊢ cons_expr e1 e2 : list_type A.
  Proof.
    intros. repeat solve_typing.
    Qed.

  (* d) Define a function head of type list A → A + 1 *)
  Definition head : val := 
    λ: "l", "l" <> (InjR #LitUnit) (λ: "h" <>, InjL "h").

  Lemma head_typed n Γ A :
    type_wf n A →
    TY n; Γ ⊢ head: (list_type A → (A + Unit)).
  Proof.
    intros. solve_typing.
    Qed.

  (* e) Define a function [tail] of type list A → list A *)
  Definition split : val :=
    λ: "l", "l" <> ((InjR #LitUnit), nil_val)
      (λ: "h" "r", match: (Fst "r") with InjL "h'" => (InjL "h", let: "r'" := Snd "r" in cons_expr "h'" "r'")
                                       | InjR <> => (InjL "h", Snd "r")
                   end).
  Definition tail : val := 
    λ: "l", Snd (split "l").

  Lemma tail_typed n Γ A :
    type_wf n A →
    TY n; Γ ⊢ tail: (list_type A → list_type A).
  Proof.
    intros. repeat solve_typing.
    Admitted.

End church_encodings.

Section free_theorems.

  (** Exercise 4 (LN Exercise 27): Free Theorems I *)
  (* a) State a free theorem for the type ∀ α, β. α → β → α × β *)
  Lemma free_thm_1 :
    ∀ f : val,
    TY 0; ∅ ⊢ f : (∀: ∀: #1 → #0 → #1 × #0) →
    ∀ (v1 v2 : val), is_closed [] v1 → is_closed [] v2 → 
      big_step (f <> <> v1 v2) (v1, v2)%V.
  Proof.
    intros f [Htycl Hty]%sem_soundness v1 v2 Hcl1 Hcl2.
    specialize (Hty ∅ δ_any). simp type_interp in Hty.
    destruct Hty as (v & Hb & Hv).
    { constructor. }

    simp type_interp in Hv. destruct Hv as (e & -> & ? & Hv).
    specialize_sem_type Hv with (λ v, v = v1) as τ1.
    { naive_solver. }
    simp type_interp in Hv. destruct Hv as (ve0 & ? & Hv).

    simp type_interp in Hv. destruct Hv as (e2 & -> & ? & Hv).
    specialize_sem_type Hv with (λ v, v = v2) as τ2.
    { naive_solver. }
    simp type_interp in Hv. destruct Hv as (ve1 & ? & Hv).

    simp type_interp in Hv. destruct Hv as (x & e0 & -> & ? & Hv).
    specialize (Hv v1). simp type_interp in Hv. simpl in Hv.
    destruct Hv as (ve2 & ? & Hv); first done.

    simp type_interp in Hv. destruct Hv as (x' & e1 & -> & ? & Hv).
    specialize (Hv v2). simp type_interp in Hv. simpl in Hv.
    destruct Hv as (ve3 & ? & Hv); first done.

    simp type_interp in Hv. destruct Hv as (v1' & v2' & -> & Hv1 & Hv2).
    simp type_interp in Hv1. simpl in Hv1. subst v1'.
    simp type_interp in Hv2. simpl in Hv2. subst v2'.

    bs_step_det.
    by rewrite subst_map_empty in Hb.
  Qed.

  (* b) State a free theorem for the type ∀ α, β. α × β → α *)
  Lemma free_thm_2 :
    ∀ f : val,
    TY 0; ∅ ⊢ f : (∀: ∀: #1 × #0 → #1) →
    ∀ (v1 v2 : val), is_closed [] v1 → is_closed [] v2 → 
      big_step (f <> <> (v1, v2)%E) v1.
  Proof.
    intros f [Htycl Hty]%sem_soundness v1 v2 Hcl1 Hcl2.
    specialize (Hty ∅ δ_any). simp type_interp in Hty.
    destruct Hty as (v & Hb & Hv).
    { constructor. }

    simp type_interp in Hv. destruct Hv as (? & -> & ? & Hv).
    specialize_sem_type Hv with (λ v, v = v1) as τ1.
    { naive_solver. }
    simp type_interp in Hv. destruct Hv as (? & ? & Hv).

    simp type_interp in Hv. destruct Hv as (? & -> & ? & Hv).
    specialize_sem_type Hv with (λ v, v = v2) as τ2.
    { naive_solver. }
    simp type_interp in Hv. destruct Hv as (? & ? & Hv).

    simp type_interp in Hv. destruct Hv as (? & ? & -> & ? & Hv).
    specialize (Hv (v1, v2)%V). simp type_interp in Hv. simpl in Hv.
    destruct Hv as (ve & ? & Hv).
    { exists v1, v2. split_and!; first done.
      all: simp type_interp; simpl; done.
    }

    simp type_interp in Hv. simpl in Hv; subst ve.
    rewrite subst_map_empty in Hb.
    bs_step_det.
  Qed.

  (* c) State a free theorem for the type ∀ α, β. α → β *)
  Lemma free_thm_3 :
    ∀ f : val,
    TY 0; ∅ ⊢ f : (∀: ∀: #1 → #0) →
    False.
  Proof.
    intros f [Htycl Hty]%sem_soundness.
    specialize (Hty ∅ δ_any). simp type_interp in Hty.
    destruct Hty as (v & Hb & Hv).
    { constructor. }

    simp type_interp in Hv. destruct Hv as (? & -> & ? & Hv).
    specialize_sem_type Hv with (λ v, v = #0) as τ1.
    { naive_solver. }
    simp type_interp in Hv. destruct Hv as (? & ? & Hv).

    simp type_interp in Hv. destruct Hv as (? & -> & ? & Hv).
    specialize_sem_type Hv with (λ v, False) as τ2.
    { naive_solver. }
    simp type_interp in Hv. destruct Hv as (? & ? & Hv).

    simp type_interp in Hv. destruct Hv as (? & ? & -> & ? & Hv).
    specialize (Hv #0). simp type_interp in Hv. simpl in Hv.
    destruct Hv as (ve & ? & Hv). { done. }

    (* Oh no! *)
    simp type_interp in Hv. simpl in Hv. done.
  Qed.


  (** Exercise 5 (LN Exercise 28): Free Theorems II *)
  Lemma free_theorem_either :
    ∀ f : val,
    TY 0; ∅ ⊢ f : (∀: #0 → #0 → #0) →
    ∀ (v1 v2 : val), is_closed [] v1 → is_closed [] v2 →
      big_step (f <> v1 v2) v1 ∨ big_step (f <> v1 v2) v2.
  Proof.
    intros f [Htycl Hty]%sem_soundness v1 v2 Hcl1 Hcl2.
    specialize (Hty ∅ δ_any). simp type_interp in Hty.
    destruct Hty as (v & Hb & Hv).
    { constructor. }

    simp type_interp in Hv. destruct Hv as (? & -> & ? & Hv).
    specialize_sem_type Hv with (λ v, v = v1 ∨ v = v2) as τ.
    { naive_solver. }
    simp type_interp in Hv. destruct Hv as (? & ? & Hv).

    simp type_interp in Hv. destruct Hv as (? & ? & -> & ? & Hv).
    specialize (Hv v1). simp type_interp in Hv. simpl in Hv.
    destruct Hv as (ve & ? & Hv). { by left. }

    simp type_interp in Hv. destruct Hv as (? & ? & -> & ? & Hv).
    specialize (Hv v2). simp type_interp in Hv. simpl in Hv.
    destruct Hv as (ve & ? & Hv). { by right. }

    simp type_interp in Hv. simpl in Hv.

    rewrite subst_map_empty in Hb.
    destruct Hv as [-> | ->]; [left | right]; bs_step_det.
  Qed.

  (** Exercise 6 (LN Exercise 29): Free Theorems III *)
  (* Hint: you might want to use the fact that our reduction is deterministic. *)
  Lemma big_step_det e v1 v2 :
    big_step e v1 → big_step e v2 → v1 = v2.
  Proof.
    induction 1 in v2 |-*; inversion 1; subst; eauto 2.
    all: naive_solver.
  Qed.

  Lemma free_theorems_magic :
    ∀ (A A1 A2 : type) (f g : val),
    type_wf 0 A → type_wf 0 A1 → type_wf 0 A2 →
    is_closed [] f → is_closed [] g →
    TY 0; ∅ ⊢ f : (∀: (A1 → A2 → #0) → #0) →
    TY 0; ∅ ⊢ g : (A1 → A2 → A) →
    ∀ v, big_step (f <> g) v →
    ∃ (v1 v2 : val), big_step (g v1 v2) v.
  Proof.
    (* Hint: you may find the following lemmas useful:
        - [sem_val_rel_cons]
        - [type_wf_closed]
        - [val_rel_is_closed]
        - [big_step_preserve_closed]
    *)
    intros A A1 A2 f g HwfA HwfA1 HwfA2 Hclf Hclg [Htyfcl Htyf]%sem_soundness [Htygcl Htyg]%sem_soundness v Hb.

    specialize (Htyf ∅ δ_any). simp type_interp in Htyf.
    destruct Htyf as (vf & Hbf & Hvf).
    { constructor. }

    specialize (Htyg ∅ δ_any). simp type_interp in Htyg.
    destruct Htyg as (vg & Hbg & Hvg).
    { constructor. }

    rewrite subst_map_empty in Hbf. rewrite subst_map_empty in Hbg.
    apply big_step_val in Hbf. apply big_step_val in Hbg.
    subst vf vg.

    (* if we know that big_step is deterministic *)
    (* We pick the interpretation [(λ v, ∃ v1 v2, big_step (g v1 v2) v)].
       Then we can equate the existential we get from Hvf with v,
        since big step is deterministic.

      We need to show that g satisfies this interpretation.
      For that, we already get v1: A1, v2:A2.
      So we use the Hvg fact to get a : A with g v1 v2 ↓ a.
      With that we can show the semantic interpretation.
     *)

    simp type_interp in Hvf. destruct Hvf as (ef & -> & ? & Hvf).
    specialize_sem_type Hvf with (λ v,
      ∃ (v1 v2 : val), is_closed [] v1 ∧ is_closed [] v2 ∧ big_step (g v1 v2) v) as τ.
    { intros v' (v1 & v2 & ? & ? & Hb').
      eapply big_step_preserve_closed.
      2: apply Hb'.
      simpl. rewrite !andb_True. done.
    }
    simp type_interp in Hvf. destruct Hvf as (ve & ? & Hvf).
    simp type_interp in Hvf. destruct Hvf as (? & ef' & -> & ? & Hvf).

    specialize (Hvf g). simp type_interp in Hvf.
    destruct Hvf as (ve2 & Hbe2 & Hvf).
    { simp type_interp in Hvg. destruct Hvg as (xg0 & eg0 & -> & ? & Hvg).
      eexists _, _. split_and!; [done | done | ].
      intros v1 Hv1.

      specialize (Hvg v1). simp type_interp in Hvg.
      destruct Hvg as (? & ? & Hvg).
      { eapply sem_val_rel_cons. rewrite type_wf_closed; done. }

      simp type_interp in Hvg. destruct Hvg as (xg1 & eg1 & -> & ? & Hvg).
      simp type_interp. eexists _. split; first done.
      simp type_interp. eexists _, _. split_and!; [done | done | ].
      intros v2 Hv2.

      specialize (Hvg v2). simp type_interp in Hvg.
      destruct Hvg as (? & ? & Hvg).
      { eapply sem_val_rel_cons. rewrite type_wf_closed; done. }

      simp type_interp. eexists. split; first done.
      simp type_interp. simpl.

      exists v1, v2. split_and!.
      - eapply val_rel_closed; done.
      - eapply val_rel_closed; done.
      - bs_step_det.
    }
    simp type_interp in Hvf. simpl in Hvf.
    destruct Hvf as (v1 & v2 & ? & ? & Hbs).
    exists v1, v2.
    assert (ve2 = v) as ->; last done.
    eapply big_step_det; last apply Hb.
    bs_step_det.
  Qed.

End free_theorems.