fix Gen.{nat,pos}_split{2,}

fixes #180

Author: gasche

src/core/QCheck.ml

@@ -248,27 +248,27 @@ module Gen = struct
     List.sort (fun (w1, _) (w2, _) -> poly_compare w1 w2) samples |> List.rev_map snd
 
   let range_subset ~size low high st =
-    if not (low <= high && size <= high - low + 1) then invalid_arg "Gen.range_subset";
+    let range_size = high - low + 1 in
+    if not (0 <= size && size <= range_size) then
+      invalid_arg "Gen.range_subset";
     (* The algorithm below is attributed to Floyd, see for example
        https://eyalsch.wordpress.com/2010/04/01/random-sample/
        https://math.stackexchange.com/questions/178690
 
-       Note: the code be made faster by checking membership in [arr]
-       directly instead of using an additional Set. None of our
-       dependencies implements dichotomic search, so using Set is
-       easier.
+       Note: the code is easier to read when drawing from [0..range_size-1]
+       rather than [low..high]. We draw in [0..bound], and shift the
+       results by adding [low] when writing them to the result array.
     *)
     let module ISet = Set.Make(Int) in
     let s = ref ISet.empty in
-    let arr = Array.make size 0 in
-    for i = high - size to high do
-      let pos = int_range high i st in
-      let choice =
-        if ISet.mem pos !s then i else pos
-      in
-      arr.(i - low) <- choice;
+    for i = range_size - size to range_size - 1 do
+      let pos = int_range 0 i st in
+      let choice = if ISet.mem pos !s then i else pos in
       s := ISet.add choice !s;
     done;
+    let arr = Array.make size 0 in
+    let idx = ref 0 in
+    ISet.iter (fun choice -> arr.(!idx) <- low + choice; incr idx) !s;
     arr
 
   let array_subset size arr st =
@@ -334,31 +334,36 @@ module Gen = struct
 
   (* nat splitting *)
 
-  let nat_split2 n st =
-    if (n < 2) then invalid_arg "nat_split2";
+  let pos_split2 n st =
+    if (n < 2) then invalid_arg "pos_split2";
     let n1 = int_range 1 (n - 1) st in
     (n1, n - n1)
 
-  let pos_split2 n st =
+  let nat_split2 n st =
     let n1 = int_range 0 n st in
     (n1, n - n1)
 
   let pos_split ~size:k n st =
-    if (k > n) then invalid_arg "nat_split";
-    (* To split n into n{0}+n{1}+..+n{k-1}, we draw distinct "boundaries"
-       b{-1}..b{k-1}, with b{-1}=0 and b{k-1} = n
-       and the k-1 intermediate boundaries b{0}..b{k-2}
-       chosen randomly distinct in [1;n-1].
-
-       Then each n{i} is defined as b{i}-b{i-1}. *)
-    let b = range_subset ~size:(k-1) 1 (n - 1) st in
-    Array.init k (fun i ->
-      if i = 0 then b.(0)
-      else if i = k-1 then n - b.(i-1)
-      else b.(i) - b.(i-1)
-    )
+    if 0 = k && 0 = n then [||]
+    else begin
+      if not (0 < k && k <= n) then invalid_arg "pos_split";
+      (* To split n into n{0}+n{1}+..+n{k-1}, we draw distinct "boundaries"
+         b{-1}..b{k-1}, with b{-1}=0 and b{k-1} = n
+         and the k-1 intermediate boundaries b{0}..b{k-2}
+         chosen randomly distinct in [1;n-1].
+
+         Then each n{i} is defined as b{i}-b{i-1}. *)
+      let b = range_subset ~size:(k-1) 1 (n - 1) st in
+      if k = 1 then [|n|]
+      else Array.init k (fun i ->
+        if i = 0 then b.(0)
+        else if i = k-1 then n - b.(i-1)
+        else b.(i) - b.(i-1)
+      )
+    end
 
   let nat_split ~size:k n st =
+    if not (0 <= k && 0 <= n) then invalid_arg "nat_split";
     pos_split ~size:k (n+k) st
     |> Array.map (fun v -> v - 1)
 

src/core/QCheck.mli

@@ -467,22 +467,22 @@ module Gen : sig
 
       This is useful to split sizes to combine sized generators.
 
-      @raise Invalid_argument unless [n >= 2].
-
       @since 0.18
   *)
 
   val pos_split2 : int -> (int * int) t
