From machine_program_logic.program_logic Require Import weakestpre.
From HypVeri Require Import lifting rules.rules_base stdpp_extra.
From HypVeri.algebra Require Import base reg mem pagetable trans mailbox base_extra.
From HypVeri.lang Require Import lang_extra mem_extra reg_extra pagetable_extra trans_extra.

Section mem_relinquish.

Context `{hypparams: HypervisorParameters}.
Context `{vmG: !gen_VMG Σ}.

Lemma p_relinquish_inv_consist σ h i j ps tt:
  inv_trans_pgt_consistent σ ->
  inv_trans_ps_disj σ ->
  σ.2 !! h = Some (Some (j, i, ps, tt, true)) ->
  j ≠ i ->
  inv_trans_pgt_consistent (update_page_table_global revoke_access (update_transaction σ h (j, i, ps, tt, false)) i ps).
Proof.
  intros Hinv_con Hinv_disj Hlk Hneq_i.
  rewrite /inv_trans_pgt_consistent /inv_trans_pgt_consistent' /=.
  rewrite map_Forall_lookup.
  intros h' meta Hlookup'.
  rewrite lookup_insert_Some in Hlookup'.
  destruct Hlookup' as [[<- <-]|[Hneq Hlookup']].
  {
    intros p Hin.
    specialize (Hinv_con h _ Hlk p Hin).
    simpl in Hinv_con.
    generalize dependent σ.1.1.1.2.
    generalize dependent σ.2.
    induction ps using set_ind_L .
    - set_solver + Hin.
    - intros tran Htran pgt Hpgt.
      simpl.
      rewrite set_fold_disj_union_strong.
      {
        rewrite set_fold_singleton.
        destruct (decide (x = p)).
        {
          subst.
          destruct tt;first done.
          rewrite Hpgt.
          apply p_upd_pgt_pgt_not_elem.
          done.
          rewrite lookup_insert_Some.
          left. split;auto.
          assert (Hrw: {[j;i]} ∖ {[i]} = ({[j]} : gset _)).
          set_solver + Hneq_i.
          rewrite /revoke_access Hrw //.
          rewrite Hpgt.
          apply p_upd_pgt_pgt_not_elem.
          done.
          rewrite lookup_insert_Some.
          left. split;auto.
          assert (Hrw: {[i]} ∖ {[i]} = (∅ : gset _)).
          set_solver +.
          rewrite /revoke_access Hrw //.
        }
        {
          destruct (pgt !! x).
          {
            feed specialize IHps.
            set_solver + Hin n.
            apply (IHps (<[h := Some (j, i, X, tt, true)]>tran));eauto.
            rewrite lookup_insert //.
            rewrite lookup_insert_ne //.
          }
          {
            feed specialize IHps.
            set_solver + Hin n.
            apply (IHps (<[h := Some (j, i, X, tt, true)]>tran));eauto.
            rewrite lookup_insert //.
          }
        }
      }
      apply upd_is_strong_assoc_comm.
      set_solver + H0.
  }
  {
    rewrite /inv_trans_pgt_consistent /inv_trans_pgt_consistent' /= in Hinv_con.
    specialize (Hinv_con h' meta Hlookup').
    simpl in Hinv_con.
    destruct meta as [[[[[sv rv] ps'] tt'] b]|];last done.
    simpl in *.
    intros p Hin.
    specialize (Hinv_con p Hin).
    assert (p ∉ ps).
    {
      intro.
      specialize (Hinv_disj h' _ Hlookup').
      simpl in Hinv_disj.
      pose proof (elem_of_pages_in_trans' p (delete h' σ.2)) as [_ Hin'].
      feed specialize Hin'.
      exists h.
      eexists.
      split.
      rewrite lookup_delete_ne //.
      done.
      set_solver + Hin H0 Hin' Hinv_disj.
    }
    destruct tt',b;auto;try apply p_upd_pgt_pgt_not_elem;auto.
  }
Qed.

Lemma mem_relinquish {tt wi sacc i j q p_tx} {ps: gset PID}
      ai r0 wh:
  (* has access to the page which the instruction is in *)
  tpa ai ≠ p_tx ->
  tpa ai ∈ sacc ->
  (* current instruction is hvc *)
  decode_instruction wi = Some(Hvc) ->
  (* the hvc call to invoke is relinquish *)
  decode_hvc_func r0 = Some(Relinquish) ->
  {SS{{(* the encoding of instruction wi is stored in location ai *)
       ▷(PC @@ i ->r ai) ∗  ▷ ai ->a wi ∗
       ▷ (R0 @@ i ->r r0) ∗
       ▷ (R1 @@ i ->r wh) ∗
       (* the pagetable, the owership ra is not required *)
       ▷ i -@A> sacc ∗
       (* the descriptor is ready in the tx page *)
       ▷ TX@ i := p_tx ∗
       (* is the receiver and the transaction has been retrieved *)
       ▷ wh -{q}>t (j, i, ps, tt) ∗ ▷ wh ->re true }}}
  ExecI @ i
  {{{ RET (false, ExecI) ;
      (* PC is incremented *)
      PC @@ i ->r (ai ^+ 1)%f ∗ ai ->a wi ∗
      (* donesn't have access to psd anymore *)
      i -@A> (sacc ∖ ps) ∗
      (* return Succ to R0 *)
      R0 @@ i ->r (encode_hvc_ret_code Succ) ∗
      R1 @@ i ->r wh ∗
      (* the same tx *)
      TX@ i := p_tx ∗
      (* the transaction is marked as unretrieved *)
      wh -{q}>t(j, i, ps, tt) ∗ wh ->re false
      }}}.
Proof.
  iIntros (Hneq_tx Hin_acc Hdecode_i Hdecode_f Φ) "(>PC & >mem_ins & >R0 & >R1 & >acc & >tx & >tran & >re) HΦ".
  iApply (sswp_lift_atomic_step ExecI);[done|].
  iIntros (n σ1) "%Hsche state".
  rewrite /scheduled /= /scheduler in Hsche.
  assert (σ1.1.1.2 = i) as Heq_cur. { case_bool_decide;last done. by apply fin_to_nat_inj. }
  clear Hsche.
  iModIntro.
  iDestruct "state" as "(Hnum & mem & regs & mb & rx_state & pgt_owned & pgt_acc & pgt_excl &
                            trans & hpool & retri & %Hwf & %Hdisj & %Hconsis)".
  (* valid regs *)
  iDestruct ((gen_reg_valid3 i PC ai R0 r0 R1 wh Heq_cur) with "regs PC R0 R1")
    as "(%Hlookup_PC & %Hlookup_R0 & %Hlookup_R1)";eauto.
  (* valid pt *)
  iDestruct (access_agree_check_true (tpa ai) i with "pgt_acc acc") as %Hcheckpg_ai;eauto.
  (* valid mem *)
  iDestruct (gen_mem_valid ai wi with "mem mem_ins") as %Hlookup_ai.
  (* valid tx rx *)
  iDestruct (mb_valid_tx i p_tx with "mb tx") as %Heq_tx.
  (* valid trans *)
  iDestruct (trans_valid_Some with "trans tran") as %[re Hlookup_tran].
  iDestruct (trans_valid_handle_Some with "tran") as %Hvalid_handle.
  iDestruct (retri_valid_Some with "retri re") as %[meta Hlookup_tran'].
  rewrite Hlookup_tran in Hlookup_tran'.
  inversion Hlookup_tran'. subst re. clear meta Hlookup_tran' H1.
  iSplit.
  - (* reducible *)
    iPureIntro.
    apply (reducible_normal i Hvc ai wi);eauto.
    rewrite Heq_tx //.
  - iModIntro.
    iIntros (m2 σ2) "vmprop_auth %HstepP".
    iFrame "vmprop_auth".
    apply (step_ExecI_normal i Hvc ai wi) in HstepP;eauto.
    2: rewrite Heq_tx //.
    remember (exec Hvc σ1) as c2 eqn:Heqc2.
    rewrite /exec /hvc Hlookup_R0 /= Hdecode_f /relinquish Hlookup_R1 /get_transaction /= Hlookup_tran Heq_cur /= in Heqc2.
    case_bool_decide;last done. clear H0. simpl in Heqc2.
    case_bool_decide;last done. clear H0.
    destruct HstepP;subst m2 σ2; subst c2; simpl.
    rewrite /gen_vm_interp.
    (* unchanged part *)
    rewrite (preserve_get_rx_gmap σ1).
    2: rewrite p_upd_pc_mb //.
    rewrite (preserve_get_mb_gmap σ1).
    2: rewrite p_upd_pc_mb //.
    rewrite p_upd_pc_mem p_upd_reg_mem p_rvk_acc_mem p_upd_tran_mem.
    iFrame "Hnum mem mb rx_state".
    (* upd regs *)
    rewrite (u_upd_pc_regs _ i ai). 2: done.
    2: { rewrite u_upd_reg_regs p_rvk_acc_current_vm p_upd_tran_current_vm.
         rewrite (preserve_get_reg_gmap σ1). rewrite lookup_insert_ne. solve_reg_lookup. done. f_equal.
    }
    rewrite u_upd_reg_regs p_rvk_acc_current_vm p_upd_tran_current_vm Heq_cur.
    rewrite (preserve_get_reg_gmap σ1);last done.
    iDestruct ((gen_reg_update2_global PC i _ (ai ^+ 1)%f R0 i _ (encode_hvc_ret_code Succ)) with "regs PC R0")
      as ">[$ [PC R0]]";eauto.
    (* upd pgt *)
    rewrite (preserve_get_own_gmap (update_page_table_global revoke_access (update_transaction σ1 wh (j, i, ps, tt, false)) i ps) (update_incr_PC _)).
    2: rewrite p_upd_pc_pgt p_upd_reg_pgt //.
    rewrite p_rvk_acc_own. rewrite (preserve_get_own_gmap σ1);last done.
    iFrame "pgt_owned".
    rewrite (preserve_get_access_gmap (update_page_table_global revoke_access (update_transaction σ1 wh (j, i, ps, tt, true)) i ps) (update_incr_PC _)).
    2: rewrite p_upd_pc_pgt p_upd_reg_pgt //.
    iDestruct (access_agree with "pgt_acc acc") as %Hlookup_pgt_acc.
    rewrite (u_rvk_acc_acc _ _ _ sacc).
    2: {
      rewrite p_upd_tran_pgt.
      intros p Hin_p.
      specialize (Hconsis wh _ Hlookup_tran p Hin_p).
      simpl in Hconsis.
      destruct tt.
      done.
      eexists;eauto.
      eexists;eauto.
    }
    2: rewrite (preserve_get_access_gmap σ1);done.
    rewrite (preserve_get_access_gmap σ1);last done.
    iDestruct (access_update (sacc∖ ps) with "pgt_acc acc") as ">[$ acc]". done.
    rewrite (preserve_get_excl_gmap (update_page_table_global revoke_access (update_transaction σ1 wh (j, i, ps, tt, false)) i ps) (update_incr_PC _)).
    2: rewrite p_upd_pc_pgt p_upd_reg_pgt //.
    rewrite (p_rvk_acc_excl _ _ j tt).
    2: {
      intros p Hin.
      specialize (Hconsis wh _ Hlookup_tran p Hin).
      destruct tt; simpl in Hconsis.
      done.
      rewrite p_upd_tran_pgt //.
      rewrite p_upd_tran_pgt //.
    }
    rewrite (preserve_get_excl_gmap σ1);last done.
    iFrame "pgt_excl".
    (* upd tran *)
    rewrite (preserve_get_trans_gmap (update_transaction σ1 wh (j, i, ps, tt, false)) (update_incr_PC _)).
    2: rewrite p_upd_pc_trans p_upd_reg_trans //.
    rewrite u_upd_tran_trans.
    rewrite insert_id.
    2: rewrite /get_trans_gmap /get_transactions_gmap lookup_fmap Hlookup_tran //=.
    iFrame "trans".
    (* upd hp *)
    rewrite (preserve_get_hpool_gset (update_transaction σ1 wh (j, i, ps, tt, false)) (update_incr_PC _)).
    2: rewrite p_upd_pc_trans p_upd_reg_trans //. rewrite p_upd_tran_hp.
    iFrame "hpool".
    (* upd retri *)
    rewrite (preserve_get_retri_gmap (update_transaction σ1 wh (j, i, ps, tt, false)) (update_incr_PC _)).
    2: rewrite p_upd_pc_trans p_upd_reg_trans //.
    2: { exists (Some (j, i, ps, tt, true)). split;eauto. }
    rewrite u_upd_tran_retri.
    iDestruct (retri_update_flip with "retri re") as ">[$ re]".
    (* inv_trans_wellformed *)
    rewrite (preserve_inv_trans_wellformed (update_transaction σ1 wh (j, i, ps, tt, false))).
    2: rewrite p_upd_pc_trans p_upd_reg_trans //.
    iAssert (⌜inv_trans_wellformed (update_transaction σ1 wh (j, i, ps, tt, false))⌝%I) as "$". iPureIntro.
    apply (p_upd_tran_inv_wf σ1 wh);eauto.
    (* inv_trans_pgt_consistent *)
    rewrite (preserve_inv_trans_pgt_consistent (update_page_table_global revoke_access (update_transaction σ1 wh (j, i, ps, tt, false)) i ps) (update_incr_PC _)).
    2: rewrite p_upd_pc_trans p_upd_reg_trans //.
    2: rewrite p_upd_pc_pgt p_upd_reg_pgt //.
    iAssert (⌜inv_trans_pgt_consistent (update_page_table_global revoke_access (update_transaction σ1 wh (j, i, ps, tt, false)) i ps)⌝%I) as "$". iPureIntro.
    apply p_relinquish_inv_consist;auto.
    { destruct Hwf as [_ [Hwf _]]. specialize (Hwf wh _ Hlookup_tran). done. }
    (* inv_trans_ps_disj *)
    rewrite (preserve_inv_trans_ps_disj (update_transaction σ1 wh (j, i, ps, tt, false))).
    2: rewrite p_upd_pc_trans p_upd_reg_trans //.
    iAssert (⌜inv_trans_ps_disj (update_transaction σ1 wh (j, i, ps, tt, false))⌝%I) as "$". iPureIntro.
    eapply p_upd_tran_inv_disj.
    apply Hdisj.
    exact Hlookup_tran.
    done.
    (* just_scheduled *)
    iModIntro.
    rewrite /just_scheduled_vms /just_scheduled.
    rewrite /scheduled /machine.scheduler /= /scheduler.
    rewrite p_upd_pc_current_vm p_upd_reg_current_vm p_rvk_acc_current_vm p_upd_tran_current_vm.
    rewrite Heq_cur.
    iSplitL "".
    set fl := (filter _ _).
    assert (fl = []) as ->.
    {
      rewrite /fl.
      induction n.
      - simpl.
        rewrite filter_nil //=.
      - rewrite seq_S.
        rewrite filter_app.
        rewrite IHn.
        simpl.
        rewrite filter_cons_False /=. rewrite filter_nil //.
        rewrite andb_negb_l //.
    }
    by iSimpl.
    (* Φ *)
    case_bool_decide;last contradiction.
    simpl. iApply "HΦ".
    rewrite /fresh_handles. iFrame.
Qed.

Lemma mem_relinquish_invalid_handle {E i wi sacc r0 r2 wh p_tx} ai:
  tpa ai ≠ p_tx ->
  (tpa ai) ∈ sacc ->
  (* the current instruction is hvc *)
  (* the decoding of wi is correct *)
  decode_instruction wi = Some(Hvc) ->
  (* the hvc call to invoke is retrieve *)
  decode_hvc_func r0 = Some(Relinquish) ->
  wh ∉ valid_handles ->
  {SS{{(* the encoding of instruction wi is stored in location ai *)
       ▷ (PC @@ i ->r ai) ∗ ▷ ai ->a wi ∗
       (* registers *)
       ▷ (R0 @@ i ->r r0) ∗
       ▷ (R1 @@ i ->r wh) ∗
       ▷ (R2 @@ i ->r r2) ∗
       ▷ i -@A> sacc ∗
       ▷ TX@ i := p_tx}}}
   ExecI @ i; E
   {{{ RET (false, ExecI) ;
       (* PC is incremented *)
       PC @@ i ->r (ai ^+ 1)%f ∗ ai ->a wi ∗
       R0 @@ i ->r (encode_hvc_ret_code Error) ∗
       R1 @@ i ->r wh ∗
       R2 @@ i ->r (encode_hvc_error InvParam) ∗
       i -@A> sacc ∗
       TX@ i := p_tx
   }}}.
Proof.
  iIntros (Hneq_tx Hin_acc Hdecode_i Hdecode_f Hnin_wh Φ)
          "(>PC & >mem_ins & >R0 & >R1 & >R2 & >acc & >tx) HΦ".
  iApply (sswp_lift_atomic_step ExecI);[done|].
  iIntros (n σ1) "%Hsche state".
  rewrite /scheduled /= /scheduler in Hsche.
  assert (σ1.1.1.2 = i) as Heq_cur. { case_bool_decide;last done. by apply fin_to_nat_inj. }
  clear Hsche.
  iModIntro.
  iDestruct "state" as "(Hnum & mem & regs & mb & rx_state & pgt_owned & pgt_acc & pgt_excl &
                            trans & hpool & retri & %Hwf & %Hdisj & %Hconsis)".
  (* valid regs *)
  iDestruct ((gen_reg_valid4 i PC ai R0 r0 R1 wh R2 r2 Heq_cur) with "regs PC R0 R1 R2")
    as "(%Hlookup_PC & %Hlookup_R0 & %Hlookup_R1 & %Hlookup_R2)";eauto.
  (* valid pt *)
  iDestruct (access_agree_check_true (tpa ai) i with "pgt_acc acc") as %Hcheckpg_ai;eauto.
  (* valid mem *)
  iDestruct (gen_mem_valid ai wi with "mem mem_ins") as %Hlookup_ai.
  (* valid tx *)
  iDestruct (mb_valid_tx i p_tx with "mb tx") as %Heq_tx.
  (* valid tran *)
  iSplit.
  - (* reducible *)
    iPureIntro.
    apply (reducible_normal i Hvc ai wi);auto.
    rewrite Heq_tx //.
  - iModIntro.
    iIntros (m2 σ2) "vmprop_auth %HstepP".
    iFrame "vmprop_auth".
    apply (step_ExecI_normal i Hvc ai wi) in HstepP;eauto.
    2: rewrite Heq_tx //.
    remember (exec Hvc σ1) as c2 eqn:Heqc2.
    rewrite /exec /hvc Hlookup_R0 /= Hdecode_f /relinquish Hlookup_R1 /= in Heqc2.
    assert (Hwh_None: get_transaction σ1 wh = None).
    {
      destruct Hwf as [_ [_ Hwf]].
      rewrite /inv_finite_handles in Hwf.
      rewrite Hwf in Hnin_wh.
      rewrite not_elem_of_dom in Hnin_wh.
      rewrite /get_transaction Hnin_wh //.
      case_bool_decide;done.
    }
    rewrite Hwh_None /= in Heqc2.
    destruct HstepP;subst m2 σ2; subst c2; simpl.
    iDestruct (hvc_error_update (E:= E) InvParam with "PC R0 R2 [$Hnum $mem $regs $mb $rx_state $pgt_owned $pgt_acc $pgt_excl $ trans $hpool $retri]")
    as ">[[$ $] ?]". 1-4: auto. iPureIntro. auto.
    rewrite /scheduled /machine.scheduler /= /scheduler.
    rewrite p_upd_pc_current_vm 2!p_upd_reg_current_vm Heq_cur.
    case_bool_decide;last contradiction.
    simpl. iApply "HΦ".
    iFrame.
    by iFrame.
Qed.

Lemma mem_relinquish_fresh_handle {E i wi sacc r0 r2 wh sh q p_tx} ai:
  tpa ai ≠ p_tx ->
  (tpa ai) ∈ sacc ->
  decode_instruction wi = Some(Hvc) ->
  decode_hvc_func r0 = Some(Relinquish) ->
  wh ∈ sh ->
  {SS{{▷ (PC @@ i ->r ai) ∗ ▷ ai ->a wi ∗
       ▷ (R0 @@ i ->r r0) ∗
       ▷ (R1 @@ i ->r wh) ∗
       ▷ (R2 @@ i ->r r2) ∗
       ▷ i -@A> sacc ∗
       ▷ TX@ i := p_tx ∗
       ▷ fresh_handles q sh}}}
   ExecI @ i; E
   {{{ RET (false, ExecI) ;
       PC @@ i ->r (ai ^+ 1)%f ∗ ai ->a wi ∗
       R0 @@ i ->r (encode_hvc_ret_code Error) ∗
       R1 @@ i ->r wh ∗
       R2 @@ i ->r (encode_hvc_error InvParam) ∗
       i -@A> sacc ∗
       TX@ i := p_tx ∗
       fresh_handles q sh
   }}}.
Proof.
  iIntros (Hneq_tx Hin_acc Hdecode_i Hdecode_f Hin_wh Φ)
          "(>PC & >mem_ins & >R0 & >R1 & >R2 & >acc & >tx & >[hp handles]) HΦ".
  iApply (sswp_lift_atomic_step ExecI);[done|].
  iIntros (n σ1) "%Hsche state".
  rewrite /scheduled /= /scheduler in Hsche.
  assert (σ1.1.1.2 = i) as Heq_cur. { case_bool_decide;last done. by apply fin_to_nat_inj. }
  clear Hsche.
  iModIntro.
  iDestruct "state" as "(Hnum & mem & regs & mb & rx_state & pgt_owned & pgt_acc & pgt_excl &
                            trans & hpool & retri & %Hwf & %Hdisj & %Hconsis)".
  (* valid regs *)
  iDestruct ((gen_reg_valid4 i PC ai R0 r0 R1 wh R2 r2 Heq_cur) with "regs PC R0 R1 R2")
    as "(%Hlookup_PC & %Hlookup_R0 & %Hlookup_R1 & %Hlookup_R2)";eauto.
  (* valid pt *)
  iDestruct (access_agree_check_true (tpa ai) i with "pgt_acc acc") as %Hcheckpg_ai;eauto.
  (* valid mem *)
  iDestruct (gen_mem_valid ai wi with "mem mem_ins") as %Hlookup_ai.
  (* valid tx *)
  iDestruct (mb_valid_tx i p_tx with "mb tx") as %Heq_tx.
  (* valid hpool *)
  iDestruct (hpool_valid with "hpool hp") as %Heq_hp.
  (* valid tran *)
  iAssert (⌜get_transaction σ1 wh = None⌝%I) as %Hwh_None.
  {
    iDestruct (big_sepS_elem_of _ _ wh with "handles") as "[tran _]".
    done.
    iDestruct (trans_valid_None with "trans tran") as %Hlookup_tran.
    iPureIntro.
    rewrite /get_transaction Hlookup_tran //.
    case_bool_decide;done.
  }
  iSplit.
  - (* reducible *)
    iPureIntro.
    apply (reducible_normal i Hvc ai wi);auto.
    rewrite Heq_tx //.
  - iModIntro.
    iIntros (m2 σ2) "vmprop_auth %HstepP".
    iFrame "vmprop_auth".
    apply (step_ExecI_normal i Hvc ai wi) in HstepP;eauto.
    2: rewrite Heq_tx //.
    remember (exec Hvc σ1) as c2 eqn:Heqc2.
    rewrite /exec /hvc Hlookup_R0 /= Hdecode_f /relinquish Hlookup_R1 /= in Heqc2.
    rewrite Hwh_None /= in Heqc2.
    destruct HstepP;subst m2 σ2; subst c2; simpl.
    iDestruct (hvc_error_update (E:= E) InvParam with "PC R0 R2 [$Hnum $mem $regs $mb $rx_state $pgt_owned $pgt_acc $pgt_excl $ trans $hpool $retri]")
    as ">[[$ $] ?]". 1-4: auto. iPureIntro. auto.
    rewrite /scheduled /machine.scheduler /= /scheduler.
    rewrite p_upd_pc_current_vm 2!p_upd_reg_current_vm Heq_cur.
    case_bool_decide;last contradiction.
    simpl. iApply "HΦ".
    iFrame.
    by iFrame.
