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.
Definition ltColor (c1 c2: color) : Prop :=
match c1, c2 with
| Red, Black => True
| _, _ => False
end.
Definition sigT_of_prod {A B : Type} (p : A * B) : {_ : A & B} :=
let (a, b) := p in existT (fun _ : A => B) a b.
Definition prod_of_sigT {A B : Type} (p : {_ : A & B}) : A * B :=
let (a, b) := p in (a, b).
Definition wf_hc (c1 c2 : (nat * color)) : Prop :=
lexprod nat (fun _ => color) lt (fun _ => ltColor) (sigT_of_prod c1) (sigT_of_prod c2).
(* begin genRBTree_height *)
Program Fixpoint genRBTree_height (hc : nat*color) {wf wf_hc hc} : G tree :=
match hc with
| (0, Red) => returnGen Leaf
| (0, Black) => oneOf [returnGen Leaf;
(do! n <- arbitrary; returnGen (Node Red Leaf n Leaf))]
| (S h, Red) => liftGen4 Node (returnGen Black) (genRBTree_height (h, Black))
arbitrary (genRBTree_height (h, Black))
| (S h, Black) => do! c' <- genColor;
let h' := match c' with Red => S h | Black => h end in
liftGen4 Node (returnGen c') (genRBTree_height (h', c'))
arbitrary (genRBTree_height (h', c')) end.
(* end genRBTree_height *)
Next Obligation.
unfold wf_hc; simpl; left; omega.
Qed.
Next Obligation.
unfold wf_hc; simpl; left; omega.
Qed.
Next Obligation.
unfold wf_hc; simpl; destruct c'; [right; apply I | left; omega].
Qed.
Next Obligation.
unfold wf_hc; simpl; destruct c'; [right; apply I | left; omega].
Qed.
Next Obligation.
abstract (apply well_founded_hc).
Defined.

Lemma genRBTree_height_eq (hc : nat*color) :
genRBTree_height hc =
match hc with
| (0, Red) => returnGen Leaf
| (0, Black) => oneOf [returnGen Leaf;
(do! n <- arbitrary; returnGen (Node Red Leaf n Leaf))]
| (S h, Red) => liftGen4 Node (returnGen Black) (genRBTree_height (h, Black))
arbitrary (genRBTree_height (h, Black))
| (S h, Black) => do! c' <- genColor;
let h' := match c' with Red => S h | Black => h end in
liftGen4 Node (returnGen c') (genRBTree_height (h', c'))
arbitrary (genRBTree_height (h', c')) end.