Require Import Coq.Sets.Ensembles.
Require Import Coq.Logic.Classical_Prop.
Require Import Coq.Classes.RelationClasses.
Require Import Coq.Relations.Relation_Definitions.
Require Import Coq.Logic.ChoiceFacts.
Require Import Coq.Logic.ClassicalChoice.
Require Import Logic.lib.RelationPairs_ext.
Require Import Logic.GeneralLogic.KripkeModel.
Require Import Logic.SeparationLogic.Model.SeparationAlgebra.
Require Import Logic.SeparationLogic.Model.OrderedSA.

(***********************************)
(* Separation Algebra Generators   *)
(***********************************)

Section trivialSA.
  Context {worlds: Type}.

  Definition trivial_Join: Join worlds:=  (fun a b c => False).

  Definition trivial_SA: @SeparationAlgebra worlds trivial_Join.
  Proof.
    constructor; intuition.
    inversion H.
  Qed.

  (* Trivial algebra is upwards closed *)
  (* Trivial is NOT downwards closed *)
  Definition trivial_uSA {R: Relation worlds}:
    @UpwardsClosedSeparationAlgebra worlds R (trivial_Join).
  Proof.
    intros until m2; intros.
    inversion H.
  Qed.

  (*Increasing*)
  Definition trivial_incrSA: @IncreasingSeparationAlgebra
                           worlds eq trivial_Join.
  Proof.
    constructor; intros.
    hnf; intros.
    inversion H.
  Qed.

  (*Trivial is NOT necessarily unital*)

  (*Trivial is NOT necessarily residual*)

End trivialSA.

Section unitSA.
  Definition unit_Join: Join unit:=  (fun _ _ _ => True).

  Definition unit_SA: @SeparationAlgebra unit unit_Join.
  Proof.
    constructor.
    + intros. constructor.
    + intros; exists tt; split; constructor.
  Qed.

  (* Unit algebra is upwards closed *)
  Definition unit_uSA:
    @UpwardsClosedSeparationAlgebra unit eq unit_Join.
  Proof.
    intros; exists tt, tt; intuition.
    + destruct m1; reflexivity.
    + destruct m2; reflexivity.
  Qed.

  (* Unit algebra is downwards closed *)
  Definition unit_dSA:
    @DownwardsClosedSeparationAlgebra unit eq unit_Join.
  Proof.
    intros; exists tt; intuition.
    destruct m; reflexivity.
  Qed.

  (*Increasing*)
  Instance unit_incrSA:
    @IncreasingSeparationAlgebra unit eq unit_Join.
  Proof.
    constructor; intros; hnf; intros.
    destruct n, n'; reflexivity.
  Qed.

  Instance unit_residual:
    @ResidualSeparationAlgebra unit eq unit_Join.
  Proof.
    constructor; intros.
    exists tt; exists tt; split.
    + constructor.
    + destruct n; reflexivity.
  Qed.

  Definition unit_unital:
    @UnitalSeparationAlgebra unit eq unit_Join.
  Proof.
    apply <- (@incr_unital_iff_residual unit eq (eq_preorder unit) unit_Join); auto.
    + apply unit_residual.
    + apply unit_incrSA.
  Qed.

End unitSA.

Section equivSA.
  Context {worlds: Type}.

  Definition equiv_Join: Join worlds:=  (fun a b c => a = c /\ b = c).

  Definition equiv_SA: @SeparationAlgebra worlds equiv_Join.
  Proof.
    constructor.
    + intros.
      inversion H.
      split; tauto.
    + intros.
      simpl in *.
      destruct H, H0.
      subst mx my mxy mz.
      exists mxyz; do 2 split; auto.
  Qed.

  (* Identity algebra is upwards closed *)
  (* Identity is NOT downwards closed *)
  Definition identity_uSA {R: Relation worlds}:
    @UpwardsClosedSeparationAlgebra worlds R equiv_Join.
  Proof.
    intros until m2; intros.
    destruct H; subst m1 m2.
    exists n, n; do 2 split; auto.
  Qed.

  Definition equiv_incrSA: @IncreasingSeparationAlgebra
                           worlds eq equiv_Join.
  Proof.
    constructor; intros.
    hnf; intros.
    inversion H; subst.
    constructor.
  Qed.

  Definition ikiM_uSA {R: Relation worlds} {po_R: PreOrder Krelation} {ikiM: IdentityKripkeIntuitionisticModel worlds} {J: Join worlds}: UpwardsClosedSeparationAlgebra worlds.
  Proof.
    intros until m2; intros.
    apply Korder_identity in H0.
    subst n.
    exists m1, m2.
    split; [| split]; auto; reflexivity.
  Qed.

  Definition ikiM_dSA {R: Relation worlds} {po_R: PreOrder Krelation} {ikiM: IdentityKripkeIntuitionisticModel worlds} {J: Join worlds}: DownwardsClosedSeparationAlgebra worlds.
  Proof.
    intros until n2; intros.
    apply Korder_identity in H0.
    apply Korder_identity in H1.
    subst n1 n2.
    exists m.
    split; auto; reflexivity.
  Qed.

  (*Identity is NOT necessarily increasing*)

  (*Identity is NOT necessarily unital*)

  (*Identity is NOT necessarily residual*)

End equivSA.

Section optionSA.
  Context (worlds: Type).

  Inductive option_join {J: Join worlds}: option worlds -> option worlds -> option worlds -> Prop :=
  | None_None_join: option_join None None None
  | None_Some_join: forall a, option_join None (Some a) (Some a)
  | Some_None_join: forall a, option_join (Some a) None (Some a)
  | Some_Some_join: forall a b c, join a b c -> option_join (Some a) (Some b) (Some c).

  Definition option_Join {SA: Join worlds}: Join (option worlds):=
    (@option_join SA).

  Definition option_SA
             {J: Join worlds}
             {SA: SeparationAlgebra worlds}:
    @SeparationAlgebra (option worlds) (option_Join).
  Proof.
    constructor.
    + intros.
      simpl in *.
      destruct H.
    - apply None_None_join.
    - apply Some_None_join.
    - apply None_Some_join.
    - apply Some_Some_join.
      apply join_comm; auto.
      + intros.
        simpl in *.
        inversion H; inversion H0; clear H H0; subst;
        try inversion H4; try inversion H5; try inversion H6; subst;
        try congruence;
        [.. | destruct (join_assoc _ _ _ _ _ H1 H5) as [? [? ?]];
           eexists; split; apply Some_Some_join; eassumption];
        eexists; split;
        try solve [ apply None_None_join | apply Some_None_join
                    | apply None_Some_join | apply Some_Some_join; eauto].
  Qed.

  (* Ordered option Upwards closed *)
  Lemma option_ord_uSA
        {R: Relation worlds}
        {J: Join worlds}
        (uSA: UpwardsClosedSeparationAlgebra worlds):
    @UpwardsClosedSeparationAlgebra (option worlds) (option01_relation Krelation) option_Join.
  Proof.
    hnf; intros.
    inversion H; subst.
    - inversion H0; subst.
      + exists None, None; repeat split; try constructor.
      + exists (Some a), None; repeat split; try constructor.
    - inversion H0; subst.
      exists None, (Some b); repeat split; try constructor; auto.
    - inversion H0; subst.
      exists (Some b), None; repeat split; try constructor; auto.
    - inversion H0; subst.
      destruct
        (uSA  _ _ _ _ H1 H3) as [n1 [n2 [HH1 [HH2 HH3]]]].
      exists (Some n1), (Some n2); repeat split; try constructor; auto.
  Qed.

  (* Downwards closed IF the algebra is increasing*)
  Lemma option_ord_dSA
        {R: Relation worlds}
        {po_R: PreOrder Krelation}
        {J: Join worlds}
        {SA: SeparationAlgebra worlds}
        (dSA: DownwardsClosedSeparationAlgebra worlds)
        {incrSA: IncreasingSeparationAlgebra worlds}:
    @DownwardsClosedSeparationAlgebra (option worlds) (option01_relation Krelation) option_Join.
  Proof.
    hnf; intros.
    inversion H0; [ | | inversion H1]; subst.
    - exists n2; inversion H; subst; split; auto;
      destruct n2; constructor.
    - exists n2; inversion H; subst; split; auto;
      destruct n2; constructor.
      + inversion H1.
      + inversion H; subst.
        apply all_increasing in H6.
        inversion H1; subst.
        transitivity b; auto.
    - exists (Some a); split; try constructor.
      inversion H; subst; auto.
    - exists (Some a); split; try constructor.
      transitivity (Some b); auto.
      inversion H; subst.
      constructor. eapply all_increasing. apply join_comm; eassumption.
    - inversion H; subst.
      destruct (dSA _ _ _ _ _ H6 H2 H5) as [n [HH1 HH2]].
      exists (Some n); split; constructor; auto.
  Qed.

  Lemma option_ord_incr_None
        {R: Relation worlds}
        {po_R: PreOrder Krelation}
        {J: Join worlds}
        {SA: SeparationAlgebra worlds}:
    @increasing (option worlds) (option01_relation Krelation) option_Join None.
  Proof.
    hnf; intros.
    inversion H; subst.
    + constructor.
    + constructor.
      reflexivity.
  Qed.

  Lemma option_ord_res_None
        {R: Relation worlds}
        {po_R: PreOrder Krelation}
        {J: Join worlds}
        {SA: SeparationAlgebra worlds}:
    forall n, @residue (option worlds) (option01_relation Krelation) option_Join n None.
  Proof.
    hnf; intros.
    exists n.
    split.
    + destruct n; constructor.
    + destruct n; constructor.
      reflexivity.
  Qed.

  Lemma option_ord_USA
        {R: Relation worlds}
        {po_R: PreOrder Krelation}
        {J: Join worlds}
        {SA: SeparationAlgebra worlds}:
    @UnitalSeparationAlgebra (option worlds) (option01_relation Krelation) option_Join.
  Proof.
    constructor.
    intros.
    exists None.
    split.
    + apply option_ord_res_None.
    + apply option_ord_incr_None.
  Qed.

  (* Disjoint option Upwards closed*)
  Lemma option_disj_uSA
        {R: Relation worlds}
        {J: Join worlds}
        (uSA: UpwardsClosedSeparationAlgebra worlds):
    @UpwardsClosedSeparationAlgebra (option worlds) (option00_relation Krelation) option_Join.
  Proof.
    hnf; intros.
    inversion H; subst.
    - inversion H0; subst.
      exists None, None; repeat split; try constructor.
    - inversion H0; subst.
      exists None, (Some b); repeat split; try constructor; auto.
    - inversion H0; subst.
      exists (Some b), None; repeat split; try constructor; auto.
    - inversion H0; subst.
      destruct
        (uSA  _ _ _ _ H1 H3) as [n1 [n2 [HH1 [HH2 HH3]]]].
      exists (Some n1), (Some n2); repeat split; try constructor; auto.
  Qed.

  (* Disjointed option Downwards *)
  Lemma option_disj_dSA
        {R: Relation worlds}
        {J: Join worlds}
        {SA: SeparationAlgebra worlds}
        (dSA: DownwardsClosedSeparationAlgebra worlds):
    @DownwardsClosedSeparationAlgebra (option worlds) (option00_relation Krelation) option_Join.
  Proof.
    hnf; intros.
    inversion H0; [ | inversion H1]; subst.
    - exists n2; inversion H; subst; split; auto;
      destruct n2; constructor.
    - exists (Some a); split; try constructor.
      inversion H; subst; auto.
    - inversion H; subst.
      destruct (dSA _ _ _ _ _ H6 H2 H5) as [n [HH1 HH2]].
      exists (Some n); split; constructor; auto.
  Qed.

  Lemma option_disj_incr_None
        {R: Relation worlds}
        {po_R: PreOrder Krelation}
        {J: Join worlds}
        {SA: SeparationAlgebra worlds}:
    @increasing (option worlds) (option00_relation Krelation) option_Join None.
  Proof.
    hnf; intros.
    inversion H; subst.
    + constructor.
    + constructor.
      reflexivity.
  Qed.

  Lemma option_disj_res_None
        {R: Relation worlds}
        {po_R: PreOrder Krelation}
        {J: Join worlds}
        {SA: SeparationAlgebra worlds}:
    forall n, @residue (option worlds) (option00_relation Krelation) option_Join n None.
  Proof.
    hnf; intros.
    exists n.
    split.
    + destruct n; constructor.
    + destruct n; constructor.
      reflexivity.
  Qed.

  Lemma option_disj_USA
        {R: Relation worlds}
        {po_R: PreOrder Krelation}
        {J: Join worlds}
        {SA: SeparationAlgebra worlds}:
    @UnitalSeparationAlgebra (option worlds) (option00_relation Krelation) option_Join.
  Proof.
    constructor.
    intros.
    exists None.
    split.
    + apply option_disj_res_None.
    + apply option_disj_incr_None.
  Qed.

  Lemma option_disj_USA'
        {R: Relation worlds}
        {po_R: PreOrder Krelation}
        {J: Join worlds}
        {SA: SeparationAlgebra worlds}:
    @UnitalSeparationAlgebra' (option worlds) (option00_relation Krelation) option_Join.
  Proof.
    constructor.
    intros.
    exists None.
    split.
    + apply option_disj_res_None.
    + hnf; intros.
      inversion H; subst.
      apply option_disj_incr_None.
  Qed.

