From iris.proofmode Require Import proofmode spec_patterns coq_tactics ltac_tactics reduction.
Require Import Eqdep_dec List.
From cap_machine Require Import classes rules macros_helpers.
From cap_machine Require Export iris_extra addr_reg_sample.
From cap_machine Require Export class_instances solve_pure solve_addr_extra.
From cap_machine Require Import NamedProp proofmode_instr_rules.
From machine_utils Require Export tactics.
From iris.bi Require Import bi.
Import bi.

From Ltac2 Require Import Ltac2.
From Ltac2 Require Option Bool Constr.
Set Default Proof Mode "Classic".

Section codefrag.
  Context {Σ:gFunctors} {memg:memG Σ} {regg:regG Σ}
          `{MP: MachineParameters}.

(* TODO: move elsewhere: to region.v? *)
Definition codefrag (a0: Addr) (cs: list Word) :=
  ([[ a0, (a0 ^+ length cs)%a ]] ↦ₐ [[ cs ]])%I.

Lemma codefrag_contiguous_region a0 cs :
  codefrag a0 cs -∗
  ⌜ContiguousRegion a0 (length cs)⌝.
Proof using.
  iIntros "Hcs". unfold codefrag.
  iDestruct (big_sepL2_length with "Hcs") as %Hl.
  set an := (a0 + length cs)%a in Hl |- *.
  unfold ContiguousRegion.
  destruct an eqn:Han; subst an; [ by eauto |]. cbn.
  exfalso. rewrite finz_seq_between_length /finz.dist in Hl.
  solve_addr.
Qed.

Lemma codefrag_lookup_acc a0 (cs: list Word) (i: nat) w:
  SimplTC (cs !! i) (Some w) →
  codefrag a0 cs -∗
  (a0 ^+ i)%a ↦ₐ w ∗ ((a0 ^+ i)%a ↦ₐ w -∗ codefrag a0 cs).
Proof.
  iIntros (Hi) "Hcs".
  iDestruct (codefrag_contiguous_region with "Hcs") as %Hub.
  rewrite /codefrag.
  destruct Hub as [? Hub].
  iDestruct (big_sepL2_lookup_acc with "Hcs") as "[Hw Hcont]"; only 2: by eauto.
  eapply finz_seq_between_lookup with (n:=length cs).
  { apply lookup_lt_is_Some_1; eauto. }
  { solve_addr. }
  iFrame.
Qed.

End codefrag.

(* Administrative reduction steps *)
Ltac wp_pure := iApply wp_pure_step_later; [ by auto | iNext ; iIntros "_" ].
(* TODO iIntros "_" fixes the lc 1 introduces in Iris 4.0.0, but I'm not sure that is the right place *)
Ltac wp_end := iApply wp_value.
Ltac wp_instr :=
  iApply (wp_bind (fill [SeqCtx]));
  (* Reduce the [fill] under the WP... *)
  let X := fresh in
  set X := (fill _);
  cbv in X;
  subst X;
  cbv beta.

Lemma z_addr_base {fb off off1 off2 : Z} {f0 f1 f2 : finz fb}: (f0 + off1)%f = Some f1 → (f0 + off2)%f = Some f2 → (off2 - off1)%Z = off → (f1 ^+ off)%f = f2.
  Proof. solve_addr. Qed.
Ltac solve_block_move :=
    first [solve_addr+ |
    lazymatch goal with
      | |- (?prev_a ^+ ?off)%f = ?cur_a =>
          match goal with
            | H: (prev_a + ?off1)%a = Some cur_a |- _ => solve_addr+H
            | H1: (?base_a + ?off1)%a = Some prev_a |- _ =>
                match goal with
                | H2 : (base_a + ?off2)%a = Some cur_a |- _ =>
                    apply (z_addr_base H1 H2); solve_addr+
                end
          end
     end ].

(* Ltac specifically meant for switching to the next block. Use `changePCto` to perform more arbitrary moves *)
Ltac changePC_next_block new_a :=
  match goal with |- context [ Esnoc _ _ (PC ↦ᵣ WCap _ _ _ ?prev_a)%I ] =>
                    rewrite (_: prev_a = new_a) ; [ | solve_block_move  ] end.
(* More powerful ltac to change the address of the pc. Might take longer to solve than the more specific alternative above.*)
Ltac changePCto0 new_a :=
  match goal with |- context [ Esnoc _ _ (PC ↦ᵣ WCap _ _ _ ?a)%I ] =>
    rewrite (_: a = new_a); [| solve_addr]
  end.
Tactic Notation "changePCto" constr(a) := changePCto0 a.

Ltac codefrag_facts h :=
  let h := constr:(h:ident) in
  match goal with |- context [ Esnoc _ h (codefrag ?a_base ?code) ] =>
    (match goal with H : ContiguousRegion a_base _ |- _ => idtac end ||
     let HH := fresh in
     iDestruct (codefrag_contiguous_region with h) as %HH;
     cbn [length map encodeInstrsW] in HH
    );
    (match goal with H : SubBounds _ _ a_base (a_base ^+ _)%a |- _ => idtac end ||
     (try match goal with H : SubBounds ?b ?e _ _ |- _ =>
            let HH := fresh in
            assert (HH: SubBounds b e a_base (a_base ^+ length code)%a) by solve_addr;
            cbn [length map encodeInstrsW] in HH
          end))
  end.

Ltac clear_codefrag_facts h :=
  let h := constr:(h:ident) in
  match goal with |- context [ Esnoc _ h (codefrag ?a_base ?code) ] =>
    try match goal with H : ContiguousRegion a_base _ |- _ => clear H end;
    try match goal with H : SubBounds _ _ a_base (a_base ^+ _)%a |- _ => clear H end
  end.

(* Classes for the code sub-block focusing tactic *)

Create HintDb proofmode_focus.
#[export] Hint Constants Opaque : proofmode_focus.
#[export] Hint Variables Opaque : proofmode_focus.

Definition App {A} (l1 l2 ll: list A) :=
  ll = l1 ++ l2.
#[export] Hint Mode App - ! ! - : proofmode_focus.

Lemma App_nil_r A (l: list A) : App l [] l.
Proof. rewrite /App app_nil_r //. Qed.
#[export] Hint Resolve App_nil_r : proofmode_focus.

Lemma App_nil_l A (l: list A) : App [] l l.
Proof. reflexivity. Qed.
#[export] Hint Resolve App_nil_l : proofmode_focus.

Lemma App_nil_default A (l1 l2: list A): App l1 l2 (l1 ++ l2).
Proof. reflexivity. Qed.
#[export] Hint Resolve App_nil_default | 100 : proofmode_focus.

Definition NthSubBlock {A} (ll: list A) (n: nat) (l1: list A) (l: list A) (l2: list A) :=
  ll = l1 ++ l ++ l2.
#[export] Hint Mode NthSubBlock + ! + - - - : proofmode_focus.

Lemma NthSubBlock_O_rest A (l1 l2: list A):
  NthSubBlock (l1 ++ l2) 0 [] l1 l2.
Proof. reflexivity. Qed.
#[export] Hint Resolve NthSubBlock_O_rest : proofmode_focus.

Lemma NthSubBlock_O_last A (l1: list A):
  NthSubBlock l1 0 [] l1 [].
Proof. unfold NthSubBlock. rewrite app_nil_r//. Qed.
#[export] Hint Resolve NthSubBlock_O_last | 100 : proofmode_focus.

Lemma NthSubBlock_S A (l0 ll l1 l2 l3 l0l1: list A) n:
  NthSubBlock ll n l1 l2 l3 →
  App l0 l1 l0l1 →
  NthSubBlock (l0 ++ ll) (S n) l0l1 l2 l3.
Proof. unfold NthSubBlock. intros -> ->. rewrite app_assoc //. Qed.
#[export] Hint Resolve NthSubBlock_S : proofmode_focus.

Section codefrag_subblock.
  Context {Σ:gFunctors} {memg:memG Σ} {regg:regG Σ}
          `{MP: MachineParameters}.

  Lemma codefrag_block0_acc a0 (l1 l2: list Word):
    codefrag a0 (l1 ++ l2) -∗
    codefrag a0 l1 ∗
    (codefrag a0 l1 -∗ codefrag a0 (l1 ++ l2)).
  Proof.
    rewrite /codefrag. iIntros "H".
    iDestruct (codefrag_contiguous_region with "H") as %Hregion.
    destruct Hregion as [an Han]. rewrite app_length in Han |- *.
    iDestruct (region_mapsto_split _ _ (a0 ^+ length l1)%a with "H") as "[H1 H2]".
    by solve_addr. by rewrite /finz.dist; solve_addr.
    iFrame. iIntros "H1".
    rewrite region_mapsto_split. iFrame. solve_addr. rewrite /finz.dist; solve_addr.
  Qed.

  Lemma codefrag_block_acc (n: nat) a0 (cs: list Word) l1 l l2:
    NthSubBlock cs n l1 l l2 →
    codefrag a0 cs -∗
    ∃ (ai: Addr), ⌜(a0 + length l1)%a = Some ai⌝ ∗
    codefrag ai l ∗
    (codefrag ai l -∗ codefrag a0 cs).
  Proof.
    unfold NthSubBlock. intros ->. rewrite /codefrag. iIntros "H".
    iDestruct (codefrag_contiguous_region with "H") as %[a1 Ha1].
    rewrite !app_length in Ha1 |- *.
    iDestruct (region_mapsto_split _ _ (a0 ^+ length l1)%a with "H") as "[H1 H2]".
    solve_addr. rewrite /finz.dist; solve_addr.
    iExists (a0 ^+ length l1)%a. iSplitR. iPureIntro; solve_addr.
    iDestruct (region_mapsto_split _ _ ((a0 ^+ length l1) ^+ length l)%a with "H2") as "[H2 H3]".
    solve_addr. rewrite /finz.dist; solve_addr. iFrame.
    iIntros "H2".
    rewrite region_mapsto_split. iFrame. 2: solve_addr. 2: rewrite /finz.dist; solve_addr.
    rewrite region_mapsto_split. iFrame. solve_addr. rewrite /finz.dist; solve_addr.
  Qed.

End codefrag_subblock.

Lemma focus_block_0_SubBounds (b e b' : Addr) k m :
  SubBounds b e b' (b' ^+ k)%a →
  ContiguousRegion b' m →
  (0 ≤ m)%Z →
  (m ≤ k)%Z →
  SubBounds b e b' (b' ^+ m)%a.
Proof. solve_addr. Qed.

(* More efficient version of codefrag_facts (avoids calling [solve_addr])
   because we have slightly more information when doing focus_block_0 *)
Ltac focus_block_0_codefrag_facts hi a0 :=
  let HCR := fresh in
  iDestruct (codefrag_contiguous_region with hi) as %HCR;
  cbn [length map encodeInstrsW] in HCR;
  try lazymatch goal with HSB : SubBounds ?b ?e a0 (a0 ^+ _)%a |- _ =>
    let HSB' := fresh in
    unshelve epose proof (focus_block_0_SubBounds _ _ _ _ _ HSB HCR _ _) as HSB';
    [solve_pure ..|];
    cbn [length map encodeInstrsW] in HSB'
  end.

