Module Impl (GenLow : GenLowInterface.Sig) :
GenHighInterface.Sig GenLow.
Definition liftGen {A B} (f: A -> B) (a : G A)
: G B := nosimpl
(bindGen a (fun x =>
returnGen (f x))).
Definition liftGen2 {A1 A2 B}
(C : A1 -> A2 -> B) (m1 : G A1) (m2  : G A2)
: G B:=
nosimpl (bindGen m1 (fun x1 => bindGen m2 (fun x2 => returnGen (C x1 x2)))).
Definition liftGen3 {A1 A2 A3 R} (F : A1 -> A2 -> A3 -> R)
(m1 : G A1) (m2 : G A2) (m3 : G A3)
: G R := nosimpl(
bindGen m1 (fun x1 =>
bindGen m2 (fun x2 =>
bindGen m3 (fun x3 =>
returnGen (F x1 x2 x3))))).
Definition liftGen4 {A1 A2 A3 A4 R}
(F : A1 -> A2 -> A3 -> A4 -> R)
(m1 : G A1) (m2 : G A2) (m3 : G A3) (m4: G A4)
: G R := nosimpl(
bindGen m1 (fun x1 =>
bindGen m2 (fun x2 =>
bindGen m3 (fun x3 =>
bindGen m4 (fun x4 =>
returnGen (F x1 x2 x3 x4)))))).
Definition liftGen5 {A1 A2 A3 A4 A5 R}
(F : A1 -> A2 -> A3 -> A4 -> A5 -> R)
(m1 : G A1) (m2 : G A2) (m3 : G A3) (m4: G A4) (m5 : G A5)
: G R := nosimpl(
bindGen m1 (fun x1 =>
bindGen m2 (fun x2 =>
bindGen m3 (fun x3 =>
bindGen m4 (fun x4 =>
bindGen m5 (fun x5 =>
returnGen (F x1 x2 x3 x4 x5))))))).
Definition sequenceGen {A : Type} (ms : list (G A)) : G (list A) :=
foldr (fun m m' => bindGen m  (fun x =>
bindGen m' (fun xs =>
returnGen (x :: xs)))) (returnGen [::]) ms.
Fixpoint foldGen {A B : Type} (f : A -> B -> G A) (l : list B) (a : A)
: G A := nosimpl(
match l with
| [::] => returnGen a
| (x :: xs) => bindGen (f a x) (foldGen f xs)
end).
Definition oneOf_ {A : Type} (def: G A) (gs : list (G A)) : G A :=
bindGen (choose (0, length gs - 1)) (nth def gs).
Definition oneof {A} :=
@deprecate (G A -> list (G A) -> G A) "oneof" "oneOf_" oneOf_.
Fixpoint pick {A : Type} (def : G A) (xs : list (nat * G A)) n : nat * G A :=
match xs with
| nil => (0, def)
| (k, x) :: xs =>
if (n < k) then (k, x)
else pick def xs (n - k)
end.
(* This should use urns! *)
Fixpoint pickDrop {A : Type} (xs : list (nat * G (option A))) n : nat * G (option A) * list (nat * G (option A)) :=
match xs with
| nil => (0, returnGen None, nil)
| (k, x) :: xs =>
if (n < k) then  (k, x, xs)
else let '(k', x', xs') := pickDrop xs (n - k)
in (k', x', (k,x)::xs')
end.
Definition sum_fst {A : Type} (gs : list (nat * A)) : nat :=
foldl (fun t p => t + (fst p)) 0 gs.
Definition freq_ {A : Type} (def : G A) (gs : list (nat * G A))
: G A :=
let tot := sum_fst gs in
bindGen (choose (0, tot-1)) (fun n =>
@snd _ _ (pick def gs n)).
Definition frequency {A}:=
@deprecate (G A -> list (nat * G A) -> G A) "frequency" "freq_" freq_.
Fixpoint backtrackFuel {A : Type} (fuel : nat) (tot : nat) (gs : list (nat * G (option A))) : G (option A) :=
match fuel with
| O => returnGen None
| S fuel' => bindGen (choose (0, tot-1)) (fun n =>
let '(k, g, gs') := pickDrop gs n in
bindGen g (fun ma =>
match ma with
| Some a => returnGen (Some a)
| None => backtrackFuel fuel' (tot - k) gs'
end ))
end.
Definition backtrack {A : Type} (gs : list (nat * G (option A))) : G (option A) :=
backtrackFuel (length gs) (sum_fst gs) gs.
Definition vectorOf {A : Type} (k : nat) (g : G A)
: G (list A) :=
foldr (fun m m' =>
bindGen m (fun x =>
bindGen m' (fun xs => returnGen (cons x xs)))
) (returnGen nil) (nseq k g).
Definition listOf {A : Type} (g : G A) : G (list A) :=
sized (fun n => bindGen (choose (0, n)) (fun k => vectorOf k g)).
Definition elems_ {A : Type} (def : A) (l : list A) :=
let n := length l in
bindGen (choose (0, n - 1)) (fun n' =>
returnGen (List.nth n' l def)).
Definition elements {A} :=
@deprecate (A -> list A -> G A) "elements" "elems_" elems_.
Definition bindGenOpt {A B}
(g : G (option A)) (f : A -> G (option B)) : G (option B) :=
bindGen g (fun ma =>
match ma with
| None => returnGen None
| Some a => f a
end).
Definition retryBody {A : Type}
(retry : nat -> G (option A) -> G (option A))
(n : nat) (g : G (option A)) : G (option A) :=
bindGen g (fun x =>
match x, n with
| Some a, _ => returnGen (Some a)
| None, O => returnGen None
| None, S n' => retry n' g
end).
(* Rerun a generator [g] until it returns a [Some], or stop after [n+1] tries. *)
Fixpoint retry {A : Type} (n : nat) (g : G (option A)) :
G (option A) :=
retryBody retry n g.
(* Filter the output of a generator [g], returning [None] when the predicate [p] is [false]. The generator is run once. *)
Definition suchThatMaybe1 {A : Type} (g : G A) (p : A -> bool) :
G (option A) :=
fmap (fun x => if p x then Some x else None) g.
(* Retry a generator [g : G A] until it returns a value satisfying the predicate, or stop after [size+1] times, where [size] is the current size value. *)
Definition suchThatMaybe {A : Type} (g : G A) (p : A -> bool) :
G (option A) :=
sized (fun n => retry n (suchThatMaybe1 g p)).
(* Retry a generator [g : G (option A)] until it returns a value satisfying the predicate, or stop after [size+1] times, where [size] is the current size value. *)
Definition suchThatMaybeOpt {A : Type} (g : G (option A))
(p : A -> bool) : G (option A) :=
sized (fun n => retry n (fmap (fun x =>
match x with
| None => None
| Some a => if p a then Some a else None
end) g)).
(* Retry a generator until it returns a value, or stop after [size+1] times. *)
Definition retrySized {A : Type} (g : G (option A)) : G (option A) :=
sized (fun n => retry n g).
(* * Semantics *)
Ltac solveLiftGenX :=
intros; split; intros;
repeat
match goal with
| [ H : exists _, _ |- _ ] => destruct H as [? [? ?]]
| [ H : semGenSize _ _ _ |- _ ] =>
try (apply semBindSize in H; destruct H as [? [? ?]]);
try (apply semReturnSize in H; subst)
end;
[ by repeat (eexists; split; [eassumption |])
| repeat (apply semBindSize; eexists; split; try eassumption);
by apply semReturnSize ].
Program Instance liftGenUnsized {A B} (f : A -> B) (g : G A)
`{Unsized _ g} : Unsized (liftGen f g).
Next Obligation.
by rewrite ! semLiftGenSize (unsized s1 s2).
Qed.
Program Instance liftGenMonotonic {A B} (f : A -> B) (g : G A)
`{SizeMonotonic _ g} : SizeMonotonic (liftGen f g).
Next Obligation.
rewrite ! semLiftGenSize. apply imset_incl. by apply monotonic.
Qed.
Program Instance liftGen2Monotonic {A1 A2 B} (f : A1 -> A2 -> B) (g1 : G A1)
`{SizeMonotonic _ g1} (g2 : G A2) `{SizeMonotonic _ g2} :
SizeMonotonic (liftGen2 f g1 g2).
Next Obligation.
rewrite ! semLiftGen2Size. rewrite ! curry_imset2l.
move => b [a1 [Ha1 [a2 [Ha2 <-]]]].
do 2 (eexists; split; first by eapply (monotonic H1); eauto).
reflexivity.
Qed.

(* CH: Made this more beautiful than the rest *)
(* CH: Should anyway use dependent types for a generic liftGenN *)
Lemma semLiftGen4Size A1 A2 A3 A4 B (f : A1 -> A2 -> A3 -> A4 -> B)
(g1 : G A1) (g2 : G A2) (g3 : G A3) (g4 : G A4) s :
semGenSize (liftGen4 f g1 g2 g3 g4) s <-->
[set b : B | exists a1 a2 a3 a4, semGenSize g1 s a1 /\ semGenSize g2 s a2 /\
semGenSize g3 s a3 /\ semGenSize g4 s a4 /\ f a1 a2 a3 a4 = b].