fair cartesian product for Gen

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 () ->

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. *)

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;
     ]