-  (** [nat_split2 n] generates pairs [(n1, n2)] of strictly positive
+  (** [pos_split2 n] generates pairs [(n1, n2)] of strictly positive
       (nonzero) natural numbers with [n1 + n2 = n].
 
+      @raise Invalid_argument unless [n >= 2].
+
       This is useful to split sizes to combine sized generators.
 
       @since 0.18
   *)
 
   val nat_split : size:int -> int -> int array t
-  (** [nat_split2 ~size:k n] generates [k]-sized arrays [n1,n2,..nk]
+  (** [nat_split ~size:k n] generates [k]-sized arrays [n1,n2,..nk]
       of natural numbers in [[0;n]] with [n1 + n2 + ... + nk = n].
 
       This is useful to split sizes to combine sized generators.
@@ -493,15 +493,15 @@ module Gen : sig
   *)
 
   val pos_split : size:int -> int -> int array t
-  (** [nat_split2 ~size:k n] generates [k]-sized arrays [n1,n2,..nk]
+  (** [pos_split ~size:k n] generates [k]-sized arrays [n1,n2,..nk]
       of strictly positive (non-zero) natural numbers with
       [n1 + n2 + ... + nk = n].
 
       This is useful to split sizes to combine sized generators.
 
       Complexity O(k log k).
 
-      @raise Invalid_argument unless [k <= n].
+      @raise Invalid_argument unless [0 < k <= n] or [0 = k = n].
 
       @since 0.18
   *)

test/core/QCheck_expect_test.ml

@@ -83,7 +83,16 @@ module Overall = struct
          true)
 end
 
-(* positive tests of the various generators *)
+(* positive tests of the various generators
+
+   Note: it is important to disable shrinking for these tests, as the
+   shrinkers will suggest inputs that are coming from the generator
+   themselves -- which we want to test -- so their reduced
+   counter-example are confusing rather than helpful.
+
+   This is achieved by using (Test.make ~print ...), without a ~shrink
+   argument.
+*)
 module Generator = struct
   open QCheck
 
@@ -127,6 +136,50 @@ module Generator = struct
       ~name:"tree_rev_is_involutive"
       QCheck.(make IntTree.gen_tree)
       (fun tree -> IntTree.(rev_tree (rev_tree tree)) = tree)
+
+  let nat_split2_spec =
+    Test.make ~name:"nat_split2 spec"
+      (make
+        ~print:Print.(pair int (pair int int))
+         Gen.(small_nat >>= fun n ->
+              pair (return n) (nat_split2 n)))
+      (fun (n, (a, b)) ->
+         0 <= a && 0 <= b && a + b = n)
+
+  let pos_split2_spec =
+    Test.make ~name:"pos_split2 spec"
+      (make
+        ~print:Print.(pair int (pair int int))
+         Gen.(small_nat >>= fun n ->
+              (* we need n > 2 *)
+              let n = n + 2 in
+              pair (return n) (pos_split2 n)))
+      (fun (n, (a, b)) ->
+         (0 < a && 0 < b && a + b = n))
+
+  let nat_split =
+    Test.make ~name:"nat_split"
+      (make
+         ~print:Print.(pair int (array int))
+         Gen.(small_nat >>= fun n ->
+              pair (return n) (nat_split ~size:n n)))
+      (fun (n, arr) ->
+         Array.length arr = n
+         && Array.for_all (fun k -> 0 <= k) arr
+         && Array.fold_left (+) 0 arr = n)
+
+  let pos_split =
+    Test.make ~name:"pos_split"
+      (make
+         ~print:Print.(triple int int (array int))
+         Gen.(pair small_nat small_nat >>= fun (m, n) ->
+              (* we need both size>0 and n>0 and size <= n *)
+              let size = 1 + min m n and n = 1 + max m n in
+              triple (return size) (return n) (pos_split ~size n)))
+      (fun (m, n, arr) ->
+         Array.length arr = m
+         && Array.for_all (fun k -> 0 < k) arr
+         && Array.fold_left (+) 0 arr = n)
 end
 
 (* negative tests that exercise shrinking behaviour *)
@@ -407,6 +460,10 @@ let _ =
     Generator.list_repeat_test;
     Generator.array_repeat_test;
     Generator.passing_tree_rev;
+    Generator.nat_split2_spec;
+    Generator.pos_split2_spec;
+    Generator.nat_split;
+    Generator.pos_split;
     (*Shrink.test_fac_issue59;*)
     Shrink.big_bound_issue59;
     Shrink.long_shrink;

test/core/qcheck_output.txt.expected

@@ -863,7 +863,7 @@ stats dist:
    4150517416584649600.. 4611686018427387903: #################                                               189
 ================================================================================
 1 warning(s)
-failure (26 tests failed, 1 tests errored, ran 66 tests)
+failure (26 tests failed, 1 tests errored, ran 70 tests)
 random seed: 153870556
 
 +++ Stats for int_dist_empty_bucket ++++++++++++++++++++++++++++++++++++++++++++++++++++++++