Module DoNotation.
Notation "'do!' X <- A ; B" :=
(bindGen A (fun X => B))
(at level 200, X ident, A at level 100, B at level 200).
End DoNotation.
(* Sized generator of simple types *)
Fixpoint gen_type_size (n : nat) : G type :=
match n with
| 0 => returnGen N
| S n' =>
do! m <- choose (0, n');
liftGen2 Arrow (gen_type_size (n' - m)) (gen_type_size (n' - (n' - m)))
end.
(* Generator of simple types *)
Definition gen_type : G type := bindGen arbitrary gen_type_size.
(* Returns the list of bindings that have type tau in e *)
Definition vars_with_type (e : env) (tau : type) : list term :=
map (fun p => Id (snd p))
(filter (fun p => proj1_sig (Sumbool.bool_of_sumbool (type_eq_dec tau (fst p))))
(combine e (seq 0 (List.length e)))).
Definition sigT_of_prod {A B : Type} (p : A * B) : {_ : A & B} :=
let (a, b) := p in existT (fun _ : A => B) a b.
Definition lt_pair (c1 c2 : (nat * type)) : Prop :=
lexprod nat (fun _ => type) lt (fun _ => lt_type) (sigT_of_prod c1) (sigT_of_prod c2).
(* Generator of app-free well-typed terms of type tau *)
Fixpoint gen_term_no_app (tau : type)  (e : env) : G term :=
match vars_with_type e tau with
| [] =>
match tau with
| N => liftGen Const arbitrary
| Arrow tau1 tau2 =>
liftGen Abs (gen_term_no_app tau2 (tau1 :: e))
end
| def :: vars =>
oneOf_ (returnGen def)
[ match tau with
| N => liftGen Const arbitrary
| Arrow tau1 tau2 =>
liftGen Abs (gen_term_no_app tau2 (tau1 :: e))
end;
elems_ def (def :: vars)]
end.
(* Generator of well-typed terms of type tau. [fst p] is the maximum number of applications *)
Program Fixpoint gen_term_size (p : nat * type) {wf lt_pair p} : env -> G term :=
fun (e : env) => (* apparently with this trick we get a more manageable term *)
match p with
| (0, tau) => gen_term_no_app tau e
| (S n', tau) =>
match vars_with_type e tau with
| [] =>
oneOf_ (gen_term_no_app tau e)
[ (do! tau' <- gen_type;
do! m <- choose (0, n');
do! m' <- choose (n' -  m, n');
liftGen2 App (@gen_term_size (n' - m, (Arrow tau' tau)) _ e)
(@gen_term_size (n' - m', tau') _ e));
(match tau with
| N => liftGen Const arbitrary
| Arrow tau1 tau2 =>
liftGen Abs (@gen_term_size (S n', tau2) _ (tau1 :: e))
end)]
| def :: vars =>
oneOf_ (gen_term_no_app tau e)
[ (do! tau' <- gen_type;
do! m <- choose (0, n');
do! m' <- choose (n' - m, n');
liftGen2 App (@gen_term_size (n' - m, (Arrow tau' tau)) _ e)
(@gen_term_size (n' - m', tau') _ e));
(match tau with
| N => liftGen Const arbitrary
| Arrow tau1 tau2 =>
liftGen Abs (@gen_term_size (S n', tau2) _ (tau1 :: e))
end);
elems_ def (def :: vars) ]
end
end.
Solve Obligations with
program_simpl; unfold lt_pair; apply left_lex; omega.
Solve Obligations with
program_simpl; unfold lt_pair; apply right_lex; unfold lt_type; simpl; omega.
Next Obligation.
unfold MR. apply wf_inverse_image. apply wf_lt_pair.
Defined.
Definition gen_term_size_unfold (p : nat * type) (e : env) : G term :=
match p with
| (0, tau) => gen_term_no_app tau e
| (S n', tau) =>
match vars_with_type e tau with
| [] =>
oneOf_ (gen_term_no_app tau e)
[ (do! tau' <- gen_type;
do! m <- choose (0, n');
do! m' <- choose (n' - m, n');
liftGen2 App (gen_term_size (n' - m, (Arrow tau' tau)) e)
(gen_term_size (n' - m', tau') e));
(match tau with
| N => liftGen Const arbitrary
| Arrow tau1 tau2 =>
liftGen Abs (@gen_term_size (S n', tau2) (tau1 :: e))
end)]
| def :: vars =>
oneOf_ (gen_term_no_app tau e)
[ (do! tau' <- gen_type;
do! m <- choose (0, n');
do! m' <- choose (n' - m, n');
liftGen2 App (gen_term_size (n' - m, (Arrow tau' tau)) e)
(@gen_term_size (n' - m', tau') e));
(match tau with
| N => liftGen Const arbitrary
| Arrow tau1 tau2 =>
liftGen Abs (gen_term_size (S n', tau2) (tau1 :: e))
end);
elems_ def (def :: vars) ]
end
end.

Lemma gen_term_size_eq (e : env) (p : nat * type) :
gen_term_size p e =
gen_term_size_unfold p e.