From Coq Require Import List ZArith.
From mathcomp Require Import ssreflect ssrnat ssrbool eqtype seq tuple.
From BitBlasting Require Import State QFBV CNF.
From ssrlib Require Import Var SsrOrder ZAriths Tactics.

Set Implicit Arguments.
Unset Strict Implicit.
Import Prenex Implicits.

Definition cnf_lit_not a := [:: [:: neg_lit a]].
Definition cnf_lit_xor a1 a2 := [:: [:: a1; a2]; [:: neg_lit a1; neg_lit a2]].
Definition cnf_lit_eq a1 a2 := [:: [:: neg_lit a1; a2]; [:: a1; neg_lit a2]].

Lemma cnf_lit_not_negb E b l :
  enc_bit E l b -> (interp_cnf E (cnf_lit_not l)) = (~~ b).
Proof.
  rewrite /enc_bit /cnf_lit_not /=. rewrite interp_lit_neg_lit orbF andbT.
    by move/eqP => ->.
Qed.

Lemma cnf_lit_xor_neqb E b1 b2 l1 l2 :
  enc_bit E l1 b1 -> enc_bit E l2 b2 ->
  (interp_cnf E (cnf_lit_xor l1 l2)) = (b1 != b2).
Proof.
  rewrite /enc_bit /cnf_lit_xor /=. rewrite !interp_lit_neg_lit.
  move=> /eqP -> /eqP ->. by case: b1; case: b2.
Qed.

Lemma cnf_lit_eq_eqb E b1 b2 l1 l2 :
  enc_bit E l1 b1 -> enc_bit E l2 b2 ->
  (interp_cnf E (cnf_lit_eq l1 l2)) = (b1 == b2).
Proof.
  rewrite /enc_bit /cnf_lit_eq /=. rewrite !interp_lit_neg_lit.
  move=> /eqP -> /eqP ->. by case: b1; case: b2.
Qed.

(* ===== Word extension ===== *)

Definition extzip_ff := extzip lit_ff lit_ff.

Lemma enc_bits_unzip1_extzip E ls1 ls2 bs1 bs2 :
  enc_bit E lit_tt b1 -> enc_bits E ls1 bs1 -> enc_bits E ls2 bs2 ->
  enc_bits E (unzip1 (extzip_ff ls1 ls2)) (unzip1 (extzip0 bs1 bs2)).
Proof.
  rewrite /extzip_ff /extzip0 !unzip1_extzip => Henc_tt Henc1 Henc2.
  rewrite (enc_bits_size Henc1) (enc_bits_size Henc2).
  rewrite (enc_bits_cat Henc1); first reflexivity.
  rewrite enc_bits_copy; first reflexivity.
  rewrite enc_bit_not. exact: Henc_tt.
Qed.

Lemma enc_bits_unzip2_extzip E ls1 ls2 bs1 bs2 :
  enc_bit E lit_tt b1 -> enc_bits E ls1 bs1 -> enc_bits E ls2 bs2 ->
  enc_bits E (unzip2 (extzip_ff ls1 ls2)) (unzip2 (extzip0 bs1 bs2)).
Proof.
  rewrite /extzip_ff /extzip0 !unzip2_extzip => Henc_tt Henc1 Henc2.
  rewrite (enc_bits_size Henc1) (enc_bits_size Henc2).
  rewrite (enc_bits_cat Henc2); first reflexivity.
  rewrite enc_bits_copy; first reflexivity.
  rewrite enc_bit_not. exact: Henc_tt.
Qed.

(* Tactics for proving newer_than_? and env_preserve *)

Ltac t_newer_hook :=
  match goal with
  | H : is_true (newer_than_lit ?g ?l) |- is_true (newer_than_lit ?g ?l)
    => assumption
  | H : is_true (newer_than_lits ?g ?l) |- is_true (newer_than_lits ?g ?l)
    => assumption
  | |- is_true (newer_than_lit (?g + _) (Pos ?g)) => exact: newer_than_lit_add_diag_r
  | |- is_true (newer_than_lit (?g + _) (Neg ?g)) => exact: newer_than_lit_add_diag_r
  | |- is_true (?g <=? ?g)%positive => exact: Pos.leb_refl
  | |- is_true (?g <=? ?g + _)%positive => exact: pos_leb_add_diag_r
  | |- is_true (?g <? ?g + _)%positive => exact: pos_ltb_add_diag_r
  | H: is_true (?g1 <? ?g)%positive |- is_true (?g1 <=? ?g)%positive
    => exact: (pos_ltb_leb_incl H)
  | H : is_true (?g1 <=? ?g)%positive |- is_true (?g2 <=? ?g)%positive
    => apply: (pos_leb_trans _ H)
  | H : is_true (?g <=? ?g1)%positive |- is_true (?g <=? ?g2)%positive
    => apply: (pos_leb_trans H)
  | H : is_true (?g1 <=? ?g)%positive |- is_true (?g2 <? ?g)%positive
    => apply: (pos_ltb_leb_trans _ H)
  | H : is_true (newer_than_lit ?g1 ?l) |- is_true (newer_than_lit ?g2 ?l)
    => apply: (newer_than_lit_le_newer H)
  | H : is_true (newer_than_lits ?g1 ?l) |- is_true (newer_than_lits ?g2 ?l)
    => apply: (newer_than_lits_le_newer H)
  | H : is_true (newer_than_cnf ?g1 ?l) |- is_true (newer_than_cnf ?g2 ?l)
    => apply: (newer_than_cnf_le_newer H)
  | H : is_true (newer_than_lit _ lit_ff) |- _ => rewrite -newer_than_lit_tt_ff in H
  | |- is_true (newer_than_lit _ lit_ff) => rewrite -newer_than_lit_tt_ff
  | |- is_true (newer_than_lits _ (_ :: _)) =>
    rewrite newer_than_lits_cons; apply/andP; split
  | H : is_true (newer_than_lits _ (_ :: _)) |- _ =>
    let H1 := fresh in
    let H2 := fresh in
    rewrite newer_than_lits_cons in H; move/andP: H; intros [H1 H2]
  | |- is_true (newer_than_lits _ (_ ++ _)) =>
    rewrite newer_than_lits_cat; apply/andP; split
  | H : is_true (newer_than_lits _ (_ ++ _)) |- _ =>
    let H1 := fresh in
    let H2 := fresh in
    rewrite newer_than_lits_cat in H; move/andP: H; intros [H1 H2]
  | |- is_true (newer_than_lits _ (catrev _ _)) =>
    rewrite newer_than_lits_catrev; apply/andP; split
  | H : is_true (newer_than_lits _ (catrev _ _)) |- _ =>
    let H1 := fresh in
    let H2 := fresh in
    rewrite newer_than_lits_catrev in H; move/andP: H; intros [H1 H2]
  | |- is_true (newer_than_cnf _ (_ :: _)) =>
    rewrite newer_than_cnf_cons; apply/andP; split
  | H : is_true (newer_than_cnf _ (_ :: _)) |- _ =>
    let H1 := fresh in
    let H2 := fresh in
    rewrite newer_than_cnf_cons in H; move/andP: H; intros [H1 H2]
  | |- is_true (newer_than_cnf _ (_ ++ _)) =>
    rewrite newer_than_cnf_cat; apply/andP; split
  | H : is_true (newer_than_cnf _ (_ ++ _)) |- _ =>
    let H1 := fresh in
    let H2 := fresh in
    rewrite newer_than_cnf_cat in H; move/andP: H; intros [H1 H2]
  | |- is_true (newer_than_cnf _ (catrev _ _)) =>
    rewrite newer_than_cnf_catrev; apply/andP; split
  | H : is_true (newer_than_cnf _ (catrev _ _)) |- _ =>
    let H1 := fresh in
    let H2 := fresh in
    rewrite newer_than_cnf_catrev in H; move/andP: H; intros [H1 H2]
  | |- is_true (newer_than_lits _ (copy _ _)) => apply: newer_than_lits_copy
  | H : is_true (newer_than_lit _ (neg_lit _)) |- _ => rewrite newer_than_lit_neg in H
  | |- is_true (newer_than_lit _ (neg_lit _)) => rewrite newer_than_lit_neg
  | |- is_true (newer_than_lits _ (unzip1 (extzip_ff _ _))) =>
    rewrite /extzip_ff unzip1_extzip
  | |- is_true (newer_than_lits _ (unzip2 (extzip_ff _ _))) =>
    rewrite /extzip_ff unzip2_extzip
  end.

Ltac t_auto_newer := t_auto with ltac:(fun _ => t_newer_hook).

Ltac t_preserve_hook :=
  match goal with
  | |- env_preserve ?E (env_upd ?E ?g _) ?g => exact: env_upd_eq_preserve
  | H : env_preserve ?E _ ?g |- env_preserve ?E _ ?g => apply: (env_preserve_trans H)
  | H : env_preserve _ ?E ?g |- env_preserve _ ?E ?g => apply: (env_preserve_trans _ H)
  | H : env_preserve ?E1 ?E2 ?g1 |- env_preserve ?E1 ?E2 ?g2 =>
    apply: (env_preserve_le H)
  end.

Ltac t_auto_preserve := t_auto with ltac:(fun _ => t_preserve_hook || t_newer_hook).

(* ===== Variable generation ===== *)

Definition generator := positive.

Definition gen (g : generator) : generator * literal := (g + 1, Pos g)%positive.

(* ===== A map from variables to literal representation ===== *)

Definition vm := SSAVM.t word.

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

Definition newer_than_vm g (m : vm) :=
  forall v rs, SSAVM.find v m = Some rs -> newer_than_lits g rs.

Lemma newer_than_vm_add_r x m y : newer_than_vm x m -> newer_than_vm (x + y) m.
Proof.
  move=> Hnew v rs Hfind. move: (Hnew v rs Hfind). exact: newer_than_lits_add_r.
Qed.

Lemma newer_than_vm_le_newer x m y :
  newer_than_vm x m -> (x <=? y)%positive -> newer_than_vm y m.
Proof.
  move=> Hnew Hle v rs Hfind. move: (Hnew v rs Hfind) => {Hnew} Hnew.
  exact: (newer_than_lits_le_newer Hnew Hle).
Qed.

Lemma newer_than_vm_add m ig og v lrs :
  (ig <=? og)%positive ->
  newer_than_lits og lrs ->
  newer_than_vm ig m ->
  newer_than_vm og (SSAVM.add v lrs m).
Proof.
  move=> Hg Hnewls Hnewvm x rs. case Hxv: (x == v).
  - rewrite (SSAVM.Lemmas.find_add_eq Hxv). case=> ?; subst. assumption.
  - move/negP: Hxv=> Hxv. rewrite (SSAVM.Lemmas.find_add_neq Hxv) => Hfind.
    move: (Hnewvm x rs Hfind) => Hnew. exact: (newer_than_lits_le_newer Hnew Hg).
Qed.

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

Definition consistent1 m E s v : Prop :=
    match SSAVM.find v m with
    | None => True
    | Some rs => enc_bits E rs (SSAStore.acc v s)
    end.
Definition consistent m E s := forall v, consistent1 m E s v.

