diff --git a/gen.ml b/gen.ml index 61373034..bf01e6a1 100644 --- a/gen.ml +++ b/gen.ml @@ -767,6 +767,45 @@ let product a b = in next +let fair_product a b = + fun () -> + let gena = a () in + let genb = b () in + let all_a = ref [] in + let all_b = ref [] in + let cur = ref `GetLeft in + let rec next () = + match !cur with + | `Stop -> raise EOG + | `GetLeft -> + let xa = try Some (gena()) with EOG -> None in + begin match xa with + | None -> cur := `GetRight + | Some a -> all_a := a :: !all_a; cur := `ProdLeft (a, !all_b) + end; + next () + | `GetRight -> + let xb = try Some (genb()) with EOG -> None in + begin match xb with + | None -> cur := `Stop; raise EOG + | Some b -> all_b := b::!all_b; cur := `ProdRight (b, !all_a) + end; + next () + | `ProdLeft (_, []) -> + cur := `GetRight; + next() + | `ProdLeft (x, y::l) -> + cur := `ProdLeft (x, l); + x, y + | `ProdRight (_, []) -> + cur := `GetLeft; + next() + | `ProdRight (y, x::l) -> + cur := `ProdRight (y, l); + x, y + in + next + (** Group equal consecutive elements together. *) let group ?(eq=(=)) enum = fun () -> diff --git a/gen.mli b/gen.mli index a2dfaf33..1faabf17 100644 --- a/gen.mli +++ b/gen.mli @@ -266,7 +266,12 @@ val intersperse : 'a -> 'a t -> 'a t (** Put the separator element between all elements of the given enum *) val product : 'a t -> 'b t -> ('a * 'b) t - (** Cartesian product. *) + (** Cartesian product. If the first sequence is infinite, some pairs + will never be generated. *) + +val fair_product : 'a t -> 'b t -> ('a * 'b) t + (** Cartesian product, in no predictable order. Contrary to {!product} this + function does eventually yield every pair *) val group : ?eq:('a -> 'a -> bool) -> 'a t -> 'a list t (** Group equal consecutive elements together. *) diff --git a/tests/test_gen.ml b/tests/test_gen.ml index e13f1499..adf9feae 100644 --- a/tests/test_gen.ml +++ b/tests/test_gen.ml @@ -110,6 +110,11 @@ let test_product () = OUnit.assert_equal [1,4; 1,5; 2,4; 2,5; 3,4; 3,5] (Gen.to_list e); () +let test_fair_product () = + let e = Gen.fair_product (Gen.repeat ()) (1--3) in + let _ = Gen.take 10 e in (* succeeds -> ok *) + () + let suite = "test_gen" >::: [ "test_singleton" >:: test_singleton; @@ -127,4 +132,5 @@ let suite = "test_interleave" >:: test_interleave; "test_intersperse" >:: test_intersperse; "test_product" >:: test_product; + "test_fair_product" >:: test_fair_product; ]