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

From Coq Require Import
     Arith
     Basics
     List
     Recdef
     ZArith
     Lia
     ssreflect
     ssrbool.
From mathcomp Require Import
     ssrnat.
From QuickChick Require Import
     Classes
     Sets
     Producer
     Enumerators
     Generators
     Tactics.

Import ListNotations QcDefaultNotation.
Open Scope qc_scope.

Set Bullet Behavior "Strict Subproofs".

(** Basic generator instances *)
#[global] Instance genBoolSized : GenSized bool :=
  {| arbitrarySized x := elems_ true [true; false] |}.

#[global]
Instance genNatSized : GenSized nat :=
  {| arbitrarySized x := choose (0,x) |}.

#[global] Instance genZSized : GenSized Z :=
  {| arbitrarySized x := let z := Z.of_nat x in
                         choose (-z, z)%Z |}.

#[global] Instance genNSized : GenSized N :=
  {| arbitrarySized x := let n := N.of_nat x in
                         choose (N0, n) |}.

#[global] Instance genListSized {A : Type} `{GenSized A} : GenSized (list A) :=
  {| arbitrarySized x := listOf (arbitrarySized x) |}.

(* [3] is a lower priority than [Classes.GenOfGenSized],
   avoiding an infinite loop in typeclass resolution. *)

#[global] Instance genList {A : Type} `{Gen A} : Gen (list A) | 3 :=
  {| arbitrary := listOf arbitrary |}.

#[global] Instance genOption {A : Type} `{Gen A} : Gen (option A) | 3 :=
  {| arbitrary := freq [ (1, ret None)
                       ; (7, liftM Some arbitrary)] |}.

#[global] Instance genPairSized {A B : Type} `{GenSized A} `{GenSized B}
: GenSized (A*B) :=
  {| arbitrarySized x :=
       liftM2 pair (arbitrarySized x)
                   (arbitrarySized x)
  |}.

#[global] Instance genPair {A B : Type} `{Gen A} `{Gen B} : Gen (A * B) :=
  {| arbitrary := liftM2 pair arbitrary arbitrary |}.