Qed.

Lemma mem_relinquish_invalid_trans {E i wi sacc r0 r2 wh meta q p_tx} ai:
  tpa ai ≠ p_tx ->
  (tpa ai) ∈ sacc ->
  decode_instruction wi = Some(Hvc) ->
  decode_hvc_func r0 = Some(Relinquish) ->
  meta.1.1.2 ≠ i ->
  {SS{{▷ (PC @@ i ->r ai) ∗ ▷ ai ->a wi ∗
       ▷ (R0 @@ i ->r r0) ∗
       ▷ (R1 @@ i ->r wh) ∗
       ▷ (R2 @@ i ->r r2) ∗
       ▷ i -@A> sacc ∗
       ▷ TX@ i := p_tx ∗
       ▷ wh -{q}>t (meta)
       }}}
   ExecI @ i; E
   {{{ RET (false, ExecI) ;
       PC @@ i ->r (ai ^+ 1)%f ∗ ai ->a wi ∗
       R0 @@ i ->r (encode_hvc_ret_code Error) ∗
       R1 @@ i ->r wh ∗
       R2 @@ i ->r (encode_hvc_error Denied) ∗
       i -@A> sacc ∗
       TX@ i := p_tx ∗
       wh -{q}>t (meta)
   }}}.
Proof.
  iIntros (Hneq_tx Hin_acc Hdecode_i Hdecode_f Hin_wh Φ)
          "(>PC & >mem_ins & >R0 & >R1 & >R2 & >acc & >tx & >tran) HΦ".
  iApply (sswp_lift_atomic_step ExecI);[done|].
  iIntros (n σ1) "%Hsche state".
  rewrite /scheduled /= /scheduler in Hsche.
  assert (σ1.1.1.2 = i) as Heq_cur. { case_bool_decide;last done. by apply fin_to_nat_inj. }
  clear Hsche.
  iModIntro.
  iDestruct "state" as "(Hnum & mem & regs & mb & rx_state & pgt_owned & pgt_acc & pgt_excl &
                            trans & hpool & retri & %Hwf & %Hdisj & %Hconsis)".
  (* valid regs *)
  iDestruct ((gen_reg_valid4 i PC ai R0 r0 R1 wh R2 r2 Heq_cur) with "regs PC R0 R1 R2")
    as "(%Hlookup_PC & %Hlookup_R0 & %Hlookup_R1 & %Hlookup_R2)";eauto.
  (* valid pt *)
  iDestruct (access_agree_check_true (tpa ai) i with "pgt_acc acc") as %Hcheckpg_ai;eauto.
  (* valid mem *)
  iDestruct (gen_mem_valid ai wi with "mem mem_ins") as %Hlookup_ai.
  (* valid tx *)
  iDestruct (mb_valid_tx i p_tx with "mb tx") as %Heq_tx.
  (* valid tran *)
  iDestruct (trans_valid_Some with "trans tran") as %[? Hlookup_tran].
  iDestruct (trans_valid_handle_Some with "tran") as %Hvalid_handle.
  iSplit.
  - (* reducible *)
    iPureIntro.
    apply (reducible_normal i Hvc ai wi);auto.
    rewrite Heq_tx //.
  - iModIntro.
    iIntros (m2 σ2) "vmprop_auth %HstepP".
    iFrame "vmprop_auth".
    apply (step_ExecI_normal i Hvc ai wi) in HstepP;eauto.
    2: rewrite Heq_tx //.
    remember (exec Hvc σ1) as c2 eqn:Heqc2.
    rewrite /exec /hvc Hlookup_R0 /= Hdecode_f /relinquish Hlookup_R1 /= in Heqc2.
    rewrite /get_transaction Hlookup_tran /= in Heqc2.
    destruct meta as [[[? ?] ?] ?].
    case_bool_decide;last contradiction. clear H0. simpl in Heqc2.
    case_bool_decide;rewrite Heq_cur // in H0.
    rewrite andb_false_r /= in Heqc2.
    destruct HstepP;subst m2 σ2; subst c2; simpl.
    iDestruct (hvc_error_update (E:= E) Denied with "PC R0 R2 [$Hnum $mem $regs $mb $rx_state $pgt_owned $pgt_acc $pgt_excl $ trans $hpool $retri]")
    as ">[[$ $] ?]". 1-4: auto. iPureIntro. auto.
    rewrite /scheduled /machine.scheduler /= /scheduler.
    rewrite p_upd_pc_current_vm 2!p_upd_reg_current_vm Heq_cur.
    case_bool_decide;last contradiction.
    simpl. iApply "HΦ".
    iFrame.
    by iFrame.
Qed.

Lemma mem_relinquish_not_retrieved{E i wi sacc r0 r2 wh q p_tx} ai:
  tpa ai ≠ p_tx ->
  (tpa ai) ∈ sacc ->
  decode_instruction wi = Some(Hvc) ->
  decode_hvc_func r0 = Some(Relinquish) ->
  {SS{{▷ (PC @@ i ->r ai) ∗ ▷ ai ->a wi ∗
       ▷ (R0 @@ i ->r r0) ∗
       ▷ (R1 @@ i ->r wh) ∗
       ▷ (R2 @@ i ->r r2) ∗
       ▷ i -@A> sacc ∗
       ▷ TX@ i := p_tx ∗
       ▷ wh -{q}>re false
       }}}
   ExecI @ i; E
   {{{ RET (false, ExecI) ;
       PC @@ i ->r (ai ^+ 1)%f ∗ ai ->a wi ∗
       R0 @@ i ->r (encode_hvc_ret_code Error) ∗
       R1 @@ i ->r wh ∗
       R2 @@ i ->r (encode_hvc_error Denied) ∗
       i -@A> sacc ∗
       TX@ i := p_tx ∗
       wh -{q}>re false
   }}}.
