Fix polymorphic generators in mutual recursive types

Author: vch9

src/ppx_deriving_qcheck/ppx_deriving_qcheck.ml

@@ -61,6 +61,12 @@ let pat ~loc ?sized s =
   let s = name ?sized s in
   A.pvar s
 
+let rec add_fuel ~loc expr =
+  match expr.pexp_desc with
+  | Pexp_fun (x, y, z, expr) ->
+    { expr with pexp_desc = Pexp_fun (x, y, z, add_fuel ~loc expr) }
+  | _ -> [%expr fun n -> [%e expr]]
+
 let rec curry_args ~loc args body =
   match args with
   | [] -> body
@@ -98,7 +104,7 @@ let sized_and_gen ~loc ?(args = []) typ gen patt =
   let gen_expr = curry_args ~loc args gen in
   let gen = [%stri let [%p patt] = [%e expr_sized]] in
   let gen_sized = [%stri let rec [%p pat_sized] = [%e gen_expr]] in
-  (gen_sized, Some gen)
+  (gen_sized, gen)
 
 
 let gen ~loc ?(env = TypeGen.empty) lg =
@@ -146,25 +152,19 @@ let sized ~loc ~env typ_name (is_rec : 'a -> bool)
 
 let mutually_recursive_gens ~loc typs gens =
   let (module A) = Ast_builder.make loc in
-  let rec aux l typs =
-    match (l, typs) with
-    | ([%stri let rec [%p? ppat] = [%e? expr]], gen) :: xs, typ :: typs
-    | ([%stri let [%p? ppat] = [%e? expr]], gen) :: xs, typ :: typs ->
-      (match gen with
-       | None ->
-         (* Current generator is not recursive. However, we want the generators
-            to share the fuel. Thus, we transform the non-recursive generator to
-            a recursive one. *)
-         let expr = [%expr fun n -> [%e expr]] in
-         let ppat' = pat ~loc ~sized:true typ in
-         let (_, gen) = sized_and_gen ~loc typ [%expr None] ppat in
-         aux (([%stri let [%p ppat'] = [%e expr]], gen) :: xs) (typ :: typs)
-       | Some x ->
-         (A.value_binding ~pat:ppat ~expr, x) :: aux xs typs)
-    | ([], []) -> []
-    | _ -> assert false
+
+  let gens =
+    List.map2 (fun x y -> match (x, y) with
+        | ([%stri let rec [%p? ppat] = [%e? expr]], `Needed call_to_size), _ ->
+          (A.value_binding ~pat:ppat ~expr, call_to_size)
+        | ([%stri let [%p? _] = [%e? expr]], `Useless call_to_size), typ ->
+          let expr = add_fuel ~loc expr in
+          let ppat = pat ~loc ~sized:true typ in
+          (A.value_binding ~pat:ppat ~expr, call_to_size)
+        | (_, _) -> assert false)
+      gens typs
   in
-  let (sized_gens, gens) = List.split (aux gens typs) in
+  let (sized_gens, gens) = List.split gens in
   let mutual_gens = A.pstr_value Recursive sized_gens in
   mutual_gens :: gens
 
@@ -343,8 +343,7 @@ and gen_from_constr ~loc ~env ty xs =
   | xs ->
     let args = List.map (gen_from_type ~loc ~env) xs in
     match gen with
-    | [%expr [%e? f] [%e? x]] ->
-      [%expr [%e apply_args ~loc args f] [%e x]]
+    | [%expr [%e? f] [%e? x]] -> [%expr [%e apply_args ~loc args f] [%e x]]
     | _ -> apply_args ~loc args gen
 
 and gen_from_constr_decl ~loc ?(env = TypeGen.empty)
@@ -490,18 +489,22 @@ let gen_from_type_declaration ~loc ?(env = TypeGen.empty) td =
       td.ptype_params
   in
 
+  let (sized, call_to_size) = sized_and_gen ~loc ~args typ gen_expr gen_pat in
   match gen_expr with
   | [%expr fun n -> [%e? _]] ->
-    sized_and_gen ~loc ~args typ gen_expr gen_pat
+    (sized, `Needed call_to_size)
   | _ ->
-    ([%stri let [%p gen_pat] = [%e curry_args ~loc args gen_expr]], None)
-
+    ([%stri let [%p gen_pat] = [%e curry_args ~loc args gen_expr]],
+      `Useless call_to_size)
 
 let derive_gen ~loc xs : structure =
   match xs with
   | _, [ x ] ->
     let (sized, gen) = gen_from_type_declaration ~loc x in
-    sized :: Option.fold ~none:[] ~some:(fun x -> [x]) gen
+    (match gen with
+     | `Needed gen ->
+       [sized; gen]
+     | _ -> [sized])
   | _, xs ->
     let (module A) = Ast_builder.make loc in
     let typs = List.map (fun td -> td.ptype_name.txt) xs in

test/ppx_deriving_qcheck/deriver/test.ml

@@ -545,14 +545,14 @@ let test_forest () =
   let expected =
     [
       [%stri
-        let rec gen_tree_sized n =
+        let rec gen_tree_sized gen_a n =
           QCheck.Gen.map
             (fun gen0 -> Node gen0)
               (QCheck.Gen.map
                  (fun (gen0, gen1) -> (gen0, gen1))
-                 (QCheck.Gen.pair QCheck.Gen.int (gen_forest_sized (n / 2))))
+                 (QCheck.Gen.pair gen_a ((gen_forest_sized gen_a)(n / 2))))
 
-        and gen_forest_sized n =
+        and gen_forest_sized gen_a n =
           match n with
           | 0 -> QCheck.Gen.pure Nil
           | n ->
@@ -564,20 +564,22 @@ let test_forest () =
                       (fun gen0 -> Cons gen0)
                       (QCheck.Gen.map
                          (fun (gen0, gen1) -> (gen0, gen1))
-                         (QCheck.Gen.pair (gen_tree_sized (n / 2)) (gen_forest_sized (n / 2))))
+                         (QCheck.Gen.pair
+                            ((gen_tree_sized gen_a) (n / 2))
+                            ((gen_forest_sized gen_a) (n / 2))))
                   );
                 ]];
-        [%stri let gen_tree = QCheck.Gen.sized @@ gen_tree_sized];
-        [%stri let gen_forest = QCheck.Gen.sized @@ gen_forest_sized];
+        [%stri let gen_tree gen_a = QCheck.Gen.sized @@ (gen_tree_sized gen_a)];
+        [%stri let gen_forest gen_a = QCheck.Gen.sized @@ (gen_forest_sized gen_a)];
     ]
   in
   let actual =
     f
     @@ extract
          [%stri
-           type tree = Node of (int * forest)
+           type 'a tree = Node of ('a * 'a forest)
 
-           and forest = Nil | Cons of (tree * forest)]
+           and 'a forest = Nil | Cons of ('a tree * 'a forest)]
   in
   check_eq ~expected ~actual "deriving forest"