(* Enumerator #[global] Instances *)

#[global] Instance enumBoolSized : EnumSized bool :=
  {| enumSized x := elems_ true [true; false] |}.

#[global] Instance enumNatSized : EnumSized nat :=
  {| enumSized x := chooseEnum (0,x) |}.

#[global] Instance enumZSized : EnumSized Z :=
  {| enumSized x := let z := Z.of_nat x in
                         chooseEnum (-z, z)%Z |}.

#[global] Instance enumNSized : EnumSized N :=
  {| enumSized x := let n := N.of_nat x in
                         chooseEnum (N0, n) |}.

#[global] Instance enumListSized {A : Type} `{Enum A} : EnumSized (list A) :=
  {| enumSized x := bindEnum (choose (0,x)) (fun x => vectorOf x enum) |}.

(* [3] is a lower priority than [Classes.EnumOfEnumSized],
   avoiding an infinite loop in typeclass resolution. *)

#[global] Instance enumList {A : Type} `{Enum A} : Enum (list A) | 3 :=
  {| enum := listOf enum |}.

#[global] Instance enumOption {A : Type} `{Enum A} : Enum (option A) | 3 :=
  {| enum := oneOf [ (ret None)
                   ; (liftM Some enum)] |}.

#[global] Instance enumPairSized {A B : Type} `{EnumSized A} `{EnumSized B}
: EnumSized (A*B) :=
  {| enumSized x :=
       liftM2 pair (enumSized x)
                   (enumSized x)
  |}.

#[global] Instance enumPair {A B : Type} `{Enum A} `{Enum B} : Enum (A * B) :=
  {| enum := liftM2 pair enum enum |}.


#[global] Instance enumOpt {A} (H : Enum A) : Enum (option A) :=
  {| enum :=
       match enum with
       | MkEnum f => MkEnum (fun n => LazyList.lcons None (fun _ => LazyList.mapLazyList Some (f n)))
       end
  |}.


(** Shrink#[global] Instances *)
#[global] Instance shrinkBool : Shrink bool :=
  {| shrink x :=
       match x with
         | false => nil
         | true  => cons false nil
       end
  |}.


#[global] Instance enumOptCorrect A `{EnumCorrect A} :
  Correct _ (@enum _ (enumOpt _)).
Proof.
  constructor.  
  intros t; split; eauto.
  - intros. exact I.
  - intros _. 
    simpl.
    inv H. inv H0. unfold semProd. simpl.
    destruct t.
    + destruct (prodCorrect a).
      destruct H3. now reflexivity.
      inv H3. simpl in *.
      eexists x. split. eassumption.
      destruct H. simpl in *. destruct enum0.
      unfold semEnumSize in *. simpl. right.
      eapply LazyList.lazy_in_map_iff.
      eexists. split. reflexivity.
      eassumption.
    + exists 0. split. reflexivity.
      destruct H. simpl. destruct enum0.
      unfold semEnumSize in *. simpl. left. reflexivity.
Qed.       

#[global] Instance enumOpt_SizeMonotonic A `{EnumMonotonic A} :
  SizeMonotonic (@enum _ (enumOpt _)).
Proof.
  destruct H. destruct enum. simpl in *. 
  intros s1 s2 Hleq [ x |]; simpl.
  - eapply H0 in Hleq. simpl in *.
    intros Hin.
    unfold semEnumSize in *. simpl in *. right.
    inv Hin; try congruence.
    eapply LazyList.lazy_in_map_iff.
    eexists. split. reflexivity.
    eapply Hleq.
    eapply LazyList.lazy_in_map_iff in H. destruct H.
    inv H. inv H2. eassumption.
  - unfold semEnumSize in *. simpl in *. firstorder.
Qed.

#[global] Instance enumOpt_SizeFP A `{Enum A} :
  SizeFP (@enum _ (enumOpt _)).
Proof.
  destruct H. destruct enum. simpl in *.
  
  intros s1 s2 Hleq Hnin. simpl.
  unfold semEnumSize in *. simpl in *. firstorder.
Qed.


#[global] Instance enumNatSized_CorrectSized :
  CorrectSized (@enumSized _ enumNatSized).
Proof.
  constructor. 
  intros t; split; eauto.
  - intros. exact I.
  - intros _. 
      assert (Hsize := @Enumerators.semChooseSizeEnumNat).
      simpl in *.
      exists t. exists t. split. exact I. eapply Hsize.
      reflexivity. ssromega.
Qed.

#[global] Instance enumNatSized_SizedMonotonic :
  SizedMonotonic (@enumSized _ enumNatSized).
Proof.
  intros s s1 s2 Hleq. simpl.
  intros x Hin.
  assert (Hsize := @Enumerators.semChooseSizeEnumNat).
  eapply Hsize in Hin; eauto.
  simpl in Hin.
  eapply Hsize.
  reflexivity.
  simpl.
  ssromega.
Qed.

#[global] Instance enumNatSized_SizeMonotonic s:
  SizeMonotonic (@enumSized _ enumNatSized s).
Proof.
  intros s1 s2 Hleq. simpl.
  intros x Hin.
  assert (Hsize := @Enumerators.semChooseSizeEnumNat).
  eapply Hsize in Hin; eauto.
  simpl in Hin.
  eapply Hsize. reflexivity. simpl.
  ssromega.
Qed.

(* TODO. These case be derived automatically. Change the default definitions *)

#[global] Instance enumListSized_SizedMonotonic A `{EnumMonotonic A} :
  SizedMonotonic (@enumSized (list A) enumListSized).
Proof.
  intros s s1 s2 Hleq x Hin.
  eapply semBindSizeEnum in Hin. destruct Hin as [a [Hin He]].
  assert (Hvec := @semVectorOfSize E _ _). eapply Hvec in He. destruct He as [Heq Hin'].
  eapply semBindSizeEnum. exists a. split.
  eapply Enumerators.semChooseSizeEnumNat in Hin. simpl in *.
  eapply Enumerators.semChooseSizeEnumNat. now eauto.
  simpl. now ssromega. now eauto.
  eapply Hvec. split; eauto.
Qed. 
 

#[global] Instance enumListSized_SizeMonotonic A `{EnumMonotonic A} s :
  SizeMonotonic (@enumSized (list A) enumListSized s).
Proof.
  eapply bindMonotonic; eauto with typeclass_instances.
Qed.


#[global] Instance enumListSized_CorrectSized A `{EnumMonotonicCorrect A} :
  CorrectSized (@enumSized (list A) enumListSized).
Proof.
  assert (Hvec := @semVectorOfSize E _ _).
  constructor. intros l; induction l; simpl.
  - split; intros Hin; simpl in *. now constructor.
    eexists 0. eexists 0. split; eauto.
    eapply semBindSizeEnum. eexists. split.
    eapply Enumerators.semChooseSizeEnum. now eauto.
    2:{ eapply Hvec. split. reflexivity. eapply sub0set. }
    now eauto.
  - split; intros Hin. now constructor.
    destruct IHl as [_ IHl].
    edestruct IHl as [x [s1 [Hin1 He]]]. now constructor.
    eapply semBindSizeEnum in He.
    destruct He as [y [Hen Hv]].
    eapply Hvec in Hv; eauto with typeclass_instances. destruct Hv as [Heq Hinl].
    subst.
    eapply Enumerators.semChooseSizeEnumNat in Hen; eauto. simpl in *.

    destruct H2. destruct H1. edestruct prodCorrect with (a := a).
    edestruct H2 as [s2 [ _ Hin'] ]. reflexivity.
      
    eexists (x + 1). eexists (s1 + s2). split. reflexivity.
    eapply semBindSizeEnum. exists (length l+1). split.
    eapply Enumerators.semChooseSizeEnumNat; eauto. simpl in *. ssromega.
    eapply Hvec. split. simpl. now ssromega.
    eapply cons_subset.

    eapply H0; [| eassumption ].  now ssromega.

    eapply subset_trans. eassumption.
    eapply H0. now ssromega. 
Qed.

#[global] Instance enumPairSized_SizedMonotonic A {_ : EnumSized A} { _ : SizedMonotonic (@enumSized A _)}
  B {_ : EnumSized B} { _ : SizedMonotonic (@enumSized B _)}:
  SizedMonotonic (@enumSized (A * B) enumPairSized).
Proof.
  intros s s1 s2 Hleq.
  simpl. rewrite !semBindSizeEnum.
  eapply incl_bigcup_compat; eauto.
  intros x.
  simpl. rewrite !semBindSizeEnum.
  eapply incl_bigcup_compat; eauto.
  intros y. eapply subset_refl.
Qed.

#[global] Instance enumPairSized_SizeMonotonic A {_ : EnumSized A} { _ : forall s, SizeMonotonic (@enumSized A _ s)}
         B {_ : EnumSized B} { _ : forall s, SizeMonotonic (@enumSized B _ s)} s :
  SizeMonotonic (@enumSized (A * B) enumPairSized s).
Proof.
  do 2 (eapply bindMonotonic; [ eauto with typeclass_instances .. | intros ? ]).
  eapply returnGenSizeMonotonic; eauto with typeclass_instances.
Qed.

#[global] Instance enumPairSized_CorrectSized
         A {_ : EnumSized A} { _ : forall s, SizeMonotonic (@enumSized A _ s)}
         { _ : SizedMonotonic (@enumSized A _)} { _ : CorrectSized (@enumSized A _)}
         B {_ : EnumSized B} { _ : forall s, SizeMonotonic (@enumSized B _ s)}
         { _ : SizedMonotonic (@enumSized B _)} { _ : CorrectSized (@enumSized B _)}:
  CorrectSized (@enumSized (A * B) enumPairSized).
Proof.
  constructor. split.
  { intros. reflexivity. }
  destruct a.
  intros _.
  edestruct H1 as [[_ Hca]].
  edestruct H4 as [[_ Hcb]].
  destruct Hca as [x1 [s1 [_ Hin1]]]. reflexivity.
  destruct Hcb as [x2 [s2 [_ Hin2]]]. reflexivity.
  eexists (x1 + x2). simpl. eexists (s1 + s2). split. reflexivity.
  simpl. eapply semBindSizeEnum.
  eexists. split.
  eapply H0. 2:{ eapply H; [| eassumption ]. ssromega. }
  ssromega.
  eapply semBindSizeEnum.
  eexists. split.
  eapply H2. 2:{ eapply H3; [| eassumption ]. ssromega. }
  ssromega.
  eapply semReturnSizeEnum. reflexivity. 
Qed.       


#[global] Instance enumOption_SizeMonotonic A {_ : Enum A} { _ : SizeMonotonic (@enum A _)} :
  SizeMonotonic (@enum (option A) enumOption).
Proof.
  simpl. eapply oneofMonotonic.
  - eauto with typeclass_instances.
  - eapply returnGenSizeMonotonic; eauto with typeclass_instances.
  - eapply cons_subset.
    eapply returnGenSizeMonotonic; eauto with typeclass_instances.
    eapply cons_subset.
    + eapply bindMonotonic; [ eauto with typeclass_instances .. | intros y ].
      eapply returnGenSizeMonotonic; eauto with typeclass_instances.
    + eapply sub0set.
Qed.

 
#[global] Instance enumOption_Correct A {_ : Enum A} { _ : Correct A (@enum A _)}:
  Correct _ (@enum (option A) enumOption).
Proof.
  constructor. split. intros. reflexivity.
  intros _.
  simpl. destruct a.
  - edestruct H as [ [ _ Hc ] ]. destruct Hc as [x [H1 H2]]. reflexivity.
    eexists x. split. reflexivity.
    eapply semOneofSize. eauto with typeclass_instances.
    eexists. split. right. left.
    eexists. eapply semBindSizeEnum. eexists. split. eassumption.
    eapply semReturnSizeEnum. reflexivity.
  - exists 0. split. reflexivity.
    eapply semOneofSize. eauto with typeclass_instances.
    eexists. split. left. reflexivity.
    eapply semReturnSizeEnum. reflexivity.
Qed.       

Lemma andb_len x1 x2 x3 x4 : 
  (x1 <=? x2)%num && (x3 <=? x4)%num <-> (x1 <= x2)%num /\ (x3 <= x4)%num.
Proof.
  destruct (x1 <=? x2)%num eqn:Heq1; simpl; split; try easy;
    destruct (x3 <=? x4)%num eqn:Heq2; simpl; try easy.
  - eapply N.leb_le in Heq1. eapply N.leb_le in Heq2.
    easy.
  - intros [H1 H2]. eapply N.leb_nle in Heq2.
    eauto.
  - intros [H1 H2]. eapply N.leb_nle in Heq1. eauto.
  - intros [H1 H2]. eapply N.leb_nle in Heq1. eauto.
Qed.   

#[global] Instance enumNSized_SizedMonotonic :
  SizedMonotonic (@enumSized _ enumNSized).
Proof. 
  intros s s1 s2 Hleq. simpl.
  intros x Hin.
  assert (Hsize := Enumerators.semChooseSizeEnum N).
  apply Hsize in Hin; [ | lia ].
  apply Hsize; lia.
Qed. 

#[global] Instance enumNSized_SizeMonotonic s:
  SizeMonotonic (@enumSized _ enumNSized s).
Proof.
  intros s1 s2 Hleq. simpl.
  intros x Hin.
  assert (Hsize := Enumerators.semChooseSizeEnum N).
  eapply Hsize in Hin; eauto; [ | lia].
  apply Hsize; lia.
Qed.

(* Lemma of_nat_bin t : *)
(*   (N.of_nat (nat_of_bin t)) = t. *)
(* Proof. *)
(*   destruct t. reflexivity. *)
(*   simpl. *)
(*   induction p; simpl. *)
(*   - admit. *)
(*   -  *)
  
#[global] Instance enumNSized_CorrectSized :
  CorrectSized (@enumSized _ enumNSized).
Proof.
  constructor. 
  intros t; split; eauto.
  - intros. exact I.
  - intros _. 
    assert (Hsize := Enumerators.semChooseSizeEnum N).
    exists (N.to_nat t). exists 0. split. exact I. eapply Hsize; lia.
Qed.
  
#[global] Instance enumBoolSized_SizeMonotonic s:
  SizeMonotonic (@enumSized _ enumBoolSized s).
Proof.
  intros s1 s2 Hleq.
  assert ( Heq := @semElementsSize E _ _).
  simpl in *. rewrite Heq.
  eapply subset_refl.
Qed.

#[global] Instance enumBoolSized_CorrectSized :
  CorrectSized (@enumSized _ enumBoolSized).
Proof.
  constructor. 
  intros t; split; eauto.
  - intros. exact I.
  - intros _. 
    destruct t; exists 0; simpl;
    eapply semElements; eauto with typeclass_instances.
    now left.
    right. now left.
Qed.       


#[global] Instance enumBoolSized_SizedMonotonic :
  SizedMonotonic (@enumSized _ enumBoolSized).
Proof.
  intros s s1 s2 Hleq. simpl.
  assert ( Heq := @semElementsSize E _ _).
  simpl in *. rewrite Heq.
  eapply subset_refl.
Qed.


(** Shrinking of nat starts to become annoying *)
Function shrinkNatAux (x : nat) {measure (fun x => x) x} : list nat :=
  match x with
    | O => nil
    | S n =>
      let x' := Nat.div x 2 in
      x' :: shrinkNatAux x'
  end.
Proof.
  move => x n Eq;
  pose proof (Nat.divmod_spec n 1 0 0) as H.
  assert (H' : (0 <= 1)%coq_nat) by lia; apply H in H';
  subst; simpl in *; clear H.
  destruct (Nat.divmod n 1 0 0) as [q u];  destruct u; simpl in *; lia.
Defined.

#[global] Instance shrinkNat : Shrink nat :=
  {| shrink := shrinkNatAux |}.

(** Shrinking of Z is even more so *)
Lemma abs_div2_pos : forall p, Z.abs_nat (Z.div2 (Z.pos p)) = Nat.div2 (Pos.to_nat p).
Proof.
  intros. destruct p.
    rewrite /Z.div2 /Pos.div2.
      rewrite /Z.abs_nat. rewrite Pos2Nat.inj_xI.
      rewrite <- Nat.add_1_r. rewrite Nat.mul_comm.
      rewrite Nat.div2_div. rewrite Nat.div_add_l; simpl; lia.
    rewrite /Z.div2 /Pos.div2.
      rewrite /Z.abs_nat. rewrite Pos2Nat.inj_xO.
      rewrite Nat.mul_comm. rewrite Nat.div2_div. rewrite Nat.div_mul. reflexivity. lia.
  reflexivity.
Qed.

Lemma neg_succ : forall p,
  Z.neg (Pos.succ p) = Z.pred (Z.neg p).
Proof.
  move => p. rewrite <- Pos.add_1_r.
  rewrite <- Pos2Z.add_neg_neg.
  rewrite <- Z.sub_1_r.
  reflexivity.
Qed.

Lemma neg_pred : forall p, (p > 1)%positive ->
  Z.neg (Pos.pred p) = Z.succ (Z.neg p).
Proof.
  move => p Hp. destruct p using Pos.peano_ind. by inversion Hp.
  rewrite Pos.pred_succ. rewrite neg_succ. rewrite Z.succ_pred.
  reflexivity.
Qed.

Lemma abs_succ_neg : forall p,
  Z.abs_nat (Z.succ (Z.neg p)) = Nat.pred (Pos.to_nat p).
Proof.
  move => p. destruct p using Pos.peano_ind. reflexivity.
  rewrite -neg_pred /Z.abs_nat. rewrite Pos2Nat.inj_pred. reflexivity.
  apply Pos.lt_1_succ. apply Pos.lt_gt; apply Pos.lt_1_succ.
Qed.

Lemma abs_succ_div2_neg : forall p,
  Z.abs_nat (Z.succ (Z.div2 (Z.neg p))) = Nat.div2 (Pos.to_nat p) \/
  Z.abs_nat (Z.succ (Z.div2 (Z.neg p))) = Nat.pred (Nat.div2 (Pos.to_nat p)).
Proof.
  intros. destruct p.
    left. rewrite /Z.div2 /Pos.div2.
      rewrite neg_succ. rewrite <- Zsucc_pred. rewrite /Z.abs_nat. rewrite Pos2Nat.inj_xI.
      rewrite <- Nat.add_1_r. rewrite Nat.mul_comm.
      rewrite Nat.div2_div. rewrite Nat.div_add_l; simpl; lia.
    right. rewrite /Z.div2 /Pos.div2.
      rewrite Pos2Nat.inj_xO. rewrite Nat.mul_comm.
      rewrite Nat.div2_div. rewrite Nat.div_mul. simpl.
      apply abs_succ_neg. lia.
  left. simpl. reflexivity.
Qed.

Function shrinkZAux (x : Z) {measure (fun x => Z.abs_nat x) x}: list Z :=
  match x with
    | Z0 => nil
    | Zpos _ => rev (cons (Z.pred x) (cons (Z.div2 x) (shrinkZAux (Z.div2 x))))
    | Zneg _ => rev (cons (Z.succ x) (cons (Z.succ (Z.div2 x)) (shrinkZAux (Z.succ (Z.div2 x)))))
  end.
Proof.
  move => ? p ?. subst. rewrite abs_div2_pos. rewrite Zabs2Nat.inj_pos.
    rewrite Nat.div2_div. apply Nat.div_lt. apply Pos2Nat.is_pos. lia.
  move => ? p ?. subst. destruct (abs_succ_div2_neg p) as [H | H].
    rewrite {}H /Z.abs_nat. rewrite Nat.div2_div.
      apply Nat.div_lt. apply Pos2Nat.is_pos. lia.
    rewrite {}H /Z.abs_nat. eapply Nat.le_lt_trans. apply Nat.le_pred_l. rewrite Nat.div2_div.
      apply Nat.div_lt. apply Pos2Nat.is_pos. lia.
Qed.

#[global] Instance shrinkZ : Shrink Z :=
  {| shrink := shrinkZAux |}.

Open Scope program_scope.

#[global] Instance shrinkN : Shrink N :=
  {| shrink := map Z.to_N ∘ shrink ∘ Z.of_N |}.

Definition shrinkListAux {A : Type} (shr : A -> list A) : list A -> list (list A) :=
  fix shrinkListAux_ (l : list A) : list (list A) :=
    match l with
    | nil => nil
    | cons x xs =>
      xs :: map (fun xs' => cons x xs') (shrinkListAux_ xs)
         ++ map (fun x'  => cons x' xs) (shr x )
    end.

#[global] Instance shrinkList {A : Type} `{Shrink A} : Shrink (list A) :=
  {| shrink := shrinkListAux shrink |}.

#[global] Instance shrinkPair {A B} `{Shrink A} `{Shrink B} : Shrink (A * B) :=
  {|
    shrink := fun (p : A * B) =>
      let (a,b) := p in
      map (fun a' => (a',b)) (shrink a) ++ map (fun b' => (a,b')) (shrink b)
  |}.

#[global] Instance shrinkOption {A : Type} `{Shrink A} : Shrink (option A) :=
  {| shrink m :=
       match m with
         | None => []
         | Some x => None :: (map Some (shrink x))
       end
  |}.

(** Arbitraries are derived automatically! *)


(**#[global] Instance correctness *)

(* Needed to add this! *)
Opaque semProdSize.

#[global]
Program Instance arbNatMon :
  @SizeMonotonic nat G ProducerGen (@arbitrary nat _).
Next Obligation.
  rewrite !semSizedSize !semChooseSizeNat // => n /andP [/leP H1 /leP H2].
  move : H => /leP => Hle. apply/andP.
  split; apply/leP; ssromega.
Qed.

(** Correctness proof about built-in generators *)

#[global]
Instance boolSizeMonotonic : SizeMonotonic (@arbitrary bool _).
Proof.
  unfold arbitrary, GenOfGenSized.
  eapply sizedSizeMonotonic; unfold arbitrarySized, genBoolSized.
  - eauto with typeclass_instances.
  - intros; eauto with typeclass_instances.
  - intros n s1 s2 Hs. eapply subset_refl.
Qed.

#[global]
Instance boolSizedMonotonic : SizedMonotonic (@arbitrarySized bool _).
Proof.
  intros n s1 s2 Hs. eapply subset_refl.
Qed.

#[global]
Instance boolCorrect : Correct bool arbitrary.
Proof.
  constructor. unfold arbitrary, GenOfGenSized.
  rewrite semSized.
  unfold arbitrarySized, genBoolSized.
  intros x. split; intros H; try now constructor.
  exists 0. split. constructor.
  eapply semElementsSize; eauto with typeclass_instances.
  destruct x; try solve [left; auto]; right; left; auto.
Qed.

Local Open Scope set_scope.

Lemma arbBool_correct:
  semProd arbitrary <--> [set: bool].
Proof.
rewrite /arbitrary /arbitrarySized /genBoolSized /=.
rewrite semSized => n; split=> // _.
exists n; split=> //.
apply semElementsSize => //=;
                           eauto with typeclass_instances.
destruct n; repeat (try solve [left; auto]; right).
Qed.

Lemma arbNat_correct:
  semProd arbitrary <--> [set: nat].
Proof.
rewrite /arbitrary /=.
rewrite semSized => n; split=> // _; exists n; split=> //.
by rewrite (semChooseSizeNat _ _ _) /=.
Qed.

#[global]
Instance ArbNatGenCorrect : Correct nat arbitrary.
Proof.
  constructor. now apply arbNat_correct.
Qed.

Lemma arbInt_correct s :
  semProdSize arbitrary s <-->
  [set z | (- Z.of_nat s <= z <= Z.of_nat s)%Z].
Proof.
rewrite semSizedSize semChooseSize.
  by move=> n; split=> [|] [? ?]. ssromega.
Qed.

Lemma arbBool_correctSize s :
  semProdSize arbitrary s <--> [set: bool].
Proof.
rewrite /arbitrary //=.
rewrite semSizedSize semElementsSize //; split=> _ //=; case a=> //=.
repeat (try solve [left; auto]; right).
repeat (try solve [left; auto]; right).
Qed.

Lemma arbNat_correctSize s :
  semProdSize arbitrary s <--> [set n : nat | (n <= s)%coq_nat].
Proof.
by rewrite semSizedSize semChooseSizeNat // => n /=; case: leP.
Qed.

Lemma arbInt_correctSize : semProd arbitrary <--> [set: Z].
Proof.
rewrite /arbitrarySized semSized => n; split=> // _; exists (Z.abs_nat n); split=> //.
simpl.
rewrite Nat2Z.inj_abs_nat (semChooseSize _ _ _); ssromega.
Qed.

Lemma arbList_correct:
  forall {A} `{H : Arbitrary A} (P : nat -> A -> Prop) s,
    (semProdSize arbitrary s <--> P s) ->
    (semProdSize arbitrary s <-->
     (fun (l : list A) => length l <= s /\ (forall x, List.In x l -> P s x))).
Proof.
  move => A G S H P s Hgen l.
  rewrite !/arbitrary /genList.
  split.
  - move => /semListOfSize [Hl Hsize]. split => // x HIn //=. apply Hgen. auto.
  - move => [Hl HP].
    apply semListOfSize; eauto with typeclass_instances.
    split => // x HIn.
    apply Hgen. auto.
Qed.

Opaque ret.
Opaque liftM.
Lemma arbOpt_correct:
  forall {A} `{H : Arbitrary A} (P : nat -> A -> Prop) s,
    (semProdSize arbitrary s <--> P s) ->
    (semProdSize arbitrary s <-->
     (fun (m : option A) =>
        match m with
          | None => true
          | Some x => P s x
        end)).
Proof.
  move => A G S Arb P s Hgen m. rewrite !/arbitrary /genOption; split.
  - move => /semFrequencySize [[w g] H2].
    move: H2 => [[H2 | [H2 | H2]] H3];
    destruct m => //=; apply Hgen => //=;
    inversion H2; subst; auto; simpl in *.
    + apply (@semReturnSize Generators.G ProducerGen ProducerSemanticsGen (option A) None s (Some a)) in H3; inversion H3.
    + apply (@semLiftProdSize _ ProducerGen _ _ _ Some (@arbitrary _ G) s (Some a)) in H3; eauto with typeclass_instances. inversion H3 as [x [H0 H1]].
      inversion H1; subst; auto.
  - destruct m eqn:Hm; simpl in *; move => HP; subst.
    + apply semFrequencySize; simpl.
      exists (7, liftM Some arbitrary); split; auto.
      * right; left; auto.
      * simpl. apply (@semLiftProdSize _ ProducerGen _ _ _ Some (@arbitrary _ G) s (Some a)); simpl;
                 eauto with typeclass_instances.
        apply imset_in; apply Hgen; auto.
    + apply semFrequencySize; simpl.
      exists (1, ret None); split; auto.
      * left; auto.
      * simpl.
        apply (@semReturnSize Generators.G ProducerGen ProducerSemanticsGen (option A) None s None). constructor.
Qed.

Lemma arbPair_correctSize
      {A B} `{Arbitrary A} `{Arbitrary B} (Sa : nat -> set A)
      (Sb : nat -> set B) s:
    (semProdSize arbitrary s <--> Sa s) ->
    (semProdSize arbitrary s <--> Sb s) ->
    (semProdSize arbitrary s <--> setX (Sa s) (Sb s)).
Proof.
  move => Hyp1 Hyp2 .
  simpl.
  rewrite semLiftProd2Size; move => [a b].
  split.
  by move => [[a' b'] [[/= /Hyp1 Ha /Hyp2 Hb] [Heq1 Heq2]]]; subst; split.
  move => [/Hyp1 Ha /Hyp2 Hb]. eexists; split; first by split; eauto.
  reflexivity.
Qed.