Ltac focus_block_0 h hi hcont :=
  let h := constr:(h:ident) in
  let hi := constr:(hi:ident) in
  let hcont := constr:(hcont:ident) in
  let x := iFresh in
  match goal with |- context [ Esnoc _ h (codefrag ?a0 _) ] =>
    iPoseProof (codefrag_block0_acc with h) as x;
    eapply tac_and_destruct with x _ hi hcont _ _ _;
    [pm_reflexivity|pm_reduce;iSolveTC|
     pm_reduce;
     lazymatch goal with
     | |- False =>
       let hi := pretty_ident hi in
       let hcont := pretty_ident hcont in
       fail "focus_block_0:" hi "or" hcont "not fresh"
     | _ => idtac
     end];
    focus_block_0_codefrag_facts hi a0
  end.

Tactic Notation "focus_block_0" constr(h) "as" constr(hi) constr(hcont) :=
  focus_block_0 h hi hcont.

Lemma focus_block_SubBounds (b e b' b'': Addr) k m n :
  SubBounds b e b' (b' ^+ k)%a →
  ContiguousRegion b'' m →
  (b' + n)%a = Some b'' →
  (0 ≤ n)%Z →
  (0 ≤ m)%Z →
  ((n + m) <= k)%Z →
  SubBounds b e b'' (b'' ^+ m)%a.
Proof. solve_addr. Qed.

(* More efficient version of codefrag_facts (avoids calling [solve_addr])
   because we have slightly more information when doing focus_block *)
Ltac focus_block_codefrag_facts hi a0 Ha_base :=
  let HCR := fresh in
  iDestruct (codefrag_contiguous_region with hi) as %HCR;
  cbn [length map encodeInstrsW] in HCR;
  try lazymatch goal with HSB : SubBounds ?b ?e a0 (a0 ^+ _)%a |- _ =>
    let HSB' := fresh in
    unshelve epose proof (focus_block_SubBounds _ _ _ _ _ _ _ HSB HCR Ha_base _ _ _) as HSB';
    [solve_pure ..|];
    cbn [length map encodeInstrsW] in HSB'
  end.

Ltac focus_block n h a_base Ha_base hi hcont :=
  let h := constr:(h:ident) in
  let hi := constr:(hi:ident) in
  let hcont := constr:(hcont:ident) in
  let x := iFresh in
  match goal with |- context [ Esnoc _ h (codefrag ?a0 _) ] =>
    iPoseProof ((codefrag_block_acc n) with h) as (a_base) x;
      [ once (typeclasses eauto with proofmode_focus) | ];
    let xbase := iFresh in
    let y := iFresh in
    eapply tac_and_destruct with x _ xbase y _ _ _;
      [pm_reflexivity|pm_reduce;iSolveTC|pm_reduce];
    iPure xbase as Ha_base;
    eapply tac_and_destruct with y _ hi hcont _ _ _;
      [pm_reflexivity|pm_reduce;iSolveTC|pm_reduce];
    focus_block_codefrag_facts hi a0 Ha_base;
    changePC_next_block a_base
  end.

Tactic Notation "focus_block" constr(n) constr(h) "as"
       ident(a_base) simple_intropattern(Ha_base) constr(hi) constr(hcont) :=
  focus_block n h a_base Ha_base hi hcont.

Ltac unfocus_block hi hcont h :=
  let hi := constr:(hi:ident) in
  let hcont := constr:(hcont:ident) in
  let h := constr:(h:ident) in
  clear_codefrag_facts hi;
  iDestruct (hcont with hi) as h.

Tactic Notation "unfocus_block" constr(hi) constr(hcont) "as" constr(h) :=
  unfocus_block hi hcont h.

(* "iApply with on-demand framing" *)

(* TODO: These proofs should probably not use the proof mode tactics *)
Lemma envs_clear_spatial_sound_rev {PROP: bi} (Δ: envs PROP) :
  envs_wf Δ →
  of_envs (envs_clear_spatial Δ) ∗ [∗] env_spatial Δ ⊢ of_envs Δ.
Proof.
  intros.
  rewrite !of_envs_eq /envs_clear_spatial /=.
  rewrite -persistent_and_sep_assoc.
  iIntros "H". iDestruct "H" as "[% [[? _] ?]]". iFrame. iSplit; eauto.
Qed.

Lemma tac_specialize_assert_delay {PROP: bi} (Δ: envs PROP) j q R P1 P2 P1' F Q :
  envs_lookup j Δ = Some (q, R) →
  IntoWand q false R P1 P2 → AddModal P1' P1 Q →
  envs_entails (envs_delete true j q Δ) (P1' ∗ F) →
  match
    envs_app false (Esnoc Enil j P2)
      (envs_clear_spatial (envs_delete true j q Δ))
  with
  | Some Δ1 =>
    envs_entails Δ1 (F -∗ Q)
  | None => False
  end → envs_entails Δ Q.
Proof.
  rewrite envs_entails_unseal. intros ??? HH.
  destruct (envs_app _ _ _) eqn:?; last done.
  intros HQ.
  rewrite envs_lookup_sound //.
  set Δ' := envs_delete true j q Δ in HH Heqo |- *.
  iIntros "[HR HΔ']".
  iAssert (⌜envs_wf Δ'⌝)%I as %?.
  { rewrite of_envs_eq. iDestruct "HΔ'" as "[? ?]". eauto. }
  rewrite envs_clear_spatial_sound.
  set Δ'i := envs_clear_spatial Δ'.
  set Δ's := env_spatial Δ'.
  iDestruct "HΔ'" as "[#HΔ'i Hs]".
  iDestruct (envs_clear_spatial_sound_rev with "[$HΔ'i $Hs]") as "HΔ'"; auto.
  iDestruct (HH with "HΔ'") as "[HP1' HF]".
  iApply add_modal. eauto. iFrame. iIntros "HP1".
  iApply (HQ with "[HR HP1] HF").
  { iApply (envs_app_singleton_sound with "[] [HR HP1]"). eauto. iApply "HΔ'i".
    iApply (into_wand with "[HR]"). eauto. eauto. eauto. }
Qed.

(* Typeclass instances to look-up framable resources in the goal, for
   [FramableMachineResource] from [machine_utils/tactics.v]
*)

Class FramableRegisterPointsto (r: RegName) (w: Word) := {}.
#[export] Hint Mode FramableRegisterPointsto + - : typeclass_instances.
Class FramableMemoryPointsto (a: Addr) (dq: dfrac) (w: Word) := {}.
#[export] Hint Mode FramableMemoryPointsto + - - : typeclass_instances.
Class FramableCodefrag (a: Addr) (l: list Word) := {}.
#[export] Hint Mode FramableCodefrag + - : typeclass_instances.

Instance FramableRegisterPointsto_default r w :
  FramableRegisterPointsto r w
| 100. Qed.

Instance FramableMemoryPointsto_default a dq w :
  FramableMemoryPointsto a dq w
| 100. Qed.

Instance FramableCodefrag_default a l :
  FramableCodefrag a l
| 100. Qed.

Instance FramableMachineResource_reg `{regG Σ} r w :
  FramableRegisterPointsto r w →
  FramableMachineResource (r ↦ᵣ w).
Qed.

Instance FramableMachineResource_mem `{memG Σ} a dq w :
  FramableMemoryPointsto a dq w →
  FramableMachineResource (a ↦ₐ{dq} w).
Qed.

Instance FramableMachineResource_codefrag `{memG Σ} a l :
  FramableCodefrag a l →
  FramableMachineResource (codefrag a l).
Qed.

(* remembering names after auto-framing done by iFrameAuto *)

Ltac2 Type hyp_table_kind := [ Reg | Mem | Codefrag ].

Ltac2 record_framed
      (table: (constr * constr * hyp_table_kind) list ref)
      (framed: constr * constr)
  :=
  let (hname, hh) := framed in
  let (lhs, kind) :=
    lazy_match! hh with
    | (?r ↦ᵣ _)%I => (r, Reg)
    | (?a ↦ₐ{_} _)%I => (a, Mem)
    | (codefrag ?a _) => (a, Codefrag)
    end in
  table.(contents) := (hname, lhs, kind) :: table.(contents).

(* iApplyCapAuto *)

(* Helpers *)

Ltac solve_to_wand tt :=
    iSolveTC ||
    let P := match goal with |- IntoWand _ _ ?P _ _ => P end in
    fail "iSpecialize:" P "not an implication/wand".

Ltac2 iTypeOf (h: constr) :=
  let Δ := match! goal with [ |- envs_entails ?Δ _ ] => Δ end in
  (* XXX should be pm_eval *)
  eval cbv in (envs_lookup $h $Δ).

Ltac2 iFresh () :=
  ltac1:(iStartProof);
  lazy_match! goal with
  | [ |- envs_entails (Envs ?Δp ?Δs ?c) ?q ] =>
    let c' := eval vm_compute in (Pos.succ $c) in
    ltac1:(Δp Δs c' q |-
           change_no_check (envs_entails (Envs Δp Δs c') q))
      (Ltac1.of_constr Δp) (Ltac1.of_constr Δs) (Ltac1.of_constr c')
      (Ltac1.of_constr q);
    '(IAnon $c)
  end.

Ltac iNamedIntro :=
  let x := iFresh in
  iIntros x; iNamed x.
Ltac2 iNamedIntro () := ltac1:(iNamedIntro).

Ltac2 on_lasts tacs :=
  Control.extend [] (fun _ => ()) tacs.

(* iApplyCapAuto_init *)

Ltac2 iSpecializeDelay (h: constr) :=
  refine '(tac_specialize_assert_delay _ $h _ _ _ _ _ _ _ _ _ _ _ _);
  Control.shelve_unifiable ();
  Control.dispatch [
    (fun _ => ltac1:(pm_reflexivity));
    (fun _ => ltac1:(solve_to_wand tt));
    (fun _ => ltac1:(class_apply @add_modal_id));
    (fun _ => ltac1:(pm_reduce));
    (fun _ => ltac1:(pm_reduce))
  ].

Ltac iApplyHypLast H :=
  eapply tac_apply with H _ _ _;
  [pm_reflexivity
  |iSolveTC
  |pm_reduce; pm_prettify].

Ltac2 iApplyCapAutoT_init0 lemma :=
  let tbl := { contents := [] } in
  let x := iFresh () in
  ltac1:(x lem |- once (iPoseProofCore lem as false (fun H => iRename H into x)))
    (Ltac1.of_constr x) (Ltac1.of_constr lemma);
  on_lasts [(fun _ =>
    iSpecializeDelay x > [|
      ltac1:(h |-
        let f := iFresh in
        iIntros f;
        iApplyHypLast h;
        iRevert f) (Ltac1.of_constr x)
    ]
  )];
  tbl.

Ltac2 iApplyCapAuto_init0 lemma :=
  let _ := iApplyCapAutoT_init0 lemma in ().

Ltac2 Notation "iApplyCapAuto_init" lem(constr) :=
  iApplyCapAuto_init0 lem.

Ltac iApplyCapAuto_init lemma :=
  let f := ltac2:(lem |- iApplyCapAuto_init0 (Option.get (Ltac1.to_constr lem))) in
  f lemma.

(* Name resources in the goal according to the table *)

Definition check_addr_eq (a b: Addr) `{FinZEq _ a b res} := res.

Ltac2 name_cap_resource (name, lhs, kind) :=
  match kind with
  | Reg =>
    match! goal with [ |- context [ (?r ↦ᵣ ?x)%I ] ] =>
      assert_constr_eq r lhs;
      ltac1:(x r name |- change (r ↦ᵣ x)%I with (name ∷ (r ↦ᵣ x))%I)
        (Ltac1.of_constr x) (Ltac1.of_constr r) (Ltac1.of_constr name)
    end
  | Mem =>
    match! goal with [ |- context [ (?a ↦ₐ{?dq} ?x)%I ] ] =>
      let is_lhs := eval unfold check_addr_eq in (@check_addr_eq $a $lhs _ _) in
      assert_constr_eq is_lhs 'true;
      ltac1:(x dq a name |- change (a ↦ₐ{dq} x)%I with (name ∷ (a ↦ₐ{dq} x))%I)
        (Ltac1.of_constr x) (Ltac1.of_constr dq) (Ltac1.of_constr a) (Ltac1.of_constr name)
    end
  | Codefrag =>
    match! goal with [ |- context [ codefrag ?a ?l ] ] =>
      let is_lhs := eval unfold check_addr_eq in (@check_addr_eq $a $lhs _ _) in
      assert_constr_eq is_lhs 'true;
      ltac1:(l a name |- change (codefrag a l) with (name ∷ (codefrag a l)))
        (Ltac1.of_constr l) (Ltac1.of_constr a) (Ltac1.of_constr name)
    end
  end.

Lemma envs_entails_rew_goal {Σ} (Δ: envs (uPredI (iResUR Σ))) P P' :
  P = P' →
  envs_entails Δ P' →
  envs_entails Δ P.
Proof. intros ->. auto. Qed.

Ltac2 reintro_cap_resources tbl :=
  (refine '(@envs_entails_rew_goal _ _ _ _ _ _);
   Control.shelve_unifiable ()) >
  [ List.iter (fun x => try (name_cap_resource x)) (tbl.(contents)); reflexivity | () ];
  iNamedIntro ().

(* cleanup *)
(* TODO: make this extensible. Remove updatePcPerm? unfolding sometimes causes issues. *)
Ltac2 iApplyCapAuto_cleanup () :=
  cbn [rules_Get.denote rules_AddSubLt.denote updatePcPerm].

(* iApplyCapAutoCore *)

Ltac iNamedAccu_fail_explain :=
  lazymatch goal with
  | |- envs_entails _ (?remaining ∗ _) =>
    fail "iApplyCapAuto: the following resources could not be found in the context:"
         remaining
  end.

Ltac2 iApplyCapAutoCore lemma :=
  let tbl := iApplyCapAutoT_init0 lemma in
  let iFrameCap := fun () => record_framed tbl (iFrameAuto ()) in
  grepeat (fun _ =>
    Control.extend [] (fun _ => try (Control.once solve_pure))
      [ (fun _ => try (iFrameCap ())); (fun _ => ()) ]);
  on_lasts [ (fun _ => ltac1:(iNamedAccu || iNamedAccu_fail_explain)); (fun _ => ()) ];
  on_lasts [ (fun _ =>
    iNamedIntro ();
    try ltac1:(iNext);
    reintro_cap_resources tbl;
    iApplyCapAuto_cleanup ()
  )].

Ltac2 Notation "iApplyCapAuto" lem(constr) := iApplyCapAutoCore lem.
Tactic Notation "iApplyCapAuto" constr(lem) :=
  let f := ltac2:(lem |- iApplyCapAutoCore (Option.get (Ltac1.to_constr lem))) in
  f lem.

(* iInstr *)

Definition as_weak_addr_incr (a b: Addr) (z: Z) `{AsWeakFinZIncr _ a b z} :=
  (b, Z.to_nat z).

Lemma addr_incr_zero (a: Addr) : (a ^+ 0)%a = a.
Proof. solve_addr. Qed.

Lemma addr_incr_zero_nat (a: Addr) : (a ^+ 0%nat)%a = a.
Proof. solve_addr. Qed.

Ltac iInstr_lookup0 hprog hi hcont :=
  let hprog := constr:(hprog:ident) in
  lazymatch goal with |- context [ Esnoc _ hprog (codefrag ?a_base _) ] =>
  lazymatch goal with |- context [ Esnoc _ ?hpc (PC ↦ᵣ (WCap _ _ _ ?pc_a))%I ] =>
    let base_off := eval unfold as_weak_addr_incr in (@as_weak_addr_incr pc_a a_base _ _) in
    lazymatch base_off with
    | (?base, ?off) =>
      iPoseProofCore (codefrag_lookup_acc _ _ off with hprog) as false (fun H =>
        eapply tac_and_destruct with H _ hi hcont _ _ _;
        [pm_reflexivity
        |pm_reduce; iSolveTC
        |pm_reduce];
        rewrite ?addr_incr_zero ?addr_incr_zero_nat
      )
     end
  end end.

Tactic Notation "iInstr_lookup" constr(hprog) "as" constr(hi) constr(hcont) :=
  iInstr_lookup0 hprog hi hcont.

Ltac iInstr_get_rule0 hi cont :=
  let hi := constr:(hi:ident) in
  once (
    (lazymatch goal with |- context [ Esnoc _ hi (_ ↦ₐ encodeInstrW ?instr)%I ] => idtac end
     + (lazymatch goal with |- context [ Esnoc _ hi (_ ↦ₐ ?instr)%I ] =>
           fail 1 "Next instruction is not of the form (encodeInstrW _):" instr
         end + fail "" hi "not found"))
  );
  lazymatch goal with |- context [ Esnoc _ hi (_ ↦ₐ encodeInstrW ?instr)%I ] =>
    dispatch_instr_rule instr cont
  end.

Tactic Notation "iInstr_get_rule" constr(hi) tactic(cont) := iInstr_get_rule0 hi cont.

Ltac iInstr_close hprog :=
  (* because of iApplyCapAuto's context shuffling, [hi] and [hcont]
     are not valid anymore... recover them. *)
  (* XXX make this a bit more robust *)
  lazymatch goal with |- context [ Esnoc _ ?hi (_ ↦ₐ encodeInstrW _)%I ] =>
  lazymatch goal with |- context [ Esnoc _ ?hcont (_ ↦ₐ encodeInstrW _ -∗ _)%I ] =>
    notypeclasses refine (tac_specialize false _ hi _ hcont _ _ _ _ _ _ _ _ _);
    [pm_reflexivity
    |pm_reflexivity
    |iSolveTC
    |pm_reduce];
    iRename hcont into hprog
  end end.

(* TODO: find a way of displaying an error message if iApplyCapAuto fails,
   displaying the rule it was called on, and without silencing iApplyCapAuto's
   own error messages? *)
Ltac iInstr hprog :=
  let hi := iFresh in
  let hcont := iFresh in
  iInstr_lookup hprog as hi hcont;
  try wp_instr;
  iInstr_get_rule hi ltac:(fun rule =>
    iApplyCapAuto rule;
    [ .. | iInstr_close hprog; try wp_pure]
  ).

Ltac2 rec iGo hprog :=
  let stop_if_at_least_two_goals () :=
    on_lasts [ (fun _ => ()); (fun _ => ()) ] in
  match Control.case (fun _ => ltac1:(hprog |- iInstr hprog) (Ltac1.of_constr hprog)) with
  | Err (_) => ()
  | Val (_) => Control.plus stop_if_at_least_two_goals (fun _ => iGo hprog)
  end.

Ltac iGo hprog :=
  let f := ltac2:(hprog |- iGo (Option.get (Ltac1.to_constr hprog))) in
  f hprog.


Require Import Language.

Lemma ptype_construction :
  forall (A B : typ) (v1 v2 : trm),
    ptype v1 A -> ptype v2 B ->
    ptype (trm_merge v1 v2) (typ_and A B).
Proof.
  intros.
  eapply ptype_merge; eauto.
Qed.


Require Import List Bool Reals String.
Require Import matrices mymatrices.
Set Implicit Arguments.

Open Scope string.
Open Scope list.

Definition var := string.
Definition heap_var := string.
Definition var_dec := string_dec.
Definition heap_var_dec := string_dec.

Inductive num_binop := Plus | Times | Minus | Divide | Min | Max.
Inductive bool_binop := bAnd | bOr.
Inductive comp_binop := bEq | bNeq | bLe.
Inductive unop :=
| Transpose | Inverse | Neg | BNot.
Inductive binop :=
| Multiply (* matrix multiplication *)
| NumBinop : num_binop -> binop
| BoolBinop : bool_binop -> binop
| CompBinop : comp_binop -> binop.

Coercion NumBinop : num_binop >-> binop.
Coercion BoolBinop : bool_binop >-> binop.
Coercion CompBinop : comp_binop >-> binop.

(* idx min max m n : return the index into a min x max-sized
   matrix stored in the heap. *)
Definition idx (M N m n : nat) := m * N + n.

Inductive type :=
| Unit | Nat | Bool
| Matrix : nat -> nat -> type
| Computation : proposition -> var -> type -> proposition -> type

with proposition :=
| Id : type -> intro_term -> intro_term -> proposition
| HId : heap -> heap -> proposition
| Indom : heap -> intro_term -> proposition
| Top | Bottom
| And : proposition -> proposition -> proposition
| Or : proposition -> proposition -> proposition
| Implies : proposition -> proposition -> proposition
| Not : proposition -> proposition

| Exists : var -> type -> proposition -> proposition
| ForAll : var -> type -> proposition -> proposition
| ExistsHeap : heap_var -> proposition -> proposition
| ForAllHeap : heap_var -> proposition -> proposition

with heap :=
| HVar : heap_var -> heap | HEmpty
| HUpdate : heap -> intro_term -> intro_term -> type -> heap

with elim_term :=
| Var : var -> elim_term | Cast : intro_term -> type -> elim_term

with intro_term :=
| ElimTerm : elim_term -> intro_term
| TT
| Boolean : bool -> intro_term
| Ref : nat -> intro_term
| Dia : computation -> intro_term
| Mat : forall m n, numMatrix m n -> intro_term
| Unop : unop -> intro_term -> intro_term
| Binop : binop -> intro_term -> intro_term -> intro_term
| Select : intro_term -> intro_term -> intro_term (* indexing into matrix *)

with computation :=
| IntroTerm : intro_term -> computation
| LetDia : var -> elim_term -> computation -> computation
| Seq : command -> computation -> computation

with command :=
| Annotation : proposition -> command
(* Allocate x M A:
   - if M:A is an mxn matrix, allocate m*n contiguous memory locations storing
     the data in M in the heap; return the pointer to the data in M
   - if M:A is a Nat or Bool, allocate a single memory location
   *)
| Alloc : var -> intro_term -> type -> command
(* Dealloc x M A :
   - if A is an mxn matrix, deallocates the entire m*n block in memory
     referenced by M.
   - if A is Nat, Bool or Unit, deallocates the single element in memory
     location M
*)
| Dealloc : var -> intro_term -> type -> command
(* Lookup x M A :
   - if A is an m x n matrix, selects the m x n matrix stored
     in the heap at memory location M. NOTE: M CAN be an offset.
     ex: to select the 2nd row from a 3x3 numerical matrix,
         DerefMatrix x (idx 3 3 2 1) (NumericalMatrix 1 3)
   - if A is a Nat, Bool or Unit, selects the data at location M which
     must be of type A
*)
| Lookup : var -> intro_term -> type -> command
| Mutate : intro_term -> intro_term -> type -> command
| If : var -> intro_term -> type -> computation -> computation -> command
| Loop : var -> proposition -> intro_term -> type -> var -> intro_term -> var -> computation -> command
| Send : type -> intro_term -> intro_term -> command
| Receive : var -> type -> intro_term -> command.

Scheme type_mut := Induction for type Sort Set
with proposition_mut := Induction for proposition Sort Set
with heap_mut := Induction for heap Sort Set
with elim_mut := Induction for elim_term Sort Set
with intro_mut := Induction for intro_term Sort Set
with computation_mut := Induction for computation Sort Set
with command_mut := Induction for command Sort Set.

Ltac type_discriminate := right; discriminate.
Ltac type_destruct := intro t2; destruct t2; auto; try type_discriminate.
Ltac type_injection := right; injection; auto; fail.
Ltac rewrite_vars :=
  repeat match goal with
  | [ H : ?v = ?v0 |- _ ] => rewrite H in *; clear H
  end; auto.

(* Seleq-related definitions *)
Definition seleq (A : type) (h : heap) (M N : intro_term)(x:heap_var) : proposition :=
  ExistsHeap x (HId h (HUpdate (HVar x) M N A)).
Definition seleq_ (A : type) (h : heap) (M : intro_term) (x:heap_var)(y:var) :=
  Exists y A (seleq A h M (ElimTerm (Var y)) x).

(* Coercions *)
Coercion ElimTerm : elim_term >-> intro_term.
Coercion Mat : numMatrix >-> intro_term.
Coercion IntroTerm : intro_term >-> computation.

(* Type Notation *)
Notation "{[ P ]} x ; A {[ Q ]}" := (Computation P x A Q) (at level 40).
Check ({[ Id _ TT TT ]} "hi" ; Unit {[Id _ TT TT]}).

(* Computation Notation *)
Notation "'Let' 'dia' x <- k 'in' e" := (LetDia x k e) (at level 40).
Notation "c ;; e" := (Seq c e) (at level 30, right associativity).

Definition maybeNumBinopDenote (b : num_binop)
           : forall m n, numMatrix m n -> forall m' n', numMatrix m' n' -> option (numMatrix m n) :=
  match b with
  | Plus => maybeNumPlus
  | Times => maybeNumTimes
  | Minus => maybeNumMinus
  | Divide => maybeNumDiv
  | Min => maybeNumMin
  | Max => maybeNumMax
  end.

Definition boolBinopDenote (b : bool_binop) (b1 b2:bool) : bool :=
  match b with
  | bAnd => andb b1 b2
  | bOr => orb b1 b2
  end.

Definition maybeNatBinopDenote (b : num_binop) (n1 n2 : nat) : option nat :=
  match b with
  | Plus => Some (n1+n2)
  | Times => Some (n1*n2)
  | Minus => Some (n1-n2)
  | Divide => None
  | Min => Some (min n1 n2)
  | Max => Some (max n1 n2)
  end.

Definition maybeNumCompBinopDenote (b : comp_binop)
           : forall m n, numMatrix m n -> forall m' n', numMatrix m' n' -> option bool :=
  match b with
  | bEq => maybe_numEqb
  | bNeq => maybe_numNeqb
  | bLe => maybe_numLeb
  end.

Fixpoint normalize (M : intro_term) : intro_term :=
  match M with
  | Unop Transpose M' => let M'' := normalize M' in
    match M'' with
    | Mat _ _ N => Mat (transpose N)
    | _ => Unop Transpose M''
    end
  | Unop Inverse M' => let M'' := normalize M' in Unop Inverse M''
  | Unop BNot M' => let M'' := normalize M' in
    match M'' with
    | Boolean b => Boolean (negb b)
    | _ => Unop BNot M''
    end
  | Unop Neg M' => let M'' := normalize M' in
    match M'' with
    | Mat _ _ N => Mat (numNeg N)
    | _ => Unop Neg M''
    end
  | Binop Multiply M1 M2 =>
    let M1' := normalize M1 in
    let M2' := normalize M2 in
    match M1', M2' with
    | Mat m n N1, Mat n' p N2 =>
      match maybeNumProd N1 N2 with
      | Some N => Mat N
      | None => Binop Multiply (Mat N1) (Mat N2)
      end
    | _, _ => Binop Multiply M1' M2'
    end
  | Binop (NumBinop b) M1 M2 =>
    let M1' := normalize M1 in
    let M2' := normalize M2 in
    match M1', M2' with
    | Mat _ _ N1, Mat _ _ N2 =>
      match maybeNumBinopDenote b N1 N2 with
      | Some N => Mat N
      | None => Binop b (Mat N1) (Mat N2)
      end
    | Ref n1, Ref n2 =>
      match maybeNatBinopDenote b n1 n2 with
      | Some n => Ref n
      | None => Binop b (Ref n1) (Ref n2)
      end
    | _, _ => Binop b M1' M2'
    end
  | Binop (BoolBinop b) M1 M2 =>
    let M1' := normalize M1 in
    let M2' := normalize M2 in
    match M1', M2' with
    | Boolean b1, Boolean b2 =>
      Boolean (boolBinopDenote b b1 b2)
    | _, _ => Binop b M1' M2'
    end
  | Binop (CompBinop b) M1 M2 =>
    let M1' := normalize M1 in
    let M2' := normalize M2 in
    match M1', M2' with
    | Mat _ _ N1, Mat _ _ N2 =>
      match maybeNumCompBinopDenote b N1 N2 with
      | Some b => Boolean b
      | None => Binop b (Mat N1) (Mat N2)
      end
    | _, _ => Binop b M1' M2'
    end
  | Select M N =>
    let M' := normalize M in
    let N' := normalize N in
    Select M' N'
  | _ => M
  end.

(* free variables *)
Fixpoint free_in_proposition x P :=
  match P with
  | Id A M1 M2 => free_in_type x A || free_in_intro_term x M1 || free_in_intro_term x M2
  | HId h1 h2 => free_in_heap x h1 || free_in_heap x h2
  | Indom h M => free_in_intro_term x M
  | Top => false
  | Bottom => false
  | And P1 P2 => free_in_proposition x P1 || free_in_proposition x P2
  | Or P1 P2 => free_in_proposition x P1 || free_in_proposition x P2
  | Implies P1 P2 => free_in_proposition x P1 || free_in_proposition x P2
  | Not P' => free_in_proposition x P'
  | Exists y A P' => if var_dec x y then false else (free_in_type x A || free_in_proposition x P')
  | ForAll y A P' => if var_dec x y then false else (free_in_type x A || free_in_proposition x P')
  | ExistsHeap h P' => free_in_proposition x P'
  | ForAllHeap h P' => free_in_proposition x P'
  end
with free_in_type x A :=
  match A with
  | Computation P y B Q => free_in_proposition x P || free_in_type x B ||
    (if var_dec x y then false else free_in_proposition x Q)
  | _ => false
  end
with free_in_heap x H :=
  match H with
  | HUpdate H' M N A => free_in_heap x H' || free_in_intro_term x M
                     || free_in_intro_term x N || free_in_type x A
  | _ => false
  end
with free_in_elim_term x K :=
  match K with
  | Var y => if var_dec x y then true else false
  | Cast M A => free_in_intro_term x M || free_in_type x A
  end
with free_in_intro_term x M :=
  match M with
  | ElimTerm K => free_in_elim_term x K
  | TT => false
  | Boolean _ => false
  | Ref _ => false
  | Dia E => free_in_computation x E
  | Mat _ _ _ => false
  | Unop _ M' => free_in_intro_term x M'
  | Binop _ M1 M2 => free_in_intro_term x M1 || free_in_intro_term x M2
  | Select M1 M2 => free_in_intro_term x M1 || free_in_intro_term x M2
  end
with free_in_computation x E :=
  match E with
  | IntroTerm M => free_in_intro_term x M
  | LetDia y K E' => free_in_elim_term x K || (if var_dec x y then false else free_in_computation x E')
  | Seq (Annotation P) E' => free_in_proposition x P || free_in_computation x E'
  | Seq (Alloc x' M A) E' => free_in_intro_term x M || free_in_type x A ||
    if var_dec x x' then false else free_in_computation x E'
  | Seq (Dealloc x' M A) E' => free_in_intro_term x M || free_in_type x A ||
    if var_dec x x' then false else free_in_computation x E'
  | Seq (Lookup x' M A) E' => free_in_intro_term x M || free_in_type x A ||
    if var_dec x x' then false else free_in_computation x E'
  | Seq (Mutate M N A) E' =>
    free_in_intro_term x M || free_in_intro_term x N || free_in_type x A ||
    free_in_computation x E'
  | Seq (Send A M N) E' =>
    free_in_type x A || free_in_intro_term x M || free_in_intro_term x N ||
    free_in_computation x E'
  | Seq (Receive x' A N) E' => free_in_type x A || free_in_intro_term x N ||
    if var_dec x x' then false else free_in_computation x E'
  | Seq (If x' M A E1 E2) E' => free_in_intro_term x M || free_in_type x A ||
    free_in_computation x E1 || free_in_computation x E2 ||
    if var_dec x x' then false else free_in_computation x E'
  | Seq (Loop x' P M A z1 N z2 E) E' =>
    free_in_intro_term x M || free_in_type x A ||
    (if var_dec x z1 then false else
     if var_dec x x' then false else
     free_in_proposition x P) ||
    (if var_dec x z1 then false else free_in_intro_term x N) ||
    (if var_dec x z2 then false else free_in_computation x E) ||
    (if var_dec x x' then false else free_in_computation x E')
  end.

Fixpoint heap_free_in_proposition x P :=
  match P with
  | Id A M1 M2 => heap_free_in_type x A || heap_free_in_intro_term x M1 || heap_free_in_intro_term x M2
  | HId h1 h2 => heap_free_in_heap x h1 || heap_free_in_heap x h2
  | Indom h M => heap_free_in_heap x h || heap_free_in_intro_term x M
  | Top => false
  | Bottom => false
  | And P1 P2 => heap_free_in_proposition x P1 || heap_free_in_proposition x P2
  | Or P1 P2 => heap_free_in_proposition x P1 || heap_free_in_proposition x P2
  | Implies P1 P2 => heap_free_in_proposition x P1 || heap_free_in_proposition x P2
  | Not P' => heap_free_in_proposition x P'
  | Exists y A P' => heap_free_in_type x A || heap_free_in_proposition x P'
  | ForAll y A P' => heap_free_in_type x A || heap_free_in_proposition x P'
  | ExistsHeap h P' => if heap_var_dec x h then false else heap_free_in_proposition x P'
  | ForAllHeap h P' => if heap_var_dec x h then false else heap_free_in_proposition x P'
  end
with heap_free_in_type x A :=
  match A with
  | Computation P y B Q => heap_free_in_proposition x P || heap_free_in_type x B
    || heap_free_in_proposition x Q
  | _ => false
  end
with heap_free_in_heap x H :=
  match H with
  | HUpdate H' M N A => heap_free_in_heap x H' || heap_free_in_intro_term x M
                     || heap_free_in_intro_term x N || heap_free_in_type x A
  | HVar h => if heap_var_dec x h then true else false
  | HEmpty => false
  end
with heap_free_in_elim_term x K :=
  match K with
  | Var y => false
  | Cast M A => heap_free_in_intro_term x M || heap_free_in_type x A
  end
with heap_free_in_intro_term x M :=
  match M with
  | ElimTerm K => heap_free_in_elim_term x K
  | TT => false
  | Boolean _ => false
  | Ref _ => false
  | Dia E => heap_free_in_computation x E
  | Mat _ _ _ => false
  | Unop _ M' => heap_free_in_intro_term x M'
  | Binop _ M1 M2 => heap_free_in_intro_term x M1 || heap_free_in_intro_term x M2
  | Select M1 M2 => heap_free_in_intro_term x M1 || heap_free_in_intro_term x M2
  end
with heap_free_in_computation x E :=
  match E with
  | IntroTerm M => heap_free_in_intro_term x M
  | LetDia y K E' => heap_free_in_elim_term x K || heap_free_in_computation x E'
  | Seq c E' => heap_free_in_command x c || heap_free_in_computation x E'
  end
with heap_free_in_command x c :=
  match c with
  | Annotation P => heap_free_in_proposition x P
  | Alloc y M A => heap_free_in_intro_term x M || heap_free_in_type x A
  | Dealloc y M A => heap_free_in_intro_term x M || heap_free_in_type x A
  | Lookup y M A => heap_free_in_intro_term x M || heap_free_in_type x A
  | Mutate M N A => heap_free_in_intro_term x M || heap_free_in_intro_term x N || heap_free_in_type x A
  | Send A M N => heap_free_in_type x A || heap_free_in_intro_term x M || heap_free_in_intro_term x N
  | Receive y A N => heap_free_in_type x A || heap_free_in_intro_term x N
  | If y M A E1 E2 => heap_free_in_intro_term x M || heap_free_in_type x A
     || heap_free_in_computation x E1 || heap_free_in_computation x E2
  | Loop y P M A z1 N z2 E => heap_free_in_proposition x P || heap_free_in_intro_term x M
    || heap_free_in_type x A || heap_free_in_intro_term x N || heap_free_in_computation x E
  end.


Require Import Blech.Defaults.

Definition Predicate A := A → Prop.

Definition subsetof {A} (P: Predicate A) := { x | P x }.

Coercion subsetof: Predicate >-> Sortclass.


Require Import Coqlib.
Require Import List. Import ListNotations.
Require Import Coq.Logic.FunctionalExtensionality.

Require Import sha.ByteBitRelations.
Require Import sha.general_lemmas.

Module Type ABSTRACT_ENTROPY.

Parameter stream: Type.

Inductive error_code: Type :=
| catastrophic_error
| generic_error.

Inductive result X: Type: Type :=
| success: X -> stream -> @result X
| error : error_code -> stream -> @result X.

Arguments success {X} _ _.
Arguments error {X} _ _.

Parameter get_bytes: nat -> stream -> result (list Z).
Parameter get_bits: nat -> stream -> result (list bool).

Parameter get_bits_length:
  forall k s b s',
    success b s' = get_bits k s ->
    length b = k.
Parameter get_bytes_length:
  forall k s b s',
    success b s' = get_bytes k s ->
    length b = k.
Parameter get_bytes_Zlength:
  forall k s b s',
    k >= 0 ->
    success b s' = get_bytes (Z.to_nat k) s ->
    Zlength b = k.
Parameter get_bytes_isbyteZ:
  forall k s b s',
    success b s' = get_bytes k s ->
    Forall isbyteZ b.
End ABSTRACT_ENTROPY.

Module OPTIONAL_ENTROPY <: ABSTRACT_ENTROPY.

Definition stream: Type := nat -> option bool.

Inductive error_code: Type :=
| catastrophic_error
| generic_error
.

Inductive result X: Type: Type :=
| success: X -> stream -> @result X
| error : error_code -> stream -> @result X
.

Arguments success {X} _ _.
Arguments error {X} _ _.

Fixpoint get_bits (k: nat) (s: stream): result (list bool) :=
  match k with
    | O => success [] s
    | S k' =>
      match get_bits k' s with (* recursive call -- get the (k-1) bits *)
      | error e s' => error  e s'
      | success b s' =>
        match s' O with (* get the first bit *)
        | None =>
          (* what does the caller do on error? call again? *)
          error catastrophic_error (fun i => match nat_compare i k' with
                                             | Lt => s i
                                             | Eq | Gt => s (1 + i)%nat
                                             end)
        | Some e => success (b ++ [e]) (fun i => s (k + i)%nat)
        (* not 1+i? (otherwise it'll be moving it 1 + 2 + 3 forward instead of 3 forward. not ::? *)
        end
      end
  end.

Example get_bits_test1:
  forall bits s output s',
    bits = [Some false; Some true; Some false; Some false] ->
    s = (fun i => nth i bits None) ->
    success output s' = get_bits (length bits) s ->
    s' = (fun i => s (length bits + i)%nat) /\ output = [false; true; false; false].
Proof.
  intros.
  subst.
  inv H1.
  split.
  extensionality i. destruct i; reflexivity.
  reflexivity.
Qed.

Example get_bits_test2:
  forall bits s output s' s'',
    bits = [Some false; None; Some true; Some false; Some false] ->
    s = (fun i => nth i bits None) ->
    error catastrophic_error s' = get_bits 4%nat s ->
    success output s'' = get_bits 4%nat s' ->
    s'' = (fun i => s (length bits + i)%nat) /\ output = [false; true; false; false].
Proof.
  intros.
  subst.
  inv H1.
  inv H2.
  split.
  extensionality i. reflexivity.
  reflexivity.
Qed.

Example get_bits_test3:
  forall bits s output s' s'',
    bits = [None; Some false; Some true; Some false; Some false] ->
    s = (fun i => nth i bits None) ->
    error catastrophic_error s' = get_bits 4%nat s ->
    success output s'' = get_bits 4%nat s' ->
    s'' = (fun i => s (length bits + i)%nat) /\ output = [false; true; false; false].
Proof.
  intros.
  subst.
  inv H1.
  inv H2.
  split.
  extensionality i. reflexivity.
  reflexivity.
Qed.

Example get_bits_test4:
  forall bits s output s' s'',
    bits = [Some false; Some true; Some false; None; Some false] ->
    s = (fun i => nth i bits None) ->
    error catastrophic_error s' = get_bits 4%nat s ->
    success output s'' = get_bits 4%nat s' ->
    s'' = (fun i => s (length bits + i)%nat) /\ output = [false; true; false; false].
Proof.
  intros.
  subst.
  inv H1.
  inv H2.
  split.
  extensionality i. reflexivity.
  reflexivity.
Qed.
(*
Fixpoint get_bits_concrete (k: nat) (s: stream) (max: nat): @result (list bool) :=
  match k with
    | O => success [] s
    | S k' =>
      match s (max - k)%nat with
        | None => error catastrophic_error (fun i => s (1 + i)%nat)
        | Some bit =>
          match get_bits_concrete k' s max with
            | error e s' => error e (fun i => match nat_compare i (max - k)%nat with
                                                  | Gt => s' i
                                                  | Eq | Lt => s i
                                     end)
            | success b s' => success (bit::b) (fun i => s (k + i)%nat)
          end
      end
  end
.
Lemma get_bits_concrete_correct:
  forall k s, get_bits k s = get_bits_concrete k s k.
Proof.
  intros k.
  induction k as [|k']; [reflexivity|].
  intros s.
  simpl.
  rewrite IHk'. clear IHk'.
  replace (k' - k')%nat with O by omega.
  remember (s O) as sO.
  destruct sO.
  {
    (* s 0 <> None *)
    (*
    remember (get_bits_concrete k' s k') as result.
    destruct result as [string s' | e s'].
    {
      remember (get_bits_concrete k' s (S k')) as result2.
      destruct result2 as [string2 s'2 | e2 s'2].
      {
        (* case where both result return success *)
        generalize dependent s.
        induction k' as [|k'']; intros.
        simpl in *. inv Heqresult. inv Heqresult2. rewrite <- HeqsO. auto.
        simpl in *. replace (k'' - k'')%nat with O in * by omega. rewrite <- HeqsO in *.
        rewrite IHk''.
      }
      unfold get_bits_concrete.
    }
*)
    induction k' as [|k''].
    {
      simpl. rewrite <- HeqsO.
      reflexivity.
    }
    simpl.
    replace (k'' - k'')%nat with O by omega.
    rewrite <- HeqsO in *.
    replace (match k'' with
             | 0%nat => S k''
             | S l => (k'' - l)%nat
             end) with 1%nat by (destruct k''; omega).

  }
  {
    destruct k' as [|k''].
    {
      simpl. rewrite <- HeqsO.
      replace (fun i : nat =>
      match nat_compare i 0 with
      | Eq => s (S i)
      | Lt => s i
      | Gt => s (S i)
      end) with (fun i => s (S i)); [reflexivity|].
      extensionality i.
      destruct i as [|i']; reflexivity.
    }
    simpl.
    replace (k'' - k'')%nat with O by omega; rewrite <- HeqsO.
    reflexivity.
  }
  simpl.
  *)

Fixpoint get_first_None (limit:nat) (s: stream): nat :=
  match limit with
    | O => O (* bogus *)
    | S limit' =>
      match s O with
        | None => O
        | Some _ => S (get_first_None limit' (fun i => s (S i)))
      end
  end
.

Lemma get_bits_stream_success:
  forall k s b s',
    success b s' = get_bits k s ->
    s' = fun i => s (k + i)%nat.
Proof.
  intros k.
  destruct k as [|k'].
  {
    simpl. intros. inv H. extensionality i; reflexivity.
  }
  simpl.
  intros.
  destruct (get_bits k' s); try solve [inversion H].
  destruct (s0 0%nat); try solve [inversion H].
  inv H.
  reflexivity.
Qed.

Lemma get_bits_stream_success_Some:
  forall k s b s',
    success b s' = get_bits k s ->
    (forall i, (i < k)%nat -> s i <> None).
Proof.
  intros k.
  induction k as [|k']; intros. omega.
  simpl in H.
  remember (get_bits k' s) as get_bits_k'_s; destruct get_bits_k'_s; try solve [inversion H].
  remember (s0 0%nat) as s0_0; destruct s0_0; try solve [inversion H].
  inv H.
  pose proof Heqget_bits_k'_s as IHsimpl.
  pose proof (Heqget_bits_k'_s) as Hstream.
  apply get_bits_stream_success in Hstream; subst.
  destruct (le_lt_eq_dec i k'). omega.
  {
    (* i < k', use inductive case *)
    eapply IHk'; eassumption.
  }
  {
    (* i = k', use known info *)
    subst.
    replace (k' + 0)%nat with k' in Heqs0_0 by omega.
    intros contra; rewrite <- Heqs0_0 in contra; inversion contra.
  }
Qed.

Lemma get_first_None_limit:
  forall k s i,
    i = get_first_None k s ->
    (i <= k)%nat.
Proof.
  induction k as [|k']; intros. simpl in H. omega.
  simpl in H.
  remember (s O) as s_O; destruct s_O.
  destruct i as [|i']; try solve [omega].
  inversion H.
  pose proof (IHk' (fun i => s (S i)) i').
  subst; omega.
  subst; omega.
Qed.

(* TODO convert this to use get_first_None_limit *)
Lemma get_bits_stream_error:
  forall k s e s',
    error e s' = get_bits k s ->
    (exists i, (i < k)%nat /\ (forall i', s' i' = match nat_compare i' i with
                                   | Lt => s i'
                                   | Eq | Gt => s (S i')
                                 end) /\ s i = None /\ (forall i', (i' < i)%nat -> s i' <> None)).
Proof.
  intros k.
  induction k as [|k']; intros. inversion H.
  simpl in H.
  remember (get_bits k' s) as get_bits_k'_s; destruct get_bits_k'_s.
  {
    (* get_bits k' s = success *)
    pose proof (Heqget_bits_k'_s) as Hstream.
    apply get_bits_stream_success in Hstream; subst.
    replace (k' + 0)%nat with k' in H by omega.
    remember (s k') as s_k'; destruct s_k'; inv H.
    exists k'.
    repeat split; auto.
    intros.
    eapply get_bits_stream_success_Some; eassumption.
  }
  {
    (* get_bits k' s = error *)
    pose proof (IHk' s e0 s0 Heqget_bits_k'_s) as IHsimpl.
    destruct IHsimpl.
    destruct (le_lt_dec k' x).
    {
      (* x >= k' *)
      destruct H0. omega.
    }
    {
      exists x.
      inv H.
      destruct H0.
      split; auto.
    }
  }
Qed.

Lemma get_bits_length:
  forall k s b s',
    success b s' = get_bits k s ->
    length b = k.
Proof.
  intros k.
  induction k as [|k'].
  {
    intros.
    inv H.
    reflexivity.
  }
  intros until s'.
  simpl.
  intros H.
  remember (get_bits k' s) as result.
  destruct result; [|inversion H].
  remember (s0 0%nat) as s0_0.
  destruct s0_0; try solve [inversion H].
  inv H.
  rewrite app_length.
  simpl. replace (length l + 1)%nat with (S (length l)) by omega.
  rewrite IHk' with (s':=s0) (s:=s); auto.
Qed.

Definition get_bytes (k: nat) (s: stream): result (list Z) :=
  match get_bits (8 * k)%nat s with
    | success bits s' => success (bitsToBytes bits) s'
    | error e s' => error e s'
  end
.

Lemma get_bytes_length:
  forall k s b s',
    success b s' = get_bytes k s ->
    length b = k.
Proof.
  unfold get_bytes; intros.
  remember (get_bits (8 * k)%nat s) as get_bits_s.
  destruct get_bits_s; inv H.
  assert (length l = (8 * k)%nat) by (apply get_bits_length with (s:=s) (s':=s0); auto).
  apply bitsToBytes_len_gen.
  omega.
Qed.

Lemma get_bytes_Zlength:
  forall k s b s',
    k >= 0 ->
    success b s' = get_bytes (Z.to_nat k) s ->
    Zlength b = k.
Proof.
  intros.
  rewrite Zlength_correct.
  erewrite get_bytes_length.
  rewrite Z2Nat.id; [reflexivity | omega].
  eauto.
Qed.

Lemma get_bytes_isbyteZ:
  forall k s b s',
    success b s' = get_bytes k s ->
    Forall isbyteZ b.
Proof.
  intros.
  unfold get_bytes in H.
  destruct (get_bits (8 * k) s); inv H.
  eapply bitsToBytes_isbyteZ; reflexivity.
Qed.

End OPTIONAL_ENTROPY.

Module ENTROPY := OPTIONAL_ENTROPY.

Definition get_entropy (security_strength min_length max_length: Z) (prediction_resistance: bool) s := ENTROPY.get_bytes (Z.to_nat min_length) s.


(* Constraint.In を \in 中置記法でかけるようにする。 *)
(* データ構造の部分のみ。 *)
(* 2019_03_31 @suharahiromichi *)
(* 変数をリテラルにする。 *)
(* 2019_04_07 @suharahiromichi *)

From mathcomp Require Import all_ssreflect.
Require Import Finite_sets_facts.

Module List.

  Lemma In_inb {A : eqType} (x : A) (s : seq A) : List.In x s <-> x \in s.
  Proof.
    elim: s.
    - done.
    - move=> a s IHs.
      split=> /=; rewrite inE.
      + case=> H.
        * by apply/orP/or_introl/eqP.
        * by apply/orP/or_intror/IHs.
      + move/orP; case.
        * move/eqP => ->. by left.
        * move=> H.
          move/IHs in H.
            by right.
  Qed.

  Lemma InP {A : eqType} (x : A) (s : seq A) : reflect (List.In x s) (x \in s).
  Proof.
    apply: (iffP idP) => H.
    - by apply/In_inb.
    - by apply/In_inb.
  Qed.

End List.

(* Test *)
Check [:: 1; 2] == [:: 1; 2].
Check 1 \in [:: 1].
Check List.InP.
Goal [:: 1; 2] == [:: 1; 2]. Proof. apply/eqP. done. Qed.
Goal 1 \in [:: 1; 2]. Proof. apply/List.InP. left. done. Qed.

Module Literal.

  Inductive Literal := a | b | c | f | g | h | x | y | z.

  Definition eqLiteral (x y : Literal) :=
    match (x, y) with
    | (a, a) => true
    | (b, b) => true
    | (c, c) => true
    | (f, f) => true
    | (g, g) => true
    | (h, h) => true
    | (x, x) => true
    | (y, y) => true
    | (z, z) => true
    | _ => false
    end.

  Lemma Literal_eqP (x y : Literal) : reflect (x = y) (eqLiteral x y).
  Proof.
    rewrite /eqLiteral.
    by apply: (iffP idP); case: x; case: y.
  Qed.

  Definition Literal_eqMixin := EqMixin Literal_eqP.
  Canonical Literal_eqType := EqType Literal Literal_eqMixin.

  Definition Literal_enum := [:: a; b; c; f; g; h; x; y; z].

  Definition Literal_pickle (x : Literal_eqType) : nat :=
    match x with
    | a => 0
    | b => 1
    | c => 2
    | f => 3
    | g => 4
    | h => 5
    | x => 6
    | y => 7
    | z => 8
    end.

  Definition Literal_unpickle (n : nat) : option Literal_eqType :=
    match n with
    | 0 => Some a
    | 1 => Some b
    | 2 => Some c
    | 3 => Some f
    | 4 => Some g
    | 5 => Some h
    | 6 => Some x
    | 7 => Some y
    | 8 => Some z
    | _ => None
    end.

  Lemma Literal_pcancel : pcancel Literal_pickle Literal_unpickle.
  Proof.
    by case.
  Qed.

  Lemma Literal_finiteP (x : Literal_eqType)  :
    (count_mem x) Literal_enum = 1.
  Proof.
    by case: x.
  Qed.

  Lemma Literal_uniq : uniq Literal_enum.
  Proof.
    done.
  Qed.
  (* UniqFinMixin に使えるはずだが。。。 *)

End Literal.

Definition Literal_eqMixin := EqMixin Literal.Literal_eqP.
Canonical Literal_eqType := EqType Literal.Literal Literal_eqMixin.
Canonical Literal_eqType' := [eqType of Literal.Literal].

Definition Literal_countMixin := CountMixin Literal.Literal_pcancel.
Definition Literal_choiceMixin := CountChoiceMixin Literal_countMixin.

Canonical Literal_choiceType := ChoiceType Literal.Literal Literal_choiceMixin.
Canonical Literal_countType := CountType Literal_choiceType Literal_countMixin.

Definition Literal_finMixin :=
  @FinMixin Literal_countType Literal.Literal_enum Literal.Literal_finiteP.
Canonical Literal_finType := FinType Literal.Literal Literal_finMixin.

(* ちゃんと定義できていことを確認する。 *)
Check Literal_eqType : eqType.
Check Literal_choiceType : choiceType.
Check Literal_countType : countType.
Check Literal_finType : finType.

Notation Literal := Literal.Literal.        (* XXXXX *)

Module Types.

  Reserved Notation "x '@' y" (at level 50, left associativity).
  Inductive Term : Set :=
  | Base
  | Var of Literal
  | Fun of Term & Term.
  Notation "x @ y" := (Fun x y).

  Notation varx := (Var (Literal.x)).
  Notation vary := (Var (Literal.y)).
  Notation varz := (Var (Literal.z)).

  (* ** *)
  (* == *)
  (* ** *)

  Fixpoint eqt (t1 t2 : Term) : bool :=
    match (t1, t2) with
    | (Base, Base) => true
    | (Var n1, Var n2) => n1 == n2
    | (t11 @ t21, t12 @ t22) => eqt t11 t12 && eqt t21 t22
    | (_ , _) => false
  end.

  Lemma Fun_eq (t11 t21 t12 t22 : Term) :
    eqt (t11 @ t21) (t12 @ t22) = eqt t11 t12 && eqt t21 t22.
  Proof.
    done.
  Qed.

  Lemma Term_eqP : forall (t1 t2 : Term), reflect (t1 = t2) (eqt t1 t2).
  Proof.
    move=> t1 t2.
    apply: (iffP idP).
    (* eqt t1 t2 -> t1 = t2 *)
    - elim: t1 t2.
      + by case.
      + by move=> n1; elim=> n2 // /eqP ->.
      + move=> t11 H1 t12 H2.
        elim=> // t21 _ t22 _ H.
        move: (H1 t21) => <-.
        move: (H2 t22) => <-.
        * done.
        * move: H.
          rewrite Fun_eq => /andP.
            by case.
        * move: H.
          rewrite Fun_eq => /andP.
            by case.
    (* t1 = t2 -> eqt t1 t2 *)
    - move=> ->.
      elim: t2 => //= u1 H1 v1 H2.
        by apply/andP.
  Qed.

  (* *** *)
  (* \in *)
  (* *** *)

  Inductive In x : Term -> Prop :=
  | In_Fun_dom : forall t1 t2, In x t1 -> In x (Fun t1 t2)
  | In_Fun_cod : forall t1 t2, In x t2 -> In x (Fun t1 t2)
  | In_Var : In x (Var x).

  Fixpoint inb (t : Term) (x : Literal) : bool :=
    match t with
    | Base => false
    | Var y => x == y
    | t1 @ t2 => [predU inb t1 & inb t2] x
    end.

  Lemma In_inb (x : Literal) (t : Term) : In x t <-> inb t x.
  Proof.
    split.
    - elim => //= t1 t2 Hp Hb; apply/orP.
      + by left.
      + by right.
    - elim: t => //=.
      + move=> y /eqP <-.
          by apply: In_Var.
      + move=> t1 Ht1 t2 Ht2 /orP.
        case=> H.
        * apply: In_Fun_dom.
            by apply: Ht1.
        * apply: In_Fun_cod.
            by apply: Ht2.
  Qed.

  Lemma InP (x : Literal) (t : Term) : reflect (In x t) (inb t x).
  Proof.
    apply: (iffP idP) => H.
    - by apply/In_inb.
    - by apply/In_inb.
  Qed.

End Types.

Notation Term := Types.Term.
Notation Var x := (Types.Var x).
Notation Base := (Types.Base).
Notation "x @ y" := (Types.Fun x y) (at level 50, left associativity).

Notation varx := (Types.Var (Literal.x)).
Notation vary := (Types.Var (Literal.y)).
Notation varz := (Types.Var (Literal.z)).

Definition Types_Term_Mixin := @EqMixin Types.Term Types.eqt Types.Term_eqP.
Canonical Types_Term_EqType := @EqType Types.Term Types_Term_Mixin.

Compute varx @ vary @ Base == varx @ vary @ Base.
Compute varx @ Base @ vary == varx @ vary @ Base.

Canonical Types_Term_predType := @mkPredType Literal Types.Term Types.inb.

Compute Types.inb (varx @ vary @ Base) Literal.x.
Compute Literal.x \in varx @ vary @ Base.
Compute Literal.y \notin varx @ vary @ Base.

(* 変数の個数 cardinal *)
(* Literal_finType を利用している。 *)
Check #| Types.inb (varx @ vary @ Base) |.
Compute #| Types.inb (varx @ vary @ Base) |.

Module Constraint.
  Definition Term := (Types_Term_EqType * Types_Term_EqType)%type.
  Definition Terms := (seq Term)%type.

  Fixpoint eqt (t1 t2 : Term) : bool :=
    (t1.1 == t2.1) && (t1.2 == t2.2).

  Lemma Term_eqP : forall (t1 t2 : Term), reflect (t1 = t2) (eqt t1 t2).
  Proof.
    move=> t1 t2.
    apply: (iffP idP).
    - case: t1 => t1_1 t2_1; case: t2 => t1_2 t2_2 /=.
      move/andP.
      case.
      move/eqP => <-.
      move/eqP => <-.
      done.
    - move=> <-.
      case: t1 => t1' t2' /=.
      apply/andP.
      done.
  Qed.

  Definition In (x : Literal) (constraints : Terms) : Prop :=
    List.Exists (fun c : Term =>
                   let: (t1, t2) := (c.1, c.2) in Types.In x t1 \/ Types.In x t2)
                constraints.

  Definition inb (s : Terms) (x : Literal) : bool :=
    has
      (fun c : Term =>
         let: (t1, t2) := (c.1, c.2) in Types.inb t1 x || Types.inb t2 x)
      s.

  Lemma In_inb (x : Literal) (s : Terms) : In x s <-> inb s x.
  Proof.
    split.
    - elim: s => /= [| a s IHs] H.
      + by inversion H.
      + inversion H; subst; clear H.
        * case: H1 => H.
          ** apply/orP/or_introl/orP/or_introl. (* left. left *)
               by apply/Types.In_inb.
          ** apply/orP/or_introl/orP/or_intror. (* left. right *)
               by apply/Types.In_inb.
        * apply/orP/or_intror.              (* right *)
            by apply: IHs.
    - elim: s => /= [| a s IHs] H.
      + done.
      + move/orP in H.
        case: H => H.
        * apply: List.Exists_cons_hd.
          move/orP in H.
          case: H => H.
          ** by apply/or_introl/Types.In_inb. (* left *)
          ** by apply/or_intror/Types.In_inb. (* right *)
        * apply: List.Exists_cons_tl.
            by move/IHs in H.
  Qed.

  Lemma InP (x : Literal) (s : Terms) : reflect (In x s) (inb s x).
  Proof.
    apply: (iffP idP) => H.
    - by apply/In_inb.
    - by apply/In_inb.
  Qed.

End Constraint.

Definition Constraint_Term_Mixin :=
  @EqMixin Constraint.Term Constraint.eqt Constraint.Term_eqP.
Canonical Constraint_Term_EqType :=
  @EqType Constraint.Term Constraint_Term_Mixin.

Check (varx @ vary, Base) : Constraint.Term.
Check (varx @ vary, Base) : Constraint_Term_EqType.

Compute (varx @ vary, Base) == (varx @ vary, Base).

Definition Constraint_Terms_EqType := (seq Constraint_Term_EqType)%type.

(* Canonical Constraint_Term_predType := mkPredType Constraint.inb. *)
(* 第二引数を省くと、うまくいかない。Constraint.Terms と解釈されるため。 *)
Canonical Constraint_Term_predType :=
  @mkPredType Literal Constraint_Terms_EqType Constraint.inb.
Set Printing All.
Print Constraint_Term_predType.
Unset Printing All.

Check [:: (varx, vary)] : Constraint.Terms.
Check [:: (varx, vary)] : seq Constraint_Term_EqType.
Check [:: (varx, vary)] : Constraint_Terms_EqType.

Definition sc := [:: (varx, vary)] : seq Constraint_Term_EqType.
Definition sc' := [:: (varx, vary)] : Constraint.Terms.
Definition sc'' := [:: (varx, vary)] : Constraint_Terms_EqType.

Compute sc == sc''.

Compute Constraint.inb sc Literal.x.
Compute Constraint.inb sc' Literal.x.
Compute Constraint.inb sc'' Literal.x.

Compute Literal.x \in sc''.                 (* ちょっと制限がある？ *)

(* 変数の個数 cardinal *)
(* Literal_finType を利用している。 *)
Check #| Constraint.inb sc |.
Compute #| Constraint.inb sc |.

Notation inb := (Constraint.inb).

Module Unify.

  Lemma Base0 : #|Base| = 0.
  Proof. by rewrite card0. Qed.

  Lemma Var1 (x : Literal) : #|Var x| = 1.
  Proof. by rewrite card1. Qed.

  Lemma Funs (x y : Term) : #|x @ y| <= #|x| + #|y|.
  Proof.
  Admitted.

  Lemma nil0 : #|inb [::]| = 0.
  Proof.
(*  rewrite /inb /Types.inb /=. *)
    by rewrite card0.
  Qed.

  Definition test_le constraints1 constraints2 :=
    #|inb constraints1| <= #|inb constraints2|.

End Unify.

(* END *)


From mathcomp Require Import all_ssreflect.

From AUChain Require Import
     GlobalState
     Messages
     MessageTuple
     Parameters.

From mathcomp Require Import finmap.

From RecordUpdate Require Import RecordSet.
Set Implicit Arguments.
Unset Strict Implicit.
Unset Printing Implicit Defensive.

(** * Network
      This file contains the basic operation for the synchronous network.
**)

(* Notation for record update *)

Notation "x [[ proj  :=  v ]]" := (set proj (fun _ => v) x)
                                    (at level 12, left associativity).

Notation "x [ proj  ::=  f ]" := (set proj f x)
                                     (at level 12, f at next level, left associativity).

(** Honest send of messages *)
Definition flood_msg (m : Message) (N : GlobalState) : GlobalState :=
  let: newMessages := [seq mkMessageTuple m r 1 | r <- exec_order N] in
  N[[msg_buff := newMessages ++ (msg_buff N)]][[history := m :: (history N)]].

(** Adversarial send of messages An adversary can for each of the
    messages choose to have them delivered in round t+1 or t+2.  *)
Definition Delay : Type := seq_sub [:: 1 ; 2].

Definition nat_of_delay (d : Delay) : nat.
  case d=> n; case: (n == 1)=> _; [exact 1| exact 2].
Defined.

Coercion nat_of_delay : Delay >-> nat.
Definition DelayMap := Party -> Delay.

Lemma delay12 (d : Delay): (d : nat) = 1 \/ (d : nat) = 2.
Proof. case d=> //= n. rewrite mem_seq2. case (_ == _) => //= _; by [constructor |constructor 2]. Qed.

Definition flood_msg_adv (m : Message) (ϕ : DelayMap) (N : GlobalState)  : GlobalState :=
  let: newMessages := [seq mkMessageTuple m r (ϕ(r)) | r <- exec_order N] in
  N[[msg_buff := newMessages ++ (msg_buff N)]][[history := m :: (history N)]].

(** Fetching of messages for a party *)
Definition collect_zero_pairs (r : Party) (N : GlobalState) : MessagePool :=
  [seq mt <- (msg_buff N) | (cd mt == 0) && (rcv mt == r)].

Definition rmv_zero_msgs (r : Party ) (N : GlobalState) : GlobalState :=
  N[[msg_buff := [seq mt <- (msg_buff N) | ~~ ((cd mt == 0) && (rcv mt == r))] ]].

Definition fetch_msgs (r : Party) (N : GlobalState): (seq Message * GlobalState) :=
  ([seq msg mt | mt <- collect_zero_pairs r N], rmv_zero_msgs r N).

(** Round tick *)
Definition round_tick (N : GlobalState) : GlobalState :=
  N[[msg_buff := [seq m[[cd := (cd m) - 1]] | m <- msg_buff N] ]].


(** (Syntactic) Typing for System F_mu_ref with tapes *)
From Autosubst Require Export Autosubst.
From stdpp Require Export stringmap fin_map_dom gmap.
From self.prob_lang Require Export lang notation.

Canonical Structure varO := leibnizO var.

Reserved Notation "⟦ τ ⟧" (at level 0, τ at level 70).
Reserved Notation "⟦ τ ⟧ₑ" (at level 0, τ at level 70).
Reserved Notation "⟦ Γ ⟧*" (at level 0, Γ at level 70).

(** * Types *)
Inductive type :=
  | TUnit : type
  | TInt : type
  | TBool : type
  | TProd : type → type → type
  | TSum : type → type → type
  | TArrow : type → type → type
  | TRec : {bind 1 of type} → type
  | TForall : {bind 1 of type} → type
  | TExists : {bind 1 of type} → type
  | TVar : var → type
  | TRef : type → type
  | TTape : type.

(** Which types are unboxed -- we can only do CAS on locations which hold unboxed types *)
Inductive UnboxedType : type → Prop :=
  | UnboxedTUnit : UnboxedType TUnit
  | UnboxedTInt : UnboxedType TInt
  | UnboxedTBool : UnboxedType TBool
  | UnboxedTRef τ : UnboxedType (TRef τ).

(** Types which support direct equality test (which coincides with ctx equiv).
    This is an auxiliary notion. *)
Inductive EqType : type → Prop :=
  | EqTUnit : EqType TUnit
  | EqTInt : EqType TInt
  | EqTBool : EqType TBool
  | EqTProd τ τ' : EqType τ → EqType τ' → EqType (TProd τ τ')
  | EqSum τ τ' : EqType τ → EqType τ' → EqType (TSum τ τ')
  (* | EQRef τ : EqType (Tref τ) *)
.

Lemma unboxed_type_ref_or_eqtype τ :
  UnboxedType τ → (EqType τ ∨ (∃ τ', τ = TRef τ') ∨ τ = TTape).
Proof. inversion 1; first [ left; by econstructor | right; naive_solver ]. Qed.

(** Autosubst instances *)
Global Instance Ids_type : Ids type. derive. Defined.
Global Instance Rename_type : Rename type. derive. Defined.
Global Instance Subst_type : Subst type. derive. Defined.
Global Instance SubstLemmas_typer : SubstLemmas type. derive. Qed.

Definition binop_int_res_type (op : bin_op) : option type :=
  match op with
  | MultOp => Some TInt | PlusOp => Some TInt | MinusOp => Some TInt
  | EqOp => Some TBool | LeOp => Some TBool | LtOp => Some TBool
  | _ => None
  end.
Definition binop_bool_res_type (op : bin_op) : option type :=
  match op with
  | XorOp => Some TBool | EqOp => Some TBool
  | _ => None
  end.
Definition unop_int_res_type (op : un_op) : option type :=
  match op with
  | MinusUnOp => Some TInt
  | _ => None
  end.
Definition unop_bool_res_type (op : un_op) : option type :=
  match op with
  | NegOp => Some TBool
  | _ => None
  end.

Declare Scope FType_scope.
Delimit Scope FType_scope with ty.
Bind Scope FType_scope with type.
Notation "()" := TUnit : FType_scope.
Notation "# x" := (TVar x) : FType_scope.
Infix "*" := TProd : FType_scope.
Notation "(*)" := TProd (only parsing) : FType_scope.
Infix "+" := TSum : FType_scope.
Notation "(+)" := TSum (only parsing) : FType_scope.
Infix "→" := TArrow : FType_scope.
Notation "(→)" := TArrow (only parsing) : FType_scope.

Notation "μ: τ" :=
  (TRec τ%ty)
  (at level 100, τ at level 200) : FType_scope.
Notation "∀: τ" :=
  (TForall τ%ty)
  (at level 100, τ at level 200) : FType_scope.
Notation "∃: τ" :=
  (TExists τ%ty)
  (at level 100, τ at level 200) : FType_scope.
Notation "'ref' τ" := (TRef τ%ty) (at level 10, τ at next level, right associativity): FType_scope.

(** * Typing judgements *)
Reserved Notation "Γ ⊢ₜ e : τ" (at level 74, e, τ at next level).
Reserved Notation "⊢ᵥ v : τ" (at level 20, v, τ at next level).

(* Shift all the indices in the context by one, *)
(*    used when inserting a new type interpretation in Δ. *)
Notation "⤉ Γ" := (Autosubst_Classes.subst (ren (+1)) <$> Γ) (at level 10, format "⤉ Γ").

(** We model type-level lambdas and applications as thunks *)
Notation "Λ: e" := (λ: <>, e)%E (at level 200, only parsing).
Notation "Λ: e" := (λ: <>, e)%V (at level 200, only parsing) : val_scope.
Notation "'TApp' e" := (App e%E #()%E) (at level 200, only parsing).

(* To unfold a recursive type, we need to take a step. We thus define the
unfold operator to be the identity function. *)
Definition rec_unfold : val := λ: "x", "x".
Definition unpack : val := λ: "x" "y", "y" "x".

Notation "'unpack:' x := e1 'in' e2" := (unpack e1%E (Lam x%binder e2%E))
  (at level 200, x at level 1, e1, e2 at level 200, only parsing) : expr_scope.

Notation "'unpack:' x := e1 'in' e2" := (unpack e1%E (LamV x%binder e2%E))
  (at level 200, x at level 1, e1, e2 at level 200, only parsing) : val_scope.

(** Operation lifts *)
Global Instance insert_binder (A : Type): Insert binder A (stringmap A) :=
  binder_insert.

(** Typing itself *)
Inductive typed : stringmap type → expr → type → Prop :=
  | Var_typed Γ x τ :
     Γ !! x = Some τ →
     Γ ⊢ₜ Var x : τ
  | Val_typed Γ v τ :
     ⊢ᵥ v : τ →
     Γ ⊢ₜ Val v : τ
  | BinOp_typed_int Γ op e1 e2 τ :
     Γ ⊢ₜ e1 : TInt → Γ ⊢ₜ e2 : TInt →
     binop_int_res_type op = Some τ →
     Γ ⊢ₜ BinOp op e1 e2 : τ
  | BinOp_typed_bool Γ op e1 e2 τ :
     Γ ⊢ₜ e1 : TBool → Γ ⊢ₜ e2 : TBool →
     binop_bool_res_type op = Some τ →
     Γ ⊢ₜ BinOp op e1 e2 : τ
  | UnOp_typed_int Γ op e τ :
     Γ ⊢ₜ e : TInt →
     unop_int_res_type op = Some τ →
     Γ ⊢ₜ UnOp op e : τ
  | UnOp_typed_bool Γ op e τ :
     Γ ⊢ₜ e : TBool →
     unop_bool_res_type op = Some τ →
     Γ ⊢ₜ UnOp op e : τ
  | UnboxedEq_typed Γ e1 e2 τ :
     UnboxedType τ →
     Γ ⊢ₜ e1 : τ → Γ ⊢ₜ e2 : τ →
     Γ ⊢ₜ BinOp EqOp e1 e2 : TBool
  | Pair_typed Γ e1 e2 τ1 τ2 :
     Γ ⊢ₜ e1 : τ1 → Γ ⊢ₜ e2 : τ2 →
     Γ ⊢ₜ Pair e1 e2 : τ1 * τ2
  | Fst_typed Γ e τ1 τ2 :
     Γ ⊢ₜ e : τ1 * τ2 →
     Γ ⊢ₜ Fst e : τ1
  | Snd_typed Γ e τ1 τ2 :
     Γ ⊢ₜ e : τ1 * τ2 →
     Γ ⊢ₜ Snd e : τ2
  | InjL_typed Γ e τ1 τ2 :
     Γ ⊢ₜ e : τ1 →
     Γ ⊢ₜ InjL e : τ1 + τ2
  | InjR_typed Γ e τ1 τ2 :
     Γ ⊢ₜ e : τ2 →
     Γ ⊢ₜ InjR e : τ1 + τ2
  | Case_typed Γ e0 e1 e2 τ1 τ2 τ3 :
     Γ ⊢ₜ e0 : τ1 + τ2 →
     Γ ⊢ₜ e1 : (τ1 → τ3) →
     Γ ⊢ₜ e2 : (τ2 → τ3) →
     Γ ⊢ₜ Case e0 e1 e2 : τ3
  | If_typed Γ e0 e1 e2 τ :
     Γ ⊢ₜ e0 : TBool →
     Γ ⊢ₜ e1 : τ →
     Γ ⊢ₜ e2 : τ →
     Γ ⊢ₜ If e0 e1 e2 : τ
  | Rec_typed Γ f x e τ1 τ2 :
     <[f:=(τ1 → τ2)%ty]>(<[x:=τ1]>Γ) ⊢ₜ e : τ2 →
     Γ ⊢ₜ (Rec f x e) : (τ1 → τ2)
  | App_typed Γ e1 e2 τ1 τ2 :
     Γ ⊢ₜ e1 : (τ1 → τ2) →
     Γ ⊢ₜ e2 : τ1 →
     Γ ⊢ₜ App e1 e2 : τ2
 | TLam_typed Γ e τ :
     ⤉ Γ ⊢ₜ e : τ →
     Γ ⊢ₜ (Λ: e) : ∀: τ
  | TApp_typed Γ e τ τ' :
     Γ ⊢ₜ e : (∀: τ) →
     Γ ⊢ₜ e #() : τ.[τ'/]
  | TFold Γ e τ :
     Γ ⊢ₜ e : τ.[(μ: τ)%ty/] →
     Γ ⊢ₜ e : (μ: τ)
  | TUnfold Γ e τ :
     Γ ⊢ₜ e : (μ: τ)%ty →
     Γ ⊢ₜ rec_unfold e : τ.[(μ: τ)%ty/]
  | TPack Γ e τ τ' :
     Γ ⊢ₜ e : τ.[τ'/] →
     Γ ⊢ₜ e : (∃: τ)
  | TUnpack Γ e1 x e2 τ τ2 :
     Γ ⊢ₜ e1 : (∃: τ) →
     <[x:=τ]>(⤉ Γ) ⊢ₜ e2 : (Autosubst_Classes.subst (ren (+1)) τ2) →
     Γ ⊢ₜ (unpack: x := e1 in e2) : τ2
  | TAlloc Γ e τ : Γ ⊢ₜ e : τ → Γ ⊢ₜ Alloc e : ref τ
  | TLoad Γ e τ : Γ ⊢ₜ e : ref τ → Γ ⊢ₜ Load e : τ
  | TStore Γ e e' τ : Γ ⊢ₜ e : ref τ → Γ ⊢ₜ e' : τ → Γ ⊢ₜ Store e e' : ()
  | TAllocTape Γ : Γ ⊢ₜ AllocTape : TTape
  | TFlip Γ e : Γ ⊢ₜ e : TTape -> Γ ⊢ₜ Flip e : TBool
  | TFlipU Γ e : Γ ⊢ₜ e : TUnit -> Γ ⊢ₜ Flip e : TBool
with val_typed : val → type → Prop :=
  | Unit_val_typed : ⊢ᵥ LitV LitUnit : TUnit
  | Int_val_typed (n : Z) : ⊢ᵥ LitV (LitInt n) : TInt
  | Bool_val_typed (b : bool) : ⊢ᵥ LitV (LitBool b) : TBool
  | Pair_val_typed v1 v2 τ1 τ2 :
     ⊢ᵥ v1 : τ1 →
     ⊢ᵥ v2 : τ2 →
     ⊢ᵥ PairV v1 v2 : (τ1 * τ2)
  | InjL_val_typed v τ1 τ2 :
     ⊢ᵥ v : τ1 →
     ⊢ᵥ InjLV v : (τ1 + τ2)
  | InjR_val_typed v τ1 τ2 :
     ⊢ᵥ v : τ2 →
     ⊢ᵥ InjRV v : (τ1 + τ2)
  | Rec_val_typed f x e τ1 τ2 :
     <[f:=TArrow τ1 τ2]>(<[x:=τ1]>∅) ⊢ₜ e : τ2 →
     ⊢ᵥ RecV f x e : (τ1 → τ2)
  | TLam_val_typed e τ :
     ∅ ⊢ₜ e : τ →
     ⊢ᵥ (Λ: e) : (∀: τ)
where "Γ ⊢ₜ e : τ" := (typed Γ e τ)
and "⊢ᵥ e : τ" := (val_typed e τ).

Lemma binop_int_typed_safe (op : bin_op) (n1 n2 : Z) τ :
  binop_int_res_type op = Some τ → is_Some (bin_op_eval op (LitV (LitInt n1)) (LitV (LitInt n2))).
Proof. destruct op; simpl; eauto. Qed.

Lemma binop_bool_typed_safe (op : bin_op) (b1 b2 : bool) τ :
  binop_bool_res_type op = Some τ → is_Some (bin_op_eval op (LitV (LitBool b1)) (LitV (LitBool b2))).
Proof. destruct op; naive_solver. Qed.

Lemma unop_int_typed_safe (op : un_op) (n : Z) τ :
  unop_int_res_type op = Some τ → is_Some (un_op_eval op (LitV (LitInt n))).
Proof. destruct op; simpl; eauto.  Qed.

Lemma unop_bool_typed_safe (op : un_op) (b : bool) τ :
  unop_bool_res_type op = Some τ → is_Some (un_op_eval op (LitV (LitBool b))).
Proof. destruct op; naive_solver. Qed.


Require Import VST.msl.seplog.
Require Import VST.msl.log_normalize.
Local Open Scope logic.

Lemma wand_frame_intro {A} {ND: NatDed A} {SL: SepLog A}: forall (P Q: A),
  Q |-- P -* P * Q.
Proof.
  intros.
  apply wand_sepcon_adjoint.
  rewrite sepcon_comm; auto.
Qed.

Lemma wand_frame_intro' {A} {ND: NatDed A} {SL: SepLog A}: forall (P Q R: A),
  (P * Q |-- R) ->
  Q |-- P -* R.
Proof.
  intros.
  apply wand_sepcon_adjoint.
  rewrite sepcon_comm; auto.
Qed.

Lemma wand_frame_elim {A} {ND: NatDed A} {SL: SepLog A}: forall (P Q: A),
  P * (P -* Q) |-- Q.
Proof.
  intros.
  rewrite sepcon_comm.
  apply wand_sepcon_adjoint; auto.
Qed.

Lemma wand_frame_elim' {A} {ND: NatDed A} {SL: SepLog A}: forall (P P' Q: A),
  (P |-- P') -> P * (P' -* Q) |-- Q.
Proof.
  intros.
  rewrite sepcon_comm.
  apply wand_sepcon_adjoint; auto.
  apply wand_derives; auto.
Qed.

Lemma wand_frame_ver {A} {ND: NatDed A} {SL: SepLog A}: forall (P Q R: A),
  (P -* Q) * (Q -* R) |-- P -* R.
Proof.
  intros.
  apply -> wand_sepcon_adjoint.
  rewrite (sepcon_comm _ P), <- sepcon_assoc.
  eapply derives_trans.
  + eapply sepcon_derives; [apply wand_frame_elim | apply derives_refl].
  + apply wand_frame_elim.
Qed.

Lemma wand_frame_hor {A} {ND: NatDed A} {SL: SepLog A}: forall (P1 P2 Q1 Q2: A),
  (P1 -* Q1) * (P2 -* Q2) |-- P1 * P2 -* Q1 * Q2.
Proof.
  intros.
  apply -> wand_sepcon_adjoint.
  rewrite <- (sepcon_assoc _ P1), (sepcon_comm _ P1), <- (sepcon_assoc P1), (sepcon_assoc _ _ P2), (sepcon_comm _ P2).
  apply sepcon_derives; apply wand_frame_elim.
Qed.

Lemma wand_frame_frame {A} {ND: NatDed A} {SL: SepLog A}: forall (P Q F: A),
  P -* Q |-- P * F -* Q * F.
Proof.
  intros.
  apply -> wand_sepcon_adjoint.
  rewrite <- sepcon_assoc, (sepcon_comm _ P).
  apply sepcon_derives; [apply wand_frame_elim | auto].
Qed.