Proof.
  iIntros (Hneq_tx Hin_acc Hdecode_i Hdecode_f Φ)
          "(>PC & >mem_ins & >R0 & >R1 & >R2 & >acc & >tx & >re) HΦ".
  iApply (sswp_lift_atomic_step ExecI);[done|].
  iIntros (n σ1) "%Hsche state".
  rewrite /scheduled /= /scheduler in Hsche.
  assert (σ1.1.1.2 = i) as Heq_cur. { case_bool_decide;last done. by apply fin_to_nat_inj. }
  clear Hsche.
  iModIntro.
  iDestruct "state" as "(Hnum & mem & regs & mb & rx_state & pgt_owned & pgt_acc & pgt_excl &
                            trans & hpool & retri & %Hwf & %Hdisj & %Hconsis)".
  (* valid regs *)
  iDestruct ((gen_reg_valid4 i PC ai R0 r0 R1 wh R2 r2 Heq_cur) with "regs PC R0 R1 R2")
    as "(%Hlookup_PC & %Hlookup_R0 & %Hlookup_R1 & %Hlookup_R2)";eauto.
  (* valid pt *)
  iDestruct (access_agree_check_true (tpa ai) i with "pgt_acc acc") as %Hcheckpg_ai;eauto.
  (* valid mem *)
  iDestruct (gen_mem_valid ai wi with "mem mem_ins") as %Hlookup_ai.
  (* valid tx *)
  iDestruct (mb_valid_tx i p_tx with "mb tx") as %Heq_tx.
  (* valid tran *)
  iDestruct (retri_valid_Some with "retri re") as %[? Hlookup_re].
  iDestruct (retri_valid_handle_Some with "re") as %Hvalid_handle.
  iSplit.
  - (* reducible *)
    iPureIntro.
    apply (reducible_normal i Hvc ai wi);auto.
    rewrite Heq_tx //.
  - iModIntro.
    iIntros (m2 σ2) "vmprop_auth %HstepP".
    iFrame "vmprop_auth".
    apply (step_ExecI_normal i Hvc ai wi) in HstepP;eauto.
    2: rewrite Heq_tx //.
    remember (exec Hvc σ1) as c2 eqn:Heqc2.
    rewrite /exec /hvc Hlookup_R0 /= Hdecode_f /relinquish Hlookup_R1 /= in Heqc2.
    rewrite /get_transaction Hlookup_re /= in Heqc2.
    destruct x as [[[? ?] ?] ?].
    assert (Heq_c2 : (m2,σ2) = (ExecI, update_incr_PC (update_reg (update_reg σ1 R0 (encode_hvc_ret_code Error)) R2 (encode_hvc_error Denied)))).
    {
      case_bool_decide;last contradiction.
      destruct HstepP;subst m2 σ2; subst c2; done.
    }
    inversion Heq_c2. clear H1 H2 Heq_c2 Heqc2.
    iDestruct (hvc_error_update (E:= E) Denied with "PC R0 R2 [$Hnum $mem $regs $mb $rx_state $pgt_owned $pgt_acc $pgt_excl $ trans $hpool $retri]")
    as ">[[$ $] ?]". 1-4: auto. iPureIntro. auto.
    rewrite /scheduled /machine.scheduler /= /scheduler.
    rewrite p_upd_pc_current_vm 2!p_upd_reg_current_vm Heq_cur.
    case_bool_decide;last contradiction.
    simpl. iApply "HΦ".
    iFrame.
    by iFrame.
Qed.

End mem_relinquish.


(* Some boilerplate *)
Fixpoint fib n := match n with
  | O => 1
  | S m => match m with
    | O => 1
    | S o => fib o + fib m end end.

Ltac sleep n :=
  try (cut (fib n = S (fib n)); reflexivity).

(* Tune that depending on your PC *)
Let time := 18.

(* Beginning of demo *)

Section Demo.

Variable i : True.

Lemma a : True.
Proof using i.
  sleep time.
  idtac.
  sleep time.
  (* Error, jump back to fix it, then Qed again *)
  exact (i i).
Qed.

Lemma b : True.
Proof using i.
  sleep time.
  idtac.
  sleep time.
  (* Here we use "a" *)
  exact a.
Qed.

Lemma work_here : True /\ True.
Proof using i.
(* Jump directly here, Coq reacts immediately *)
split.
  exact b.  (* We can use the lemmas above *)
exact a.
Qed.

End Demo.


Require Import Min.

Require Import msl.msl_standard.
Require Import Maps.
Require Import FuncListMachine.
Require Import lemmas.

Open Scope pred.