End optionSA.

Section exponentialSA.
  Definition fun_Join (A B: Type) {J_B: Join B}: Join (A -> B) :=
    (fun a b c => forall x, join (a x) (b x) (c x)).

  Definition fun_SA
             (A B: Type)
             {Join_B: Join B}
             {SA_B: SeparationAlgebra B}: @SeparationAlgebra (A -> B) (fun_Join A B).
  Proof.
    constructor.
    + intros.
      simpl in *.
      intros x; specialize (H x).
      apply join_comm; auto.
    + intros.
      simpl in *.
      destruct (choice (fun x fx => join (my x) (mz x) fx /\ join (mx x) fx (mxyz x) )) as [myz ?].
    - intros x; specialize (H x); specialize (H0 x).
      apply (join_assoc _ _ _ _ _ H H0); auto.
    - exists myz; firstorder.
  Qed.

  (* Exponential is upwards closed *)
  Lemma fun_uSA
        (A B: Type)
        {R_B: Relation B}
        {J_B: Join B}
        (uSA_B: UpwardsClosedSeparationAlgebra B):
    @UpwardsClosedSeparationAlgebra (A -> B) (pointwise_relation A R_B) (fun_Join A B).
  Proof.
    hnf; intros.
    unfold join, fun_Join in H.
    unfold Krelation, pointwise_relation in H0.
    destruct (choice (fun x nn => join (fst nn) (snd nn) (n x) /\
                               Krelation (m1 x) (fst nn) /\
                               Krelation (m2 x) (snd nn)))
      as [nn H1].
    intros x.
    destruct (uSA_B (m x) (n x) (m1 x) (m2 x) (H x) (H0 x)) as [x1 [x2 ?]];
      exists (x1, x2); auto.
    exists (fun x => fst (nn x)), (fun x => snd (nn x)).
    simpl; repeat split; intros x; specialize (H1 x); destruct H1 as [H1 [H2 H3]];
    assumption.
  Qed.

  (* Exponential is downwards closed *)
  Lemma fun_dSA
        (A B: Type)
        {R_B: Relation B}
        {J_B: Join B}
        (dSA_B: DownwardsClosedSeparationAlgebra B):
    @DownwardsClosedSeparationAlgebra (A -> B) (pointwise_relation A R_B) (fun_Join A B).
  Proof.
    hnf; intros.
    unfold join, fun_Join in H.
    unfold Krelation, pointwise_relation in H0.
    destruct (choice (fun x n => join (n1 x) (n2 x) (n) /\
                               Krelation (n) (m x)))
      as [n H2].
    intros x.
    destruct (dSA_B (m1 x) (m2 x) (m x) (n1 x) (n2 x) (H x) (H0 x)) as [x1 [x2 ?]]; auto.
    + apply H1.
    + exists x1; auto.
    + exists n; split; hnf; intros x; specialize (H2 x); destruct H2; auto.
  Qed.

  (* Exponential is increasing *)
  Lemma fun_incrSA
        (A B: Type)
        {R_B: Relation B}
        {J_B: Join B}
        (incr_B: IncreasingSeparationAlgebra B):
    @IncreasingSeparationAlgebra (A -> B) (pointwise_relation A R_B) (fun_Join A B).
  Proof.
    constructor; intros.
    hnf; intros.
    hnf; intros.
    specialize (H a).
    eapply all_increasing; eauto.
  Qed.

  (* Exponential is Unital*)
  Lemma fun_unitSA
        (A B: Type)
        {R_B: Relation B}
        {J_B: Join B}
        (USA_B: UnitalSeparationAlgebra B):
    @UnitalSeparationAlgebra (A -> B) (pointwise_relation A R_B) (fun_Join A B).
  Proof.
    constructor; intros.
    destruct (choice (fun x mx => residue (n x) mx /\ increasing mx)) as [M HH].
    { intros;
      specialize (incr_exists (n x)); intros [y HH];
      exists y; auto. }
    exists M; split.
    + cut (forall x, residue (n x) (M x)).
      - clear; unfold residue; intros.
        apply choice in H; destruct H as [n' H].
        exists n'; split; hnf; intros x;
        specialize (H x); destruct H; auto.
      - intros x; specialize (HH x); destruct HH; auto.
    + unfold increasing; intros.
      unfold join, fun_Join in H.
      hnf; intros x.
      specialize (HH x); destruct HH as [ _ HH].
      apply HH.
      auto.
  Qed.

  (* Exponential is Unital' *)
  Lemma fun_unitSA'
        (A B: Type)
        {R_B: Relation B}
        {J_B: Join B}
        (USA'_B: UnitalSeparationAlgebra' B):
    @UnitalSeparationAlgebra' (A -> B) (pointwise_relation A R_B) (fun_Join A B).
  Proof.
    constructor; intros.
    destruct (choice (fun x mx => residue (n x) mx /\ increasing' mx)) as [M HH].
    { intros;
      specialize (incr'_exists (n x)); intros [y HH];
      exists y; auto. }
    exists M; split.
    + cut (forall x, residue (n x) (M x)).
      - clear; unfold residue; intros.
        apply choice in H; destruct H as [n' H].
        exists n'; split; hnf; intros x;
        specialize (H x); destruct H; auto.
      - intros x; specialize (HH x); destruct HH; auto.
    + unfold increasing', increasing; intros.
      unfold join, fun_Join in H.
      hnf; intros x.
      specialize (HH x); destruct HH as [ _ HH].
      eapply (HH _); eauto.
      apply H.
  Qed.

End exponentialSA.

Section sumSA.

  Inductive sum_worlds {worlds1 worlds2}: Type:
    Type:=
  | lw (w:worlds1): sum_worlds
  | rw (w:worlds2): sum_worlds.

  Inductive sum_join {A B: Type} {J1: Join A} {J2: Join B}:
    @sum_worlds A B ->
    @sum_worlds A B ->
    @sum_worlds A B-> Prop :=
  | left_join a b c:
      join a b c ->
      sum_join (lw a) (lw b) (lw c)
  | right_join a b c:
      join a b c ->
      sum_join (rw a) (rw b) (rw c).

  Definition sum_Join (A B: Type) {Join_A: Join A} {Join_B: Join B}: Join (@sum_worlds A B) :=
    (@sum_join A B Join_A Join_B).

  Definition sum_SA (A B: Type) {Join_A: Join A} {Join_B: Join B} {SA_A: SeparationAlgebra A} {SA_B: SeparationAlgebra B}: @SeparationAlgebra (@sum_worlds A B) (sum_Join A B).
  Proof.
    constructor.
    - intros; inversion H;
      constructor; apply join_comm; auto.
    - intros.
      inversion H; subst;
      inversion H0;
      destruct (join_assoc _ _ _ _ _ H1 H3) as [myz [HH1 HH2]].
      + exists (lw myz); split; constructor; auto.
      + exists (rw myz); split; constructor; auto.
  Qed.

End sumSA.

Section productSA.

  Definition prod_Join (A B: Type) {Join_A: Join A} {Join_B: Join B}: Join (A * B) :=
    (fun a b c => join (fst a) (fst b) (fst c) /\ join (snd a) (snd b) (snd c)).

  Definition prod_SA (A B: Type) {Join_A: Join A} {Join_B: Join B} {SA_A: SeparationAlgebra A} {SA_B: SeparationAlgebra B}: @SeparationAlgebra (A * B) (prod_Join A B).
  Proof.
    constructor.
    + intros.
      simpl in *.
      destruct H; split;
      apply join_comm; auto.
    + intros.
      simpl in *.
      destruct H, H0.
      destruct (join_assoc _ _ _ _ _ H H0) as [myz1 [? ?]].
      destruct (join_assoc _ _ _ _ _ H1 H2) as [myz2 [? ?]].
      exists (myz1, myz2).
      do 2 split; auto.
  Qed.

  Lemma prod_uSA
        (A B: Type)
        {R_A: Relation A}
        {R_B: Relation B}
        {Join_A: Join A}
        {Join_B: Join B}
        {dSA_A: UpwardsClosedSeparationAlgebra A}
        {dSA_B: UpwardsClosedSeparationAlgebra B}:
    @UpwardsClosedSeparationAlgebra (A * B) (RelProd R_A R_B) (@prod_Join _ _ Join_A Join_B).
  Proof.
    intros until m2; intros.
    destruct H, H0.
    destruct (join_Korder_up _ _ _ _ H H0) as [fst_n1 [fst_n2 [? [? ?]]]].
    destruct (join_Korder_up _ _ _ _ H1 H2) as [snd_n1 [snd_n2 [? [? ?]]]].
    exists (fst_n1, snd_n1), (fst_n2, snd_n2).
    do 2 split; simpl; auto;
    constructor; auto.
  Qed.

  Lemma prod_dSA
        (A B: Type)
        {R_A: Relation A}
        {R_B: Relation B}
        {Join_A: Join A}
        {Join_B: Join B}
        {uSA_A: DownwardsClosedSeparationAlgebra A}
        {uSA_B: DownwardsClosedSeparationAlgebra B}:
    @DownwardsClosedSeparationAlgebra (A * B) (RelProd R_A R_B) (@prod_Join _ _ Join_A Join_B).
  Proof.
    intros until n2; intros.
    destruct H, H0, H1.
    destruct (join_Korder_down _ _ _ _ _ H H0 H1) as [fst_n [? ?]].
    destruct (join_Korder_down _ _ _ _ _ H2 H3 H4) as [snd_n [? ?]].
    exists (fst_n, snd_n).
    do 2 split; simpl; auto.
  Qed.

  Lemma prod_incr
        (A B: Type)
        {R_A: Relation A}
        {R_B: Relation B}
        {Join_A: Join A}
        {Join_B: Join B}:
    forall (a: A) (b: B),
      increasing a -> increasing b ->
      @increasing _
                   (RelProd R_A R_B)
                   (@prod_Join _ _ Join_A Join_B) (a,b).
  Proof.
    intros. hnf; intros.
    destruct n, n'.
    inversion H1; simpl in *.
    hnf; intros; split.
    apply H; auto.
    apply H0; auto.
  Qed.

  Lemma prod_incrSA
        (A B: Type)
        {R_A: Relation A}
        {R_B: Relation B}
        {Join_A: Join A}
        {Join_B: Join B}
        {incrSA_A: IncreasingSeparationAlgebra A}
        {incrSA_B: IncreasingSeparationAlgebra B}:
    @IncreasingSeparationAlgebra (A * B) (RelProd R_A R_B) (@prod_Join _ _ Join_A Join_B).
  Proof.
    constructor; intros.
    destruct x; apply prod_incr; auto.
    + apply incrSA_A.
    + apply incrSA_B.
  Qed.

  Lemma prod_residualSA
        (A B: Type)
        {R_A: Relation A}
        {R_B: Relation B}
        {Join_A: Join A}
        {Join_B: Join B}
        {residualSA_A: ResidualSeparationAlgebra A}
        {residualSA_B: ResidualSeparationAlgebra B}:
    @ResidualSeparationAlgebra (A * B) (RelProd R_A R_B) (@prod_Join _ _ Join_A Join_B).
  Proof.
    constructor; intros.
    destruct n as [a b].
    inversion residualSA_A;
      inversion residualSA_B.
    destruct (residue_exists a) as [a' [a'' [Ha1 Ha2]]].
    destruct (residue_exists0 b) as [b' [b'' [Hb1 Hb2]]].
    exists (a', b'); hnf; intros.
    exists (a'', b''); hnf; intros;
    split; hnf; intros;
    split; simpl; auto.
  Qed.

  Lemma prod_unitalSA
        (A B: Type)
        {R_A: Relation A}
        {R_B: Relation B}
        {Join_A: Join A}
        {Join_B: Join B}
        {unitalSA_A: UnitalSeparationAlgebra A}
        {unitalSA_B: UnitalSeparationAlgebra B}:
    @UnitalSeparationAlgebra (A * B) (RelProd R_A R_B) (@prod_Join _ _ Join_A Join_B).
  Proof.
    inversion unitalSA_A.
    inversion unitalSA_B.
    constructor; intros.
    - destruct n as [a b].
      destruct (incr_exists a) as [a' [Ha1 Ha2]].
      destruct (incr_exists0 b) as [b' [Hb1 Hb2]].
      exists (a', b'); split; hnf; intros.
      + destruct Ha1 as [a'' [Ha1 Ha3]].
        destruct Hb1 as [b'' [Hb1 Hb3]].
        exists (a'',b''); split; hnf; hnf; intros; try constructor; simpl; auto.
      + destruct n, n'.
        inversion H; simpl in *.
        apply Ha2 in H0.
        apply Hb2 in H1.
        split; auto.
  Qed.

End productSA.

Class SeparationAlgebra_unit (worlds: Type) {J: Join worlds} := {
                                                                 unit: worlds;
                                                                 unit_join: forall n, join n unit n;
                                                                 unit_spec: forall n m, join n unit m -> n = m
                                                               }.

(***********************************)
(* Preorder examples               *)
(***********************************)

(***********************************)
(* dSA uSA examples                *)
(***********************************)

(***********************************)
(* More examples                   *)
(***********************************)

(*
Program Definition nat_le_kiM: KripkeIntuitionisticModel nat :=
  Build_KripkeIntuitionisticModel nat (fun a b => a <= b) _.
Next Obligation.
  constructor; hnf; intros.
  + apply le_n.
  + eapply NPeano.Nat.le_trans; eauto.
Qed.

(* TODO: Probably don't need this one. *)
Program Definition SAu_kiM (worlds: Type) {J: Join worlds} {SA: SeparationAlgebra worlds} {SAu: SeparationAlgebra_unit worlds} : KripkeIntuitionisticModel worlds :=
  Build_KripkeIntuitionisticModel worlds (fun a b => exists b', join b b' a) _.
Next Obligation.
  constructor; hnf; intros.
  + exists unit; apply unit_join.
  + destruct H as [? ?], H0 as [? ?].
    destruct (join_assoc _ _ _ _ _ H0 H) as [? [? ?]].
    exists x2; auto.
Qed.

Definition Heap (addr val: Type): Type := addr -> option val.

Instance Heap_Join (addr val: Type): Join (Heap addr val) :=
  @fun_Join _ _ (@option_Join _ (equiv_Join _)).

Instance Heap_SA (addr val: Type): SeparationAlgebra (Heap addr val) :=
  @fun_SA _ _ _ (@option_SA _ _ (equiv_SA _)).

Instance mfHeap_kiM (addr val: Type): KripkeIntuitionisticModel (Heap addr val) :=
  identity_kiM _.

Instance gcHeap_kiM (addr val: Type): KripkeIntuitionisticModel (Heap addr val) :=
  @fun_kiM _ _ (@option_ord_kiM _ (identity_kiM _)).

Definition Stack (LV val: Type): Type := LV -> val.

Definition StepIndex_kiM (worlds: Type) {po_R: PreOrder Krelation}: KripkeIntuitionisticModel (nat * worlds) := @prod_kiM _ _ nat_le_kiM kiM.

Definition StepIndex_Join (worlds: Type) {J: Join worlds}: Join (nat * worlds) :=
  @prod_Join _ _ (equiv_Join _) J.

Definition StepIndex_SA (worlds: Type) {J: Join worlds} {SA: SeparationAlgebra worlds}:
  @SeparationAlgebra (nat * worlds) (StepIndex_Join worlds) := @prod_SA _ _ _ _ (equiv_SA _) SA.

Definition StepIndex_dSA (worlds: Type) {po_R: PreOrder Krelation}
           {J: Join worlds} {dSA: UpwardsClosedSeparationAlgebra worlds}:
  @UpwardsClosedSeparationAlgebra (nat * worlds) (StepIndex_Join worlds) (StepIndex_kiM worlds):= @prod_dSA _ _ _ _ _ _ (@identity_dSA _ nat_le_kiM) dSA.

 *)


(** Generated by coq-of-ocaml *)
Require Import OCaml.OCaml.

Local Set Primitive Projections.
Local Open Scope string_scope.
Local Open Scope Z_scope.
Local Open Scope type_scope.
Import ListNotations.

Unset Positivity Checking.
Unset Guard Checking.

Inductive nat : Set :=
| O : nat
| S : nat -> nat.

Inductive natural : Set :=
| Zero : natural
| Succ : natural -> natural.

Fixpoint plus (plus_arg0 : natural) (plus_arg1 : natural) {struct plus_arg0}
  : natural :=
  match plus_arg0 with
  | Zero => plus_arg1
  | Succ n => Succ (plus n plus_arg1)
  end.

Definition synth (lf2 : natural) : natural := plus (plus Zero Zero) lf2.


Require Import Coq.Bool.Bool.
Require Import Coq.ZArith.ZArith.
Require Import Coq.micromega.Lia.
Require Import Crypto.Util.ZUtil.Notations.
Require Import Crypto.Util.ZUtil.Definitions.
Require Import Crypto.Util.ZUtil.Hints.Core.
Require Import Crypto.Util.ZUtil.Testbit.

Local Open Scope bool_scope. Local Open Scope Z_scope.

Module Z.
  Lemma lor_m1'_r x : Z.lor x (-1) = -1.
  Proof. apply Z.lor_m1_r. Qed.
#[global]
  Hint Rewrite lor_m1'_r : zsimplify_const zsimplify zsimplify_fast.

  Lemma lor_m1'_l x : Z.lor (-1) x = -1.
  Proof. apply Z.lor_m1_l. Qed.
#[global]
  Hint Rewrite lor_m1'_l : zsimplify_const zsimplify zsimplify_fast.

  Lemma lor_add a b (Hand : a &' b = 0) : a |' b = a + b.
  Proof. rewrite <- Z.lxor_lor, Z.add_nocarry_lxor by assumption; reflexivity. Qed.

  Lemma lor_small_neg a b
        (Ha : - 2^b <= a < 0)
        (Hb : 0 < b) :
    2^b |' a = a.
  Proof.
    apply Z.bits_inj_iff; red; intros; rewrite Z.lor_spec.
    destruct (Z.eqb_spec b n); subst.
    - now rewrite (Testbit.Z.testbit_small_neg a), orb_true_r.
    - now rewrite Z.pow2_bits_false, orb_false_l. Qed.
End Z.


From mathcomp Require Import ssreflect ssrfun ssrbool ssrnat seq path order eqtype.
From okasaki Require Import ssrlia.
Require Import Psatz.
Import Order.TTheory.
Notation ordType := (orderType tt).
Import Order.NatOrder.
From Equations Require Import Equations.

Open Scope order_scope.

Lemma snd_true3 a b : a || true || b.
Proof. by case: a. Qed.

Lemma trd_true3 a b : b || a || true.
Proof. by case: a; case b. Qed.

Lemma snd_true2 a : a || true.
Proof. by case: a. Qed.

Lemma add_lt_le a b c d: (a < c)%N -> (b <= d)%N -> (a + b < c + d)%N.
Proof. by move=> *; rewrite [a + b]addnC -ltn_subRL -addnBA // ltn_addr. Qed.

Lemma add_le_lt a b c d: (a <= c)%N -> (b < d)%N -> (a + b < c + d)%N.
Proof. move=> *; rewrite [a + b]addnC [c + d]addnC. by apply: add_lt_le. Qed.

Hint Resolve trd_true3 snd_true3 snd_true2 lexx ltxx : core.

Module leftistheap.
(* Leftist heaps are heap-ordered binary trees that satisfy the                         *)
(* leftist property: the rank of any left child is at least as large as the rank of its *)
(* right sibling. The rank of a node is defined to be the length of its right spine     *)
(* (i.e., the rightmost path from the node in question to an empty node).               *)
Section LeftistDef.
Variables (T: ordType).

Inductive heap :=
| Emp : heap
| Node : nat -> T-> heap -> heap -> heap.

End LeftistDef.
Arguments Node {T}.
Arguments Emp {T}.
Notation "[ tl | n , x | tr ]" := (Node n x tl tr) (at level 0).
Notation "'[||]'" := Emp.

(* Definition of arbitrary measure on Heaps *)
Module Measure.
Structure mixin_of {T : ordType} (measure : heap T -> nat) :=
  Mixin {
    measure1            : nat;
    f                   : nat -> nat -> nat;
    measureNode         : forall tl tr n x, measure [tl| x, n |tr] = f (measure tl) (measure tr);
    measure_NodexEE     : f (measure [||]) (measure [||]) = measure1;
    measure_Node_E      : forall h1 h2 n x, (measure [||] < measure [h1 | n, x | h2])%N
  }.

Notation class_of := mixin_of (only parsing).

Section classdef.

Structure type T := Pack { sort; _ : @class_of T sort }.

Local Coercion sort : type >-> Funclass.

Variables (T : ordType) (cT : (type T)).

(** Projection out of [type] *)
Definition class :=
  let: Pack _ c := cT return class_of cT in c.

End classdef.
Module Exports.

Coercion sort : type >-> Funclass.
(** Some shorthands *)
Notation measureType := type.

Notation MeasureMixin := Mixin.

End Exports.
End Measure.

Section Specifications.
Export Measure.Exports.
Variables (T : ordType).
Implicit Type h : heap T.

Definition empty h :=
if h is Emp then true else false.

Theorem emptyP h : reflect (h = Emp) (empty h).
Proof. case h; by constructor. Qed.

(* Definition of decidable equality on Heaps  *)
Fixpoint eqheap h1 h2 :=
  match h1, h2 with
  | Emp, Emp => true
  | Node n1 x1 tl1 tr1, Node n2 x2 tl2 tr2 => [&& (n1 == n2), (x1 == x2), eqheap tl1 tl2 & eqheap tr1 tr2]
  | _, _ => false
  end.

Lemma eqheapP : Equality.axiom eqheap.
Proof.
elim=> [|??? IHh1 ? IHh2] [|*] /=; try by constructor.
by apply: (iffP and4P); case=> [/eqP-> /eqP-> /IHh1-> /IHh2->].
Qed.

Canonical heap_eqMixin := EqMixin eqheapP.
Canonical heap_eqType := Eval hnf in EqType (heap T) heap_eqMixin.

Lemma eqheapE : eqheap = eq_op. Proof. by []. Qed.

Fixpoint mem_heap h :=
  if h is Node _ y tl tr then xpredU1 y (xpredU (mem_heap tl) (mem_heap tr)) else xpred0.

Definition heap_eqclass := heap T.
Identity Coercion heap_of_eqclass : heap_eqclass >-> heap.
Coercion pred_of_heap (h : heap_eqclass) : {pred T} := mem_heap h.

Canonical heap_predType := PredType pred_of_heap.
(* The line below makes mem_heap a canonical instance of topred. *)
Canonical mem_heap_predType := PredType mem_heap.

Lemma in_node x y n tl tr:
  x \in [tl | n, y | tr] = [|| x == y, x \in tl | x \in tr].
Proof. by []. Qed.

(* Definition of rank and size and simple properties *)
Section Mesures.

Fixpoint rank h : nat := if h is Node _ _ _ b then (rank b).+1 else O.

Definition f1 (x y : nat) := y.+1.

Lemma rank__NodexEE : f1 (rank [||]) (rank  [||]) = 1%N.
Proof. by []. Qed.

Lemma rankNode tl tr n x : rank [tl| x, n |tr] = f1 (rank tl) (rank tr).
Proof. by []. Qed.

Lemma rank_Node_E h1 h2 n x : (rank  [||] < rank [h1 | n, x | h2])%N.
Proof. by []. Qed.

Fixpoint size h : nat :=
if h is Node _ _ a b then (size a) + (size b) + 1 else O.

Definition f2 (x y : nat) := x + y + 1.

Lemma sizeNode tl tr n x : size [tl| x, n |tr] = f2 (size tl) (size tr).
Proof. by []. Qed.

Lemma size__NodexEE : f2 (size  [||]) (size  [||]) = 1%N.
Proof. by []. Qed.
Lemma size_Node_E h1 h2 n x : (size  [||] < size [h1 | n, x | h2])%N.
Proof. by rewrite addn_gt0 leqnn. Qed.

End Mesures.

(* Definition of count, all and simple properties *)
Section HeapCount.
Variables a : pred T.

Fixpoint count h : nat := if h is [tl| _, x |tr] then a x + count
 tl + count tr else 0.

Lemma count_Node tl tr y n: count [tl| n, y |tr] = a y + count tl + count tr.
Proof. by []. Qed.

Lemma count_E : count [||] = 0.
Proof. by []. Qed.

Lemma count_NodexEE n y : count [[||]| n, y | [||]] = a y.
Proof. by rewrite /= !addn0. Qed.

Fixpoint all h :=
if h is [tl| _, x |tr] then [&& a x, (all tl) & all tr] else true.

End HeapCount.

Lemma count_a_predT a h : (count a h <= count predT h)%N.
Proof. elim: h=> //= n s *; case: (a s); by repeat apply: leq_add. Qed.
Hint Resolve count_a_predT: core.

Lemma size_count h : size =1 count predT.
Proof. by elim=> //= n x h1 -> h2->; ssrnatlia. Qed.

Lemma allE h a : reflect (count a h = size h) (all a h).
Proof.
rewrite size_count //.
elim: h => /= [| _ s h1 /(sameP eqP) <- h2 /(sameP eqP) <-]; first by constructor.
apply/introP => [/and3P[-> /eqP-> /eqP->]|] // H.
suffices: ((a s) + count a h1 + count a h2 < 1 + count predT h1 + count predT h2)%N.
- by move=> LE EQ; move: EQ LE=>->; rewrite ltnn.
move: H=> /nandP[/negPf->|/nandP[] H] //=.
- apply: add_lt_le=> //; apply : add_lt_le=> //.
1: have: (count a h1 < count predT h1)%N=>*.
3: have: (count a h2 < count predT h2)%N=>*.
1,3: rewrite ltn_neqAle; apply/andP; by split.
- apply: add_lt_le=> //; apply add_le_lt=> //; by case (a s).
- apply: add_le_lt=> //. apply leq_add=> //; by case (a s).
Qed.

Lemma in_count x h: x \in h = (0 < count (xpred1 x) h)%N.
Proof.
elim h=> // ??? IHh1 ? IHh2; by rewrite in_node /= !addn_gt0 IHh1 IHh2 eq_sym lt0b orbA.
Qed.

(* if h is heap oredered (see below) LE x h is same as   *)
(* x is less or equal then all elements in heaps         *)
Definition LE x h : bool :=
if h is Node n y a b then (x <= y) else true.

Lemma LE_trans y x h: LE y h -> (x <= y) -> LE x h.
Proof. by case: h=> //= ???? H /le_trans/(_ H). Qed.

(* Definition of heap ordered heaps and simpl properties *)
Fixpoint heap_ordered h : bool :=
if h is Node n x tl tr then
  [&& (LE x tl), (LE x tr), (heap_ordered tl) & (heap_ordered tr)]
else true.

Lemma LE_correct h x y: heap_ordered h -> x \in h -> LE y h -> y <= x.
Proof.
by elim: h x y=> [| n x h1 IHh1 h2 IHh2 x' y' /= /and4P[L1 L2 H1 H2]] //
/or3P[/eqP ->|/(IHh1 _ _ H1)/(_ L1)|/(IHh2 _ _ H2)/(_ L2)] // H /le_trans/(_ H).
Qed.

Lemma LE_spec x h : heap_ordered h -> LE x h -> all (>= x) h.
Proof.
elim h=> // n y h1 IHh1 h2 IHh2 /= /and4P[???? XY].
by rewrite XY IHh1 // ?IHh2 // ?(LE_trans y).
Qed.

Fact heap_ordered_NodexEE n x : heap_ordered [ [||]| n, x |  [||]].
Proof. by []. Qed.
Hint Resolve heap_ordered_NodexEE: core.
Fact heap_ordered_E : heap_ordered [||]. Proof. by []. Qed.

(* Some basic definitions for reasoning about spines in heaps *)
Variant side := R | L.

Definition spine := seq side.

Fixpoint spine_in p h : bool :=
match p, h with
| R :: p, Node _ _ _ tr => spine_in p tr
| L :: p, Node _ _ tl _ => spine_in p tl
| [::],   Emp          => true
| _,      _          => false
end.

Fixpoint right p : bool :=
match p with
| R :: p' => right p'
| L :: _  => false
| [::]    => true
end.

Lemma rigth_correct x s: right (x :: s) -> ((x = R) * right s).
Proof. by case: s; case x. Qed.

Lemma right_spine_ex h: exists s, right s && spine_in s h.
Proof.
elim: h=> [|n s hl _ hr [] s']; by [exists nil| exists (R :: s')].
Qed.

Lemma spine_in_E s : spine_in s Emp -> s = [::].
Proof. by case: s=> [|[]] //=. Qed.

(* Theory for heaps with leftist measure *)

Section Measure.
Variables measure : measureType T.
Definition measure1 : nat := Measure.measure1 _ (Measure.class T measure).
Definition measure_NodexEE := Measure.measure_NodexEE _ (Measure.class T measure).
Definition measureNode := Measure.measureNode _ (Measure.class T measure).
Definition measure_Node_E := Measure.measure_Node_E _ (Measure.class T measure).
Definition f := Measure.f _ (Measure.class T measure).
Hint Resolve measure_Node_E : core.

Fixpoint measure_inv h : bool :=
match h with
| [tl| n, x |tr] => [&& (n == measure h), (measure_inv tl) & (measure_inv tr)]
| [||]           => true
end.

Lemma measure_inv_NodexEE x : measure_inv [ [||]| measure1, x | [||]].
Proof. by rewrite /= measureNode measure_NodexEE eq_refl. Qed.
Hint Resolve measure_inv_NodexEE : core.

Fact measure_E : measure_inv [||]. Proof. by []. Qed.

(* leftist invariant for arbitrary mesure and defonition of leftist heap *)
Fixpoint leftist_inv h : bool :=
if h is Node n x tl tr then
  [&& (measure tr <= measure tl)%N, (leftist_inv tl) & (leftist_inv tr)]
else true.

Fact leftist_inv_NodexEE n x : leftist_inv [[||]| n, x |  [||]].
Proof. move=> /=; by rewrite leqnn. Qed.
Hint Resolve leftist_inv_NodexEE : core.

Fact leftist_inv_E : leftist_inv [||]. Proof. by []. Qed.

Definition leftist_measure_inv h := leftist_inv h && measure_inv h.

Lemma case_leftist_measure_inv_r n x tl tr :
leftist_measure_inv (Node n x tl tr) -> leftist_measure_inv tr.
Proof.
by rewrite /leftist_measure_inv /= => /and3P[/and3P[_ _ -> _ /andP[]]].
Qed.

Lemma case_leftist_measure_inv_rl n x tl tr:
leftist_measure_inv (Node n x tl tr) ->
[&& leftist_measure_inv tr, leftist_measure_inv tl & (measure tr <= measure tl)%N].
Proof.
by rewrite /leftist_measure_inv => /andP[] /= => /and3P[->->-> /and3P[_ ->->]].
Qed.

Lemma case_leftist_measure_inv n x tl tr:
leftist_measure_inv (Node n x tl tr) -> (n == f (measure tl) (measure tr)) &&
[&& leftist_measure_inv tr, leftist_measure_inv tl & (measure tr <= measure tl)%N].
Proof.
move=> LI. move: LI (LI)=> /case_leftist_measure_inv_rl -> /andP[_] /= /and3P[];
by rewrite measureNode=>->.
Qed.

Definition leftistheap h :=
[&& leftist_inv h, measure_inv h & heap_ordered h].

Fact leftistheap_NodexEE x : leftistheap [[||]| measure1, x |[||]].
Proof. by rewrite /leftistheap measure_inv_NodexEE heap_ordered_NodexEE leftist_inv_NodexEE. Qed.
Hint Resolve leftistheap_NodexEE : core.

Fact leftistheap_E : leftistheap  [||].
Proof. by rewrite /leftistheap. Qed.
Hint Resolve leftistheap_E : core.

(* makeT takes two heaps and element and make heap from them *)
(* with correct mesure                                       *)
Definition makeT (x : T) h1 h2 :=
let m1 := measure h1 in
  let m2 := measure h2 in
      if (m2 <= m1)%N then [h1| f m1 m2, x |h2] else [h2| f m2 m1, x |h1] .

Lemma makeT_LI_inv x tl tr : leftist_inv tl -> leftist_inv tr ->
leftist_inv (makeT x tl tr).
Proof.
rewrite /makeT; case: ifP=> H /=->->; first by rewrite H.
- suffices: (measure tl <= measure tr)%N => [->|] //; move: H.
move=> /neq0_lt0n; by ssrnatlia.
Qed.

Lemma makeT_rk_inv x tl tr : measure_inv tl -> measure_inv tr ->
measure_inv (makeT x tl tr).
Proof.
by rewrite /makeT; case: ifP=> /= _ ->->; rewrite measureNode eq_refl.
Qed.

Lemma makeT_peserve_HO_inv x tl tr :
heap_ordered tl -> heap_ordered tr -> LE x tl -> LE x tr ->
 heap_ordered (makeT x tl tr).
Proof. rewrite /makeT; by case: ifP => _ /=->->->->. Qed.

Lemma makeT_spec h1 h2 x a:
count a (makeT x h1 h2) = a x + count a h1 + count a h2.
Proof. rewrite /makeT; case: ifP=> /=; ssrnatlia. Qed.

Lemma makeT_in_spec h1 h2 x y :
x \in (makeT y h1 h2) = ((x == y) || (x \in h1) || (x \in h2)).
Proof.
by rewrite !in_count makeT_spec eq_sym; case : (x == y)=> /=; rewrite !addn_gt0.
Qed.

(* fuction merge makes from two heaps nwe one that contains *)
(* same elements (see merge_spec)                           *)
(* Here your can see special total realization of merge     *)
Fixpoint merge' a :=
if a is Node n x a1 b1 then
let fix merge_a b :=
  if b is Node m y a2 b2 then
    if x <= y then
      makeT x a1 (merge' b1 b)
    else
      makeT y a2 (merge_a b2)
  else a in
merge_a
else id.

Arguments merge' !tl !tr : rename.
Definition merge := nosimpl merge'.

Ltac merge_cases := match goal with
| H : (?x <= ?y) = true |- _ => move: (H); rewrite /merge /= =>->; rewrite ?H
| H : (?x <= ?y) = false |- _ => move: (H)=> /merge_a /= ->; rewrite ?H
end.

Ltac merge_casesxy x y := case H : (x <= y); merge_cases.

Lemma merge_E_h h: merge Emp h = h.
Proof. by []. Qed.

Lemma merge_h_E h: merge h Emp = h.
Proof. by case: h. Qed.

Lemma merge_a nl nr x y tll tlr trr trl : (x <= y) = false ->
merge (Node nl x tll tlr) (Node nr y trl trr) =
makeT y trl (merge  (Node nl x tll tlr) trr).
Proof. rewrite /merge /= => ->. by elim: trr. Qed.

Lemma merge_measure_inv h1 h2:
measure_inv h1 -> measure_inv h2 -> measure_inv (merge h1 h2) .
Proof.
elim: h1 h2=> // ? x ??? IHhr. elim=> // ? y ??? IH'hr.
merge_casesxy x y => /and3P[E M M' /and3P[EQ M1 M2]];
by rewrite makeT_rk_inv // (IHhr, IH'hr) //= (EQ, E) (M1, M) (M2, M').
Qed.

Lemma merge_LE h1 h2 x: LE x h1 -> LE x h2 -> LE x (merge h1 h2).
Proof.
elim: h1 h2 x => // ? x ????. elim=> // ? y *.
by merge_casesxy x y; rewrite /makeT; case: ifP.
Qed.

Lemma merge_HO_inv h1 h2:
heap_ordered h1 -> heap_ordered h2 -> heap_ordered (merge h1 h2).
Proof.
elim: h1 h2=> // ? x ??? IHhr. elim=> // ? y ??? IH'hr H1 H2.
move : (H1) (H2)=> /and4P[???? /and4P[????]]. by merge_casesxy x y;
rewrite makeT_peserve_HO_inv ?IHhr ?IH'hr ?merge_LE //=; move: H; case: ltgtP.
Qed.

Lemma merge_LI_inv h1 h2:
 leftist_inv h1 -> leftist_inv h2 -> leftist_inv (merge h1 h2).
Proof.
elim: h1 h2=> // ? x ??? IHhr. elim=> // ? y ??? IH'hr H1 H2.
move : (H1) (H2) => /and3P[??? /and3P[*]]. merge_casesxy x y;
by rewrite makeT_LI_inv ?IHhr ?IH'hr.
Qed.

Theorem merge_LH h1 h2 :
leftistheap h1 -> leftistheap h2 -> leftistheap (merge h1 h2).
Proof.
move=> /and3P[???/and3P[*]].
by rewrite /leftistheap merge_LI_inv ?merge_measure_inv ?merge_HO_inv.
Qed.

Theorem merge_spec h1 h2 a:  count a (merge h1 h2) = count a h1 + count a h2.
Proof.
elim: h1 h2=> [?|? x ??? IHh2] //. elim=> [/=|? y ??? IHh22]; first by ssrnatlia.
by merge_casesxy x y; rewrite makeT_spec ?IHh2 ?IHh22 /=; ssrnatlia.
Qed.

Lemma merge_size h1 h2 : size (merge h1 h2) = size h1 + size h2.
Proof. by rewrite !size_count ?merge_spec. Qed.

Theorem merge_in_spec  h1 h2 x : x \in (merge h1 h2) = ((x \in h1) || (x \in h2)).
Proof. by rewrite !in_count merge_spec !addn_gt0. Qed.

(* Two defenition of insert, their equality and some properties *)
Definition insert (x : T) h :=  merge [[||]| measure1, x | [||]] h.

Fixpoint insert' (x : T) h :=
if h is Node n y a b then
let h' := Node n y a b in
  if x <= y then
    Node (f (measure h') (measure [||])) x h' Emp
  else makeT y a (insert' x b)
else Node measure1 x Emp Emp.

Lemma insert_measure_inv h x: measure_inv h -> measure_inv (insert x h).
Proof. by move=> *; apply: merge_measure_inv. Qed.

Lemma insert_HO_inv h x: heap_ordered h -> heap_ordered (insert x h).
Proof. by move=> *; apply: merge_HO_inv. Qed.

Lemma insert_LI_inv h x: leftist_inv h -> leftist_inv (insert x h).
Proof. by move=> *; apply: merge_LI_inv. Qed.

Theorem insert_LH h x: leftistheap h -> leftistheap (insert x h).
Proof. by move=> *; apply: merge_LH. Qed.

Theorem insert_spec h x a : count a (insert x h) = a x + count a h.
Proof. by rewrite merge_spec /=; ssrnatlia. Qed.

Lemma insert_size h x: size (insert x h) = (size h).+1.
Proof. by rewrite !size_count ?merge_spec. Qed.

Theorem insertE x h : measure_inv h -> leftist_inv h -> insert' x h = insert x h.
Proof.
rewrite /insert. elim h=> // n y h1 IHh1 h2 IHh2 /= /and3P[??? /and3P[*]].
merge_casesxy x y=> //; rewrite ?IHh2 // /makeT.
have: (measure [h1 | n, y | h2] <= measure [||] = false)%N=> [|->] //.
by apply/negbTE; rewrite -ltnNge.
Qed.

Theorem insert_in_spec h x y : x \in (insert y h) = ((x == y) || (x \in h)).
Proof.
rewrite merge_in_spec /= !in_node; case (x == y); by case (x \in h).
Qed.

(* Properties of findmin function *)
Definition findmin h :=
if h is Node _ x _ _ then Some x else None.

Theorem findmin_None h: None = findmin h <-> h = [||].
Proof. split=> [|-> //]. by case : h. Qed.

Theorem findmin_Some h z: heap_ordered h ->
Some z = findmin h -> all (>= z) h.
Proof.
case: h=> //= _ ??? /and3P[?? /andP[?? [->]]]; by rewrite lexx ?LE_spec.
Qed.

Lemma findmin_cases h: (exists x, Some x = findmin h) \/ (h = Emp).
Proof. case: h=> [|n z h1 h2 /=]; first by right. by left; exists z. Qed.

Theorem findmin_spec x h: heap_ordered h ->
((x \in h) && LE x h) <-> (Some x = findmin h).
Proof.
split=> [/andP[]|]; move: h H=> [] //.
- move=> n y h1 h2 /= /and4P[L1 L2 H1 H2]
  /or3P[/eqP-> //|/(LE_correct _ _ _ H1)/(_ L1)|/(LE_correct _ _ _ H2)/(_ L2)] XY YX;
suffices: x = y=> [-> //|]; apply: le_anti; rewrite XY YX //.
move=> ????? [->]. by rewrite in_node eq_refl /=.
Qed.

(* Properties of findmin function *)
Definition deletemin h := if h is Node _ _ a b then merge a b else Emp.

Lemma case_leftistheap n x h1 h2 :
leftistheap (Node n x h1 h2) -> leftistheap h1 && leftistheap h2.
Proof.
by rewrite/leftistheap=>/and3P[/=/and3P[_->->/=/and3P[_->->/and4P[_ _->->]]]].
Qed.

Lemma deletemin_rk_inv h: measure_inv h -> measure_inv (deletemin h).
Proof. case: h=> //= ???? /and3P[*]; by rewrite merge_measure_inv. Qed.

Lemma deletemin_ho_inv h: heap_ordered h -> heap_ordered (deletemin h).
Proof. case: h=> //= ???? /and4P[*]. by rewrite merge_HO_inv. Qed.

Lemma deletemin_LI_inv h: leftist_inv h -> leftist_inv (deletemin h).
Proof. case: h=> //= ???? /and3P[*]. by rewrite merge_LI_inv. Qed.

Lemma deletemin_correct h:
leftistheap h -> leftistheap (deletemin h).
Proof.
by case: h=> //=???? H; apply merge_LH; move: H=> /case_leftistheap/andP[].
Qed.

Theorem deletemin_spec h x:
  Some x = findmin h <-> (count^~ (insert x (deletemin h)) =1 count^~ h).
Proof.

case: h=> /=; first (split=> // /(_ (pred1 x)) /=; by rewrite eq_refl).
move=> _ s??; split=> [[-> a]|/(_ (pred1 x))];
rewrite insert_spec merge_spec addnA // -2?addnA=> /eqP;
rewrite eqn_add2r/= eq_refl=> /eqP/esym. case E: (s == x)=> //. by rewrite (eqP E).
Qed.

Lemma deletemin_size h : size (deletemin h) = (size h).-1.
Proof. case: h=> // n x tl tr; by rewrite ?size_count //= merge_spec. Qed.

Theorem deletemin_in_spec h x y:
Some x = findmin h -> (y \in (insert x (deletemin h))) = (y \in h).
Proof.
case: h=> // ???? [->].
by rewrite insert_in_spec /deletemin merge_in_spec // orbA.
Qed.

(* In this section we define heapsort and prove it's properties *)
(* heapsort fist translates sequense of elements into heap and  *)
(* than takes the minimal element of the heap while it is not   *)
(* empty                                                        *)
Section Heapsort.
Implicit Type hh : seq (option (heap T)).

Fixpoint seq_to_seqheap (st : seq T) :=
if st is h :: t then cons [ [||]| measure1, h | [||]] (seq_to_seqheap t)
else [::].

Fixpoint count_sseq a hh :=
  if hh is sh :: hh' then
    if sh is some h then count a h + count_sseq a hh'
    else count_sseq a hh'
  else 0.

Fixpoint count_seq a (sh : seq (heap T)) :=
  if sh is h :: sh' then count a h + count_seq a sh' else 0.

Theorem seq_to_seqheap_spec s: count_seq^~ (seq_to_seqheap s) =1 seq.count^~ s.
Proof. move=> a; elim: s=> //= ??->; ssrnatlia. Qed.

(* Now we want to define fromseq function that will translate  *)
(* sequense of elements into heap in O(n) time                 *)
(* Our implementation is based on divide and conqure principle *)
(* adapted to total languages                                  *)
Fixpoint fromseqheap_push h1 hh :=
  match hh with
  | None :: hh' | [::] as hh' => (some h1) :: hh'
  | (some h2) :: hh' => None :: fromseqheap_push (merge h2 h1) hh'
  end.

Fixpoint fromseqheap_pop h1 hh :=
  if hh is sh2 :: hh' then
    let h2 := if sh2 is some h then h else [||] in
   fromseqheap_pop (merge h2 h1) hh' else h1.

Fixpoint fromseqheap_rec hh sh  :=
  match sh with
  | [:: x1, x2 & h'] => let h1 := merge x1 x2 in
    fromseqheap_rec (fromseqheap_push h1 hh) h'
  | [:: h] => fromseqheap_pop h hh
  | [::] => fromseqheap_pop [||] hh
  end.

Definition fromseqheap := (fromseqheap_rec [::]).

Fixpoint fromseqheap_rec1 hh sh :=
  if sh is x :: h then fromseqheap_rec1 (fromseqheap_push x hh) h else fromseqheap_pop (  [||]) hh.

Lemma fromseqheapE sh : fromseqheap sh = fromseqheap_rec1 [::] sh.
Proof.
transitivity (fromseqheap_rec1 [:: None] sh); last by case: sh.
rewrite /fromseqheap; move: [::] {2}_.+1 (ltnSn (seq.size sh)./2) => hh n.
elim: n => // n IHn in hh sh *. case: sh => [|x [|y s]] //=; by rewrite ?merge_h_E=> //= /IHn->.
Qed.

Definition fromseq := fromseqheap \o seq_to_seqheap.

Lemma fromseqheap_pop_spec a h hh :
  count a (fromseqheap_pop h hh) = count_sseq a hh + count a h.
Proof.
elim: hh => [|[?|]? IHhh] //= in h *; by rewrite IHhh merge_spec; ssrnatlia.
Qed.

Lemma fromseqheap_push_spec h hh a:
  count_sseq a (fromseqheap_push h hh )= count a h + count_sseq a hh.
Proof.
elim: hh=> [|[?|]? IHhh] //= in h *. by rewrite IHhh merge_spec; ssrnatlia.
Qed.

Lemma fromseqheap_rec1_spec sh hh a:
  count a (fromseqheap_rec1 hh sh) = count_sseq a hh + count_seq a sh.
Proof.
elim: sh => [|?? IHsh] /= in hh *;
by rewrite ?fromseqheap_pop_spec // IHsh fromseqheap_push_spec; ssrnatlia.
Qed.

Lemma fromseqheap_spec sh :
 count^~ (fromseqheap sh) =1 count_seq^~ sh.
Proof. move=> a; by rewrite fromseqheapE fromseqheap_rec1_spec. Qed.

Theorem fromseq_cspec (s : seq T):
  seq.count^~ s =1 count^~ (fromseq s).
Proof.
move=> a; by rewrite /fromseq /= fromseqheap_spec seq_to_seqheap_spec.
Qed.

(* In this section we prove that fromseq s statisfy any invariant *)
(* with spesial properties                                        *)
Section Invariant.
Definition invariant := heap T -> bool.

Variables inv : invariant.

Hypothesis merge_invariat : forall h1 h2, inv h1 -> inv h2 -> inv (merge h1 h2).
Hypothesis inv_E : inv [||].
Hypothesis inv_NodexEE : forall x, inv [[||]| measure1, x |  [||]].
Hint Resolve inv_E : core.

Definition some_inv sh := if sh is some h then inv h else true.

Lemma inv_fromseqheap_pop h sh:
  (seq.all some_inv sh) -> inv h -> inv (fromseqheap_pop h sh).
Proof.
elim: sh=> [|[?|]? IHhh] //= in h *=> /andP[*]; by rewrite IHhh ?merge_invariat.
Qed.

Lemma all_inv_fromseqheap_push h hh: inv h -> (seq.all some_inv hh) ->
seq.all some_inv (fromseqheap_push h hh) = (inv h) && (seq.all some_inv hh).
Proof.
by elim: hh=> [|[?|]? IHhh] //= in h *=> HH /andP[SH ?];
rewrite IHhh ?merge_invariat // HH SH.
Qed.

Lemma inv_fromseqheap_rec1 hh sh :
(seq.all some_inv  hh) -> (seq.all inv sh) -> (inv (fromseqheap_rec1 hh sh)).
Proof.
elim: sh => [|h sh IHsh] /= in hh *=> ?; rewrite ?inv_fromseqheap_pop //.
move => /andP[HH ?]. by rewrite IHsh ?all_inv_fromseqheap_push // HH.
Qed.

Lemma inv_fromseqheap sh : seq.all inv sh -> inv (fromseqheap sh).
Proof. move=> ASH; by rewrite fromseqheapE inv_fromseqheap_rec1. Qed.

Lemma all_inv_seq_toseqheap s : seq.all inv (seq_to_seqheap s).
Proof. elim: s=> //= a h->; by rewrite inv_NodexEE. Qed.

Lemma inv_fromseq s : inv (fromseq s).
Proof. by rewrite /fromseq /= inv_fromseqheap // all_inv_seq_toseqheap. Qed.
End Invariant.

Lemma measure_inv_E : measure_inv [||].
Proof. by []. Qed.

Definition rank_rk_fromseq : forall s, measure_inv (fromseq s) :=
  inv_fromseq measure_inv merge_measure_inv measure_inv_E measure_inv_NodexEE.

Definition heap_ordered_fromseq : forall s, heap_ordered (fromseq s) :=
  inv_fromseq heap_ordered merge_HO_inv heap_ordered_E (heap_ordered_NodexEE measure1).

Definition leftist_inv_fromseq : forall s, leftist_inv (fromseq s) :=
  inv_fromseq leftist_inv merge_LI_inv leftist_inv_E (leftist_inv_NodexEE measure1).

Definition leftistheap_fromseq : forall s, leftistheap (fromseq s) :=
  inv_fromseq leftistheap merge_LH leftistheap_E leftistheap_NodexEE.

Equations fromheap h : seq T by wf (size h) lt:=
fromheap [||] := [::];
fromheap [tl| n, x |tr] := x :: (fromheap (deletemin [tl| n, x |tr])).

Next Obligation. rewrite merge_size; ssrnatlia. Qed.

Lemma fromheap_spec1 h: count^~ h =1 seq.count^~ (fromheap h).
Proof.
move=> p. apply_funelim (fromheap h)=> // ???? /= <-. by rewrite merge_spec; ssrnatlia.
Qed.

Lemma fromheap_srec1_in x h: (x \in h) = (x \in (fromheap h)).
Proof. by rewrite !in_count -has_pred1 has_count -fromheap_spec1. Qed.

Lemma fromheap_spec2 h: (heap_ordered h) -> sorted <=%O (fromheap h).
Proof.
apply_funelim (fromheap h)=> //= _ x h1 h2 IHh /and4P[L1 L2 H1 H2];
rewrite path_sortedE ?IHh ?all_count -?count_predT -?fromheap_spec1 ?merge_spec ?merge_HO_inv //.
- by rewrite (allE _ _ (LE_spec _ _ H1 L1)) (allE _ _ (LE_spec _ _ H2 L2)) !size_count // eq_refl.
by exact le_trans.
Qed.

Definition heapsort s := fromheap (fromseq s).

Theorem sorted_heapsort s : (sorted <=%O (heapsort s)).
Proof. by rewrite fromheap_spec2 // heap_ordered_fromseq. Qed.

Lemma count_heapsort s : seq.count^~ (heapsort s) =1 seq.count^~ s.
Proof. move=> p. by rewrite /heapsort -fromheap_spec1 fromseq_cspec. Qed.

Theorem perm_heapsort s : perm_eql (heapsort s) s.
Proof. apply/permPl/permP; by apply: count_heapsort. Qed.
End Heapsort.
End Measure.

Definition rank_measureMixin : Measure.mixin_of rank := Measure.Mixin T rank 1%N f1 rankNode rank__NodexEE rank_Node_E.
Canonical rank_measureType : measureType T := Measure.Pack T rank rank_measureMixin.

Definition size_measureMixin : Measure.mixin_of size := Measure.Mixin T size 1%N f2 sizeNode size__NodexEE size_Node_E.
Canonical size_measureType : measureType T := Measure.Pack T size size_measureMixin.

Definition leftist_rank_inv := (leftist_measure_inv rank_measureType).
Definition case_leftist_rank_inv := (case_leftist_measure_inv rank_measureType).
Definition case_leftist_rank_inv_r := (case_leftist_measure_inv_r rank_measureType).
Definition case_leftist_rank_inv_rl := (case_leftist_measure_inv_rl rank_measureType).

Section Grank.

(** The (general) rank of a tree is the length of the shortest path from the root to leaves *)
Fixpoint grank h : nat :=
  if h is Node _ _ tl tr then
    (minn (grank tl) (grank tr)).+1
  else 0.

Theorem grank_rk h : leftist_rank_inv  h -> grank h = rank h.
Proof.
elim: h=> // ??? IHh1 ? IHh2 /case_leftist_rank_inv/andP[].
rewrite /f/=/f1=>_ /and3P[??].
by rewrite IHh1 // IHh2 // minnC ?Order.NatOrder.minnE=> ->.
Qed.

End Grank.

(* We can obtain that right spine is the shoterst for heaps with *)
(* rank measure                                                  *)
Section Spine.

Lemma length_right_spine h s:
leftist_rank_inv h -> right s -> spine_in s h -> length s = rank h.
Proof.
move=> /andP[_ RC]. elim: h s RC=>[[|[]]|????? IHtr [_ _|[?/=/and3P[_ *]|]]] //.
apply/eqnP=> /=; apply/eqP. by apply: IHtr.
Qed.

Theorem rigth_spine_shortest H s1 s2:
 right s1 -> leftist_rank_inv H -> spine_in s1 H -> spine_in s2 H ->
(length s1 <= length s2)%nat.
Proof.
elim: H s1 s2=> [????/spine_in_E->|?? tl IHtl tr IHtr [//|a s1 [???
|[] s2 /rigth_correct [] ->]]] //= => [? /case_leftist_rank_inv_r ?|
? /case_leftist_rank_inv_rl /and3P[??/=?]]=>*.
- rewrite -addn1 -[(length s2).+1]addn1 leq_add2r IHtr //.
rewrite (length_right_spine tr s1) //. case: (right_spine_ex tl)=> s /andP[*].
suffices: (rank tl <= length s2)%N; first by ssrnatlia.
by rewrite -(length_right_spine tl s) // IHtl.
Qed.
End Spine.
End Specifications.
End leftistheap.


Require Export Vector VectorDef.

Delimit Scope vector_scope with vector_scope.

Notation "[||]" := (nil _) : vector_scope.
Notation "h ':::' t" := (cons _ h _ t) (at level 60, right associativity) : vector_scope.

Notation " [| x |] " := ((x ::: [||])%vector_scope) : vector_scope.
Notation " [| x ; y ; .. ; z |] " := (cons _ x _ (cons _ y _ .. (cons _ z _ (nil _)) ..)) : vector_scope.
Notation "v [@ p ]" := (nth v p) (at level 1, format "v [@ p ]").

Lemma Vector_replace_nth X n (v : Vector.t X n) i (x : X) :
  (Vector.replace v i x) [@i] = x.
Proof.
  induction i.
  - cbn. revert x. pattern v. revert n v. eapply Vector.caseS.
    cbn. reflexivity.
  - cbn. revert x i  IHi. pattern v. revert n v. eapply Vector.caseS.
    intros. cbn. eapply IHi.
Qed.

Lemma Vector_replace_nth2 X n (v : Vector.t X n) i j (x : X) :
  i <> j -> (Vector.replace v i x) [@j] = v[@j].
Proof.
  revert v. pattern i, j. revert n i j.
  eapply Fin.rect2; intros; try congruence.
  - revert f H. pattern v. revert n v.
    eapply Vector.caseS.
    cbn. reflexivity.
  - revert f H. pattern v. revert n v.
    eapply Vector.caseS.
    cbn. reflexivity.
  - revert g f H H0. pattern v. revert n v.
    eapply Vector.caseS. firstorder congruence.
Qed.


(* infotheo: information theory and error-correcting codes in Coq             *)
(* Copyright (C) 2020 infotheo authors, license: LGPL-2.1-or-later            *)
From mathcomp Require Import all_ssreflect ssralg fingroup finalg perm zmodp.
From mathcomp Require Import matrix.
From mathcomp Require boolp.
From mathcomp Require Import Rstruct.
Require Import Reals.
Require Import ssrR Reals_ext ssr_ext ssralg_ext logb Rbigop.
Require Import fdist proba entropy num_occ channel_code channel typ_seq.

(******************************************************************************)
(*          Elements of the theory of types (in information theory)           *)
(*                                                                            *)
(* P_n(A)                  == type                                            *)
(* type_counting           == Upper-bound of the number of types              *)
(* T_{P}                   == typed tuples, tuples that are representative of *)
(*                            a type                                          *)
(* tuple_dist_type_entropy == probability of tuples representative of a type  *)
(*                            using the entropy                               *)
(* card_typed_tuples       == Upper-bound of the number of tuples             *)
(*                            representative of a type using the entropy      *)
(*                                                                            *)
(******************************************************************************)

Reserved Notation "'P_' n '(' A ')'" (at level 9, n, A at next level).
Reserved Notation "'T_{' P '}'" (at level 9).
Reserved Notation "P '.-typed_code' c" (at level 50, c at next level).

Declare Scope types_scope.

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

Local Open Scope entropy_scope.
Local Open Scope num_occ_scope.
Local Open Scope R_scope.

Module type.

Section type_def.

Variable A : finType.
Variable n : nat.

Record type : predArgType := mkType {
  d :> fdist A ;
  f : {ffun A -> 'I_n.+1} ;
  d_f : forall a, d a = INR (f a) / INR n }.

End type_def.

End type.

Coercion type_coercion := type.d.

Notation "'P_' n '(' A ')'" := (type.type A n) : types_scope.

Local Open Scope types_scope.

Definition ffun_of_type A n (P : P_ n ( A )) := let: type.mkType _ f _ := P in f.

Lemma type_fun_type A n (_ : n != O) (P : P_ n ( A )) a :
  ((type.f P) a)%:R = n%:R * P a.
Proof.
case: P => /= d f d_f; by rewrite d_f mulRCA mulRV ?INR_eq0' // mulR1.
Qed.

Lemma INR_type_fun A n (P : P_ n ( A )) a : ((type.f P) a)%:R / n%:R = P a.
Proof. destruct P as [d f d_f] => /=. by rewrite d_f. Qed.

Lemma no_0_type A (d : fdist A) (t : {ffun A -> 'I_1}) :
  (forall a, d a = (t a)%:R / 0%:R) -> False.
Proof.
move=> H; apply R1_neq_R0.
rewrite -(FDist.f1 d).
transitivity (\sum_(a | a \in A) INR (t a) / 0); first exact/eq_bigr.
rewrite -big_distrl /= -big_morph_natRD.
rewrite (_ : (\sum_(a in A) _)%nat = O) ?mul0R //.
transitivity (\sum_(a in A) 0)%nat; first by apply eq_bigr => a _; rewrite (ord1 (t a)).
by rewrite big_const iter_addn.
Qed.

Definition type_of_tuple (A : finType) n (ta : n.+1.-tuple A) : P_ n.+1 ( A ).
set f := [ffun a => N(a | ta)%:R / n.+1%:R].
assert (H1 : forall a, (0 <= f a)%R).
  move=> a; rewrite ffunE; apply divR_ge0; by [apply leR0n | apply ltR0n].
have H2 : \sum_(a in A) f a = 1%R.
  under eq_bigr do rewrite ffunE /=.
  by rewrite -big_distrl /= -big_morph_natRD sum_num_occ_alt mulRV // INR_eq0'.
have H : forall a, (N(a | ta) < n.+2)%nat.
  move=> a; rewrite ltnS; by apply num_occ_leq_n.
refine (@type.mkType _ n.+1 (FDist.make H1 H2)
  [ffun a => @Ordinal n.+2 (N(a | ta)) (H a)] _).
by move=> a /=; rewrite !ffunE.
Defined.

Lemma type_ext A n (t1 t2 : P_ n ( A )) : type.f t1 = type.f t2 -> t1 = t2.
Proof.
case: t1 t2 => d1 f1 H1 /= [] d2 f2 H2 /= f1f2.
subst f2.
suff ? : d1 = d2 by subst d2; congr type.mkType; exact: boolp.Prop_irrelevance.
apply fdist_ext => /= a; by rewrite H1 H2.
Qed.

Definition type_eq A n (t1 t2 : P_ n ( A )) :=
  match t1, t2 with
    | type.mkType _ f1 _, type.mkType _ f2 _ => f1 == f2
  end.

Lemma type_eqP A n : Equality.axiom (@type_eq A n).
Proof.
case=> d1 f1 H1 [] d2 f2 H2 /=.
apply: (iffP idP) => [/eqP H|[] _ -> //].
subst f2.
suff ? : d1 = d2 by subst d2; congr type.mkType; exact: boolp.Prop_irrelevance.
apply fdist_ext => /= a; by rewrite H1 H2.
Qed.

Definition type_eqMixin A n := EqMixin (@type_eqP A n).
Canonical type_eqType A n := Eval hnf in EqType _ (@type_eqMixin A n).

Lemma type_ffunP A n (P Q : P_ n.+1 ( A )) :
  (forall c, P c = Q c) -> P = Q.
Proof.
move=> H.
destruct P as [d1 f1 H1].
destruct Q as [d2 f2 H2].
rewrite /= in H.
apply/type_eqP => /=.
apply/eqP/ffunP => a.
apply/val_inj/INR_eq.
move: {H}(H a); rewrite H1 H2 eqR_mul2r //.
apply/invR_neq0; by rewrite INR_eq0.
Qed.

Definition nneg_fun_of_ffun (A : finType) n (f : {ffun A -> 'I_n.+2}) : nneg_finfun A.
set d := [ffun a : A => INR (f a) / INR n.+1].
refine (@mkNNFinfun _ d _); apply/forallP_leRP => a.
by rewrite ffunE; apply divR_ge0; [apply leR0n | apply ltR0n].
Defined.

Definition fdist_of_ffun (A : finType) n (f : {ffun A -> 'I_n.+2})
  (Hf : (\sum_(a in A) f a)%nat == n.+1) : fdist A.
set pf := nneg_fun_of_ffun f.
have H : \sum_(a in A) pf a == 1 :> R.
  rewrite /pf; under eq_bigr do rewrite ffunE /=.
  rewrite /Rdiv -big_distrl /= -big_morph_natRD.
  by move/eqP : Hf => ->; rewrite mulRV // INR_eq0'.
exact:(FDist.mk H).
Defined.

Lemma fdist_of_ffun_prop (A : finType) n (f : {ffun A -> 'I_n.+2})
  (Hf : (\sum_(a in A) f a)%nat == n.+1) :
forall a : A, (fdist_of_ffun Hf) a = INR (f a) / INR n.+1.
Proof. by move=> a; rewrite ffunE. Qed.

Definition type_choice_f (A : finType) n (f : {ffun A -> 'I_n.+1}) : option (P_ n ( A )).
destruct n; first by exact None.
refine (match Sumbool.sumbool_of_bool (\sum_(a in A) f a == n.+1)%nat with
          | left H => Some (@type.mkType _ _ (fdist_of_ffun H) f (fdist_of_ffun_prop H))
          | right _ => None
        end).
Defined.

Lemma ffun_of_fdist A n (d : fdist A) (t : {ffun A -> 'I_n.+2})
  (H : forall a : A, d a = INR (t a) / INR n.+1) : (\sum_(a in A) t a)%nat == n.+1.
Proof.
suff : INR (\sum_(a in A) t a) == INR n.+1 * \sum_(a | a \in A) d a.
  by move/eqP; rewrite (FDist.f1 d) mulR1 => /INR_eq/eqP.
apply/eqP.
transitivity (INR n.+1 * (\sum_(a|a \in A) INR (t a) / INR n.+1)).
  by rewrite -big_distrl -big_morph_natRD mulRCA mulRV ?mulR1 // INR_eq0'.
congr (_ * _); exact/eq_bigr.
Qed.

Lemma type_choice_pcancel A n : pcancel (@type.f A n) (@type_choice_f A n).
Proof.
case=> d t H /=.
destruct n.
  by move: (no_0_type H).
rewrite /type_choice_f /=; f_equal.
move: (ffun_of_fdist H) => H'.
destruct Sumbool.sumbool_of_bool as [e|e]; last first.
  by rewrite H' in e.
congr Some.
set d1 := fdist_of_ffun _.
suff ? : d1 = d by subst d; congr type.mkType; apply boolp.Prop_irrelevance.
apply fdist_ext => /= a; by rewrite ffunE H.
Qed.

Lemma type_choiceMixin A n : choiceMixin (P_ n ( A )).
Proof. apply (PcanChoiceMixin (@type_choice_pcancel A n)). Qed.

Canonical type_choiceType A n := Eval hnf in ChoiceType _ (type_choiceMixin A n).

Definition type_pickle A n (P : P_ n (A)) : nat.
destruct P as [d f H].
exact: (pickle f).
(*destruct (finfun_countMixin A [finType of 'I_n.+1]) as [pi unpi Hcan].
apply (pi f).*)
Defined.

Definition type_unpickle A n (m : nat) : option (P_ n ( A )).
destruct n.
  exact None.
pose unpi : option {ffun A -> 'I_n.+2} := unpickle m.
case: unpi; last first.
  exact None.
move=> f.
refine (match Sumbool.sumbool_of_bool ((\sum_(a in A) f a)%nat == n.+1) with
          | left H => Some (@type.mkType _ _ (fdist_of_ffun H) f (fdist_of_ffun_prop H))
          | right _ => None
        end).
Defined.

Lemma type_count_pcancel A n : pcancel (@type_pickle A n) (@type_unpickle A n).
Proof.
destruct n.
  case=> d t H /=; by move: (no_0_type H).
case=> d t H /=.
rewrite pickleK.
move: (ffun_of_fdist H) => H'.
destruct Sumbool.sumbool_of_bool as [e|e]; last first.
  by rewrite H' in e.
congr Some.
set d1 := fdist_of_ffun _.
suff ? : d1 = d by subst d; congr type.mkType; apply boolp.Prop_irrelevance.
apply/fdist_ext => a; by rewrite ffunE H.
Qed.

Definition type_countMixin A n := CountMixin (@type_count_pcancel A n).
Canonical type_countType A n :=
  Eval hnf in CountType (P_ n ( A )) (@type_countMixin A n).

Definition type_enum_f (A : finType) n (f : { f : {ffun A -> 'I_n.+1} | (\sum_(a in A) f a)%nat == n} ) : option (P_ n ( A )).
destruct n.
  apply None.
refine (Some (@type.mkType _ _ (fdist_of_ffun (proj2_sig f)) (sval f) (fdist_of_ffun_prop (proj2_sig f)))).
Defined.

Definition type_enum A n := pmap (@type_enum_f A n)
  (enum [finType of {f : {ffun A -> 'I_n.+1} | (\sum_(a in A) f a)%nat == n}]).

Lemma type_enumP A n : Finite.axiom (@type_enum A n).
Proof.
destruct n.
  case=> d t H /=; by move: (no_0_type H).
case=> d t H /=.
move: (ffun_of_fdist H) => H'.
have : Finite.axiom (enum [finType of { f : {ffun A -> 'I_n.+2} | (\sum_(a in A) f a)%nat == n.+1}]).
  rewrite enumT; by apply enumP.
move/(_ (@exist {ffun A -> 'I_n.+2} (fun f => \sum_(a in A) f a == n.+1)%nat t H')) => <-.
rewrite /type_enum /= /type_enum_f /= count_map.
by apply eq_count.
Qed.

Definition type_finMixin A n := Eval hnf in FinMixin (@type_enumP A n).
Canonical type_finType A n := Eval hnf in FinType _ (@type_finMixin A n).

Section type_facts.
Variable A : finType.
Local Open Scope nat_scope.

Lemma type_counting n : #| P_ n ( A ) | <= expn (n.+1) #|A|.
Proof.
rewrite -(card_ord n.+1) -card_ffun /=.
rewrite cardE /enum_mem.
apply (@leq_trans (size (map (@ffun_of_type A n) (Finite.enum (type_finType A n))))).
  by rewrite 2!size_map.
rewrite cardE.
apply: uniq_leq_size.
  rewrite map_inj_uniq //.
    move: (enum_uniq (type_finType A n)).
    by rewrite enumT.
  case=> d f Hd [] d2 f2 Hd2 /= ?; subst f2.
  have ? : d = d2 by apply/fdist_ext => a; rewrite Hd Hd2.
  subst d2; congr type.mkType; exact: boolp.Prop_irrelevance.
move=> /= f Hf; by rewrite mem_enum.
Qed.

Lemma type_card_neq0 n : 0 < #|A| -> 0 < #|P_ n.+1(A)|.
Proof.
case/card_gt0P => a _.
apply/card_gt0P.
have [f Hf] : [finType of {f : {ffun A -> 'I_n.+2} | \sum_(a in A) f a == n.+1}].
  exists [ffun a1 => if pred1 a a1 then Ordinal (ltnSn n.+1) else Ordinal (ltn0Sn n.+1)].
  rewrite (bigD1 a) //= big1; first by rewrite ffunE eqxx addn0.
  move=> p /negbTE Hp; by rewrite ffunE Hp.
exists (@type.mkType _ _ (fdist_of_ffun Hf) _ (fdist_of_ffun_prop Hf)).
by rewrite inE.
Qed.

Lemma type_empty1 n : #|A| = 0 -> #|P_ n(A)| = 0.
Proof.
move=> A0; apply eq_card0; case=> d ? ?.
move: (fdist_card_neq0 d); by rewrite A0.
Qed.

Lemma type_empty2 : #|P_ 0(A)| = 0.
Proof.
apply eq_card0; case=> d f Hf.
exfalso.
by move/no_0_type in Hf.
Qed.

End type_facts.

Section typed_tuples.
Variables (A : finType) (n : nat) (P : P_ n ( A )).

Local Open Scope nat_scope.

Definition typed_tuples :=
  [set t : n.-tuple A | [forall a, P a == (INR N(a | t) / INR n)%R] ].

End typed_tuples.

Notation "'T_{' P '}'" := (typed_tuples P) : types_scope.

Section typed_tuples_facts.
Variables (A : finType) (n' : nat).
Let n := n'.+1.
Variable P : P_ n ( A ).

Lemma type_numocc ta (Hta : ta \in T_{P}) a : N(a | ta) = type.f P a.
Proof.
move: Hta.
rewrite in_set.
move/forallP/(_ a)/eqP.
destruct P as [d f H] => /= Htmp.
apply/INR_eq/esym; move: Htmp.
rewrite H eqR_mul2r //.
by apply/invR_neq0; rewrite INR_eq0.
Qed.

Lemma typed_tuples_not_empty' : exists x : seq A,
  exists Hx : size x == n, Tuple Hx \in T_{P}.
Proof.
exists (flatten (map (fun x0 => nseq (type.f P x0) x0) (enum A))).
have Hx : size (flatten [seq nseq (type.f P x0) x0 | x0 <- enum A]) == n.
  rewrite size_flatten /shape -map_comp sumn_big_addn big_map.
  case: (P) => P' f HP' /=.
  apply/eqP.
  transitivity (\sum_(a in A) f a)%nat; last first.
     apply/eqP; by apply ffun_of_fdist with P'.
  apply congr_big => //.
  by rewrite enumT.
  move=> a /= _.
  by rewrite size_nseq.
exists Hx.
rewrite inE.
apply/forallP => a.
rewrite /num_occ /= -size_filter.
rewrite filter_flatten size_flatten /shape -!map_comp sumn_big_addn big_map.
rewrite (bigD1 a) // big1 /= => [|a' Ha'].
- rewrite addn0 -(INR_type_fun P).
  apply/eqP.
  do 2 f_equal.
  rewrite -{1}(_ : size (nseq (type.f P a) a) = type.f P a); last by rewrite size_nseq.
  congr (size _).
  apply/esym/all_filterP/all_pred1P.
  by rewrite size_nseq.
- transitivity (size (@List.nil A)) => //.
  congr (size _).
  apply/eqP/negPn; rewrite -has_filter; apply/hasPn => l Hl.
  by case/nseqP : Hl => ->.
Qed.

Lemma typed_tuples_not_empty : { t | t \in T_{P} }.
Proof.
apply sigW.
case: typed_tuples_not_empty' => x [Hx H].
by exists (Tuple Hx).
Qed.

End typed_tuples_facts.

Section typed_tuples_facts_continued.
Variables (A : finType) (n : nat).
Hypothesis Hn : n != O.
Variable P : P_ n ( A ).

Lemma typed_tuples_not_empty_alt : {t : n.-tuple A | t \in T_{P}}.
Proof. destruct n => //. apply typed_tuples_not_empty. Qed.

Local Open Scope fdist_scope.
Local Open Scope tuple_ext_scope.
Local Open Scope vec_ext_scope.

Lemma tuple_dist_type t : tuple_of_row t \in T_{P} ->
  P `^ n t = \prod_(a : A) P a ^ (type.f P a).
Proof.
move=> Hx.
rewrite fdist_rVE.
rewrite (_ : \prod_(i < n) P (t ``_ i) =
  \prod_(a : A) (\prod_(i < n) (if a == t ``_ i then P t ``_ i else 1))); last first.
  rewrite exchange_big; apply eq_big ; first by [].
  move=> i _.
  rewrite (bigID (fun y => y == t ``_ i)) /=.
  rewrite -/(INR n.+1) big_pred1_eq eqxx big1 ?mulR1 //.
  by move=> i0 /negbTE ->.
apply eq_bigr => a _.
rewrite -big_mkcond /= -/(INR n.+1).
transitivity (\prod_(i < n | t ``_ i == a) (INR (type.f P a) / INR n)).
  apply eq_big => // i.
  move/eqP => ->.
  by rewrite INR_type_fun.
rewrite big_const iter_mulR INR_type_fun.
congr (_ ^ _).
rewrite /typed_tuples inE in Hx.
move/forallP/(_ a)/eqP : Hx.
rewrite -INR_type_fun eqR_mul2r; last by apply/invR_neq0; rewrite INR_eq0; exact/eqP.
move/INR_eq => ->.
rewrite num_occ_alt cardsE /=.
apply eq_card => /= n0.
by rewrite /in_mem /= tnth_mktuple.
Qed.

Local Close Scope tuple_ext_scope.

Lemma tuple_dist_type_entropy t : tuple_of_row t \in T_{P} ->
  P `^ n t = exp2 (- INR n * `H P).
Proof.
move/(@tuple_dist_type t) => ->.
rewrite (_ : \prod_(a : A) P a ^ (type.f P) a =
             \prod_(a : A) exp2 (P a * log (P a) * INR n)); last first.
  apply eq_bigr => a _.
  case/boolP : (0 == P a) => H; last first.
    have {}H : 0 < P a.
      have := FDist.ge0 P a.
      case/Rle_lt_or_eq_dec => // abs.
      by rewrite abs eqxx in H.
    rewrite -{1}(logK H) -exp2_pow.
    congr exp2.
    rewrite -mulRA [X in _ = X]mulRC -mulRA mulRC.
    congr (_ * _).
    by rewrite type_fun_type.
  - move/eqP : (H) => <-.
    rewrite -(_ : O = type.f P a); first by rewrite !mul0R exp2_0 /pow.
    apply INR_eq.
    rewrite {1}/INR.
    rewrite -(@eqR_mul2r ( / INR n)); last by apply/invR_neq0; rewrite INR_eq0; exact/eqP.
    by rewrite type_fun_type // -(eqP H) mulR0.
rewrite -(big_morph _ morph_exp2_plus exp2_0) -(big_morph _ (morph_mulRDl _) (mul0R _)).
by rewrite /entropy Rmult_opp_opp mulRC.
Qed.

Local Open Scope typ_seq_scope.

Lemma typed_tuples_are_typ_seq : (@row_of_tuple A n @: T_{ P }) \subset `TS P n 0.
Proof.
apply/subsetP => t Ht.
rewrite /set_typ_seq inE /typ_seq tuple_dist_type_entropy; last first.
  case/imsetP : Ht => x Hx ->.
  by rewrite row_of_tupleK.
by rewrite addR0 subR0 !leRR'.
Qed.

Lemma card_typed_tuples : INR #| T_{ P } | <= exp2 (INR n * `H P).
Proof.
rewrite -(invRK (exp2 (INR n * `H P))%R); last exact/eqP.
rewrite -exp2_Ropp -mulNR.
set aux := - INR n * `H P.
rewrite -div1R leR_pdivl_mulr // {}/aux.
case/boolP : [exists x, x \in T_{P}] => x_T_P.
- case/existsP : x_T_P => ta Hta.
  rewrite -(row_of_tupleK ta) in Hta.
  rewrite -(tuple_dist_type_entropy Hta).
  rewrite [X in X <= _](_ : _ = Pr P `^ n (@row_of_tuple A n @: T_{P})).
    by apply Pr_1.
  symmetry.
  rewrite /Pr.
  transitivity (\sum_(a| (a \in [finType of 'rV[A]_n]) && [pred x in (@row_of_tuple A n @: T_{P})] a)
      exp2 (- INR n * `H P)).
    apply eq_big => // ta'/= Hta'.
    rewrite -(@tuple_dist_type_entropy ta') //.
    case/imsetP : Hta' => x Hx ->. by rewrite row_of_tupleK.
  rewrite big_const iter_addR tuple_dist_type_entropy //.
  do 2 f_equal.
  by rewrite card_imset //; exact: row_of_tuple_inj.
- rewrite (_ : (INR #| T_{P} | = 0)%R); first by rewrite mul0R; exact/Rle_0_1.
  rewrite (_ : 0%R = INR 0) //; congr INR; apply/eqP.
  rewrite cards_eq0; apply/negPn.
  by move: x_T_P; apply contra => /set0Pn/existsP.
Qed.

Lemma card_typed_tuples_alt : INR #| T_{P} | <= exp2 (INR n * `H P).
Proof.
apply (@leR_trans (INR #| `TS P n 0 |)).
  apply/le_INR/leP.
  apply: leq_trans; last first.
    by apply subset_leq_card; exact: typed_tuples_are_typ_seq.
  by rewrite card_imset //; exact: row_of_tuple_inj.
by apply: (leR_trans (TS_sup _ _ _)); rewrite addR0; exact/leRR.
Qed.

Lemma perm_tuple_in_Ttuples ta (s : 'S_n) :
  perm_tuple s ta \in T_{P} <-> ta \in T_{P}.
Proof.
rewrite 2!in_set.
split => /forallP H; apply/forallP => a; move: H => /(_ a)/eqP => ->; by rewrite num_occ_perm.
Qed.

End typed_tuples_facts_continued.

Section enc_pre_img_partition.
Variables (A B M : finType) (n' : nat).
Let n := n'.+1.
Variable c : code A B M n.

Definition enc_pre_img (P : P_ n ( A )) := [set m | tuple_of_row (enc c m) \in T_{P}].

Lemma enc_pre_img_injective (P Q : P_ n ( A )) :
  enc_pre_img P != set0 -> enc_pre_img P = enc_pre_img Q ->
  forall a, P a = Q a.
Proof.
rewrite /enc_pre_img.
case/set0Pn => m.
rewrite in_set => HmP /setP HPQ.
have HmQ : tuple_of_row (enc c m) \in T_{Q} by move:(HPQ m) ; rewrite 2!in_set => <-.
move=> a {HPQ}.
move: HmP ; rewrite in_set => /forallP/(_ a)/eqP => ->.
move: HmQ ; rewrite in_set => /forallP/(_ a)/eqP => <-.
reflexivity.
Qed.

Definition enc_pre_img_partition :=
  enc_pre_img @: [set P in P_ n ( A ) | enc_pre_img P != set0].

Lemma cover_enc_pre_img : cover enc_pre_img_partition = [set: M].
Proof.
rewrite /cover /enc_pre_img_partition.
apply/setP => m.
rewrite in_set.
apply/bigcupP.
exists (enc_pre_img (type_of_tuple (tuple_of_row (enc c m)))).
- apply/imsetP; exists (type_of_tuple (tuple_of_row (enc c m))) => //.
  rewrite in_set.
  apply/andP; split => //.
  apply/set0Pn.
  exists m.
  rewrite 2!in_set.
  by apply/forallP => a; rewrite ffunE.
- rewrite 2!in_set.
  by apply/forallP => a; rewrite ffunE.
Qed.

Lemma trivIset_enc_pre_img : trivIset enc_pre_img_partition.
Proof.
apply/trivIsetP => S1 S2 /imsetP ; case => P1 _ HP1. case/imsetP => P2 _ HP2 HP12.
subst S1 S2.
rewrite /disjoint.
apply/pred0P => m /=.
apply/negP/negP.
move: m.
apply/forallP; rewrite -negb_exists; apply/negP; case/existsP => m /andP [H1 H2]; contradict HP12.
apply/negP/negPn/eqP/setP => m'.
case/boolP : (m' \in enc_pre_img P2) => [|/negbTE] Hcase.
- apply/negPn/negPn.
  rewrite 2!in_set; apply/forallP => a.
  move: H1; rewrite 2!in_set => /forallP/(_ a)/eqP ->.
  move: Hcase; rewrite 2!in_set => /forallP/(_ a)/eqP <-.
  move: H2; rewrite 2!in_set => /forallP/(_ a)/eqP <-.
  by rewrite eqxx.
- apply/negP/negPn; move: Hcase => /negP/negPn; apply contra => Hcase.
  rewrite 2!in_set; apply/forallP => a.
  move: H2; rewrite 2!in_set => /forallP/(_ a)/eqP ->.
  move: Hcase; rewrite 2!in_set => /forallP/(_ a)/eqP <-.
  move: H1; rewrite 2!in_set => /forallP/(_ a)/eqP <-.
  by rewrite eqxx.
Qed.

End enc_pre_img_partition.

Section sum_messages_types.
Variables (A B M : finType) (n' : nat).
Let n := n'.+1.
Variable c : code A B M n.

Lemma sum_messages_types' f :
  \sum_(P : P_ n ( A )) (\sum_(m |m \in enc_pre_img c P) f m) =
  \sum_ (S | S \in enc_pre_img_partition c) \sum_(m in S) f m.
Proof.
rewrite (bigID (fun P => [exists m, m \in enc_pre_img c P] )).
rewrite (_ : forall a b, addR_comoid a b = a + b) //.
rewrite Rplus_comm big1 ; last first.
  move=> P ; rewrite andTb negb_exists => HP.
  apply big_pred0 => m /=.
  apply/negP/negPn; by move:HP => /forallP/(_ m) ->.
rewrite /= add0R big_imset.
  apply eq_big => [P|P _] //=.
  rewrite in_set.
  case: set0Pn => [/existsP //| ?]; exact/existsP.
move=> P Q; rewrite 2!in_set => HP HQ HPQ /=.
move: (enc_pre_img_injective HP HPQ) => {HP HQ} {}HPQ.
case: P HPQ => /= Pd Pf HP HPQ.
case: Q HPQ => /= Qd Qf HQ HPQ.
apply/type_eqP => /=.
apply/eqP.
apply ffunP => a.
apply/val_inj/INR_eq.
move: {HPQ}(HPQ a); rewrite HP HQ eqR_mul2r //.
apply/invR_neq0; by rewrite INR_eq0.
Qed.

Lemma sum_messages_types f :
  \sum_(P : P_ n ( A )) (\sum_(m |m \in enc_pre_img c P) f m) = \sum_ (m : M) (f m).
Proof.
transitivity (\sum_ (m in [set: M]) (f m)); last by apply eq_bigl => b; rewrite in_set.
rewrite -(cover_enc_pre_img c) /enc_pre_img_partition sum_messages_types'.
symmetry.
by apply big_trivIset, trivIset_enc_pre_img.
Qed.

End sum_messages_types.

Section typed_code_def.
Variables (A B M : finType) (n : nat).
Variable P : P_ n ( A ).

Record typed_code := mkTypedCode {
  untyped_code :> code A B M n ;
  typed_prop : forall m, tuple_of_row (enc untyped_code m) \in T_{P} }.

End typed_code_def.

Section typed_code_of_code.
Variables (A B M : finType) (n' : nat).
Let n := n'.+1.
Variable P : P_ n ( A ).
Variable c : code A B M n.

Definition def := row_of_tuple (sval (typed_tuples_not_empty P)).
Definition Hdef := proj2_sig (typed_tuples_not_empty P).

Definition tcode_untyped_code := mkCode
  [ffun m => if tuple_of_row (enc c m) \in T_{P} then enc c m else def] (dec c).

Lemma tcode_typed_prop (m : M) : tuple_of_row ((enc tcode_untyped_code) m) \in T_{P}.
Proof.
rewrite /= ffunE; case: ifP => [//| _]; rewrite /def row_of_tupleK; exact Hdef.
Qed.

Definition tcode : typed_code B M P := mkTypedCode tcode_typed_prop.

End typed_code_of_code.

Notation "P '.-typed_code' c" := (tcode P c) : types_scope.


Require Import AutoSep WordLemmas.

(* ============================================================================
 * specification with overflow safeguard
 * - the result is correct and not overflowing if the input is 12 or less
 * ========================================================================= *)

Fixpoint fact (n : nat) :=
  match n with
    | 0 => 1
    | S n' => n * fact n'
  end.

Definition factS : spec := SPEC("n") reserving 1
  PRE[V]  [| (wordToNat (V "n") <= 12)%nat |]
  POST[R] [| wordToNat R = fact (wordToNat (V "n")) |].

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

Definition m := bmodule "factorial" {{
  bfunction "fact" ("n", "r" ) [factS]
    "r" <- 1;;
    [ PRE[V]  [| goodSize (wordToNat (V "r") * fact (wordToNat (V "n"))) |]
      POST[R] [| wordToNat R
                 = (wordToNat (V "r") * fact (wordToNat (V "n")))%nat |] ]
    While ("n" > 1) {
      "r" <- "r" * "n";;
      "n" <- "n" - 1
    };;
    Return "r"
  end
  }}.

(* ============================================================================
 * factorial up to 12 is not overflowing
 * - fact12 is too big to calculate, instead we calcuate factN 12 and deduce
 * ========================================================================= *)

Require Import NArith.

Fixpoint factN (n : nat) :=
  match n with
    | 0 => N.of_nat 1
    | S n' => Nmult (N.of_nat n) (factN n')
  end.

Lemma factN_12_lt_pow2_32 : (factN 12 < Npow2 32)%N.
  red; simpl; unfold Pos.compare; simpl; auto.
Qed.

Theorem factN_fact : forall n, N.to_nat (factN n) = fact n.
  induction n; simpl; auto.
  destruct (factN n) eqn:?; simpl in *.
  rewrite <- IHn; simpl; auto.
  rewrite <- IHn.
  rewrite Pnat.Pos2Nat.inj_mul, Pnat.SuccNat2Pos.id_succ; simpl; auto.
Qed.

Lemma fact_non_decreasing : forall x y, (x <= y)%nat -> (fact x <= fact y)%nat.
  induction 1; auto; simpl.
  destruct m0, x; simpl in *; auto;
  eapply Le.le_trans; try eassumption; apply Plus.le_plus_l.
Qed.
Local Hint Resolve fact_non_decreasing.

Lemma goodSize_def : forall x, (N.of_nat x < Npow2 32)%N -> goodSize x.
  auto.
Qed.

Lemma fact_bound : forall n, (n <= 12)%nat -> goodSize (fact n).
  intros.
  assert (goodSize (fact 12)).
  {
    apply goodSize_def.
    rewrite <- factN_fact, N2Nat.id.
    apply factN_12_lt_pow2_32.
  }
  eapply goodSize_weaken; eassumption || auto.
Qed.

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

Lemma fact_le_1 : forall r n : W, n <= natToW 1
          -> wordToNat r = (wordToNat r * fact (wordToNat n))%nat.
  intros; destruct_words; roundtrip.
  repeat (destruct w0; simpl; try omega).
Qed.
Hint Resolve fact_le_1.

Lemma fact_gt_0 : forall x, (0 < fact x)%nat.
  induction x; simpl; auto.
  rewrite Mult.mult_comm; simpl.
  generalize (fact x * x); intros; omega.
Qed.

Lemma fact_ge_1 : forall x, (1 <= fact x)%nat.
  induction x; simpl; auto.
  rewrite Mult.mult_comm; simpl.
  generalize (fact x * x); intros; omega.
Qed.
Local Hint Resolve fact_ge_1.

Lemma fact_ge : forall x, (x <= fact x)%nat.
  destruct x; simpl; auto.
  change (S x <= S x * fact x)%nat.
  rewrite <- Mult.mult_1_l at 1.
  rewrite Mult.mult_comm at 1.
  apply Mult.mult_le_compat_l; auto.
Qed.
Local Hint Resolve fact_ge.

Lemma rw1 : forall r n, natToW 1 < n
                 -> goodSize (wordToNat r * fact (wordToNat n))
                 -> wordToNat (r ^* n) * fact (wordToNat (n ^- natToW 1))
                    = wordToNat r * fact (wordToNat n).
  intros; destruct_words; roundtrip.
  destruct w0; simpl; try omega.
  replace (w0 - 0) with w0 by omega.
  assert (w * S w0 <= w * fact (S w0))%nat by (apply Mult.mult_le_compat_l; auto).
  rewrite wordToNat_wmult by (roundtrip; goodsize).
  roundtrip.
  rewrite <- Mult.mult_assoc; f_equal.
Qed.

Lemma rw2 : forall n : W, wordToNat (natToW 1) * fact (wordToNat n)
                          = fact (wordToNat n).
  intros; roundtrip; omega.
Qed.

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

Ltac finish :=
  match goal with
    | _ => solve [auto]
    | H: Regs _ Rv = _ |- _ => rewrite H
  end.

Hint Rewrite rw1 : sepFormula.
Hint Rewrite rw2 : sepFormula.
Hint Resolve fact_bound.

Theorem ok : moduleOk m.
  vcgen; sep_auto; repeat finish.
Qed.