Lemma consistent_env_upd_newer m s E g g' b :
  newer_than_vm g m -> consistent m E s -> (g <=? g')%positive ->
  consistent m (env_upd E g' b) s.
Proof.
  move=> Hnew Hcon Hle. move: (newer_than_vm_le_newer Hnew Hle) => Hnew'. move=> v.
  move: (Hnew' v) (Hcon v) => {Hnew Hnew' Hcon}. rewrite /consistent1.
  case: (SSAVM.find v m); last by done. move=> rs Hnew Henc.
  move: (Hnew rs (Logic.eq_refl (Some rs))) => {Hnew} Hnew.
  rewrite (newer_than_lits_enc_bits_env_upd _ _ _ Hnew). exact: Henc.
Qed.

Lemma env_preserve_consistent m E E' g s :
  newer_than_vm g m -> env_preserve E E' g -> consistent m E s -> consistent m E' s.
Proof.
  move=> Hnew_gm Hpre Hcon. move=> x; rewrite /consistent1.
  case Hfind: (SSAVM.find x m); last by done. move: (Hnew_gm x _ Hfind) => Hnew_glsx.
  move: (Hcon x); rewrite /consistent1. rewrite Hfind.
  exact: (env_preserve_enc_bits Hpre Hnew_glsx).
Qed.

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

Definition vm_preserve (m m' : vm) : Prop :=
  forall v ls, SSAVM.find v m = Some ls -> SSAVM.find v m' = Some ls.

Lemma vm_preserve_consistent m m' s E :
  vm_preserve m m' -> consistent m' E s -> consistent m E s.
Proof.
  move=> Hpre Hcon v. rewrite /consistent1. case Hfind: (SSAVM.find v m) => //=.
  move: (Hpre _ _ Hfind) => Hfind'. move: (Hcon v). rewrite /consistent1.
  rewrite Hfind'. by apply.
Qed.

Lemma vm_preserve_refl m : vm_preserve m m.
Proof. done. Qed.

Lemma vm_preserve_trans m1 m2 m3 :
  vm_preserve m1 m2 -> vm_preserve m2 m3 -> vm_preserve m1 m3.
Proof. move=> H12 H23 v ls Hfind1. apply: H23. apply: H12. assumption. Qed.

Lemma vm_preserve_add_diag m v ls :
  SSAVM.find v m = None -> vm_preserve m (SSAVM.add v ls m).
Proof.
  move=> Hfind x xls Hfindx. case Hxv: (x == v).
  - rewrite (eqP Hxv) Hfind in Hfindx. discriminate.
  - move/negP: Hxv => Hxv. rewrite (SSAVM.Lemmas.find_add_neq Hxv). assumption.
Qed.

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

Ltac fresh_name t :=
  match t with
  | vm => fresh "m"
  | env => fresh "E"
  | generator => fresh "g"
  | cnf => fresh "cs"
  | _.-tuple literal => fresh "ls"
  | literal => fresh "l"
  | _ => fresh
  end.

Ltac dcase_hyps :=
  match goal with
  | H : context f [let '_ := ?r in _] |- _ =>
    move: H;
    let H1 := fresh in
    let H2 := fresh in
    lazymatch type of r with
    | prod (prod (prod (prod ?t1 ?t2) ?t3) ?t4) ?t5 =>
      let t1 := fresh_name t1 in
      let t2 := fresh_name t2 in
      let t3 := fresh_name t3 in
      let t4 := fresh_name t4 in
      let t5 := fresh_name t5 in
      dcase r; move=> [[[[t1 t2] t3] t4] t5] H1 H2; dcase_hyps
    | prod (prod (prod ?t1 ?t2) ?t3) ?t4 =>
      let t1 := fresh_name t1 in
      let t2 := fresh_name t2 in
      let t3 := fresh_name t3 in
      let t4 := fresh_name t4 in
      dcase r; move=> [[[t1 t2] t3] t4] H1 H2; dcase_hyps
    | prod (prod ?t1 ?t2) ?t3 =>
      let t1 := fresh_name t1 in
      let t2 := fresh_name t2 in
      let t3 := fresh_name t3 in
      dcase r; move=> [[t1 t2] t3] H1 H2; dcase_hyps
    | prod ?t1 ?t2 =>
      let t1 := fresh_name t1 in
      let t2 := fresh_name t2 in
      dcase r; move=> [t1 t2] H1 H2; dcase_hyps
    end
  | H : context f [if ?c then _ else _] |- _ =>
    move: H;
    let H1 := fresh in
    let H2 := fresh in
    dcase c=> H1 H2; dcase_hyps
  | H : (_, _, _, _, _) = (_, _, _, _, _) |- _ =>
    case: H;
    let H1 := fresh in
    let H2 := fresh in
    let H3 := fresh in
    let H4 := fresh in
    let H5 := fresh in
    move=> H1 H2 H3 H4 H5; dcase_hyps
  | H : (_, _, _, _) = (_, _, _, _) |- _ =>
    case: H;
    let H1 := fresh in
    let H2 := fresh in
    let H3 := fresh in
    let H4 := fresh in
    move=> H1 H2 H3 H4; dcase_hyps
  | H : (_, _, _) = (_, _, _) |- _ =>
    case: H;
    let H1 := fresh in
    let H2 := fresh in
    let H3 := fresh in
    move=> H1 H2 H3; dcase_hyps
  | H : (_, _) = (_, _) |- _ =>
    case: H;
    let H1 := fresh in
    let H2 := fresh in
    move=> H1 H2; dcase_hyps
  | H : context f [match ?t with | None => _ | Some _ => _ end] |- _ =>
    move: H;
    let H1 := fresh in
    let H2 := fresh in
    dcase t; case => H1 H2; dcase_hyps
  | H : context f [match ?t with | Some _ => _ | None => _ end] |- _ =>
    move: H;
    let H1 := fresh in
    let H2 := fresh in
    dcase t; case => H1 H2; dcase_hyps
  | |- _ => idtac
  end.

Ltac dcase_goal :=
  match goal with
  | |- context f [let '_ := ?r in _] =>
    let H := fresh in
    lazymatch type of r with
    | prod (prod (prod (prod ?t1 ?t2) ?t3) ?t4) ?t5 =>
      let t1 := fresh_name t1 in
      let t2 := fresh_name t2 in
      let t3 := fresh_name t3 in
      let t4 := fresh_name t4 in
      let t5 := fresh_name t5 in
      dcase r; move=> [[[[t1 t2] t3] t4] t5] H; dcase_goal
    | prod (prod (prod ?t1 ?t2) ?t3) ?t4 =>
      let t1 := fresh_name t1 in
      let t2 := fresh_name t2 in
      let t3 := fresh_name t3 in
      let t4 := fresh_name t4 in
      dcase r; move=> [[[t1 t2] t3] t4] H; dcase_goal
    | prod (prod ?t1 ?t2) ?t3 =>
      let t1 := fresh_name t1 in
      let t2 := fresh_name t2 in
      let t3 := fresh_name t3 in
      dcase r; move=> [[t1 t2] t3] H; dcase_goal
    | prod ?t1 ?t2 =>
      let t1 := fresh_name t1 in
      let t2 := fresh_name t2 in
      dcase r; move=> [t1 t2] H; dcase_goal
    end
  | |- (_, _, _, _, _) = (_, _, _, _, _) =>
    let H1 := fresh in
    let H2 := fresh in
    let H3 := fresh in
    let H4 := fresh in
    let H5 := fresh in
    case=> H1 H2 H3 H4 H5; dcase_goal
  | |- (_, _, _, _) = (_, _, _, _) =>
    let H1 := fresh in
    let H2 := fresh in
    let H3 := fresh in
    let H4 := fresh in
    case=> H1 H2 H3 H4; dcase_goal
  | |- (_, _, _) = (_, _, _) =>
    let H1 := fresh in
    let H2 := fresh in
    let H3 := fresh in
    case=> H1 H2 H3; dcase_goal
  | |- (_, _) = (_, _) =>
    let H1 := fresh in
    let H2 := fresh in
    case=> H1 H2; dcase_goal
  | |- _ => idtac
  end.

Ltac dite_hyps :=
  match goal with
  | H : context f [if ?c then _ else _] |- _ =>
    move: H;
    let H1 := fresh in
    let H2 := fresh in
    dcase c; case=> H1 H2; dite_hyps
  | |- _ => idtac
  end.

(* split_andb *)
Ltac split_andb_hyps :=
  repeat (match goal with
          | H : is_true (andb ?l ?r) |- _ => move/andP: H;
                                             let H1 := fresh in
                                             let H2 := fresh in
                                             move=> [H1 H2]
          end).

Ltac split_andb_goal :=
   repeat (match goal with
          | |- ?l /\ ?r => split
          | |- is_true (andb ?l ?r) => apply /andP
          end).


From discprob.basic Require Import base sval order monad bigop_ext nify.
From mathcomp Require Import ssrfun ssreflect eqtype ssrbool seq fintype choice.
From discprob.prob Require Import prob countable finite stochastic_order rearrange.
Require Import Reals Psatz.
From discprob.monad.idxval Require Import ival ival_dist.

Record ival_couplingP {A1 A2} (Is1: ival A1) (Is2: ival A2) (P: A1 → A2 → Prop) : Type :=
  mkICoupling { ic_witness :> ival {xy: A1 * A2 | P (fst xy) (snd xy)};
     ic_proj1: eq_ival Is1 (x ← ic_witness; mret (fst (sval x)));
     ic_proj2: eq_ival Is2 (x ← ic_witness; mret (snd (sval x)));
     }.

Record idist_couplingP {A1 A2} (Is1: ivdist A1) (Is2: ivdist A2) (P: A1 → A2 → Prop) : Type :=
  mkIdCoupling { idc_witness :> ivdist {xy: A1 * A2 | P (fst xy) (snd xy)};
     idc_proj1: eq_ivd Is1 (x ← idc_witness; mret (fst (sval x)));
     idc_proj2: eq_ivd Is2 (x ← idc_witness; mret (snd (sval x)));
     }.
From mathcomp Require Import bigop.

Lemma ic_coupling_to_id {A1 A2} (I1: ivdist A1) (I2: ivdist A2) P:
  ival_couplingP I1 I2 P →
  idist_couplingP I1 I2 P.
Proof.
  intros [Ic Hproj1 Hproj2].
  assert (Hsum1: is_series (countable_sum (ival.val Ic)) 1%R).
  {
    replace 1%R with R1 by auto.
    eapply (eq_ival_series _ _) in Hproj1; last eapply val_sum1.
    unshelve (eapply (countable_series_rearrange_covering_match _ _); last eassumption).
    { intros ic. rewrite //=. exists (existT (sval ic) tt).
      { abstract (rewrite Rmult_1_r => //=; destruct ic => //=). }
    }
    - apply val_nonneg.
    - apply val_nonneg.
    - intros (n1&?) (n2&?) Heq. apply sval_inj_pred => //=.
      inversion Heq. auto.
    - intros ((n1&[])&Hpf) => //=.
      unshelve (eexists).
      { exists n1.  rewrite //= Rmult_1_r // in Hpf. }
      apply sval_inj_pred => //=.
    - intros (n1&Hpf) => //=. nra.
  }
  exists {| ivd_ival := Ic; val_sum1 := Hsum1 |}; auto.
Qed.

Local Open Scope R_scope.
Record ival_coupling_nondepP {A1 A2} (Is1: ival A1) (Is2: ival A2) (P: A1 → A2 → Prop) : Type :=
  mkNonDepICoupling { ic_nondep_witness :> ival (A1 * A2);
               ic_nondep_support :
                 ∀ xy, (∃ i, ind ic_nondep_witness i = xy
                            ∧ val ic_nondep_witness i > 0) → P (fst xy) (snd xy);
     ic_nondep_proj1: eq_ival Is1 (x ← ic_nondep_witness; mret (fst x));
     ic_nondep_proj2: eq_ival Is2 (x ← ic_nondep_witness; mret (snd x));
     }.

Lemma ival_coupling_nondep_suffice {A1 A2} Is1 Is2 (P: A1 → A2 → Prop):
      ival_coupling_nondepP Is1 Is2 P →
      ival_couplingP Is1 Is2 P.
Proof.
  intros [Ic Isupp Hp1 Hp2].
  exists (ival_equip Ic _ Isupp).
  - setoid_rewrite Hp1.
    setoid_rewrite (ival_bind_P Ic).
    eapply ival_bind_congr; first reflexivity.
    intros (a1&a2) => //=; reflexivity.
  - setoid_rewrite Hp2.
    setoid_rewrite (ival_bind_P Ic).
    eapply ival_bind_congr; first reflexivity.
    intros (a1&a2) => //=; reflexivity.
Qed.

Lemma ival_coupling_refl {A} (I: ival A) : ival_couplingP I I (λ x y, x = y).
Proof.
  unshelve (eexists).
  { refine (x ← I; mret (exist _ (x, x) _)).
    done. }
  - setoid_rewrite ival_bind_mret_mret. setoid_rewrite ival_right_id. reflexivity.
  - setoid_rewrite ival_bind_mret_mret. setoid_rewrite ival_right_id. reflexivity.
Qed.

Lemma ival_coupling_sym {A1 A2} (Is1: ival A1) (Is2: ival A2) P
      (Ic: ival_couplingP Is1 Is2 P): ival_couplingP Is2 Is1 (λ x y, P y x).
Proof.
  destruct Ic as [Ic Hp1 Hp2].
  unshelve (eexists).
  { refine (mbind _ Ic). intros ((x&y)&HP).
    refine (mret (exist _ (y, x) _)); auto. }
  * setoid_rewrite Hp2. setoid_rewrite ival_assoc. apply ival_bind_congr; first by reflexivity.
    intros ((x&y)&Hpf). setoid_rewrite ival_left_id => //=. reflexivity.
  * setoid_rewrite Hp1. setoid_rewrite ival_assoc; apply ival_bind_congr; first by reflexivity.
    intros ((x&y)&Hpf). setoid_rewrite ival_left_id => //=. reflexivity.
Qed.

Lemma ival_coupling_eq {A} (Is1 Is2: ival A)
      (Ic: ival_couplingP Is1 Is2 (λ x y, x = y)): eq_ival Is1 Is2.
Proof.
  destruct Ic as [Ic Hp1 Hp2].
  setoid_rewrite Hp1. setoid_rewrite Hp2.
  apply ival_bind_congr; first by reflexivity.
  intros ((x&y)&Heq) => //=. rewrite //= in Heq; subst; reflexivity.
Qed.

Lemma ival_coupling_bind {A1 A2 B1 B2} (f1: A1 → ival B1) (f2: A2 → ival B2)
      Is1 Is2 P Q (Ic: ival_couplingP Is1 Is2 P):
  (∀ x y, P x y → ival_couplingP (f1 x) (f2 y) Q) →
  ival_couplingP (mbind f1 Is1) (mbind f2 Is2) Q.
Proof.
  intros Hfc.
  destruct Ic as [Ic Hp1 Hp2].
  unshelve (eexists).
  { refine (xy ← Ic; _). destruct xy as ((x&y)&HP).
    exact (Hfc _ _ HP).
  }
  - setoid_rewrite Hp1; setoid_rewrite ival_assoc;
      apply ival_bind_congr; first by reflexivity.
    intros ((x&y)&HP). setoid_rewrite ival_left_id => //=.
    destruct (Hfc x y); auto.
  - setoid_rewrite Hp2; setoid_rewrite ival_assoc;
      apply ival_bind_congr; first by reflexivity.
    intros ((x&y)&HP). setoid_rewrite ival_left_id => //=.
    destruct (Hfc x y); auto.
Qed.

Lemma ival_coupling_mret {A1 A2} x y (P: A1 → A2 → Prop):
  P x y →
  ival_couplingP (mret x) (mret y) P.
Proof.
  intros HP. exists (mret (exist _ (x, y) HP)); setoid_rewrite ival_left_id => //=; reflexivity.
Qed.

Lemma ival_coupling_conseq {A1 A2} (P1 P2: A1 → A2 → Prop) (I1: ival A1) (I2: ival A2):
  (∀ x y, P1 x y → P2 x y) →
  ival_couplingP I1 I2 P1 →
  ival_couplingP I1 I2 P2.
Proof.
  intros HP [Ic Hp1 Hp2].
  unshelve (eexists).
  { refine (x ← Ic; mret _).
    destruct x as (x&P).
    exists x. auto. }
  - setoid_rewrite Hp1.
    setoid_rewrite ival_assoc.
    eapply ival_bind_congr; first reflexivity; intros (?&?).
    setoid_rewrite ival_left_id; reflexivity.
  - setoid_rewrite Hp2.
    setoid_rewrite ival_assoc.
    eapply ival_bind_congr; first reflexivity; intros (?&?).
    setoid_rewrite ival_left_id; reflexivity.
Qed.

Lemma ival_coupling_plus' {A1 A2} p
      (P : A1 → A2 → Prop) (I1 I1': ival A1) (I2 I2': ival A2) :
  ival_couplingP I1 I2 P →
  ival_couplingP I1' I2' P →
  ival_couplingP (iplus (iscale p I1) (iscale (1 - p) I1'))
    (iplus (iscale p I2) (iscale (1 - p) I2')) P.
Proof.
  intros [Ic Hp1 Hp2] [Ic' Hp1' Hp2'].
  exists (ival.iplus (ival.iscale p Ic) (ival.iscale (1 - p) Ic')).
  - setoid_rewrite ival.ival_plus_bind.
    setoid_rewrite ival.ival_scale_bind.
    setoid_rewrite Hp1.
    setoid_rewrite Hp1'.
    reflexivity.
  - setoid_rewrite ival.ival_plus_bind.
    setoid_rewrite ival.ival_scale_bind.
    setoid_rewrite Hp2.
    setoid_rewrite Hp2'.
    reflexivity.
Qed.

Lemma ival_coupling_plus {A1 A2} p Hpf Hpf'
      (P : A1 → A2 → Prop) (I1 I1': ivdist A1) (I2 I2': ivdist A2) :
  ival_couplingP I1 I2 P →
  ival_couplingP I1' I2' P →
  ival_couplingP (ivdplus p Hpf I1 I1') (ivdplus p Hpf' I2 I2') P.
Proof.
  rewrite //=. apply ival_coupling_plus'.
Qed.

Lemma ival_coupling_idxOf {A1 A2} I1 I2 (P: A1 → A2 → Prop):
  ival_couplingP I1 I2 P →
  { Q :  {Q: idx I1 → idx I2 → Prop | (∀ i1 i1' i2, Q i1 i2 → Q i1' i2 → i1 = i1')} &
  ival_couplingP (idxOf I1) (idxOf I2) (λ i1 i2, P (ind I1 i1) (ind I2 i2) ∧
                                                 val I1 i1 > 0 ∧ val I2 i2 > 0 ∧
                                                 (sval Q) i1 i2  )}.
Proof.
  intros [Ic Hp1 Hp2].
  symmetry in Hp1. symmetry in Hp2.
  apply ClassicalEpsilon.constructive_indefinite_description in Hp1.
  destruct Hp1 as (h1&Hp1).
  apply ClassicalEpsilon.constructive_indefinite_description in Hp1.
  destruct Hp1 as (h1'&Hp1).
  apply ClassicalEpsilon.constructive_indefinite_description in Hp2.
  destruct Hp2 as (h2&Hp2).
  apply ClassicalEpsilon.constructive_indefinite_description in Hp2.
  destruct Hp2 as (h2'&Hp2).
  assert (Hgt_coerce: ∀ xy : support (val Ic),
             Rgt_dec (val Ic (projT1 (@existT _ (λ _, unit) (sval xy) tt)) * 1) 0).
  { abstract (intros (?&?); rewrite Rmult_1_r => //=; destruct Rgt_dec; auto). }
  unshelve (eexists).
  exists (λ i1 i2, ∃ ic1 ic2, sval (h1 ic1) = i1 ∧ sval (h2 ic2) = i2 ∧ sval ic1 = sval ic2).
  {
    abstract (intros i1 i1' i2 (ic1&ic2&Heq1&Heq2&Hcoh) (ic1'&ic2'&Heq1'&Heq2'&Hcoh');
    assert (ic2' = ic2);
    first ( abstract (destruct Hp2 as (Hbij2&?);
      rewrite -(Hbij2 ic2');
      rewrite -(Hbij2 ic2);
      f_equal; apply sval_inj_pred; congruence));
    abstract (subst; do 2 f_equal; apply sval_inj_pred; auto; congruence)). }
  unshelve (eexists).
  - refine (xy ← supp_idxOf Ic; mret _).
    unshelve (refine (exist _ ((sval (h1 _)), sval (h2 _)) _)).
      * exact (exist _ (existT (sval xy) tt) (Hgt_coerce xy)).
      * exact (exist _ (existT (sval xy) tt) (Hgt_coerce xy)).
      *
        abstract(
        split; [ by
        (abstract  (rewrite //=;
                   destruct Hp1 as (?&?&->&?) => //=;
                   destruct Hp2 as (?&?&->&?) => //=;
                   destruct (ind Ic (sval xy)) => //=)) |];
        split; [ by
        (abstract (rewrite //=; try destruct (h1 _); try destruct (h2 _); rewrite //=;
                   destruct Rgt_dec => //=)) |];
        split; [ by
        (abstract (rewrite //=; try destruct (h1 _); try destruct (h2 _); rewrite //=;
                   destruct Rgt_dec => //=)) |];
        abstract (
        rewrite //=;
        exists ((exist _ (existT (sval xy) tt) (Hgt_coerce xy)));
        exists ((exist _ (existT (sval xy) tt) (Hgt_coerce xy))); done)
        ).
  - setoid_rewrite ival_bind_mret_mret.
    symmetry.
    unshelve (eexists).
    {
      simpl. intros ((ic&?)&Hgt).
      simpl in h1.
      rewrite Rmult_1_r in Hgt.
      destruct ic as (ic&Hgtc).
      apply h1.
      exists (existT ic tt).
      simpl in *. rewrite Rmult_1_r. done.
    }
    unshelve (eexists).
    { simpl. intros ix.
      destruct (h1' ix) as ((i1&?)&Hgt).
      simpl in Hgt.
      unshelve (eexists).
      { unshelve (refine (existT _ tt)).
        exists i1. rewrite Rmult_1_r in Hgt. done.
      }
      rewrite //=.
    }
    rewrite //=.
    repeat split.
    *  intros (((a&Hgt)&[])&?). destruct Hp1 as (Hinv1&Hinv1'&?).
      rewrite //=. destruct (Rmult_1_r _). rewrite /eq_rect_r //=.
      rewrite Hinv1 => //=.
      apply sval_inj_pred => //=.
      f_equal.
      apply sval_inj_pred => //=.
    * rewrite //= => a. destruct Hp1 as (Hinv1&Hinv1'&?).
      specialize (Hinv1' a).
      destruct (h1' a) as ((?&[])&?).
      rewrite //=.
      rewrite -Hinv1'. destruct (Rmult_1_r (val Ic x)) => //=.
    * rewrite //=. intros (((a&Hgt)&[])&?). destruct Hp1 as (Hinv1&Hinv1'&?).
      rewrite //=. destruct (Rmult_1_r _). rewrite /eq_rect_r //=.
      f_equal. f_equal.
      apply sval_inj_pred => //=.
    * rewrite //=. intros (((a&Hgt)&[])&?). destruct Hp1 as (Hinv1&Hinv1'&?&Hval).
      rewrite //=. destruct (Rmult_1_r _). rewrite /eq_rect_r //=.
      rewrite Hval //= !Rmult_1_r //.
  -  setoid_rewrite ival_bind_mret_mret.
    symmetry.
    unshelve (eexists).
    {
      simpl. intros ((ic&?)&Hgt).
      simpl in h1.
      rewrite Rmult_1_r in Hgt.
      destruct ic as (ic&Hgtc).
      apply h2.
      exists (existT ic tt).
      simpl in *. rewrite Rmult_1_r. done.
    }
    unshelve (eexists).
    { simpl. intros ix.
      destruct (h2' ix) as ((i1&?)&Hgt).
      simpl in Hgt.
      unshelve (eexists).
      { unshelve (refine (existT _ tt)).
        exists i1. rewrite Rmult_1_r in Hgt. done.
      }
      rewrite //=.
    }
    rewrite //=.
    repeat split.
    * rewrite //=. intros (((a&Hgt)&[])&?). destruct Hp2 as (Hinv1&Hinv1'&?).
      rewrite //=. destruct (Rmult_1_r _). rewrite /eq_rect_r //=.
      rewrite Hinv1 => //=.
      apply sval_inj_pred => //=.
      f_equal.
      apply sval_inj_pred => //=.
    * rewrite //= => a. destruct Hp2 as (Hinv1&Hinv1'&?).
      specialize (Hinv1' a).
      destruct (h2' a) as ((?&[])&?).
      rewrite //=.
      rewrite -Hinv1'. destruct (Rmult_1_r (val Ic x)) => //=.
    * rewrite //=. intros (((a&Hgt)&[])&?). destruct Hp1 as (Hinv1&Hinv1'&?).
      rewrite //=. destruct (Rmult_1_r _). rewrite /eq_rect_r //=.
      f_equal. f_equal.
      apply sval_inj_pred => //=.
    * rewrite //=. intros (((a&Hgt)&[])&?). destruct Hp2 as (Hinv1&Hinv1'&?&Hval).
      rewrite //=. destruct (Rmult_1_r _). rewrite /eq_rect_r //=.
      rewrite Hval //= !Rmult_1_r //.
Qed.

Lemma ival_coupling_equip {X} I1 (P: X → Prop) Hpf:
  ival_couplingP I1 (ival_equip I1 P Hpf) (λ x y, x = (sval y)).
Proof.
  unshelve eexists.
  refine ((x ← ival_equip I1 P Hpf; mret (exist _ (sval x, x) _))); auto.
  - setoid_rewrite ival_bind_mret_mret.
    rewrite /sval.
    etransitivity; first (symmetry; apply ival_right_id).
    apply ival_bind_P.
  - setoid_rewrite ival_bind_mret_mret.
    rewrite /sval.
    etransitivity; first (symmetry; apply ival_right_id).
    eapply ival_bind_congr; first reflexivity.
    intros (?&?) => //=. reflexivity.
Qed.

Lemma ival_idist_wit_sum1 {X Y} (I1 : ivdist X) (I2: ivdist Y) (P: X → Y → Prop) Hwit:
  eq_ival I1 (Hwit ≫= (λ x : {xy : X * Y | P xy.1 xy.2}, mret (sval x).1)) →
  is_series (countable_sum (val Hwit)) 1.
Proof.
  intros Hproj1.
    eapply eq_ival_series in Hproj1; last apply (val_sum1).
    rewrite //= in Hproj1.
    setoid_rewrite Rmult_1_r in Hproj1.
    unshelve (eapply countable_series_rearrange_covering_match; last eapply Hproj1).
    { intros x. exists (existT (proj1_sig x) tt). exact (proj2_sig x). }
    { intros x; apply val_nonneg. }
    { intros x; apply val_nonneg. }
    { rewrite //=. intros (?&?) (?&?). inversion 1; subst. apply sval_inj_pi => //=. }
    { intros ((n&[])&Hpf). exists (exist _ n Hpf). apply sval_inj_pi => //=. }
    { intros (n&Hpf) => //=. }
Qed.

Lemma ival_idist_couplingP {X Y} (I1 : ivdist X) (I2: ivdist Y) (P: X → Y → Prop):
  ival_couplingP I1 I2 P →
  idist_couplingP I1 I2 P.
Proof.
  intros [Hwit Hproj1 Hproj2].
  unshelve (econstructor).
  { exists Hwit. eapply ival_idist_wit_sum1; eauto. }
  - rewrite /eq_ivd //=.
  - rewrite /eq_ivd //=.
Qed.

Lemma ival_coupling_support {X Y} I1 I2 (P: X → Y → Prop)
  (Ic: ival_couplingP I1 I2 P) :
  ival_couplingP I1 I2 (λ x y, ∃ Hpf: P x y,  In_isupport x I1 ∧ In_isupport y I2 ∧
                        In_isupport (exist _ (x, y) Hpf) Ic).
Proof.
  destruct Ic as [Ic Hp1 Hp2].
  cut (ival_couplingP I1 I2 (λ x y, ∃ (Hpf: P x y), In_isupport (exist _ (x, y) Hpf) Ic)).
  { intros. eapply ival_coupling_conseq; last eauto.
    intros x y (Hpf&Hin).
    destruct Hin as (ic&Hind&Hval).
    exists Hpf.
    repeat split; auto.
      - rewrite //=. symmetry in Hp1.
        destruct Hp1 as (h1&?&?&?&Hind1&Hval1).
        unshelve (eexists).
        { refine (sval (h1 _ )).
          exists (existT ic tt). rewrite //= Rmult_1_r. destruct Rgt_dec => //=.  }
        split; rewrite //=.
        * rewrite Hind1 //= Hind //=.
        * rewrite Hval1 //= Rmult_1_r //.
      - rewrite //=. symmetry in Hp2.
        destruct Hp2 as (h1&?&?&?&Hind1&Hval1).
        unshelve eexists.
        { refine (sval (h1 _)).
          exists (existT ic tt). rewrite //= Rmult_1_r. destruct Rgt_dec => //=.  }
        split; rewrite //=.
        * rewrite Hind1 //= Hind //=.
        * rewrite Hval1 //= Rmult_1_r //.
      - rewrite //=. exists ic; split => //=.
  }
  unshelve (eexists).
  refine (z ← ival_equip Ic (λ xy, In_isupport xy Ic) _; mret _).
  { destruct z as ((xy&HP)&Hin).  exists xy. abstract (exists HP; destruct xy; eauto). }
  { auto. }
  - setoid_rewrite Hp1.
    setoid_rewrite ival_bind_mret_mret.
    eapply ival_coupling_eq.
    eapply ival_coupling_bind; first eapply ival_coupling_equip.
    intros (xy&HP) ((xy'&HP')&Hin).
    rewrite //=.
    inversion 1; subst; auto. apply ival_coupling_refl.
  - setoid_rewrite Hp2.
    setoid_rewrite ival_bind_mret_mret.
    eapply ival_coupling_eq.
    eapply ival_coupling_bind; first eapply ival_coupling_equip.
    intros (xy&HP) ((xy'&HP')&Hin).
    rewrite //=.
    inversion 1; subst; auto. apply ival_coupling_refl.
Qed.

Lemma ival_coupling_support_wit {X Y} (I1: ivdist X) (I2: ivdist Y) (P: X → Y → Prop)
      (Ic: ival_couplingP I1 I2 P):
       { xy : X * Y | ∃ Hpf : P (fst xy) (snd xy),
           In_isupport (fst xy) I1 ∧ In_isupport (snd xy) I2 ∧ In_isupport (exist _ xy Hpf) Ic}.
Proof.
  specialize (ival_coupling_support I1 I2 P Ic) => Ic'.
  specialize (ival_idist_couplingP I1 I2 _ Ic') => Ic''.
  destruct Ic'' as [Ic'' ? ?].
  destruct (ivd_support_idx Ic'') as (i&?).
  destruct (ind Ic'' i) as ((x&y)&?); eauto.
Qed.

Lemma ival_coupling_proper {X Y} I1 I1' I2 I2' (P: X → Y → Prop) :
  eq_ival I1 I1' →
  eq_ival I2 I2' →
  ival_couplingP I1 I2 P →
  ival_couplingP I1' I2' P.
Proof.
  intros Heq1 Heq2 [Ic Hp1 Hp2].
  exists Ic; etransitivity; eauto; by symmetry.
Qed.


(* First, we need some imports from the Coq library *)
Set Warnings "-notation-overridden,-parsing".
From Coq Require Import Bool.Bool.
From Coq Require Import Init.Nat.
From Coq Require Import Arith.Arith.
From Coq Require Import Arith.EqNat.
From Coq Require Import omega.Omega.
From Coq Require Import Lists.List.
From Coq Require Import Strings.String.
From Coq Require Import Logic.FunctionalExtensionality.
Import ListNotations.

(* We don't use the Z scope here because we need many lemmas
   for mathematical things that are defined for natural numbers but not for
   integers. This means we will not be using the Num type. *)
(*----------------------Redefinition_Basic_Framework-------------------------*)
Declare Scope ns_scope.
Open Scope ns_scope.

Definition w : string := "w".
Definition x : string := "x".
Definition y : string := "y".
Definition z : string := "z".

Definition total_map (A : Type) := string -> A.
Definition t_empty {A : Type} (v : A) : total_map A :=
  (fun _ => v).
Definition eqb_string (x y : string) : bool :=
  if string_dec x y then true else false.
Definition t_update {A : Type} (m : total_map A)
                    (x : string) (v : A) :=
  fun x' => if eqb_string x x' then v else m x'.
Notation "x '!->' v ',' m" := (t_update m x v)
        (at level 100, v at next level, right associativity).

Definition State := total_map nat.

(* Syntax Aexp*)
Inductive Aexp : Type :=
  | ANum (n : nat)
  | AId (x : string)
  | APlus (a1 a2 : Aexp)
  | AMinus (a1 a2 : Aexp)
  | AMult (a1 a2 : Aexp).

(* Semantics Aexp *)
Fixpoint Aeval (st : State) (a : Aexp) : nat :=
  match a with
  | ANum n => n
  | AId x => st x
  | APlus a1 a2 => (Aeval st a1) + (Aeval st a2)
  | AMinus a1 a2  => (Aeval st a1) - (Aeval st a2)
  | AMult a1 a2 => (Aeval st a1) * (Aeval st a2)
  end.

(* Syntax Bexp *)
Inductive Bexp : Type :=
  | BTrue
  | BFalse
  | BEq (a1 a2 : Aexp)
  | BLe (a1 a2 : Aexp)
  | BNot (b : Bexp)
  | BAnd (b1 b2 : Bexp).

(* Semantics Bexp *)
Fixpoint Beval (st : State) (b : Bexp) : bool :=
  match b with
  | BTrue       => true
  | BFalse      => false
  | BEq a1 a2   => (Aeval st a1) =? (Aeval st a2)
  | BLe a1 a2   => (Aeval st a1) <=? (Aeval st a2)
  | BNot b1     => negb (Beval st b1)
  | BAnd b1 b2  => andb (Beval st b1) (Beval st b2)
  end.

Coercion AId : string >-> Aexp.
Coercion ANum : nat >-> Aexp.

Definition bool_to_bexp (b : bool) : Bexp :=
  if b then BTrue else BFalse.
Coercion bool_to_bexp : bool >-> Bexp.

Declare Scope imp_scope.
Bind Scope imp_scope with Aexp.
Bind Scope imp_scope with Bexp.
Delimit Scope imp_scope with imp.

Notation "x + y" := (APlus x y) (at level 50, left associativity) : imp_scope.
Notation "x - y" := (AMinus x y) (at level 50, left associativity) : imp_scope.
Notation "x * y" := (AMult x y) (at level 40, left associativity) : imp_scope.
Notation "x <= y" := (BLe x y) (at level 70, no associativity) : imp_scope.
Notation "x = y" := (BEq x y) (at level 70, no associativity) : imp_scope.
Notation "x && y" := (BAnd x y) (at level 40, left associativity) : imp_scope.
Notation "'~' b" := (BNot b) (at level 75, right associativity) : imp_scope.

Notation "'_' '!->' v" := (t_empty v)
  (at level 100, right associativity).

Definition empty_st := (_ !-> 0).

Notation "a '!->' x" := (t_update empty_st a x) (at level 100).

Notation " 'A' '[[' a ']]' st " := (Aeval st a)
        (at level 90, left associativity) : ns_scope.
Notation " 'B' '[[' b ']]' st " := (Beval st b)
        (at level 90, left associativity) : ns_scope.

Inductive Stm : Type :=
  | Sass (x : string) (a : Aexp)
  | Sskip
  | Scomp (s1 s2 : Stm)
  | Sif (b : Bexp) (s1 s2 : Stm)
  | Swhile (b : Bexp) (s : Stm).

Bind Scope ns_scope with Stm.
Notation "x '::=' a" :=
  (Sass x a) (at level 60) : ns_scope.
Notation "'SKIP'" :=
   Sskip : ns_scope.
Notation "s1 ; s2" :=
  (Scomp s1 s2) (at level 80, right associativity) : ns_scope.
Notation "'WHILE' b 'DO' s " :=
  (Swhile b s) (at level 80, right associativity) : ns_scope.
Notation "'IF_' b 'THEN' s1 'ELSE' s2 " :=
  (Sif b s1 s2) (at level 80, right associativity) : ns_scope.

Reserved Notation "'<<' S ',' st '>>' '-->' st'"
                  (at level 40).

Inductive Seval : Stm -> State -> State -> Prop :=
  | ass_ns  : forall st a1 n x,
      Aeval st a1 = n ->
      Seval (x ::= a1) st (t_update st x n)
  | skip_ns : forall st,
      Seval SKIP st st
  | comp_ns : forall s1 s2 st st' st'',
      Seval s1 st st' ->
      Seval s2 st' st'' ->
      Seval (s1;s2) st st''
  | if_tt_ns : forall st st' b s1 s2,
      Beval st b = true ->
      Seval s1 st st' ->
      Seval (IF_ b THEN s1 ELSE s2) st st'
  | if_ff_ns : forall st st' b s1 s2,
      Beval st b = false ->
      Seval s2 st st' ->
      Seval (IF_ b THEN s1 ELSE s2) st st'
  | while_tt_ns : forall st st' st'' b s,
      Beval st b = true ->
      Seval s st st' ->
      Seval (WHILE b DO s) st' st'' ->
      Seval (WHILE b DO s) st st''
  | while_ff_ns : forall b st s,
      Beval st b = false ->
      Seval (WHILE b DO s) st st.

Notation "<< S , st >> --> st'" := (Seval S st st')
                                   (at level 40).

(*-------------------------------Hoare_Logic---------------------------------*)
Declare Scope hoare_spec_scope.
Open Scope hoare_spec_scope.

(* We first define assertions, these are the core of Axiomatic Semantics. *)
Definition Assertion := State -> Prop.

Definition as2 : Assertion := fun st => st x <= st y.

(* P and Q are predicates and S is a statement as usual. We say that P is
   the precondition and Q is the postcondition. *)
Definition hoare_triple
  (P : Assertion) (S : Stm) (Q : Assertion) : Prop :=
  forall st st',
     << S , st >> --> st'  ->
     P st  ->
     Q st'.

Notation "{{ P }}  c  {{ Q }}" :=
  (hoare_triple P c Q) (at level 90, c at next level)
  : hoare_spec_scope.

(* This is the informal definition of assertions: If predicate P holds in the
   starting state and statement S terminates on that starting state, then
   predicate Q must hold in the resulting final state. *)

(* Assertions can be true or false *)
Theorem hoare_post_true : forall (P Q : Assertion) c,
  (forall st, Q st) ->
  {{P}} c {{Q}}.
Proof.
  intros P Q c H.
  unfold hoare_triple.
  intros st st' Heval HP.
  apply H.
Qed.

Theorem hoare_pre_false : forall (P Q : Assertion) c,
  (forall st, ~ (P st)) ->
  {{P}} c {{Q}}.
Proof.
  intros P Q c H.
  unfold hoare_triple.
  intros st st' Heval HP.
  unfold not in H.
  apply H in HP.
  inversion HP.
Qed.

(* We can define notation to use our normal logical operators with predicates*)
Definition assert_implies (P Q : Assertion) : Prop :=
  forall st, P st -> Q st.

Notation "P =>> Q" := (assert_implies P Q)
  (at level 80) : hoare_spec_scope.

Definition assn_sub x a P : Assertion :=
  fun (st : State) => P (x !-> Aeval st a , st).

Notation "P [ x |-> a ]" := (assn_sub x a P)
  (at level 10, x at next level).

Definition bassn b : Assertion :=
  fun st => (Beval st b = true).

Notation " 'B' '[[' b ']]'" := (bassn b)
        (at level 90, left associativity) : ns_scope.

Definition bassn2 b : Assertion :=
  fun st => ~(Beval st b = true).

Notation " '~B' '[[' b ']]'" := (bassn2 b)
        (at level 90, left associativity) : ns_scope.

Definition and_sub (P1 P2 : Assertion) : Assertion :=
  fun (st : State) => P1 st /\ P2 st.

Notation "P1 /\p P2" := (and_sub P1 P2)
  (at level 10).

Definition or_sub (P1 P2 : Assertion) : Assertion :=
  fun (st : State) => P1 st \/ P2 st.

Notation "P1 \/p P2" := (or_sub P1 P2)
  (at level 10).

(* We can also define the rules that belong to the notation. The thing about
   axiomatic semantics is that we can also proof all the rules are correct. *)
Theorem assp : forall Q x a,
  {{Q [x |-> a]}} x ::= a {{Q}}.
Proof.
  unfold hoare_triple.
  intros Q x a st st' HE HQ.
  inversion HE. subst.
  unfold assn_sub in HQ. assumption.  Qed.

Theorem consp_pre : forall (P P' Q : Assertion) S,
  {{P'}} S {{Q}} ->
  P =>> P' ->
  {{P}} S {{Q}}.
Proof.
  intros P P' Q c Hhoare Himp.
  intros st st' Hc HP. apply (Hhoare st st').
  assumption. apply Himp. assumption. Qed.

Theorem consp_post : forall (P Q Q' : Assertion) c,
  {{P}} c {{Q'}} ->
  Q' =>> Q ->
  {{P}} c {{Q}}.
Proof.
  intros P Q Q' c Hhoare Himp.
  intros st st' Hc HP.
  apply Himp.
  apply (Hhoare st st').
  assumption. assumption. Qed.

Theorem consp : forall (P P' Q Q' : Assertion) S,
  {{P'}} S {{Q'}} ->
  P =>> P' ->
  Q' =>> Q ->
  {{P}} S {{Q}}.
Proof.
  intros P P' Q Q' c Hht HPP' HQ'Q.
  apply consp_pre with (P' := P').
  apply consp_post with (Q' := Q').
  assumption. assumption. assumption.  Qed.

Theorem skipp : forall P,
     {{P}} SKIP {{P}}.
Proof.
  intros P st st' H HP. inversion H. subst.
  assumption.  Qed.

Theorem compp : forall P Q R S1 S2,
     {{Q}} S2 {{R}} ->
     {{P}} S1 {{Q}} ->
     {{P}} S1; S2 {{R}}.
Proof.
  intros P Q R c1 c2 H1 H2 st st' H12 Pre.
  inversion H12; subst.
  apply (H1 st'0 st'); try assumption.
  apply (H2 st st'0); assumption. Qed.

(* We can also define useful facts about bassn *)
Lemma Bexp_eval_true : forall b st,
  Beval st b = true -> (bassn b) st.
Proof.
  intros b st Hbe.
  unfold bassn. assumption.  Qed.

Lemma Bexp_eval_false : forall b st,
  Beval st b = false -> ~(bassn b) st.
Proof.
  intros b st Hbe.
  unfold bassn.
  rewrite Hbe.
  unfold not.
  intros.
  discriminate H.
Qed.

Theorem ifp : forall P Q b S1 S2,
  {{B[[b]] /\p P}} S1 {{Q}} ->
  {{~B[[b]] /\p P}} S2 {{Q}} ->
  {{P}} IF_ b THEN S1 ELSE S2 {{Q}}.
Proof.
  intros P Q b S1 S2 HTrue HFalse st st' HE HP.
  inversion HE; subst.
  - (* b is true *)
    apply (HTrue st st').
    assumption.
    split.
    apply Bexp_eval_true.
    assumption.
    assumption.
  - (* b is false *)
    apply (HFalse st st').
    assumption.
    split.
    apply Bexp_eval_false.
    assumption.
    assumption.
Qed.

Theorem whilep : forall P b c,
  {{B[[b]] /\p P}} c {{P}} ->
  {{P}} WHILE b DO c {{~B[[b]] /\p P}}.
Proof.
  intros P b c Hhoare st st' He HP.
  (* We need to use induction to talk about the complete loop *)
  remember (WHILE b DO c) as wcom eqn:Heqwcom.
  induction He;
  try (inversion Heqwcom);
  subst;
  clear Heqwcom.
  - (* while true *)
    apply IHHe2.
    reflexivity.
    apply (Hhoare st st').
    assumption.
    split.
    apply Bexp_eval_true.
    assumption.
    assumption.
  - (* while false *)
    split.
    apply Bexp_eval_false.
    assumption.
    assumption.
Qed.

(*----------------------------Annotated_Programs-----------------------------*)
(* The derivation trees for axiomatic semantics tend to get very wide. We can
   improve the readability of the proofs by making annotated programs. This is
   quite a complex definition, it is not needed to understand how this works.*)

Module Annotated.

(* syntax annotated programs of statements with embedded assertions *)
Inductive AStm : Type :=
  | Aass : string -> Aexp ->  Assertion -> AStm
  | Askip :  Assertion -> AStm
  | Acomp : AStm -> AStm -> AStm
  | Aif : Bexp ->  Assertion -> AStm ->  Assertion -> AStm
           -> Assertion-> AStm
  | Awhile : Bexp -> Assertion -> AStm -> Assertion -> AStm
  | Apre : Assertion -> AStm -> AStm
  | Apost : AStm -> Assertion -> AStm.

Inductive Annotated : Type :=
  | annotated : Assertion -> AStm -> Annotated.

(* Notations to make it easier to write these programs *)
Declare Scope annot.
Open Scope annot.
Notation "x '::=' a {{ P }}"
      := (Aass x a P)
      (at level 60, a at next level) : annot.
Notation "'SKIP' {{ P }}"
      := (Askip P)
      (at level 10) : annot.
Notation " S1 ; S2 "
      := (Acomp S1 S2)
      (at level 80, right associativity)  : annot.
Notation "'IF_' b 'THEN' {{ P }} d 'ELSE' {{ P' }} d' {{ Q }}"
      := (Aif b P d P' d' Q)
      (at level 80, right associativity)  : annot.
Notation "'WHILE' b 'DO' {{ Pbody }} d {{ Ppost }}"
      := (Awhile b Pbody d Ppost)
      (at level 80, right associativity) : annot.
Notation "'=>>' {{ P }} d"
      := (Apre P d)
      (at level 90, right associativity)  : annot.
Notation "d '=>>' {{ P }}"
      := (Apost d P)
      (at level 80, right associativity)  : annot.
Notation "{{ P }} d"
      := (annotated P d)
      (at level 90)  : annot.

(* We can erase all the annotations from an annotated program to get only the
   statements *)
Fixpoint extract (a : AStm) : Stm :=
  match a with
  | Aass x y _        => x ::= y
  | Askip _           => SKIP
  | Acomp a1 a2       => (extract a1 ; extract a2)
  | Aif b _ a1 _ a2 _ => IF_ b THEN extract a1 ELSE extract a2
  | Awhile b _ a _    => WHILE b DO extract a
  | Apre _ a          => extract a
  | Apost a _         => extract a
  end.

Definition extract_ann (ann : Annotated) : Stm :=
  match ann with
  | annotated P a => extract a
  end.

(* AStm S together with a precondition P determines a Hoare triple
   {{P}} (extract S) {{post S}}, where post is defined as follows: *)
Fixpoint post (a : AStm) : Assertion :=
  match a with
  | Aass x y Q             => Q
  | Askip P                => P
  | Acomp a1 a2            => post a2
  | Aif  _ _ a1 _ a2 Q     => Q
  | Awhile b Pbody c Ppost => Ppost
  | Apre _ a               => post a
  | Apost c Q              => Q
  end.

(* Extracting precondition and the postcondition *)
Definition pre_ann (ann : Annotated) : Assertion :=
  match ann with
  | annotated P a => P
  end.

Definition post_ann (ann : Annotated) : Assertion :=
  match ann with
  | annotated P a => post a
  end.

(* An annotated program is correct if the following holds *)
Definition ann_correct (ann : Annotated) :=
  {{pre_ann ann}} (extract_ann ann) {{post_ann ann}}.

(* The function verification_conditions takes a AStm a together with a
   precondition P and returns a proposition that, if it can be proved,
   implies that the triple {{P}} (extract S) {{post S}} is valid. *)
Fixpoint verification_cond (P : Assertion) (a : AStm) : Prop :=
  match a with
  | Aass x y Q =>
      (P =>> Q [x |-> y])
  | Askip Q =>
      (P =>> Q)
  | Acomp a1 a2 =>
      verification_cond P a1
      /\ verification_cond (post a1) a2
  | Aif b P1 a1 P2 a2 Q =>
      ((fun st => bassn b st /\ P st) =>> P1)
      /\ ((fun st => (bassn2 b st) /\ P st) =>> P2)
      /\ (post a1 =>> Q) /\ (post a2 =>> Q)
      /\ verification_cond P1 a1
      /\ verification_cond P2 a2
  | Awhile b Pbody a Ppost =>
      (P =>> post a)
      /\ ((fun st => bassn b st /\ post a st ) =>> Pbody)
      /\ ((fun st => (bassn2 b st) /\ post a st) =>> Ppost)
      /\ verification_cond Pbody a
  | Apre P' a =>
      (P =>> P') /\ verification_cond P' a
  | Apost a Q =>
      verification_cond P a /\ (post a =>> Q)
  end.

Lemma Bev : forall (b:Bexp) (st:State),
  (~B[[b]] st) = (~ (B[[b]]) st).
Proof.
  intros.
  unfold bassn2.
  unfold bassn.
  reflexivity.
Qed.

(* This theorem proves that verification_conditions works correctly *)
Theorem verification_correct : forall a P,
  verification_cond P a -> {{P}} (extract a) {{post a}}.
Proof.
  induction a; intros P H; simpl in *.
  - (* ass *)
    eapply consp_pre.
    apply assp.
    assumption.
  - (* skip *)
    eapply consp_pre.
    apply skipp.
    assumption.
  - (* comp *)
    destruct H as [H1 H2].
    eapply compp.
    apply IHa2.
    apply H2.
    apply IHa1.
    apply H1.
  - (* if *)
    destruct H as [HPre1 HPre2].
    destruct HPre2 as [HPre2 Hd1].
    destruct Hd1 as [Hd1 Hd2].
    destruct Hd2 as [Hd2 HThen].
    destruct HThen as [HThen HElse].
    apply IHa1 in HThen. clear IHa1.
    apply IHa2 in HElse. clear IHa2.
    apply ifp.
      + eapply consp_post with (Q':=post a2); eauto.
         eapply consp_pre; eauto.
      + eapply consp_post with (Q':=post a4); eauto.
         eapply consp_pre; eauto.
  - (* while *)
    destruct H as [HPre Hbody1].
    destruct Hbody1 as [Hbody1 Hpost1].
    destruct Hpost1 as [Hpost1 Hd].
    eapply consp_pre; eauto.
    eapply consp_post; eauto.
    apply whilep.
    eapply consp_pre; eauto.
  - (* pre *)
    destruct H as [HP Hd].
    eapply consp_pre.
    apply IHa.
    apply Hd.
    assumption.
  - (* post *)
    destruct H as [Hd HQ].
    eapply consp_post.
    apply IHa.
    apply Hd.
    assumption.
Qed.

(* Now we can verify an entire annotated program *)
Definition verification_cond_ann (ann : Annotated) : Prop :=
  match ann with
  | annotated P a => verification_cond P a
  end.

Lemma verification_correct_ann : forall ann,
  verification_cond_ann ann -> ann_correct ann.
Proof.
  intros [P a].
  apply verification_correct.
Qed.

(* verification_conditions generates many often trivial subgoals. We can make
   a tactic that can solve these subgoals on itself by using lemmas. We first
   need to define these lemmas. Some of these lemmas have already been used
   in other Coq files but we will redefine them for clarity here. Other lemmas
   used are in the Coq standard library. *)
Theorem eqb_string_true_iff : forall x y : string,
    eqb_string x y = true <-> x = y.
Proof.
   intros x y.
   unfold eqb_string.
   destruct (string_dec x y) as [|Hs].
   - subst.
     split.
     reflexivity.
     reflexivity.
   - split.
     + intros contra.
       discriminate contra.
     + intros H.
       rewrite H in Hs.
       destruct Hs.
       reflexivity.
Qed.

Theorem eqb_string_false_iff : forall x y : string,
    eqb_string x y = false <-> x <> y.
Proof.
  intros x y.
  rewrite <- eqb_string_true_iff.
  rewrite not_true_iff_false.
  reflexivity.
Qed.

Lemma t_update_eq : forall (m : total_map nat) x v,
    (x !-> v , m) x = v.
Proof.
  unfold t_update.
  intros.
  destruct (eqb_string x0 x0) eqn:E.
  - reflexivity.
  - apply eqb_string_false_iff in E.
    unfold not in E.
    destruct E.
    reflexivity.
Qed.

Theorem t_update_neq : forall (m : total_map nat) x1 x2 v,
    x1 <> x2 ->
    (x1 !-> v , m) x2 = m x2.
Proof.
  unfold t_update.
  intros.
  destruct (eqb_string x1 x2 ) eqn:E.
  - unfold not in H.
    rewrite eqb_string_true_iff in E.
    destruct H.
    apply E.
  - reflexivity.
Qed.

(* This is the tactic, taken from Software Foundations chapter Hoare2 *)
Tactic Notation "verify" :=
  apply verification_correct;
  repeat split;
  simpl; unfold assert_implies;
  unfold bassn in *;unfold bassn2 in *;
  unfold Beval in *; unfold Aeval in *;
  unfold assn_sub; intros;
  repeat rewrite t_update_eq;
  repeat (rewrite t_update_neq; [| (intro X; inversion X)]);
  simpl in *;
  repeat match goal with [H : _ /\ _ |- _] => destruct H end;
  repeat rewrite not_true_iff_false in *;
  repeat rewrite not_false_iff_true in *;
  repeat rewrite negb_true_iff in *;
  repeat rewrite negb_false_iff in *;
  repeat rewrite eqb_eq in *;
  repeat rewrite eqb_neq in *;
  repeat rewrite leb_iff in *;
  repeat rewrite leb_iff_conv in *;
  try subst;
  repeat
    match goal with
      [st : State |- _] =>
        match goal with
          [H : st _ = _ |- _] => rewrite -> H in *; clear H
        | [H : _ = st _ |- _] => rewrite <- H in *; clear H
        end
    end;
  try eauto; try omega.

End Annotated.

(*--------------------------------Soundness----------------------------------*)
(* Valid Hoare triple: For all states s and s', if P s = tt and <S, s> -> s'
     then Q s' = tt.
   Provable Hoare triple: There exists an inference tree with {P} S {Q} as
     conclusion.
   Soundness: The inference system of the axiomatic semantics is sound,
     that is for every partial correctness formula {P} S {Q} we have
     (provable) {P} S {Q} implies (valid) {P} S {Q}.
   We already defined when a Hoare triple is valid in the previous chapter,
   we can already define what it means for a Hoare triple to be provable. *)

(* Provability *)
Inductive hoare_proof : Assertion -> Stm -> Assertion -> Type :=
  | ass_p : forall Q x a,
      hoare_proof (assn_sub x a Q) (x ::= a) Q
  | skip_p : forall P,
      hoare_proof P (SKIP) P
  | comp_p  : forall P S1 Q S2 R,
      hoare_proof P S1 Q ->
      hoare_proof Q S2 R ->
      hoare_proof P (S1;S2) R
  | if_p : forall P Q b S1 S2,
    hoare_proof (fun st => P st /\ bassn b st) S1 Q ->
    hoare_proof (fun st => P st /\ ~(bassn b st)) S2 Q ->
    hoare_proof P (IF_ b THEN S1 ELSE S2 ) Q
  | while_p : forall P b S,
    hoare_proof (fun st => P st /\ bassn b st) S P ->
    hoare_proof P (WHILE b DO S ) (fun st => P st /\ ~ (bassn b st))
  | cons_p  : forall (P Q P' Q' : Assertion) S,
    hoare_proof P' S Q' ->
    (forall s, P s -> P' s) ->
    (forall s, Q' s -> Q s) ->
    hoare_proof P S Q.

(* Now we can prove that axiomatic semantics is sound *)
Theorem hoare_proof_sound : forall P S Q,
  hoare_proof P S Q -> {{P}} S {{Q}}.
Proof.
  intros.
  induction X.
  - unfold hoare_triple.
    intros.
    inversion H; subst.
    apply H0.
  - unfold hoare_triple.
    intros.
    inversion H; subst.
    apply H0.
  - unfold hoare_triple.
    intros.
    inversion H; subst.
    apply (IHX2 st'0 st').
    apply H6.
    apply (IHX1 st st'0).
    apply H3.
    apply H0.
  - unfold hoare_triple.
    intros.
    inversion H; subst.
    + apply (IHX1 st st').
      apply H7.
      split.
      apply H0.
      apply H6.
    + apply (IHX2 st st').
      apply H7.
      split.
      apply H0.
      apply Bexp_eval_false.
      apply H6.
  - unfold hoare_triple.
    intros.
    remember (WHILE b DO S ) as wStm eqn:HeqwStm.
    induction H.
    + inversion HeqwStm.
    + inversion HeqwStm.
    + inversion HeqwStm.
    + inversion HeqwStm.
    + inversion HeqwStm.
    + inversion HeqwStm; subst.
      apply IHSeval2.
      reflexivity.
      apply (IHX st st').
      apply H1.
      split.
      apply H0.
      apply Bexp_eval_true.
      apply H.
    + inversion HeqwStm; subst.
      split.
      apply H0.
      apply Bexp_eval_false.
      apply H.
  - unfold hoare_triple.
    intros.
    apply q.
    apply (IHX st st').
    apply H.
    apply p.
    apply H0.
Qed.

(*--------------------------------Completeness-------------------------------*)
(* Completeness: The inference system of the axiomatic semantics is complete,
   that is for every partial correctness formula {P} S {Q} we have
     (valid) {P} S {Q} implies (provable) {P} S {Q}.
   We can prove this using the weakest liberal precondition.
   wlp(S, Q) s = tt if and only if for all states s', if <S, s> -> s' then
     Q s' = tt. *)

Definition wlp (S : Stm) (Q : Assertion) : Assertion :=
  fun s => forall s', << S , s >> --> s' -> Q s'.

(* wlp has two important properties,
   property 1: (valid) {wlp(S, Q)} S {Q}
   property 2: If (valid p) {P} S {Q} then P => wlp(S, Q). *)
Lemma wlp_property1:
  forall S Q,
  {{wlp S Q}} S {{Q}}.
Proof.
  intros.
  unfold wlp.
  unfold hoare_triple.
  intros.
  apply H0.
  apply H.
Qed.

Lemma wlp_property2:
  forall P S Q,
      {{P}} S {{Q}}
    ->
      forall s, P s -> wlp S Q s.
Proof.
  intros.
  unfold wlp.
  unfold hoare_triple in H.
  intros.
  apply (H s).
  - apply H1.
  - apply H0.
Qed.

(* Now we can prove that Axiomatic Semantics is complete *)
Theorem hoare_proof_complete: forall P S Q,
  {{P}} S {{Q}} -> hoare_proof P S Q.
Proof.
  intros P S.
  generalize dependent P.
  induction S; intros P Q HT.
  - eapply cons_p.
    + eapply ass_p.
    + intro s.
      apply HT.
      econstructor.
      reflexivity.
    + intros.
      apply H.
  - eapply cons_p.
    + eapply skip_p.
    + intros.
      eapply H.
    + intro st.
      apply HT.
      apply skip_ns.
  - apply comp_p with (wlp S2 Q).
    + eapply IHS1.
      unfold wlp.
      intros s s' E1 H.
      intros s'' E2.
      eapply HT.
      * econstructor.
        eapply E1.
        eapply E2.
      * apply H.
    + eapply IHS2.
      intros s s' E1 H.
      apply H.
      apply E1.
  - eapply if_p.
     + apply IHS1.
       intros s s' H [H1 H2].
       eapply HT.
       * apply if_tt_ns.
         apply H2.
         apply H.
       * apply H1.
     + apply IHS2.
       intros s s' H [H1 H2].
       eapply HT.
       * apply if_ff_ns.
         unfold bassn in H2.
         apply not_true_is_false in H2.
         apply H2.
         apply H.
       * apply H1.
  - eapply cons_p with (P' := wlp (WHILE b DO S ) Q).
    + apply while_p.
      apply IHS.
      intros s s' H [H2 H3] s'' H4.
      assert(<<(WHILE b DO S ) , s >> --> s'').
      * eapply while_tt_ns.
        apply H3.
        apply H.
        apply H4.
      * eapply wlp_property1.
        { apply H0. }
        { apply H2. }
    + apply wlp_property2.
      apply HT.
    + intros s [H1 H2].
      eapply wlp_property1.
      * assert(<<(WHILE b DO S ),  s>> -->  s).
        { apply while_ff_ns.
           unfold bassn in H2.
           apply not_true_is_false in H2.
           apply H2.
        }
        { apply H. }
      * apply H1.
Qed.


Require Import Platform.AutoSep Platform.Arrays8.

Definition copyS := SPEC("dst", "dstPos", "src", "srcPos", "len") reserving 1
  Al bsD, Al bsS, Al lenD : W, Al lenS : W,
  PRE[V] array8 bsD (V "dst") * array8 bsS (V "src")
    * [| wordToNat (V "dstPos") + wordToNat (V "len") <= length bsD |]%nat
    * [| wordToNat (V "srcPos") + wordToNat (V "len") <= length bsS |]%nat
    * [| length bsD = wordToNat lenD |] * [| length bsS = wordToNat lenS |]
  POST[_] Ex bsD', array8 bsD' (V "dst") * array8 bsS (V "src")
    * [| length bsD' = length bsD |].

Definition equalS := SPEC("arr1", "pos1", "arr2", "pos2", "len") reserving 2
  Al bs1, Al bs2, Al len1 : W, Al len2 : W,
  PRE[V] array8 bs1 (V "arr1") * array8 bs2 (V "arr2")
    * [| wordToNat (V "pos1") + wordToNat (V "len") <= length bs1 |]%nat
    * [| wordToNat (V "pos2") + wordToNat (V "len") <= length bs2 |]%nat
    * [| length bs1 = wordToNat len1 |] * [| length bs2 = wordToNat len2 |]
  POST[_] array8 bs1 (V "arr1") * array8 bs2 (V "arr2").

Definition m := bmodule "array8" {{
  bfunction "copy"("dst", "dstPos", "src", "srcPos", "len", "tmp") [copyS]
    [Al bsD, Al bsS, Al lenD : W, Al lenS : W,
      PRE[V] array8 bsD (V "dst") * array8 bsS (V "src")
        * [| wordToNat (V "dstPos") + wordToNat (V "len") <= length bsD |]%nat
        * [| wordToNat (V "srcPos") + wordToNat (V "len") <= length bsS |]%nat
        * [| length bsD = wordToNat lenD |] * [| length bsS = wordToNat lenS |]
      POST[_] Ex bsD', array8 bsD' (V "dst") * array8 bsS (V "src")
        * [| length bsD' = length bsD |] ]
    While ("len" > 0) {
      Assert [Al bsD, Al bsS, Al lenD : W, Al lenS : W,
        PRE[V] [| V "len" > natToW 0 |]%word * [| V "srcPos" < natToW (length bsS) |]%word
          * [| V "dstPos" < natToW (length bsD) |]%word
          * array8 bsD (V "dst") * array8 bsS (V "src")
          * [| wordToNat (V "dstPos") + wordToNat (V "len") <= length bsD |]%nat
          * [| wordToNat (V "srcPos") + wordToNat (V "len") <= length bsS |]%nat
          * [| length bsD = wordToNat lenD |] * [| length bsS = wordToNat lenS |]
        POST[_] Ex bsD', array8 bsD' (V "dst") * array8 bsS (V "src")
          * [| length bsD' = length bsD |] ];;

      "tmp" <-*8 "src" + "srcPos";;
      "dst" + "dstPos" *<-8 "tmp";;
      "dstPos" <- "dstPos" + 1;;
      "srcPos" <- "srcPos" + 1;;
      "len" <- "len" - 1
    };;

    Return 0
  end with bfunction "equal"("arr1", "pos1", "arr2", "pos2", "len", "tmp1", "tmp2") [equalS]
    [Al bs1, Al bs2, Al len1 : W, Al len2 : W,
      PRE[V] array8 bs1 (V "arr1") * array8 bs2 (V "arr2")
        * [| wordToNat (V "pos1") + wordToNat (V "len") <= length bs1 |]%nat
        * [| wordToNat (V "pos2") + wordToNat (V "len") <= length bs2 |]%nat
        * [| length bs1 = wordToNat len1 |] * [| length bs2 = wordToNat len2 |]
      POST[_] array8 bs1 (V "arr1") * array8 bs2 (V "arr2") ]
    While ("len" > 0) {
      Assert [Al bs1, Al bs2, Al len1 : W, Al len2 : W,
        PRE[V] [| V "len" > natToW 0 |]%word * [| V "pos1" < natToW (length bs1) |]%word
          * [| V "pos2" < natToW (length bs2) |]%word
          * array8 bs1 (V "arr1") * array8 bs2 (V "arr2")
          * [| wordToNat (V "pos1") + wordToNat (V "len") <= length bs1 |]%nat
          * [| wordToNat (V "pos2") + wordToNat (V "len") <= length bs2 |]%nat
          * [| length bs1 = wordToNat len1 |] * [| length bs2 = wordToNat len2 |]
        POST[_] array8 bs1 (V "arr1") * array8 bs2 (V "arr2") ];;

      "tmp1" <-*8 "arr1" + "pos1";;
      "tmp2" <-*8 "arr2" + "pos2";;
      If ("tmp1" = "tmp2") {
        "pos1" <- "pos1" + 1;;
        "pos2" <- "pos2" + 1;;
        "len" <- "len" - 1
      } else {
        Return 0
      }
    };;

    Return 1
  end
}}.

Lemma rt1 : wordToNat (natToW 1) = 1.
  reflexivity.
Qed.

Hint Rewrite rt1 : N.

Lemma updown : forall (x : W) y,
  natToW 0 < x
  -> y + 1 + wordToNat (x ^- natToW 1) = y + wordToNat x.
  intros; rewrite wordToNat_wminus; nomega.
Qed.

Hint Rewrite updown using assumption : sepFormula.

Lemma natToW_wordToNat : forall w : W, natToW (wordToNat w) = w.
  apply natToWord_wordToNat.
Qed.

Hint Rewrite natToW_wordToNat : N.

Lemma goodBound : forall (x y u : W) n,
  (wordToNat x + wordToNat y <= n)%nat
  -> natToW 0 < y
  -> n = wordToNat u
  -> x < natToW n.
  intros; subst; nomega.
Qed.

Local Hint Immediate goodBound.

Theorem ok : moduleOk m.
  vcgen; abstract (sep_auto; (descend; eauto)).
Qed.


From Test Require Import tactic.

Section FOFProblem.

Variable Universe : Set.
Variable UniverseElement : Universe.

Variable wd_ : Universe -> Universe -> Prop.
Variable col_ : Universe -> Universe -> Universe -> Prop.

Variable col_swap1_1 : (forall A B C : Universe, (col_ A B C -> col_ B A C)).
Variable col_swap2_2 : (forall A B C : Universe, (col_ A B C -> col_ B C A)).
Variable col_triv_3 : (forall A B : Universe, col_ A B B).
Variable wd_swap_4 : (forall A B : Universe, (wd_ A B -> wd_ B A)).
Variable col_trans_5 : (forall P Q A B C : Universe, ((wd_ P Q /\ (col_ P Q A /\ (col_ P Q B /\ col_ P Q C))) -> col_ A B C)).

Theorem pipo_6 : (forall A B C D M : Universe, ((wd_ A C /\ (wd_ B D /\ (wd_ M D /\ (wd_ M A /\ (wd_ A B /\ (wd_ B C /\ (wd_ M C /\ (wd_ B M /\ (col_ A M C /\ (col_ B M D /\ (col_ D A C /\ (col_ M A D /\ col_ M A B)))))))))))) -> col_ A B C)).
Proof.
  time tac.
Qed.

End FOFProblem.


From Test Require Import tactic.

Section FOFProblem.

Variable Universe : Set.
Variable UniverseElement : Universe.

Variable wd_ : Universe -> Universe -> Prop.
Variable col_ : Universe -> Universe -> Universe -> Prop.

Variable col_swap1_1 : (forall A B C : Universe, (col_ A B C -> col_ B A C)).
Variable col_swap2_2 : (forall A B C : Universe, (col_ A B C -> col_ B C A)).
Variable col_triv_3 : (forall A B : Universe, col_ A B B).
Variable wd_swap_4 : (forall A B : Universe, (wd_ A B -> wd_ B A)).
Variable col_trans_5 : (forall P Q A B C : Universe, ((wd_ P Q /\ (col_ P Q A /\ (col_ P Q B /\ col_ P Q C))) -> col_ A B C)).

Theorem pipo_6 : (forall A B C Aprime Bprime Cprime O : Universe, ((wd_ B O /\ (wd_ C O /\ (wd_ Bprime O /\ (wd_ Cprime O /\ (wd_ A O /\ (wd_ Aprime O /\ (wd_ A B /\ (wd_ A Aprime /\ (col_ O A B /\ (col_ O B C /\ (col_ O Aprime Bprime /\ (col_ O Bprime Cprime /\ col_ O A C)))))))))))) -> col_ O Aprime Cprime)).
Proof.
  time tac.
Qed.

End FOFProblem.


From Test Require Import tactic.

Section FOFProblem.

Variable Universe : Set.
Variable UniverseElement : Universe.

Variable wd_ : Universe -> Universe -> Prop.
Variable col_ : Universe -> Universe -> Universe -> Prop.

Variable col_swap1_1 : (forall A B C : Universe, (col_ A B C -> col_ B A C)).
Variable col_swap2_2 : (forall A B C : Universe, (col_ A B C -> col_ B C A)).
Variable col_triv_3 : (forall A B : Universe, col_ A B B).
Variable wd_swap_4 : (forall A B : Universe, (wd_ A B -> wd_ B A)).
Variable col_trans_5 : (forall P Q A B C : Universe, ((wd_ P Q /\ (col_ P Q A /\ (col_ P Q B /\ col_ P Q C))) -> col_ A B C)).

Theorem pipo_6 : (forall A B C D P : Universe, ((wd_ A D /\ (wd_ A B /\ (wd_ B C /\ (wd_ A C /\ (wd_ A D /\ (wd_ C D /\ (wd_ P D /\ (wd_ P A /\ (wd_ B D /\ (wd_ B P /\ (wd_ P C /\ (col_ C P D /\ (col_ A D D /\ col_ B P C))))))))))))) -> col_ B C D)).
Proof.
  time tac.
Qed.

End FOFProblem.


Set Warnings "-notation-overridden,-parsing".

From mathcomp Require Import ssreflect ssrfun ssrbool eqtype ssrnat div ssralg.
From Tweetnacl.High Require Import mc.
From Tweetnacl.High Require Import mcgroup.
From Tweetnacl.High Require Import ladder.
From Tweetnacl.High Require Import montgomery.

Import GRing.Theory.

Local Open Scope ring_scope.

Set Implicit Arguments.
Unset Strict Implicit.

Section OptimizedLadder.
  Variable K : finECUFieldType.
  Variable M : mcuType K.

  Definition cswap {A : Type} (c : nat) (a b : A) :=
    if c == 1%N then (b, a) else (a, b).

  Fixpoint opt_montgomery_rec (n m : nat) (x a b c d : K) : K :=
    if m is m.+1 then
      let (a, b) := cswap (bitn n m) a b in
      let (c, d) := cswap (bitn n m) c d in
      let e := a + c in
      let a := a - c in
      let c := b + d in
      let b := b - d in
      let d := e^+2 in
      let f := a^+2 in
      let a := c * a in
      let c := b * e in
      let e := a + c in
      let a := a - c in
      let b := a^+2 in
      let c := d - f in
      let a := c * ((M#a - 2%:R) / 4%:R) in
      let a := a + d in
      let c := c * a in
      let a := d * f in
      let d := b * x in
      let b := e^+2 in
      let (a, b) := cswap (bitn n m) a b in
      let (c, d) := cswap (bitn n m) c d in
      opt_montgomery_rec n m x a b c d
    else
      a / c.

  Definition opt_montgomery (n m : nat) (x : K) : K :=
    opt_montgomery_rec n m x 1 x 0 1.

  Hypothesis mcu_no_square : forall x : K, x^+2 != (M#a)^+2 - 4%:R.

  Local Notation "p '#x0'" := (point_x0 p) (at level 30).
  Local Notation "p '#x'" := (point_x p) (at level 30).

  (* In montgomery.v we avoided relying on the fact that in ssreflect,
   * a / 0 = 0. However, the implementation in TweetNaCl considers the inverse
   * of 0 to be 0. Hence, we replicate this behaviour.
   *)
  Lemma point_x0_point_x_simple a b (p : point K) :
    inf_div a b = p#x -> a / b = p#x0.
  Proof.
    move=> H.
    case: (boolP (b == 0)) => [/eqP|] Hb.
    + rewrite {}Hb /= in H *.
      rewrite invr0 mulr0.
      case: p => [|p_x p_y] in H *; first by [].
      have : 0 != 0 :> K by apply: inf_div_K_Fin; exact: H.
      by rewrite eq_refl.
    + apply/eqP; rewrite eq_sym point_x0_point_x //; first by apply/eqP.
      case: p => [|p_x p_y] in H *; last by [].
      have : b == 0 by rewrite -(inf_div_K_Inf a); apply/eqP.
      by rewrite /= -(negbTE Hb).
  Qed.

  Local Notation "\- x"   := (@MCGroup.neg _ x).
  Local Notation "x \+ y" := (@MCGroup.add _ M x y).
  Local Notation "x \- y" := (x \+ (\- y)).

  Local Lemma opt_montgomery_rec_small (n m k : nat) (x a b c d : K) : k >= m ->
    opt_montgomery_rec (n + 2^k) m x a b c d = opt_montgomery_rec n m x a b c d.
  Proof.
    move=> Hm.
    elim: m => [|m IHm /=] in Hm a b c d *; first by [].
    rewrite bitn_large //.
    by case/orP: (bitnV n m) => /eqP-> /=; rewrite IHm 1?ltnW.
  Qed.

  Local Lemma opt_montgomery_rec_ok (n m : nat) (p q : mc M) (a b c d : K) :
    n < 2^m ->
    p#x0 != 0 ->
    hom_ok a c -> hom_ok b d ->
    q#x = inf_div a c -> (p + q)#x = inf_div b d ->
    opt_montgomery_rec n m (p#x0) a b c d = (p *+ n + q *+ 2^m)#x0.
  Proof.
    move=> Hn Hp Hac_ok Hbd_ok /eqP Hac /eqP Hbd.
    move/andP: {Hac_ok Hac} (conj Hac_ok Hac) => Hac.
    move/andP: {Hbd_ok Hbd} (conj Hbd_ok Hbd) => Hbd.
    rewrite -montgomery_rec_ok //.
    elim: m => [|m IHm] in q a b c d n Hn Hac Hbd *.
      by rewrite /=; apply: point_x0_point_x_simple; move/andP: Hac => [_ /eqP].
    move/andP: Hac => [Hac_ok /eqP Hac].
    move/andP: Hbd => [Hbd_ok /eqP Hbd].
    have H1 : p#x = inf_div (p#x0) 1.
      rewrite /inf_div oner_eq0 divr1 point_x_K_Fin_point_x0 //.
      by apply: point_x0_neq0_fin.
    pose E := (oncurve_mc, oner_neq0, andbT).
    case: (boolP (n < 2^m)) => [{Hn} Hn|].
      rewrite /= bitn_small0 //= (IHm (q*+2)) //.
      + by rewrite [q + _]addrC -addrA -mulr2n.
      + rewrite [X in _ * (X + _)]mulrC [_ + (a + c)^+2]addrC mulr2n.
        apply/andP; split.
        - by apply: montgomery_hom_eq_ok.
        - by apply/eqP; apply: montgomery_hom_eq; rewrite ?E //.
      + rewrite -[_^+2]mul1r [_ * (a - c)]mulrC [_ * (a + c)]mulrC.
        rewrite [_ * p#x0]mulrC mulr2n addrCA.
        apply: montgomery_hom_neq; rewrite ?E //.
        by rewrite -[q \- _]/(point_of_mc (q-_)) addrC opprD addrNK point_x_neg.
    rewrite -leqNgt => Hngeq.
    rewrite -[n](@subnK (2^m)) //.
    have H3 : n - 2 ^ m < 2 ^ m.
      by rewrite -(@ltn_add2r (2^m)) subnK // addnn -mul2n -expnS.
    rewrite /= bitn_small1 //=.
    rewrite opt_montgomery_rec_small ?montgomery_rec_small ?leqnn //.
    rewrite [(p + q)*+2]mulr2n -addrA [q + (p + q)]addrC.
    rewrite (IHm (p + q + q)) //.
    + rewrite -[_^+2]mul1r [_ * (b - d)]mulrC [_ * (b + d)]mulrC.
      rewrite [_ * p#x0]mulrC.
      apply: montgomery_hom_neq; rewrite ?E //.
      by rewrite -[_ \- q]/(point_of_mc (_-q)) addrK.
    + rewrite [X in _*(X + _)]mulrC [_ + (b+d)^+2]addrC [p+q]addrC -addrA addrA.
      apply/andP; split.
      - by apply: montgomery_hom_eq_ok.
      - by apply/eqP; apply: montgomery_hom_eq; rewrite ?E //.
  Qed.

  Lemma opt_montgomery_x (n m : nat) (x : K) :
    n < 2^m -> x != 0 ->
    forall (p : mc M), p#x0 = x -> opt_montgomery n m x = (p *+ n)#x0.
  Proof.
    move=> Hn x_neq0 p p_x_eqx.
    rewrite /opt_montgomery -{1}p_x_eqx -[p*+n]addr0 -[in _+0](@mul0rn _ (2^m)).
    apply: opt_montgomery_rec_ok; first by [].
    + by rewrite p_x_eqx.
    + by rewrite /hom_ok oner_eq0.
    + by rewrite /hom_ok oner_eq0 orbT.
    + by apply/eqP; rewrite eq_sym inf_div_K_Inf eq_refl.
    + apply/eqP; rewrite addr0 -point_x0_point_x.
      - by apply/eqP; rewrite divr1.
      - by exact: oner_neq0.
      - by apply: point_x0_neq0_fin; rewrite p_x_eqx.
  Qed.

(*   Lemma neg_x : forall (p : mc M),
    p#x0 = (-p)#x0.
  Proof.
    move=> [[|xp yp] Hp] //=.
  Qed.

  Lemma neg_x_n : forall (n:nat) (p : mc M),
    (p *+ n)#x0 = ((-p) *+ n)#x0.
  Proof.
    move => n p.
    rewrite GRing.mulNrn.
    apply neg_x.
  Qed.
 *)
Local Ltac ring_simplify_this :=
  repeat match goal with
  | _ => rewrite expr2
  | _ => rewrite exprS
  | _ => rewrite GRing.mul1r
  | _ => rewrite GRing.mulr1
  | _ => rewrite GRing.mul0r
  | _ => rewrite GRing.mulr0
  | _ => rewrite GRing.add0r
  | _ => rewrite GRing.oppr0
  | _ => rewrite GRing.addr0
  | _ => done
end.

Lemma opt_montgomery_rec_0 : forall (m n: nat) a b d,
  opt_montgomery_rec n m 0 a b 0 d = 0.
Proof.
  elim => [| m IHm] n a b d /=.
  + rewrite GRing.invr0 mulr0 //.
  + have/orP := bitnV n m.
    move => []/eqP => ->.
    all: rewrite /cswap => /=.
    all: ring_simplify_this.
    all: rewrite addrN.
    all: ring_simplify_this.
Qed.

Lemma opt_montgomery_0 (n m: nat):
  opt_montgomery n m 0 = 0.
Proof.
  rewrite /opt_montgomery.
  apply opt_montgomery_rec_0.
Qed.

Local Lemma p_x0_0_impl : forall (p: mc M),
  p #x0 = 0%:R ->
  p = MC (oncurve_0_0 M) \/ p = MC (oncurve_inf M).
Proof.
  move => [] [Hp|xp yp Hp] => /= Hxp.
  + by right ; apply mc_eq.
  + left; apply mc_eq.
  move: Hxp Hp -> => /=.
  rewrite ?exprS ?expr0 ?mulr1 ?mulr0 ?addr0.
  have/eqP := M.(pB) => Hb.
  rewrite mulf_eq0.
  move /orP => [].
  by move/eqP.
  by rewrite mulf_eq0; move /orP => []/eqP ->.
Qed.

Lemma p_x0_0_eq_0 : forall (n:nat) (p: mc M),
  p #x0 = 0%:R ->
  (p *+ n) #x0 = 0%R.
Proof.
  elim => [| n IHn] => p.
  move => _ ; rewrite mulr0n => //=.
  rewrite mulrS => Hp.
  have /IHn/p_x0_0_impl := Hp.
  move => [] ->.
  all: move/p_x0_0_impl:Hp.
  all: move => [] ->.
  all: rewrite /GRing.add //=.
  rewrite eq_refl => //=.
Qed.

Lemma opt_montgomery_ok (n m : nat) (x : K) :
    n < 2^m ->
    forall (p : mc M), p#x0 = x -> opt_montgomery n m x = (p *+ n)#x0.
Proof.
  have /orP[Hx|Hx] := orNb (x == 0) => Hn.
  apply opt_montgomery_x => //.
  move/eqP: Hx => -> => p.
  move/(p_x0_0_eq_0 n) => ->.
  apply opt_montgomery_0.
Qed.

End OptimizedLadder.