Program Definition guards (n:nat) (X:pred world) (stk:stack) : pred world :=
  fun prt =>
    match prt with (p,(r,t)) =>

      forall p',
        necR p p' ->
        X (p',(r,t)) ->
        safe t p' r stk /\
        (level p' >= n -> eventually_halts t p' r stk)
  end.
Next Obligation.
  unfold hereditary; intros.
  destruct a as [p [r' t]].
  destruct a' as [p' [r'' t']].
  hnf in H.
  simpl in H.
  unfold prog in *.
  case_eq (age1 p); intros.
  rewrite H1 in H.
  inv H.
  spec H0 p'0.
  spec H0.
  apply rt_trans with p'; auto.
  apply rt_step; auto.
  spec H0; auto.
  rewrite H1 in H; discriminate.
Qed.

Program Definition term_guards (t:store -> option nat) (n:nat) (X:pred world) (stk:stack) : pred world :=
  fun prt =>
    match prt with (p,(r,t')) =>

      forall p',
        necR p p' ->
        X (p',(r,t')) ->
        exists n',
          t r = Some n' /\
          safe t' p' r stk /\
          (level p' >= (n'+n) -> eventually_halts t' p' r stk)
  end.
Next Obligation.
  unfold hereditary; intros.
  destruct a as [? [? ?]].
  destruct a' as [? [? ?]].
  intros.
  hnf in H.
  unfold prog in *.
  simpl in H.
  case_eq (age1 p); intros; rewrite H3 in H; inv H.
  spec H0 p'. spec H0.
  apply rt_trans with p0; auto.
  apply rt_step; auto.
  spec H0; auto.
Qed.

Program Definition funptr (l:label) (A:Type) (P Q: A -> pred world) : pred world :=
  fun prt =>
    match prt with (p,(_,t)) =>
      exists i,  prog_lookup p l = Some i /\
        (forall stk p' n t' (x:A),
          termMeasure_incr t t' ->
           laterR p p' ->
           (forall r', guards n (Q x) stk (p',(r' ,t'))) ->
           (forall r , term_guards (t l) n (P x) ((i;;instr_nil _)::stk) (p',(r,t'))))
    end.
Next Obligation.
  unfold hereditary; intros.
  destruct a as [p [r' t]].
  destruct a' as [p' [r'' t']].
  hnf in H; simpl in H.
  case_eq (age1 p); intros.
  unfold prog in *.
  rewrite H1 in H.
  inv H.
  rename t' into t.
  destruct H0 as [i [? ?]].
  rewrite K.knot_age1 in H1.
  case_eq (K.unsquash p); intros.
  rewrite H2 in H1.
  destruct n; try discriminate.
  inv H1.
  exists (fmap_instr (K.approx n) i).
  split.
  unfold prog_lookup.
  rewrite K.unsquash_squash; simpl.
  unfold KnotInput.fmap.
  apply fmap_eqn.
  unfold prog_lookup in H.
  rewrite H2 in H.
  auto.
  intros.
  spec H0 stk p' n0 t' x.
  spec H0; auto.
  spec H0.
  apply t_trans with (K.squash (n,f)); auto.
  apply t_step.
  hnf. simpl.
  rewrite K.knot_age1. rewrite H2. auto.
  spec H0. auto.
  hnf; intros.
  spec H0 r p'0 H5 H6.
  destruct H0 as [n' [? ?]].
  exists n'; split; auto.
  destruct H7; split.
  revert H7; apply safe_approx with n; eauto.
  replace n with (level (K.squash (n,f))).
  apply later_level.
  apply Rt_Rft_trans with p'; auto.
  rewrite K.knot_level.
  rewrite K.unsquash_squash; auto.
  hnf; eauto.
  hnf; auto.
  constructor; auto.
  apply (intro_instr_approx_rel _ _ (i;;instr_nil _)); auto.
  apply stack_approx_rel_refl.
  intros; spec H8; auto.
  revert H8.
  apply ev_halts_approx with n.
  replace n with (level (K.squash (n,f))).
  apply later_level.
  apply Rt_Rft_trans with p'; auto.
  rewrite K.knot_level.
  rewrite K.unsquash_squash; auto.
  hnf; eauto.
  hnf; auto.
  constructor; auto.
  apply (intro_instr_approx_rel _ _ (i;;instr_nil _)); auto.
  apply stack_approx_rel_refl.

  unfold prog in *.
  rewrite H1 in H.
  discriminate.
Qed.

Lemma boxy_funptr : forall l A P Q,
  boxy K.expandM (funptr l A P Q).
Proof.
  intros. apply boxy_i; auto.
  simpl; intros.
  destruct w; destruct w'.
  destruct o; destruct o0.
  rewrite K.expandM_spec in H.
  destruct H.
  assert (k=k0).
  hnf in H.
  apply K.unsquash_inj.
  destruct (K.unsquash k). destruct (K.unsquash k0).
  destruct H; f_equal; auto.
  subst k0.
  destruct H1. simpl in *.
  destruct H0 as [i [? ?]].
  exists i; split; auto.
  intros.
  spec H3 stk p' n t' x. spec H3; auto.
  do 2 intro. destruct (H2 l0 r0); auto.
  rewrite H9; auto.
  spec H3; auto.
  spec H3; auto.
  spec H3 r p'0 H7 H8.
  destruct H3 as [n' [? ?]].
  exists n'; split; auto.
  destruct (H2 l r); auto.
  rewrite H3 in H10; discriminate.
  congruence.
Qed.

(* Definition of the hoare relation and rules *)

Definition hoare (t:terminationMeasure) (n':nat) (G R P:pred world) (c:instruction) (Q:pred world) :=
  forall t', termMeasure_incr t t' ->

  forall p n k stk,
    (forall s, G (p,(s,t'))) ->
    (forall r', guards n R stk       (p,(r',t'))) ->
    (forall r', guards n Q (k::stk)  (p,(r',t'))) ->
    (forall r,  guards (n'+n) P ((c ;; k)::stk) (p,(r,t'))).

Lemma hoare_weaken : forall t t' x x' (G G' R R' P P' Q Q':pred world) i,
  termMeasure_incr t t' ->
  x <= x' ->
  G' |-- G  ->
  R  |-- R' ->
  P' |-- P  ->
  Q  |-- Q' ->
  hoare t  x  G  R  P  i Q ->
  hoare t' x' G' R' P' i Q'.
Proof.
  intros until i; repeat intro.
  hnf in H5.
  spec H5 t'0.
  spec H5; auto.
  do 2 intro. destruct (H l r0); auto.
  rewrite H12; auto.
  spec H5 p n k stk.
  spec H5; auto.
  spec H5. repeat intro. eapply H8; eauto.
  spec H5. repeat intro. eapply H9; eauto.
  spec H5 r p' H10.
  spec H5; auto.
  destruct H5; split; auto.
  intro. spec H12. omega.
  auto.
Qed.

Lemma hoare_weaken_time : forall t x x' G R P i Q,
  x <= x' ->
  hoare t x G R P i Q ->
  hoare t x' G R P i Q.
Proof.
  intros until Q. intro. apply hoare_weaken; auto; hnf; auto.
Qed.

Lemma hoare_weaken_pre : forall t x G R P P' i Q,
  P' |-- P ->
  hoare t x G R P i Q ->
  hoare t x G R P' i Q .
Proof.
  intros until Q. intro.
  apply hoare_weaken; auto; hnf; auto.
Qed.

Lemma hoare_weaken_post : forall t x G R P i Q Q',
  Q |-- Q' ->
  hoare t x G R P i Q ->
  hoare t x G R P i Q'.
Proof.
  intros until Q'. intro.
  apply hoare_weaken; auto; hnf; auto.
Qed.

Lemma hoare_ex_pre : forall t x G R T P i Q,
  (forall z:T, hoare t x G R (P z) i Q) ->
  hoare t x G R (exp P) i Q.
Proof.
  repeat intro.
  destruct H5.
  spec H x0.
  eapply H; eauto.
Qed.

Lemma hoare_fact_pre : forall t x G R (X:Prop) P i Q,
  (X -> hoare t x G R P i Q) ->
  hoare t x G R (!!X && P) i Q.
Proof.
  repeat intro.
  destruct H5. hnf in H5.
  spec H H5. eapply H; eauto.
Qed.

Lemma hoare_return : forall t (G R:pred world),
  hoare t 0 G R R instr_return FF.
Proof.
  repeat intro.
  spec H1 r p' H3 H4.
  destruct H1. split.
  repeat intro.
  inv H6. econstructor.
  eapply step_return.
  inv H7. apply H1; auto.
  intro. spec H5. omega.
  destruct H5 as [pz [rz ?]].
  exists pz. exists rz.
  econstructor. eapply step_return.
  auto.
Qed.

Lemma hoare_getlabel : forall t l v (G R P:pred world),
  hoare t 0 G R
    (box (setM v (value_label l)) P)
    (instr_getlabel l v)
    P.
Proof.
  repeat intro.
  spec H2 (r#v <- (value_label l)).
  spec H2 p' H3.
  spec H2. apply H4. constructor; auto.
  destruct H2. split.
  repeat intro.
  inv H6. econstructor. econstructor.
  inv H7. apply H2; auto.
  intro; spec H5. omega.
  destruct H5 as [pz [rz ?]].
  exists pz; exists rz.
  econstructor. econstructor.
  auto.
Qed.

Lemma hoare_step_fetch0 : forall t v1 v2 x1 x2 (G R P:pred world),
  hoare t 0 G R
    (store_op (fun r => r#v1 = Some (value_cons x1 x2)) && box (setM v2 x1) P)
    (instr_fetch_field v1 0 v2)
    P.
Proof.
  repeat intro.
  destruct H4.
  destruct H4 as [_ [? _]].
  spec H2 (r#v2 <- x1).
  spec H2 p' H3.
  spec H2. apply H5. constructor. auto.
  destruct H2. split.
  repeat intro.
  inv H7.
  econstructor. econstructor; eauto.
  inv H8.
  apply H2; auto.
  replace x1 with a0 by congruence. auto.
  intro. spec H6. omega.
  destruct H6 as [pz [rz ?]].
  exists pz; exists rz.
  econstructor. econstructor. eauto.
  auto.
Qed.

Lemma hoare_step_fetch1 : forall t v1 v2 x1 x2 (G R P:pred world),
  hoare t 0 G R
    (store_op (fun r => r#v1 = Some (value_cons x1 x2)) && box (setM v2 x2) P)
    (instr_fetch_field v1 1 v2)
    P.
Proof.
  repeat intro.
  destruct H4.
  destruct H4 as [_ [? _]].
  spec H2 (r#v2 <- x2).
  spec H2 p' H3.
  spec H2. apply H5. constructor. auto.
  destruct H2. split.
  repeat intro.
  inv H7.
  econstructor. econstructor; eauto.
  inv H8.
  apply H2; auto.
  replace x2 with a1 by congruence. auto.
  intro. spec H6. omega.
  destruct H6 as [pz [rz ?]].
  exists pz; exists rz.
  econstructor. econstructor. eauto.
  auto.
Qed.

Lemma hoare_cons : forall t v1 v2 v3 x1 x2 (G R P:pred world),
  hoare t 0 G R
     (store_op (fun r => r#v1 = Some x1 /\ r#v2 = Some x2) &&
        box (setM v3 (value_cons x1 x2)) P)
     (instr_cons v1 v2 v3)
     P.
Proof.
  repeat intro.
  destruct H4.
  destruct H4 as [_ [[? ?] _]].
  spec H2 (r#v3 <- (value_cons x1 x2)).
  spec H2 p' H3.
  spec H2.
  apply H5. constructor. auto.
  destruct H2. split.
  repeat intro.
  inv H8.
  econstructor. econstructor; eauto.
  inv H9.
  apply H2.
  replace x1 with a1 by congruence.
  replace x2 with a2 by congruence.
  auto.
  intro. spec H7. omega.
  destruct H7 as [pz [rz ?]].
  exists pz; exists rz.
  econstructor. econstructor; eauto.
  auto.
Qed.

Lemma hoare_if : forall t x v s1 s2 G R P Q,
  P |-- store_op (fun r => match r#v with
                           | Some (value_label l) => l = L 0
                           | Some (value_cons _ _) => True
                           | _ => False
                           end) ->
  hoare t x G R (P && store_op (fun r => r#v = Some (value_label (L 0)))) s1 Q ->
  hoare t x G R (P && store_op (fun r => exists x1, exists x2, r#v = Some (value_cons x1 x2))) s2 Q ->
  hoare t x G R P (instr_if_nil v s1 s2) Q.
Proof.
  repeat intro.
  generalize H7; intros.
  apply H in H8.
  simpl in H8. destruct H8 as [_ [? _]].
  case_eq (r#v); intros; rewrite H9 in H8;
    try tauto.
  destruct v0.

  (* nil case *)
  subst l.
  spec H0 t' H2 p n k.
  spec H0 stk.
  do 3 (spec H0; auto).
  spec H0 r p' H6.
  spec H0.
  split; auto.
  simpl. intuition.
  destruct H0; split.
  repeat intro.
  inv H10.
  econstructor 1.
  eapply step_if_nil1. auto.
  inv H11.
  eapply H0; eauto.
  rewrite H22 in H9. discriminate.
  intro; spec H8; auto.
  destruct H8 as [pz [rz ?]].
  exists pz; exists rz.
  econstructor.
  eapply step_if_nil1; auto.
  auto.

  (* cons case *)
  spec H1 t' H2 p n k.
  spec H1 stk.
  do 3 (spec H1; auto).
  spec H1 r p' H6.
  spec H1. split. auto.
  simpl; intuition. eauto.
  destruct H1. split.
  repeat intro.
  inv H11.
  econstructor 1. eapply step_if_nil2; eauto.
  eapply H1; eauto.
  inv H12. rewrite H23 in H9; discriminate.
  rewrite H23 in H9. inv H9. auto.
  intro; spec H10; auto.
  destruct H10 as [pz [rz ?]].
  exists pz; exists rz.
  econstructor.
  eapply step_if_nil2; eauto.
  auto.
Qed.

Lemma hoare_call : forall t x G R v Q,
  let wp :=
    EX l:label, EX A:Type, EX lP:(A->pred world), EX lQ:(A -> pred world), EX n':nat, EX a:A,
      store_op (fun r => r#v = Some (value_label l) /\ t l r = Some n' /\ n' < x) &&
      (G --> funptr l A lP lQ) &&
      lP a && (closed (lQ a --> Q))
 in
  hoare t x G R wp (instr_call v) Q.
Proof.
  repeat intro.
  unfold wp in H4.
  destruct H4 as [l [A [lP [lQ [n' [a ?]]]]]].
  destruct H4 as [[[? ?] ?] ?].
  destruct H4 as [_ [? _]].
  destruct H4 as [? [? ?]].
  destruct x. exfalso; omega.
  assert (funptr l A lP lQ (p',(r,t'))).
  apply H5. auto.
  apply pred_nec_hereditary with (p,(r,t')); auto.
  rewrite worldNec_unfold; intuition.
  case_eq (age1 p'); intros.
  destruct H10 as [i [? ?]].
  spec H12 (k::stk) p0 n t' a.
  spec H12. hnf; auto.
  spec H12. apply t_step. auto.
  spec H12.
    intro r'. spec H2 r'.
    repeat intro.
    spec H2 p'0.
    spec H2.
    apply rt_trans with p'; auto.
    apply rt_trans with p0; auto.
    apply rt_step; auto.
    spec H2; auto.
    spec H7 (p',(r',t')).
    spec H7. simpl; hnf; auto.
    apply H7.
    rewrite worldNec_unfold; intuition.
    apply rt_trans with p0; auto.
    apply rt_step; auto.
    auto.

  spec H12 r p0. spec H12; auto.
  spec H12. apply pred_nec_hereditary with (p',(r,t')); auto.
  rewrite worldNec_unfold; intuition.
  destruct H12 as [n2 [? ?]].
  destruct (H l r). congruence.
  rewrite <- H14 in H12.
  rewrite H8 in H12. inv H12. clear H14.

  destruct H13. split.
  repeat intro.
  inv H14.
  econstructor 1.
  eapply step_call with (i':=i); eauto.
  inv H15.
  assert (p0 = p'1) by congruence. subst p'1.
  assert (l = l0) by congruence. subst l0.
  assert (i = i'0) by congruence. subst i'0.
  apply H12; auto.
  apply (af_level2 age_facts) in H11.
  congruence.
  intro. spec H13.
  apply (af_level2 age_facts) in H11. omega.
  destruct H13 as [pz [rz ?]].
  exists pz; exists rz.
  econstructor.
  econstructor; eauto.
  auto.

  split.
  clear -H11.
  repeat intro.
  apply stepstar_stepN in H. destruct H as [n H].
  revert p' H11 p'0 s' k k' stk H.
  induction n; simpl; intros.
  inv H.
  econstructor 1.
  eapply step_call_apocalypse.
  rewrite <- (af_level1 age_facts); auto.
  inv H.
  inv H1.
  hnf in H12. rewrite H12 in H11. discriminate.
  eapply IHn; eauto.
  rewrite (af_level1 age_facts) in H11.
  intro. exfalso. omega.
Qed.

Lemma hoare_assert : forall t (G R P:pred world) (Q:K.assert),
  boxy K.expandM G ->
  G && P |-- (proj1_sig Q) ->
  hoare t 0 G R P (instr_assert Q) P.
Proof.
  repeat intro.
  split.
  hnf; intros.
  inv H7.
  econstructor 1.
  econstructor.
  apply H0. split; auto.
  eapply pred_nec_hereditary.
  2: apply H2.
  instantiate (1:=s').
  rewrite worldNec_unfold. intuition.
  inv H8.
  destruct (H4 s'0 p'1); auto.
  repeat intro.
  destruct (H4 r p'); auto.
  spec H9. omega.
  destruct H9 as [p'' [r' ?]].
  exists p''; exists r'.
  econstructor 2.
  econstructor.
  apply H0.
  split; auto.
  eapply pred_nec_hereditary.
  2: apply H2.
  instantiate (1:=r).
  rewrite worldNec_unfold; intuition.
  auto.
Qed.

Lemma hoare_seq : forall t x y G R P Q S i1 i2,
  hoare t x G R P i1 Q ->
  hoare t y G R Q i2 S ->
  hoare t (x+y) G R P (i1 ;; i2) S.
Proof.
  intros. hnf; intros.
  assert (guards (x+(y+n)) P ((i1;;i2;;k)::stk) (p,(r,t'))).
  apply H; auto.
  repeat intro.
  spec H3 r' p' H5 H6.
  intuition.
  repeat intro.
  spec H5 p' H6 H7.
  destruct H5; split.
  repeat intro.
  inv H9.
  econstructor 1. econstructor.
  inv H10.
  eapply H5; eauto.
  intro.
  spec H8. omega.
  destruct H8 as [pz [rz ?]].
  exists pz; exists rz.
  econstructor. econstructor.
  auto.
Qed.

(* Verifying function bodies and entire programs *)

Definition verify_prog (t:terminationMeasure) (G:pred world) (psi:program K.assert) (G':pred world) :=
  forall n r,
     agedfrom (K.squash (n,psi),(r,t)) && |>(box K.expandM (closed G))
       |-- box K.expandM (closed G').

Lemma verify_complete : forall G G' psi t,
  G' |-- G ->
  verify_prog t G psi G' ->
  forall n r, G' (K.squash (n,psi),(r,t)).
Proof.
  intros.
  spec H0 n r.
  cut (agedfrom (K.squash (n,psi),(r,t)) |-- (box K.expandM (closed G'))).
  intros.
  eapply H1.
  hnf; apply rt_refl.
  apply K.expandM_refl.
  hnf; simpl; auto.
  apply goedel_loeb.
  apply derives_trans with
    (agedfrom (K.squash (n,psi),(r,t)) && |>(box K.expandM (closed G))); auto.
  intros a Ha; destruct Ha; split; auto.
  clear H1.
  revert a H2.
  do 3 apply box_positive; auto.
Qed.

Lemma verify_func : forall psi l (G:pred world) (A:Type) (P Q:A -> pred world) i t,
  psi#l = Some i ->

  (forall a p r t', termMeasure_incr t t' -> proj1_sig (P a) (p,(r,t')) -> exists n, t l r = Some n) ->

  (forall a n,
    let Pr  := P a && store_op (fun r' => t l r' = Some n) in
    let Pr' a' := P a' && store_op (fun r' => exists n', t l r' = Some n' /\ n' < n) in
    hoare t n (funptr l A Pr' Q && G) (Q a) Pr i FF) ->

  verify_prog t  G psi G ->
  verify_prog t (funptr l A P Q && G) psi (funptr l A P Q && G).
Proof.
  repeat intro.
  destruct H3.
  unfold closed in H6.
  do 3 rewrite box_and in H6.
  destruct H6.
  split. 2: eapply H2; eauto; split; eauto.
  destruct a. destruct a'. destruct a'0.
  destruct H5.
  destruct p. destruct p1. destruct o.
  simpl in H5, H8. subst p0 t1.
  rewrite K.expandM_spec in H4.
  assert (k = k0 /\ r = s /\ t = t0).
  simpl in H3. rewrite worldNec_unfold in H3.
  intuition.
  hnf in H3.
  apply K.unsquash_inj.
  destruct (K.unsquash k). destruct (K.unsquash k0).
  f_equal; intuition.
  destruct H5 as [? [? ?]]; subst s t0 k0.
  destruct H4 as [_ ?].
  destruct H4. simpl in H4, H5.
  assert (Hlookup :
    prog_lookup k l =
      Some (fmap_instr (K.approx (level k)) i)).
  replace (fmap_instr (K.approx (level k)) i)
    with (fmap_instr (K.approx (level k)) (fmap_instr (K.approx n) i)).
  apply nec_prog_lookup with (K.squash (n,psi)); auto.
  simpl in H3. rewrite worldNec_unfold in H3; intuition.
  unfold prog_lookup. rewrite K.unsquash_squash.
  simpl. unfold KnotInput.fmap.
  apply fmap_eqn. auto.
  rewrite collapse_instr_approx.
  rewrite min_l; auto.
  replace n with (level (K.squash (n,psi))).
  apply nec_level. simpl in H3. rewrite worldNec_unfold in H3; intuition.
  rewrite K.knot_level. rewrite K.unsquash_squash; auto.

  simpl.
  exists (fmap_instr (K.approx (level k)) i).
  split. auto.

  clear s0.
  intros stk p' n0 t' x G1 G2 G3 r0.
  revert stk p' n0 t' x G1 G2 G3.
  set (R r1 r2 :=
           match t l r1, t l r2 with
           | Some n1, Some n2 => n1 < n2
           | _, _ => False
           end).
  assert (well_founded R).
  clear. hnf; intro a.
  case_eq (t l a); intros.
  revert a H.
  induction n using (well_founded_induction lt_wf); intros.
  constructor; intros.
  hnf in H1.
  case_eq (t l y); intros.
  rewrite H2 in H1. rewrite H0 in H1.
  apply H with n0; auto.
  rewrite H2 in H1. elim H1.
  constructor; intros.
  hnf in H0. destruct (t l y).
  rewrite H in H0. elim H0. elim H0.

  move r0 after H8.
  induction r0 using (well_founded_induction H8).
  intros.
  destruct (H0 x p'0 r0 t') as [n' ?]; auto.
  eapply termMeasure_incr_trans; eauto.
  exists n'; split; auto.
  destruct (H5 l r0); congruence.

  spec H1 x n' t'.
  spec H1. eapply termMeasure_incr_trans; eauto.
  spec H1 p' n0 (instr_nil K.assert) stk.
  spec H1.

    repeat intro.
    split.
    exists (fmap_instr (K.approx (level p')) i).
    split.
    replace (fmap_instr (K.approx (level p')) i)
      with (fmap_instr (K.approx (level p')) (fmap_instr (K.approx (level k)) i)).
    apply nec_prog_lookup with k; auto.
    apply Rt_Rft; auto.
    rewrite collapse_instr_approx.
    rewrite min_l; auto.
    apply later_level in G2.
    unfold prog in *. omega.

    clear H1.
    repeat intro.
    destruct H16.
    simpl in H17.
    destruct H17 as [_ [? _]].
    destruct H17 as [n'' [? ?]].
    exists n''. split.
    destruct (H5 l r1). congruence.
    destruct (G1 l r1); congruence.
    spec H9 r1. spec H9.
    hnf. rewrite H17. rewrite H12. auto.
    spec H9 stk0 p'1 n1 t'0 x0.
    spec H9.
    eapply termMeasure_incr_trans; eauto.
    spec H9. apply t_trans with p'; auto.
    spec H9. auto.
    spec H9 p'2. spec H9. auto.
    spec H9. auto.
    destruct H9 as [n2 [? ?]].
    assert (n'' = n2).
    destruct (H5 l r1); congruence.
    subst n''.
    destruct H19; split.
    revert H19.
    apply safe_approx with (level p').
    apply later_level.
    apply Rt_Rft_trans with p'1; auto.
    hnf; auto.
    hnf; auto.
    constructor. hnf. simpl. f_equal.
    rewrite collapse_instr_approx.
    rewrite collapse_instr_approx_same.
    rewrite min_l; auto.
    apply later_level in G2.
    unfold prog in *; omega.
    apply stack_approx_rel_refl.
    intro. spec H20. auto.
    revert H20.
    apply ev_halts_approx with (level p').
    apply later_level. apply Rt_Rft_trans with p'1; auto.
    hnf; auto.
    hnf. auto.
    constructor. hnf. simpl. f_equal.
    rewrite collapse_instr_approx.
    rewrite collapse_instr_approx_same.
    rewrite min_l; auto.
    apply later_level in G2.
    unfold prog in *; omega.
    apply stack_approx_rel_refl.
    eapply H7.
    instantiate (1:=(p',(r,t))).
    simpl.
    rewrite worldLater_unfold. intuition.
    instantiate (1:= (p',(r,t'))).
    rewrite K.expandM_spec.
    split. hnf. destruct (K.unsquash p'); split; hnf; auto.
    split; simpl; auto. hnf; auto.
    eapply termMeasure_incr_trans; eauto.
    split; simpl; auto.

  spec H1. auto.
  spec H1. repeat intro. elim H14.
  spec H1 r0.
  spec H1 p'0.
  spec H1; auto.
  spec H1. split; auto.
  simpl. auto.
  destruct H1; split.
  revert H1.
  apply safe_approx with (level k).
  apply later_level. apply Rt_Rft_trans with p'; auto.
  hnf; auto.
  hnf; auto.
  constructor.
  hnf. simpl. f_equal.
  rewrite collapse_instr_approx_same; auto.
  apply stack_approx_rel_refl.
  intro; spec H13; auto.
  revert H13.
  apply ev_halts_approx with (level k).
  apply later_level. apply Rt_Rft_trans with p'; auto.
  hnf; auto.
  hnf; auto.
  constructor.
  hnf. simpl. f_equal.
  rewrite collapse_instr_approx_same; auto.
  apply stack_approx_rel_refl.
Qed.

Lemma end_assert_lemma : forall Q t n p p' r r' stk,
    stepN n t p p' r (stk++(instr_assert Q ;; instr_return ;; instr_nil _) :: nil) r' nil ->
    stepstar t p p' r (stk++(instr_return ;; instr_nil _) :: nil) r' nil /\ proj1_sig Q (p',(r',t)).
Proof.
  induction n; simpl; intros; inv H.
  destruct stk; discriminate.
  destruct stk. simpl in H1.
  inv H1. simpl.
  split; auto.
  apply stepN_stepstar with n; auto.
  inv H2. inv H.
  inv H0. auto.
  inv H.

  simpl in H1.
  inv H1.
  apply IHn with (stk:= i0::stk) in H2.
  intuition. econstructor. econstructor; auto. auto.
  apply IHn with (stk:= i0::stk) in H2.
  intuition. econstructor. econstructor; auto. auto.
  apply IHn with (stk:= i0::stk) in H2.
  intuition. econstructor. econstructor; eauto. auto.
  apply IHn with (stk:= i0::stk) in H2.
  intuition. econstructor. econstructor; eauto. auto.
  apply IHn with (stk:= i0::stk) in H2.
  intuition. econstructor. econstructor; eauto. auto.
  apply IHn with (stk:= (i1;;i2;;i3)::stk) in H2.
  intuition. econstructor. econstructor; eauto. auto.
  apply IHn with (stk:= (i1;;i0)::stk) in H2.
  intuition. econstructor. econstructor; eauto. auto.
  apply IHn with (stk:= (i2;;i0)::stk) in H2.
  intuition. econstructor. eapply step_if_nil2; eauto. auto.
  apply IHn with (stk:= (i'0;;instr_nil _)::i0::stk) in H2.
  intuition. econstructor. econstructor; eauto. auto.
  apply IHn with (stk:= (instr_call v;;i0)::stk) in H2. intuition.
  apply IHn with (stk:=stk) in H2.
  intuition. econstructor. econstructor. auto.
Qed.

(* Fundamental liveness theorem.  A verified function,
   when started in a state satisfying its precondition,
   will halt in a state satisfying its postcondition.
 *)

Lemma verify_totally_correct : forall t G A P Q psi l x
  (HQ:forall x, boxy K.expandM (Q x)),
  verify_prog t G psi G ->
  G |-- funptr l A P Q ->
  forall r,
    (forall n, P x (K.squash (n,psi),(r#(L 0) <- (value_label l),t))) ->
    exists n, exists p', exists r',
      stepstar t (K.squash (n,psi)) p'
        r ((instr_getlabel l (L 0) ;; instr_call (L 0) ;; instr_return ;; instr_nil _)::nil)
        r' nil /\
      Q x (p',(r',t)).
Proof.
  intros.
  assert (forall n r, G (K.squash (n,psi),(r,t))).
  apply verify_complete with G; auto.
  hnf; auto.
  set (r' :=(r#(L 0) <- (value_label l))).
  assert (exists n, t l r' = Some n).
  assert (funptr l A P Q (K.squash (1,psi),(r,t))) by auto.
  destruct H3 as [i [? ?]].
  spec H4 (nil : stack) (K.squash (0,psi)) 0 t x.
  spec H4. hnf; auto.
  spec H4. apply t_step. hnf.
  rewrite K.knot_age1. rewrite K.unsquash_squash.
  f_equal. apply K.unsquash_inj.
  repeat rewrite K.unsquash_squash.
  f_equal.
  change ((KnotInput.fmap (K.approx 0) oo KnotInput.fmap (K.approx 1)) psi =
    KnotInput.fmap (K.approx 0) psi).
  rewrite KnotInput.fmap_comp.
  replace 1 with (1+0) by omega.
  rewrite <- (K.approx_approx1 1 0); auto.
  spec H4.
  repeat intro. split; repeat intro.
  inv H7. constructor 2.
  inv H8.
  exists p'. exists r'0. apply stepstar_O.
  spec H4 r' (K.squash (0,psi)).
  spec H4. auto.
  spec H4. auto.
  destruct H4 as [n' [? ?]]. eauto.
  destruct H3 as [n ?].
  assert (funptr l A P Q (K.squash (S n,psi),(r,t))) by auto.
  destruct H4 as [i [? ?]].
  spec H5 ((instr_assert (exist _ (Q x) (HQ x)) ;; instr_return ;; instr_nil K.assert)::nil).
  spec H5 (K.squash (n,psi)) 0 t x.
  spec H5. hnf; auto.
  spec H5.
  apply t_step.
  hnf. rewrite K.knot_age1.
  rewrite K.unsquash_squash; simpl.
  f_equal.
  apply K.unsquash_inj.
  repeat rewrite K.unsquash_squash.
  f_equal.
  transitivity (KnotInput.fmap (K.approx n oo K.approx (S n)) psi).
  rewrite <- KnotInput.fmap_comp; auto.
  replace (K.approx n oo K.approx (S n)) with (K.approx n); auto.
  extensionality.
  unfold compose.
  change (K.approx n x0 = (K.approx n oo K.approx (1+n)) x0).
  rewrite <- K.approx_approx1; auto.
  hnf; auto.
  spec H5.
  repeat intro.
  split.
  hnf.
  intros.
  inv H8.
  econstructor 1.
  econstructor.
  auto.
  inv H9. simpl in *.
  inv H10.
  econstructor 1.
  econstructor.
  inv H8.
  inv H9.
  constructor.
  inv H8.
  intros.
  econstructor.
  econstructor.
  econstructor 2.
  econstructor.
  simpl; auto.
  econstructor. econstructor.
  econstructor 1.
  spec H5 r' (K.squash (n,psi)).
  spec H5; auto.
  spec H5. auto.
  destruct H5 as [n' [? ?]].
  assert (n = n') by congruence. subst n'.
  destruct H6. spec H7; auto.
  rewrite K.knot_level.
  rewrite K.unsquash_squash; simpl. omega.
  exists (S n).
  destruct H7 as [pz [rz ?]].
  exists pz; exists rz.
  cut
    (stepstar t (K.squash (n, psi)) pz r'
         ((i;; instr_nil K.assert)
          :: (instr_return;; instr_nil K.assert) :: nil) rz nil
         /\ Q x (pz,(rz,t))).
  intros [? ?]; split; auto.
  econstructor 2.
  econstructor.
  econstructor 2.
  econstructor; eauto.
  rewrite get_set_same; auto.
  hnf; simpl.
  rewrite K.knot_age1.
  rewrite K.unsquash_squash.
  reflexivity.
  match goal with [ |- stepstar _ ?X _ _ _ _ _ ] =>
    replace X with (K.squash (n,psi))
  end.
  auto.
  apply K.unsquash_inj.
  repeat rewrite K.unsquash_squash.
  f_equal.
  symmetry.
  transitivity (KnotInput.fmap (K.approx n oo K.approx (S n)) psi).
  rewrite <- KnotInput.fmap_comp; auto.
  change (S n) with (1+n).
  rewrite <- K.approx_approx1; auto.

  apply stepstar_stepN in H7. destruct H7.
  eapply end_assert_lemma with (Q:=exist _ (Q x) (HQ x))
    (stk:=(i;;instr_nil _)::nil). eauto.
Qed.

(* Weaker corallary: a program with a safe entry point
   will eventually halt.
 *)
Corollary verify_halts : forall t G psi l,
  verify_prog t G psi G ->
  G |-- funptr l unit (fun _ => TT) (fun _ => TT) ->
  forall r, exists n,
    eventually_halts t (K.squash (n,psi)) r
        ((instr_getlabel l (L 0) ;; instr_call (L 0) ;; instr_return ;; instr_nil _)::nil).
Proof.
  intros.
  destruct (verify_totally_correct t G unit (fun _ => TT) (fun _ => TT) psi l tt) with (r:=r)
    as [n [p' [r' [? ?]]]]; auto.
  exists n. hnf. eauto.
Qed.


Set Warnings "-notation-overridden,-parsing,-deprecated-hint-without-locality".
From Coq Require Import Strings.String.
From PLF Require Import Maps.
From PLF Require Import Types.
From PLF Require Import Smallstep.

Inductive ty : Type :=
  | Ty_Top : ty
  | Ty_Bool : ty
  | Ty_Base : string -> ty
  | Ty_Arrow : ty -> ty -> ty
  | Ty_Unit : ty
  | Ty_Prod : ty -> ty -> ty.

Inductive tm : Type :=
  | tm_var : string -> tm
  | tm_app : tm -> tm -> tm
  | tm_abs : string -> ty -> tm -> tm
  | tm_true : tm
  | tm_false : tm
  | tm_if : tm -> tm -> tm -> tm
  | tm_unit : tm
  | tm_pair : tm -> tm -> tm
  | tm_fst : tm -> tm
  | tm_snd : tm -> tm.

Declare Custom Entry stlc.
Notation "<{ e }>" := e (e custom stlc at level 99).
Notation "( x )" := x (in custom stlc, x at level 99).
Notation "x" := x (in custom stlc at level 0, x constr at level 0).
Notation "S -> T" := (Ty_Arrow S T) (in custom stlc at level 50, right associativity).
Notation "x y" := (tm_app x y) (in custom stlc at level 1, left associativity).
Notation "\ x : t , y" :=
  (tm_abs x t y) (in custom stlc at level 90, x at level 99,
                     t custom stlc at level 99,
                     y custom stlc at level 99,
                     left associativity).
Coercion tm_var : string >-> tm.
Notation "'Bool'" := Ty_Bool (in custom stlc at level 0).
Notation "'if' x 'then' y 'else' z" :=
  (tm_if x y z) (in custom stlc at level 89,
                    x custom stlc at level 99,
                    y custom stlc at level 99,
                    z custom stlc at level 99,
                    left associativity).
Notation "'true'" := true (at level 1).
Notation "'true'" := tm_true (in custom stlc at level 0).
Notation "'false'" := false (at level 1).
Notation "'false'" := tm_false (in custom stlc at level 0).
Notation "'Unit'" :=
  (Ty_Unit) (in custom stlc at level 0).
Notation "'unit'" := tm_unit (in custom stlc at level 0).
Notation "'Base' x" := (Ty_Base x) (in custom stlc at level 0).
Notation "'Top'" := (Ty_Top) (in custom stlc at level 0).

Notation "X * Y" :=
  (Ty_Prod X Y) (in custom stlc at level 2, X custom stlc, Y custom stlc at level 0).
Notation "( x ',' y )" := (tm_pair x y) (in custom stlc at level 0,
                                                x custom stlc at level 99,
                                                y custom stlc at level 99).
Notation "t '.fst'" := (tm_fst t) (in custom stlc at level 0).
Notation "t '.snd'" := (tm_snd t) (in custom stlc at level 0).

Reserved Notation "'[' x ':=' s ']' t" (in custom stlc at level 20, x constr).
Fixpoint subst (x : string) (s : tm) (t : tm) : tm :=
  match t with
  | tm_var y =>
      if eqb_string x y then s else t
  | <{\y:T, t1}> =>
      if eqb_string x y then t else <{\y:T, [x:=s] t1}>
  | <{t1 t2}> =>
      <{([x:=s] t1) ([x:=s] t2)}>
  | <{true}> =>
      <{true}>
  | <{false}> =>
      <{false}>
  | <{if t1 then t2 else t3}> =>
      <{if ([x:=s] t1) then ([x:=s] t2) else ([x:=s] t3)}>
  | <{unit}> =>
      <{unit}>
  | <{ ( l , r ) }> => <{ ( [x:=s]l , [x:=s]r ) }>
  | <{ t.fst }> => <{ ([x:=s]t).fst }>
  | <{ t.snd }> => <{ ([x:=s]t).snd }>
  end
where "'[' x ':=' s ']' t" := (subst x s t) (in custom stlc).

Inductive value : tm -> Prop :=
  | v_abs : forall x T2 t1,
      value <{\x:T2, t1}>
  | v_true :
      value <{true}>
  | v_false :
      value <{false}>
  | v_unit :
      value <{unit}>
  | v_pair : forall v1 v2,
      value v1 ->
      value v2 ->
      value <{(v1, v2)}>.

Hint Constructors value : core.
Reserved Notation "t '-->' t'" (at level 40).
Inductive step : tm -> tm -> Prop :=
  | ST_AppAbs : forall x T2 t1 v2,
         value v2 ->
         <{(\x:T2, t1) v2}> --> <{ [x:=v2]t1 }>
  | ST_App1 : forall t1 t1' t2,
         t1 --> t1' ->
         <{t1 t2}> --> <{t1' t2}>
  | ST_App2 : forall v1 t2 t2',
         value v1 ->
         t2 --> t2' ->
         <{v1 t2}> --> <{v1 t2'}>
  | ST_IfTrue : forall t1 t2,
      <{if true then t1 else t2}> --> t1
  | ST_IfFalse : forall t1 t2,
      <{if false then t1 else t2}> --> t2
  | ST_If : forall t1 t1' t2 t3,
      t1 --> t1' ->
      <{if t1 then t2 else t3}> --> <{if t1' then t2 else t3}>
  | ST_Pair1 : forall t1 t1' t2,
       t1 --> t1' ->
       <{ (t1 , t2) }> --> <{ (t1' , t2) }>
  | ST_Pair2 : forall v1 t2 t2',
       value v1 ->
       t2 --> t2' ->
       <{ (v1 , t2) }> --> <{ (v1 , t2') }>
  | ST_PairFstStep : forall t t',
       t --> t' ->
       <{ t.fst }> --> <{ t'.fst }>
  | ST_PairFst : forall v1 v2,
       <{ (v1 , v2).fst }> --> <{ v1 }>
  | ST_PairSndStep : forall t t',
       t --> t' ->
       <{ t.snd }> --> <{ t'.snd }>
  | ST_PairSnd : forall v1 v2,
       <{ (v1 , v2).snd }> --> <{ v2 }>
where "t '-->' t'" := (step t t').
Hint Constructors step : core.

Reserved Notation "T '<:' U" (at level 40).
Inductive subtype : ty -> ty -> Prop :=
  | S_Refl : forall T,
      T <: T
  | S_Trans : forall S U T,
      S <: U ->
      U <: T ->
      S <: T
  | S_Top : forall S,
      S <: <{Top}>
  | S_Arrow : forall S1 S2 T1 T2,
      T1 <: S1 ->
      S2 <: T2 ->
      <{S1->S2}> <: <{T1->T2}>
  | S_Pair : forall S1 S2 T1 T2,
      S1 <: T1 ->
      S2 <: T2 ->
      <{S1 * S2}> <: <{T1 * T2}>
where "T '<:' U" := (subtype T U).

Hint Constructors subtype : core.
Module Examples.
Open Scope string_scope.
Notation x := "x".
Notation y := "y".
Notation z := "z".
Notation A := <{Base "A"}>.
Notation B := <{Base "B"}>.
Notation C := <{Base "C"}>.
Notation String := <{Base "String"}>.
Notation Float := <{Base "Float"}>.
Notation Integer := <{Base "Integer"}>.

Example subtyping_example_0 :
  <{C->Bool}> <: <{C->Top}>.
Proof. auto. Qed.

Definition Person : ty := <{ (String * Top)}>.

Definition Student : ty := <{ (String * Float)}>.

Definition Employee : ty := <{ (String * Integer)}>.

Example sub_student_person :  Student <: Person.
Proof. apply S_Pair. apply S_Refl. apply S_Top. Qed.

Example sub_employee_person :  Employee <: Person.
Proof. apply S_Pair. apply S_Refl. apply S_Top. Qed.

Example subtyping_example_1 :
  <{Top->Student}> <: <{(C->C)->Person}>.
Proof with eauto.
  apply S_Arrow...
  apply sub_student_person. Qed.

Example subtyping_example_2 :
  <{Top->Person}> <: <{Person->Top}>.
Proof with eauto.
  apply S_Arrow... Qed.

End Examples.

Definition context := partial_map ty.
Reserved Notation "Gamma '|-' t '\in' T" (at level 40, t custom stlc, T custom stlc at level 0).
Inductive has_type : context -> tm -> ty -> Prop :=
  (* Same as before: *)
  (* pure STLC *)
  | T_Var : forall Gamma x T1,
      Gamma x = Some T1 ->
      Gamma |- x \in T1
  | T_Abs : forall Gamma x T1 T2 t1,
      (x |-> T2 ; Gamma) |- t1 \in T1 ->
      Gamma |- \x:T2, t1 \in (T2 -> T1)
  | T_App : forall T1 T2 Gamma t1 t2,
      Gamma |- t1 \in (T2 -> T1) ->
      Gamma |- t2 \in T2 ->
      Gamma |- t1 t2 \in T1
  | T_True : forall Gamma,
       Gamma |- true \in Bool
  | T_False : forall Gamma,
       Gamma |- false \in Bool
  | T_If : forall t1 t2 t3 T1 Gamma,
       Gamma |- t1 \in Bool ->
       Gamma |- t2 \in T1 ->
       Gamma |- t3 \in T1 ->
       Gamma |- if t1 then t2 else t3 \in T1
  | T_Unit : forall Gamma,
      Gamma |- unit \in Unit
  (* New rule of subsumption: *)
  | T_Sub : forall Gamma t1 T1 T2,
      Gamma |- t1 \in T1 ->
      T1 <: T2 ->
      Gamma |- t1 \in T2
  (* Pair *)
  | T_Pair : forall Gamma t1 t2 T1 T2,
      Gamma |- t1 \in T1 ->
      Gamma |- t2 \in T2 ->
      Gamma |- (t1 , t2) \in (T1 * T2)
  | T_PairFst : forall Gamma t T1 T2,
      Gamma |- t \in (T1 * T2) ->
      Gamma |- t.fst \in T1
  | T_PairSnd : forall Gamma t T1 T2,
      Gamma |- t \in (T1 * T2) ->
      Gamma |- t.snd \in T2
where "Gamma '|-' t '\in' T" := (has_type Gamma t T).
Hint Constructors has_type : core.

Module Examples2.
Import Examples.

Example e21: forall A B, empty |- ((\z:A, z) , (\z:B, z)) \in ((A->A) * (B->B)).
Proof. intros A B. eapply (T_Pair empty <{ \z:A, z }> <{ \z:B, z }> <{ A->A }> <{ B->B }>).
  - apply T_Abs. apply T_Var. auto.
  - apply T_Abs. apply T_Var. auto.
Qed.

Example e22: forall A B, empty |- (\x:(Top * (B->B)), x.snd) ((\z:A, z), (\z:B, z)) \in (B -> B).
Proof. intros A B.
  eapply T_App.
  - apply T_Abs. eapply T_PairSnd. apply T_Var. unfold "|->". rewrite t_update_eq. reflexivity.
  - apply T_Pair.
    + apply T_Sub with <{ A -> A }>. (* Anything subtype of Top *)
      * apply T_Abs. apply T_Var. unfold "|->". rewrite t_update_eq. reflexivity.
      * apply S_Top.
    + apply T_Abs. apply T_Var. unfold "|->". rewrite t_update_eq. reflexivity.
Qed.

Example e23: forall A B C, empty |- (\z:(C->C)->(Top * (B->B)), (z (\x:C, x)).snd)
              (\z:C->C, ((\z:A, z), (\z:B, z)))
         \in (B->B).
Proof. intros A B C.
  apply T_App with <{ (C -> C) -> Top * (B -> B) }>.
  - apply T_Abs. eapply T_PairSnd. eapply T_App.
    + apply T_Var. unfold "|->". rewrite t_update_eq. reflexivity.
    + apply T_Abs. apply T_Var. unfold "|->". rewrite t_update_eq. reflexivity.
  - apply T_Abs. apply T_Pair.
    + apply T_Sub with <{ A -> A }>.
      * apply T_Abs. apply T_Var. unfold "|->". rewrite t_update_eq. reflexivity.
      * apply S_Top.
    + apply T_Abs. apply T_Var. unfold "|->". rewrite t_update_eq. reflexivity.
Qed.

End Examples2.

Lemma sub_inversion_Bool : forall U,
     U <: <{Bool}> -> U = <{Bool}>.
Proof with auto.
  intros U Hs.
  remember <{Bool}> as V.
  induction Hs; subst; try auto; try (inversion HeqV).
  - rewrite -> IHHs1; apply IHHs2; reflexivity.
Qed.

Lemma sub_inversion_arrow : forall U V1 V2,
     U <: <{ V1 -> V2 }> ->
     exists U1 U2, U = <{ U1 -> U2 }> /\ V1 <: U1 /\ U2 <: V2.
Proof with eauto.
  intros U V1 V2 Hs.
  remember <{ V1 -> V2 }> as V.
  generalize dependent V2. generalize dependent V1.
  induction Hs; intros V1 V2 Ev; subst; try (inversion Ev).
  - exists V1, V2. split; try reflexivity. split; apply S_Refl.
  - destruct (IHHs2 V1 V2) as [ U1 [ U2  [Eu [ Hv1s Hu2s] ] ] ]. reflexivity.
    subst.
    destruct (IHHs1 U1 U2) as [ U1' [ U2'  [Eu' [ Hv1s' Hu2s'] ] ] ]. reflexivity.
    subst.
    exists U1', U2'.
    split. reflexivity.
    split.
    + apply (S_Trans _ _ _ Hv1s Hv1s').
    + apply (S_Trans _ _ _ Hu2s' Hu2s).
  - exists S1, S2.
    split. reflexivity.
    split.
    + rewrite <- H0. assumption.
    + rewrite <- H1. assumption.
Qed.

Lemma bool_if_bottom_for_true: forall G T, G |- true \in T -> <{ Bool }> <: T.
Proof.
  intros G T H.
  remember <{ true }> as t eqn:Et.
  induction H; try (discriminate Et).
  - constructor.
  - subst. apply S_Trans with T1.
    + apply IHhas_type. reflexivity.
    + assumption.
Qed.

Lemma true_isnt_arrow: forall G T1 T2, ~(G |- true \in (T1 -> T2)).
Proof.
  intros G T1 T2 H.
  inversion H; subst.
  apply bool_if_bottom_for_true in H.
  apply sub_inversion_arrow in H.
  destruct H as [U1 [U2 [HT0 [HU1 HU2]]]].
  discriminate HT0.
Qed.

Lemma bool_if_bottom_for_false: forall G T, G |- false \in T -> <{ Bool }> <: T.
Proof.
  intros G T H.
  remember <{ false }> as t eqn:Et.
  induction H; try (discriminate Et).
  - constructor.
  - subst. apply S_Trans with T1.
    + apply IHhas_type. reflexivity.
    + assumption.
Qed.

Lemma false_isnt_arrow: forall G T1 T2, ~(G |- false \in (T1 -> T2)).
Proof.
  intros G T1 T2 H.
  inversion H; subst.
  apply bool_if_bottom_for_false in H.
  apply sub_inversion_arrow in H.
  destruct H as [U1 [U2 [HT0 [HU1 HU2]]]].
  discriminate HT0.
Qed.

Lemma unit_is_bottom_for_unit: forall G T, G |- unit \in T -> <{ Unit }> <: T.
Proof.
  intros G T H.
  remember <{ unit }> as t eqn:Et.
  induction H; try (discriminate Et).
  - constructor.
  - subst. apply S_Trans with T1.
    + apply IHhas_type. reflexivity.
    + assumption.
Qed.

Lemma unit_isnt_arrow: forall G T1 T2, ~(G |- unit \in (T1 -> T2)).
Proof.
  intros G T1 T2 H.
  inversion H; subst.
  apply unit_is_bottom_for_unit in H.
  apply sub_inversion_arrow in H.
  destruct H as [U1 [U2 [HT0 [HU1 HU2]]]].
  discriminate HT0.
Qed.

Lemma pair_is_bottom_for_pair: forall G v1 v2 T, G |- (v1, v2) \in T -> exists T1 T2, <{ T1 * T2 }> <: T.
Proof.
  intros G v1 v2 T H.
  remember <{ (v1, v2) }> as t eqn:Et.
  induction H; try (discriminate Et).
  - subst.
    destruct IHhas_type as [T2' [T3' H']]. reflexivity.
    exists T2', T3'.
    apply S_Trans with T1.
    + apply H'.
    + assumption.
  - exists T1, T2. constructor.
Qed.

Lemma pair_isnt_arrow: forall G v1 v2 T1 T2, ~(G |- (v1, v2) \in (T1 -> T2)).
Proof.
  intros G v1 v2 T1 T2 H.
  inversion H; subst.
  apply pair_is_bottom_for_pair in H.
  destruct H as [T' [T'' H']].
  apply sub_inversion_arrow in H'.
  destruct H' as [U1 [U2 [HT0 [HU1 HU2]]]].
  discriminate HT0.
Qed.

Lemma canonical_forms_of_arrow_types : forall Gamma s T1 T2,
  Gamma |- s \in (T1 -> T2) ->
  value s ->
  exists x S1 s2, s = <{ \x:S1, s2 }>.
Proof with eauto.
  intros G s T1 T2 Ht Hv.
  generalize dependent G.
  generalize dependent T1.
  generalize dependent T2.
  induction Hv.
  - intros T1 T2' G Ht. exists x, T2, t1. reflexivity.
  - intros T2 T1 G Ht. exfalso. apply (true_isnt_arrow G T1 T2 Ht).
  - intros T2 T1 G Ht. exfalso. apply (false_isnt_arrow G T1 T2 Ht).
  - intros T2 T1 G Ht. exfalso. apply (unit_isnt_arrow G T1 T2 Ht).
  - intros T2 T1 G Ht. exfalso. apply (pair_isnt_arrow G v1 v2 T1 T2 Ht).
Qed.

Lemma abs_is_bottom_for_abs: forall G x xT b T, G |- (\x:xT, b) \in T -> exists T1 T2, <{ T1 -> T2 }> <: T.
Proof.
  intros G x xT b T H.
  remember <{ (\x:xT, b) }> as f eqn:Ef.
  induction H; try (discriminate Ef).
  - subst.
    injection Ef as Ex Et Eb. subst.
    exists xT, T1. constructor.
  - subst.
    destruct IHhas_type as [T2' [T3' H']]. reflexivity.
    exists T2', T3'.
    apply S_Trans with T1.
    + apply H'.
    + assumption.
Qed.

Lemma typing_inversion_abs : forall Gamma x S1 t2 T,
     Gamma |- (\x:S1, t2) \in T ->
     exists S2, <{ S1 -> S2 }> <: T  /\  (x |-> S1 ; Gamma) |- t2 \in S2.
Proof.
  intros G x S1 t2 T H.
  remember <{ (\x:S1, t2) }> as f eqn:Ef.
  induction H; try (discriminate Ef).
  - injection Ef as Ex Et Eb. subst.
    exists T1. split.
      constructor.
      assumption.
  - subst.
    destruct IHhas_type as [T2' [T3' H']]. reflexivity.
    exists T2'.
    split.
    * apply S_Trans with T1; assumption.
    * assumption.
Qed.

Lemma sub_inversion_pair : forall U V1 V2,
     U <: <{ V1 * V2 }> ->
     exists U1 U2, U = <{ U1 * U2 }> /\ U1 <: V1 /\ U2 <: V2.
Proof with eauto.
  intros U V1 V2 Hs.
  remember <{ V1 * V2 }> as V.
  generalize dependent V2. generalize dependent V1.
  induction Hs; intros V1 V2 Ev; subst; try (inversion Ev).
  - exists V1, V2. split; try reflexivity. split; apply S_Refl.
  - destruct (IHHs2 V1 V2) as [ U1 [ U2  [Eu [ Hv1s Hu2s] ] ] ]. reflexivity.
    subst.
    destruct (IHHs1 U1 U2) as [ U1' [ U2'  [Eu' [ Hv1s' Hu2s'] ] ] ]. reflexivity.
    subst.
    exists U1', U2'.
    split. reflexivity.
    split.
    + apply (S_Trans _ _ _ Hv1s' Hv1s).
    + apply (S_Trans _ _ _ Hu2s' Hu2s).
  - exists S1, S2.
    split. reflexivity.
    split.
    + rewrite <- H0. assumption.
    + rewrite <- H1. assumption.
Qed.

Lemma canonical_forms_of_pair_types : forall Gamma s T1 T2,
  Gamma |- s \in (T1 * T2) ->
  value s ->
  exists v1 v2, s = <{ (v1, v2) }>.
Proof with eauto.
  intros G s T1 T2 Ht Hv.
  generalize dependent G.
  generalize dependent T1.
  generalize dependent T2.
  induction Hv.
  - intros T1 T' G Ht. exfalso.
    apply abs_is_bottom_for_abs in Ht.
    destruct Ht as [T2' [T3' Ht]].
    apply sub_inversion_pair in Ht.
    destruct Ht as [U1' [U2' [Ht [Hu1 Hu2]]]].
    discriminate Ht.
  - intros T2 T1 G Ht. exfalso.
    apply bool_if_bottom_for_true in Ht.
    apply sub_inversion_pair in Ht.
    destruct Ht as [U1' [U2' [Ht [Hu1 Hu2]]]].
    discriminate Ht.
  - intros T2 T1 G Ht. exfalso.
    apply bool_if_bottom_for_false in Ht.
    apply sub_inversion_pair in Ht.
    destruct Ht as [U1' [U2' [Ht [Hu1 Hu2]]]].
    discriminate Ht.
  - intros T2 T1 G Ht. exfalso.
    apply unit_is_bottom_for_unit in Ht.
    apply sub_inversion_pair in Ht.
    destruct Ht as [U1' [U2' [Ht [Hu1 Hu2]]]].
    discriminate Ht.
  - intros T2 T1 G Ht. exists v1, v2. reflexivity.
Qed.

Lemma abs_is_bottom_for_fun: forall G x T2 t1 T, G |- \ x : T2, t1 \in T -> exists T1, <{ T2 -> T1 }> <: T.
Proof.
  intros G x T2 t1 T H.
  remember <{ (\x : T2, t1) }> as t eqn:Et.
  induction H; try (discriminate Et).
  - injection Et as Et1 Et2 Et3. subst. exists T1. constructor.
  - subst.
    destruct IHhas_type as [T3' H'].
      reflexivity.
    exists T3'. apply S_Trans with T1.
    + apply H'.
    + assumption.
Qed.

Lemma typing_inversion_var : forall Gamma (x:string) T,
  Gamma |- x \in T ->
  exists S, Gamma x = Some S /\ S <: T.
Proof with eauto.
  intros G x T H.
  remember (tm_var x) as t eqn:Et.
  induction H; try (discriminate Et).
  - exists T1. split.
    + injection Et as Et. rewrite Et in H. assumption.
    + constructor.
  - destruct IHhas_type as [S' [Hl' Hr']]...
Qed.

Lemma typing_inversion_app : forall Gamma t1 t2 T2,
  Gamma |- t1 t2 \in T2 ->
  exists T1,  Gamma |- t1 \in (T1 -> T2) /\ Gamma |- t2 \in T1.
Proof with eauto.
  intros G t1 t2 T2 H.
  remember <{ t1 t2 }> as t eqn:Et.
  induction H; try (discriminate Et).
  - exists T2.
    injection Et as Et1 Et2. subst.
    split; assumption.
  - destruct IHhas_type as [t2' [Hl' Hr']]...
Qed.

Lemma abs_arrow : forall x S1 s2 T1 T2,
  empty |- (\x:S1, s2) \in (T1 -> T2) ->
     T1 <: S1 /\ (x |-> S1 ; empty) |- s2 \in T2.
Proof with eauto.
  intros x S1 s2 T1 T2 Hty.
  apply typing_inversion_abs in Hty.
  destruct Hty as [S2 [Hsub Hty1]].
  apply sub_inversion_arrow in Hsub.
  destruct Hsub as [U1 [U2 [Heq [Hsub1 Hsub2]]]].
  injection Heq as Heq; subst... Qed.

Lemma canonical_forms_of_Bool : forall Gamma s,
  Gamma |- s \in Bool ->
  value s ->
  s = tm_true \/ s = tm_false.
Proof. intros G s H.
  inversion H; intro Hv; subst; try (inversion Hv; subst).
  - left. reflexivity.
  - right. reflexivity.
  - exfalso.
    apply abs_is_bottom_for_fun in H.
    destruct H as [T' H].
    apply sub_inversion_Bool in H.
    discriminate H.
  - left. reflexivity.
  - right. reflexivity.
  - exfalso.
    apply unit_is_bottom_for_unit in H.
    apply sub_inversion_Bool in H.
    discriminate H.
  - exfalso.
    apply pair_is_bottom_for_pair in H.
    destruct H as [T1' [T2' H]].
    apply sub_inversion_Bool in H.
    discriminate H.
Qed.

Theorem progress : forall t T,
     empty |- t \in T   ->   value t \/ (exists t', t --> t').
Proof with eauto.
  intros t T Ht.
  remember empty as Gamma.
  induction Ht; subst Gamma; auto.
  - (* T_Var *)
    discriminate.
  - (* T_App *)
    right.
    destruct IHHt1; subst...
    + (* t1 is a value *)
      destruct IHHt2; subst...
      * (* t2 is a value *)
        eapply canonical_forms_of_arrow_types in Ht1; [|assumption].
        destruct Ht1 as [x [S1 [s2 H1]]]. subst.
        exists (<{ [x:=t2]s2 }>)...
      * (* t2 steps *)
        destruct H0 as [t2' Hstp]. exists <{ t1 t2' }>...
    + (* t1 steps *)
      destruct H as [t1' Hstp]. exists <{ t1' t2 }>...
  - (* T_Test *)
    right.
    destruct IHHt1.
    + (* t1 is a value *) eauto.
    + apply canonical_forms_of_Bool in Ht1; [|assumption].
      destruct Ht1; subst...
    + destruct H. rename x into t1'. eauto.
  - (* T_Pair *)
    destruct IHHt1; subst...
    + (* t1 is a value *)
      destruct IHHt2; subst...
      * (* t2 steps *)
        destruct H0 as [t2' Hstp]. right. exists <{ (t1 , t2') }>...
    + (* t1 steps *)
      destruct H as [t1' Hstp]. right. exists <{ (t1' , t2) }>...
  - (* T_PairFst *)
    destruct IHHt...
    + (* value pair *)
      destruct (canonical_forms_of_pair_types _ _ _ _ Ht) as [v1 [v2 Ht']]; auto.
      right. subst. exists v1. constructor.
    + destruct H as [t' H]. right. exists (tm_fst t')...
  - (* T_PairSnd *)
    destruct IHHt...
    + (* value pair *)
      destruct (canonical_forms_of_pair_types _ _ _ _ Ht) as [v1 [v2 Ht']]; auto.
      right. subst. exists v2. constructor.
    + destruct H as [t' H]. right. exists (tm_snd t')...
Qed.

Lemma weakening : forall Gamma Gamma' t T,
     inclusion Gamma Gamma' ->
     Gamma |- t \in T ->
     Gamma' |- t \in T.
Proof.
  intros Gamma Gamma' t T H Ht.
  generalize dependent Gamma'.
  induction Ht; eauto using inclusion_update.
Qed.

Lemma weakening_empty : forall Gamma t T,
     empty |- t \in T ->
     Gamma |- t \in T.
Proof.
  intros Gamma t T.
  eapply weakening.
  discriminate.
Qed.

Lemma substitution_preserves_typing : forall Gamma x U t v T,
   (x |-> U ; Gamma) |- t \in T ->
   empty |- v \in U ->
   Gamma |- [x:=v]t \in T.
Proof.
Proof.
  intros Gamma x U t v T Ht Hv.
  remember (x |-> U; Gamma) as Gamma'.
  generalize dependent Gamma.
  induction Ht; intros Gamma' G; simpl; eauto.
  - destruct (eqb_stringP x x0).
    + subst. rewrite update_eq in H. injection H as H. subst.
      apply weakening_empty. assumption.
    + subst. rewrite update_neq in H.
        constructor. assumption.
        assumption.
  - destruct (eqb_stringP x x0); subst.
    + constructor.
      rewrite update_shadow in Ht.
      assumption.
    + constructor.
      apply IHHt.
      rewrite update_permute.
        reflexivity.
        assumption.
Qed.

Lemma pair_destr: forall v1 v2 T1 T2,
  empty |- (v1, v2) \in (T1 * T2) ->
  empty |- v1 \in T1  /\  empty |- v2 \in T2.
Proof.
  intros v1 v2 T1 T2 H.
  remember <{(v1, v2)}> as p eqn:Ep.
  remember <{(T1 * T2)}> as t eqn:Et.
  generalize dependent T1.
  generalize dependent T2.
  induction H; try (discriminate Ep); subst.
  - intros T2' T1' Hv. subst.
    apply sub_inversion_pair in H0.
    destruct H0 as [U1 [U2 [H' [Hu1 Hu2]]]]. subst.
    assert (HG: Gamma |- v1 \in U1 /\ Gamma |- v2 \in U2). {
      apply IHhas_type.
        reflexivity.
        reflexivity.
    }
    clear IHhas_type.
    destruct HG as [Hvu1 Hvu2].
    split.
    + eapply T_Sub. apply Hvu1. apply Hu1.
    + eapply T_Sub. apply Hvu2. apply Hu2.
  - intros T1' T2' Ets. injection Ets as Ets1 Ets2. subst.
    injection Ep as tv1 tv2. subst.
    split.
      assumption.
      assumption.
Qed.

Theorem preservation : forall t t' T,
     empty |- t \in T ->
     t --> t' ->
     empty |- t' \in T.
Proof with eauto.
  intros t t' T HT. generalize dependent t'.
  remember empty as Gamma.
  induction HT;
       intros t' HE; subst;
       try solve [inversion HE; subst; eauto].
  - (* T_App *)
    inversion HE; subst...
    (* Most of the cases are immediate by induction,
       and eauto takes care of them *)
    + (* ST_AppAbs *)
      destruct (abs_arrow _ _ _ _ _ HT1) as [HA1 HA2].
      apply substitution_preserves_typing with T0...
  - (* T_PairFst *)
    inversion HE; subst...
    apply pair_destr in HT.
    destruct HT as [Ht1 Ht2].
    assumption.
  - (* T_PairSnd *)
  inversion HE; subst...
  apply pair_destr in HT.
  destruct HT as [Ht1 Ht2].
  assumption.
Qed.


Require Import Coq.Reals.Rdefinitions.
Require Import Coq.Reals.RIneq.
Require Import Coq.micromega.Psatz.
Require Import SMTC.Tactic.
Require Import Logic.Syntax.
Require Import Logic.Semantics.
Require Import Logic.Lib.
Require Import Logic.ProofRules.
Require Import Logic.Automation.
Require Import Coq.Lists.List.
Require Import Coq.Strings.String.

(* Some useful tactics for our examples. *)

Ltac cut_setoid_rewrite_r X :=
  let xx := fresh "xx" in
  cut X; [ intro xx; rewrite xx; clear xx | ].

Ltac cut_setoid_rewrite_l X :=
  let xx := fresh "xx" in
  cut X; [ intro xx; rewrite <- xx; clear xx | ].

Ltac rewrite_rename_equiv F m :=
  cut_setoid_rewrite_r (F -|- Rename m F);
  [ | rewrite <- Rename_ok by eauto with rw_rename; reflexivity ].

Ltac decompose_hyps :=
  repeat first [ rewrite Lemmas.land_lor_distr_R
               | rewrite Lemmas.land_lor_distr_L ] ;
  repeat apply lorL.

Ltac destruct_ite :=
  match goal with
  | [ |- context [ if ?e then _ else _ ] ]
    => destruct e
  end.

Fixpoint get_vars_term (t : Term) : list Var :=
  match t with
  | VarNextT t | VarNowT t => t :: nil
  | NatT _ | RealT _ => nil
  | PlusT a b | MinusT a b | MultT a b | MaxT a b
  | Binop _ a b =>
                             get_vars_term a ++
                                           get_vars_term b
  | InvT a | CosT a | SinT a | SqrtT a | ArctanT a
  | ExpT a | Unop _ a =>
                                         get_vars_term a
  end.

Definition get_image_vars (m:list (Var*Term)) :=
  List.flat_map (fun p => get_vars_term (snd p)) m.

Fixpoint get_witness (m : list (Var*Term)) : state -> state :=
  match m with
  | nil => fun st => st
  | (x,VarNowT y) :: m =>
    fun st z =>
      if String.string_dec z y
      then st x else get_witness m st z
  | _ => fun st => st
  end.

Fixpoint get_next_vars_term (t : Term) : list Var :=
  match t with
  | VarNextT t => t :: nil
  | VarNowT _ | NatT _ | RealT _ => nil
  | PlusT a b | MinusT a b | MultT a b | MaxT a b
  | Binop _ a b =>
                             get_next_vars_term a ++
                             get_next_vars_term b
  | InvT a | CosT a | SinT a | SqrtT a | ArctanT a
  | ExpT a | Unop _ a => get_next_vars_term a
  end.

Fixpoint get_next_vars_formula (f : Formula) : list Var :=
  match f with
  | Always a | Eventually a | Enabled a =>
                              get_next_vars_formula a
  | And a b | Or a b | Imp a b =>
                       get_next_vars_formula a ++
                       get_next_vars_formula b
  | Rename _ a => get_next_vars_formula a
  | Comp l r _ => get_next_vars_term l ++
                                     get_next_vars_term r
  | _ => nil
  end.

Fixpoint remove_dup (ls : list Var) : list Var :=
  match ls with
  | nil => nil
  | l :: ls => let ls' := remove_dup ls in
               if in_dec string_dec l ls'
               then ls' else l :: ls'
  end.

Ltac enable_ex_st :=
  match goal with
  | |- lentails _ (Enabled ?X) =>
    let vars := eval compute in
    (remove_dup (get_next_vars_formula X)) in
        let rec foreach ls :=
            match ls with
            | @cons _ ?l ?ls => eapply (ex_state l); simpl;
                                foreach ls
            | _ => idtac
            end
        in
        eapply Enabled_action; simpl; intros;
        foreach vars
  end; try (eapply ex_state_any; (let st := fresh in
                                  intro st; clear st)).

Ltac enable_ex_st' :=
  match goal with
  | |- _ |-- Enabled ?X =>
        let vars := eval compute in (remove_dup (get_next_vars_formula X)) in
        let rec foreach ls :=
         (match ls with
          | ?l :: ?ls => eapply (ex_state l); simpl; foreach ls
          | _ => idtac
          end) in
        eapply Enabled_action'; [ tlaIntuition | ]; simpl; intros; foreach vars
  end; try (eapply ex_state_any; (let st := fresh in
                                  intro st; clear st)).

Ltac smart_repeat_eexists :=
  repeat match goal with
           |- exists x, _ => eexists
         end.

(* The old tactic, very slow. *)
(*
Ltac enable_ex_st :=
  eapply Enabled_action; intros; eapply ex_state_flow_any;
  auto; simpl; intros;
  repeat match goal with
         | |- context [ ?X ] =>
           match type of X with
           | Var => idtac
           | String.string => idtac
           end ;
             try match goal with
                 | X := _ |- _ => unfold X
                 end;
             eapply (@ex_state X) ; simpl ;
             match goal with
             | |- exists x (y : _), (@?F x) => fail 1
             | |- _ => idtac
             end
         end;
  try (eapply ex_state_any ;
       let st := fresh in intro st ; clear st).
*)

Lemma reason_action : forall P Q,
    (forall a b tr,
        eval_formula
          P
          (Stream.Cons a
                       (Stream.Cons b tr)) ->
        eval_formula
          Q (Stream.Cons a (Stream.Cons b tr))) ->
    (P |-- Q).
Proof.
  red. red. red. intros. destruct tr.
  destruct tr. auto.
Qed.

Ltac reason_action_tac :=
  eapply reason_action; simpl;
  let pre := fresh "pre" in
  let post := fresh "post" in
  let tr := fresh "tr" in
  intros pre post tr;
    breakAbstraction; simpl; unfold eval_comp;
    simpl; intros.

(* This solves linear real arithmetic goals.
   It should be complete. *)
Ltac solve_linear :=
  breakAbstraction; intros; unfold eval_comp in *;
  simpl in *; intuition; try psatzl R.

(* This tries to solve nonlinear real
   arithmetic goals. It is not complete
   and can be incredibly inefficient. *)
Ltac solve_nonlinear :=
  breakAbstraction; intros; unfold eval_comp in *;
  simpl in *; intuition; try psatz R.

Ltac zero_deriv_tac v :=
  eapply ContinuousProofRules.zero_deriv
  with (x:=v); [ charge_tauto | solve_linear | ].

Ltac always_imp_tac :=
  match goal with
  | [ |- ?H |-- _ ]
    => match H with
       | context[ Always ?HH ] =>
         tlaAssert (Always HH);
           [ charge_tauto |
             apply Lemmas.forget_prem; apply Always_imp ]
       end
  end.

Ltac specialize_arith_hyp H :=
  repeat match type of H with
         | ?G -> _ =>
           let HH := fresh "H" in
           assert G as HH by solve_linear;
             specialize (H HH); clear HH
         end.

Ltac specialize_arith :=
  repeat match goal with
         | [ H : ?G -> _ |- _ ] =>
           specialize_arith_hyp H
         end.

(* This simplifies real arithmetic goals.
   It sometimes is useful to run this before
   sending things to solve_nonlinear. *)
Ltac R_simplify :=
  unfold state, Value; field_simplify;
  unfold Rdiv;
  repeat rewrite RMicromega.Rinv_1;
  repeat
    match goal with
    | H:_ |- _ =>
      unfold state, Value in H; field_simplify in H;
      unfold Rdiv in H;
      repeat rewrite RMicromega.Rinv_1 in H;
      revert H
    end; intros.

(* Doesn't change the goal but runs
   z3 on real arithmetic goals. At the
   moment, you have to look in the *coq*
   buffer for the output. *)
Ltac z3_prepare :=
  intros.

Ltac z3_solve :=
  z3_prepare; smt solve.

Ltac z3_quick :=
  z3_prepare; smt solve.

(* rewrites the values of variables in the next
   state into hypothesis and goals. *)
Ltac rewrite_next_st :=
  repeat match goal with
           | [ H : eq (Stream.hd (Stream.tl _) _)  _ |- _ ]
             => rewrite H in *
         end.

(* Gets rid of arithmetic expressions of the
   form 0+_, _+0, 0*_, and _*0, _-0, 0-_. *)
Ltac rewrite_real_zeros :=
  repeat (first
            [ rewrite Rmult_0_r in *
            | rewrite Rmult_0_l in *
            | rewrite Rplus_0_r in *
            | rewrite Rplus_0_l in *
            | rewrite Rminus_0_r in *
            | rewrite Rminus_0_l in * ]).

Ltac simpl_Rmax :=
  repeat first [rewrite Rbasic_fun.Rmax_left in * by solve_linear |
                rewrite Rbasic_fun.Rmax_right in * by solve_linear ].

Local Open Scope HP_scope.

(* I'm not sure what the following three
   tactics do *)
(*
Ltac find_zeros eqs :=
  match eqs with
    | nil => constr:(@nil Var)
    | cons (DiffEqC ?y (ConstC (NatC O))) ?eqs =>
      let rest := find_zeros eqs in
      constr:(cons y rest)
    | cons _ ?eqs =>
      let rest := find_zeros eqs in
      rest
  end.
*)

(*
Ltac extract_unchanged eqs :=
  let xs := find_zeros eqs in
  let rec aux l :=
      match l with
        | nil => idtac
        | cons ?y ?l => apply zero_deriv
                        with (cp:=eqs) (x:=y);
                        try (aux l)
      end in
  aux xs.
*)

Ltac get_var_inv F x :=
  match F with
    | And ?F1 _ =>
      get_var_inv F1 x
    | And _ ?F2 =>
      get_var_inv F2 x
    | Comp (next_term x) (next_term ?e) Eq =>
      constr:(Comp x e Eq)
  end.

(* Applies differential induction with
   a known differential invariant *)
Ltac prove_diff_inv known :=
  match goal with
      |- context [ Continuous ?eqs ] =>
      match goal with
          |- (|-- _ -->> Comp (next_term ?t1)
                   (next_term ?t2) ?op) =>
          apply diff_ind with
          (Hyps:=known) (G:=Comp t1 t2 op) (cp:=eqs)
      end
  end.

(* Removes ! from variables in a Term *)
Fixpoint unnext_term (t:Term) : Term :=
  match t with
    | VarNowT x => VarNowT x
    | VarNextT x => VarNowT x
    | RealT r => RealT r
    | NatT n => NatT n
    | PlusT t1 t2 =>
      PlusT (unnext_term t1) (unnext_term t2)
    | MinusT t1 t2 =>
      MinusT (unnext_term t1) (unnext_term t2)
    | MultT t1 t2 =>
      MultT (unnext_term t1) (unnext_term t2)
    | InvT t => InvT (unnext_term t)
    | CosT t => CosT (unnext_term t)
    | SinT t => SinT (unnext_term t)
    | SqrtT t => SqrtT (unnext_term t)
    | ArctanT t => ArctanT (unnext_term t)
    | ExpT t => ExpT (unnext_term t)
    | MaxT t1 t2 => MaxT (unnext_term t1) (unnext_term t2)
    | Unop f t => Unop f (unnext_term t)
    | Binop f t1 t2 => Binop f (unnext_term t1)
                             (unnext_term t2)
  end.

(* Removes ! from variables in a Formula *)
Fixpoint unnext (F:Formula) : Formula :=
  match F with
    | Comp t1 t2 op =>
      Comp (unnext_term t1) (unnext_term t2) op
    | And F1 F2 => And (unnext F1) (unnext F2)
    | Or F1 F2 => Or (unnext F1) (unnext F2)
    | Imp F1 F2 => Imp (unnext F1) (unnext F2)
    | Syntax.Exists T f => Syntax.Exists T (fun t => unnext (f t))
    | Syntax.Forall T f => Syntax.Forall T (fun t => unnext (f t))
    | _ => F
  end.

(* Tries to prove (discrete) inductive goals in our examples.
   Only works for linear arithmetic. Leaves unsolved subgoals
   unchanged. *)
Ltac prove_inductive :=
  repeat apply or_next; repeat apply and_right;
  match goal with
    | [ |- context [Continuous ?deqs] ] =>
      match goal with
(*        | [ |- (|-- _ -->> (?HH -->> ?GG))] =>
          abstract (apply diff_ind_imp
                    with (eqs:=deqs) (H:=unnext HH)
                                     (G:=unnext GG);
                    solve [reflexivity |
                           simpl; intuition;
                           solve_linear])*)
(*        | [ |- _ ] =>
          abstract
            (apply unchanged_continuous with (eqs:=deqs);
             solve_linear)*)
        | [ |- (|-- _ -->> ?GG) ] =>
          abstract (eapply diff_ind
                    with (cp:=deqs) (G:=unnext GG)
                                    (Hyps:=TRUE);
                    try solve [reflexivity |
                               simpl; intuition;
                               solve_linear] )
      end
    | [ |- _ ] =>
      try abstract (solve_linear)
  end.

Ltac rewrite_projT2_L s :=
  let H := fresh in
  pose proof (projT2 s) as H;
    cbv beta in H; rewrite <- H; clear H.

Ltac rewrite_projT2_R s :=
  let H := fresh in
  pose proof (projT2 s) as H;
    cbv beta in H; rewrite H; clear H.


Set Implicit Arguments.

Require Import fcf.FCF.
Require Import fcf.CompFold.
Require Import fcf.PRF.
Require Import fcf.PRG_NA.

Local Open Scope list_scope.
Section HMAC_DRBG_PRG_NA.

(* HMAC-DRBG spec *)

(* The security parameter eta determines the size of PRF keys *)
Variable eta : nat.
Variable eta_nz : eta <> O.
(* The function f models HMAC *)
Variable f : Bvector eta -> Blist -> Bvector eta.

Definition RndK : Comp (Bvector eta) := {0,1}^eta.
Definition RndV : Comp (Bvector eta) := {0,1}^eta.
Definition KV : Set := (Bvector eta * Bvector eta)%type.

(* The Instantiate function *)
(* NOTE: does not reflect NIST spec *)
Definition Instantiate : Comp KV :=
  k <-$ RndK;
  v <-$ RndV;
  ret (k, v).

(* The Generate function *)
Definition to_list (A : Type) (n : nat) (v : Vector.t A n) := Vector.to_list v.

Fixpoint Gen_loop (k : Bvector eta) (v : Bvector eta) (n : nat)
  : list (Bvector eta) * Bvector eta :=
  match n with
  | O => (nil, v)
  | S n' =>
    let v' := f k (to_list v) in
    let (bits, v'') := Gen_loop k v' n' in
    (v' :: bits, v'')
  end.

(* Spec says "V || 0x00"; here we will use a list of 8 bits of 0 (a byte) *)
Fixpoint replicate {A} (n : nat) (a : A) : list A :=
  match n with
  | O => nil
  | S n' => a :: replicate n' a
  end.

Definition zeroes : list bool := replicate 8 false.

Definition Generate (state : KV) (n : nat) :
  Comp (list (Bvector eta) * KV) :=
  [k, v] <-2 state;
  [bits, v'] <-2 Gen_loop k v n;
  k' <- f k (to_list v' ++ zeroes);
  v'' <- f k' (to_list v');
  ret (bits, (k', v'')).

(* The adversary against the PRG *)
Variable blocksPerCall : nat.       (* blocks generated by GenLoop *)
Variable blocksPerCall_gt_0 : blocksPerCall > O.
Variable numCalls : nat.        (* number of calls to Generate *)
Variable numCalls_gt_0: numCalls > O.
Definition requestList : list nat := replicate numCalls blocksPerCall.
Variable A : list (list (Bvector eta)) -> Comp bool.
Variable A_wf : forall ls, well_formed_comp (A ls).

(* The constructed adversary against the PRF *)
Fixpoint oracleCompMap_inner {D R OracleIn OracleOut : Set}
           (e1 : EqDec ((list R) * (nat * KV)))
           (e2 : EqDec (list R))
           (* this is an oracleComp, not an oracle *)
           (* the oracle has type (D * R) -> D -> Comp (R, (D * R)) *)
           (oracleComp : (nat * KV) -> D -> OracleComp OracleIn OracleOut (R * (nat * KV)))
           (state : (nat * KV)) (* note this state type -- it is EXPLICITLY being passed around *)
           (inputs : list D) : OracleComp OracleIn OracleOut (list R * (nat * KV)) :=
  match inputs with
  | nil => $ ret (nil, state)
  | input :: inputs' =>
    [res, state'] <--$2 oracleComp state input;
    [resList, state''] <--$2 oracleCompMap_inner _ _ oracleComp state' inputs';
    $ ret (res :: resList, state'')
  end.

Definition oracleCompMap_outer {D R OracleIn OracleOut : Set}
           (e1 : EqDec ((list R) * (nat * KV)))
           (e2 : EqDec (list R))
           (oracleComp : (nat * KV) -> D -> OracleComp OracleIn OracleOut (R * (nat * KV)))
           (inputs : list D) : OracleComp OracleIn OracleOut (list R) :=
  [k, v] <--$2 $ Instantiate;   (* generate state inside, instead of being passed state *)
  [bits, _] <--$2 oracleCompMap_inner _ _ oracleComp (O, (k, v)) inputs;
  (* the "_" here has type (nat * KV) *)
  $ ret bits.

Definition Generate_v_PRF_oc (state : KV) (n : nat) :
  OracleComp (list bool) (Bvector eta) (list (Bvector eta) * KV) :=
  [k, v] <-2 state;
  v' <- f k (to_list v);
  [bits, v''] <-2 Gen_loop k v' n;
  k' <- f k (to_list v'' ++ zeroes);
  $ ret (bits, (k', v'')).

Fixpoint Gen_loop_oc (v : Bvector eta) (n : nat)
  : OracleComp (list bool) (Bvector eta) (list (Bvector eta) * Bvector eta) :=
  match n with
  | O => $ ret (nil, v)
  | S n' =>
    v' <--$ (OC_Query _ (to_list v)); (* ORACLE USE *)
    [bits, v''] <--$2 Gen_loop_oc v' n';
    $ ret (v' :: bits, v'')
  end.

Definition Generate_v_oc (state : KV) (n : nat) :
  OracleComp (list bool) (Bvector eta) (list (Bvector eta) * KV) :=
  [k, v_0] <-2 state;
  v <--$ (OC_Query _ (to_list v_0)); (* ORACLE USE *)
  [bits, v'] <--$2 Gen_loop_oc v n;
  (* TODO what's the state type here? and the global Generate_v_oc return type? *)
  k' <--$ (OC_Query _ (to_list v' ++ zeroes)); (* ORACLE USE *)
  $ ret (bits, (k', v')).

Definition Generate_noV_oc (state : KV) (n : nat) :
  OracleComp (list bool) (Bvector eta)  (list (Bvector eta) * KV) :=
  [k, v] <-2 state;
  [bits, v'] <--$2 Gen_loop_oc v n;
  (* TODO what's the state type here? and the global Generate_v_oc return type? *)
  k' <--$ (OC_Query _ (to_list v' ++ zeroes)); (* ORACLE USE *)
  $ ret (bits, (k', v')).

Fixpoint Gen_loop_rb_intermediate (k : Bvector eta) (v : Bvector eta) (n : nat)
  : Comp (list (Bvector eta) * Bvector eta) :=
  match n with
  | O => ret (nil, v)
  | S n' =>
    v' <-$ {0,1}^eta;
    [bits, v''] <-$2 Gen_loop_rb_intermediate k v' n';
    ret (v' :: bits, v'')
  end.

Definition Generate_rb_intermediate_oc (state : KV) (n : nat)
  : OracleComp (list bool) (Bvector eta) (list (Bvector eta) * KV) :=
  [k, v] <-2 state;
  v' <--$ $ {0,1}^eta;
  [bits, v''] <--$2 $ Gen_loop_rb_intermediate k v' n;    (* promote comp to oraclecomp, then remove from o.c. *)
  $ ret (bits, (k, v'')).

Definition Oi_oc' (i : nat) (sn : nat * KV) (n : nat)
  : OracleComp Blist (Bvector eta) (list (Bvector eta) * (nat * KV)) :=
  [callsSoFar, state] <-2 sn;
  let Generate_v_choose :=
      (* this behavior (applied with f_oracle) needs to match that of choose_Generate's *)
      if lt_dec callsSoFar i (* callsSoFar < i (override all else) *)
           then Generate_rb_intermediate_oc (* this implicitly has no v to update *)
      else if beq_nat callsSoFar O (* use oracle on 1st call w/o updating v *)
           then Generate_noV_oc
      else if beq_nat callsSoFar i (* callsSoFar = i *)
           then Generate_v_oc    (* uses provided oracle (PRF or RF) *)
      else Generate_v_PRF_oc in        (* uses PRF with (k,v) updating *)
  [bits, state'] <--$2 Generate_v_choose state n;
  $ ret (bits, (S callsSoFar, state')).

Definition PRF_Adversary (i : nat) : OracleComp Blist (Bvector eta) bool :=
  bits <--$ oracleCompMap_outer _ _ (Oi_oc' i) requestList;
  $ A bits.
(* End constructed adversary definition *)

Definition Pr_collisions := (S blocksPerCall)^2 / 2^eta.
Definition PRF_Advantage_Game i : Rat :=
  PRF_Advantage RndK ({0,1}^eta) f _ _ (PRF_Adversary i).
Fixpoint argMax(f : nat -> Rat) (n : nat) :=
  match n with
    | O => O
    | S n' => let p := (argMax f n') in
              if (le_Rat_dec (f (S n')) (f p)) then p else (S n')
                                                             end.

Definition PRF_Advantage_Max := PRF_Advantage_Game (argMax PRF_Advantage_Game numCalls).
Definition Gi_Gi_plus_1_bound := PRF_Advantage_Max + Pr_collisions.

(* The desired security property: HMAC_DRBG is a non-adaptive PRG *)
Definition HMAC_DRBG_PRG_NA :=
  PRG_Nonadaptive_Advantage _ _ ({0,1}^eta) Instantiate Generate
  (ret requestList) A <= (numCalls / 1) * Gi_Gi_plus_1_bound.

(* The proof of this property is imported and applied below *)
Require Import HMAC_DRBG_nonadaptive.

Theorem Generate_rb_eq_ideal :
  forall b x,
  evalDist (Generate_ideal (Bvector_EqDec eta) ({ 0 , 1 }^ eta) b) x ==
  evalDist (Generate_rb eta b) x.

  unfold Generate_rb, Generate_ideal in *.
  Local Opaque evalDist.
  induction b; intuition; simpl in *.
  - fcf_simp.
    reflexivity.

  - fcf_inline_first.
    fcf_skip.
    fcf_inline_first.
    fcf_skip.
    rewrite IHb.
    rewrite evalDist_right_ident.
    reflexivity.
    rewrite evalDist_right_ident.
    reflexivity.
Qed.

Theorem PRG_Advantage_eq :
  PRG_Nonadaptive_Advantage _ _ ({0,1}^eta) HMAC_DRBG_PRG_NA.Instantiate HMAC_DRBG_PRG_NA.Generate
  (ret HMAC_DRBG_PRG_NA.requestList) A ==
  | Pr [G_real f _ A blocksPerCall numCalls ] -
       Pr [G_ideal A blocksPerCall numCalls ] |.

  unfold PRG_Nonadaptive_Advantage.
  apply ratDistance_eqRat_compat.
  - unfold PRG_G1, G_real.
    fcf_skip.
    reflexivity.
    fcf_simp.
    reflexivity.

  - unfold PRG_G2, G_ideal.
    fcf_simp.
    fcf_skip.
    eapply compMap_eq; intuition.
    apply list_pred_eq.
    subst.
    apply Generate_rb_eq_ideal.
Qed.

Theorem HMAC_DRBG_PRG_NA_true : HMAC_DRBG_PRG_NA.

  unfold HMAC_DRBG_PRG_NA, PRG_Nonadaptive_Advantage.
  eapply leRat_trans.
  apply eqRat_impl_leRat.
  apply PRG_Advantage_eq.
  apply G1_G2_close; intuition.

Qed.