Add function generation support

CHANGES.md

@@ -1,5 +1,7 @@
 # Unreleased
 
+- Add support for generating function values
+  [\#277](https://github.com/ocaml-gospel/ortac/pull/277)
 - Extend QCheck-STM plugin tests
   [\#271](https://github.com/ocaml-gospel/ortac/pull/271)
 - Remove unnecessary stack access from precondition

dune-project

@@ -157,7 +157,7 @@
   (ppxlib (>= 0.26.0))
   (mdx (and :with-test (>= 2.3.0)))
   (gospel (= 0.3.0))
-  (qcheck-stm (>= 0.4))
+  (qcheck-stm (>= 0.5))
   (ortac-core (= :version))
   (ortac-runtime-qcheck-stm (= :version)))
  (conflicts
@@ -193,7 +193,7 @@
   "Nikolaus Huber <[email protected]>")
  (depends
   (ocaml (>= 4.12.0))
-  (qcheck-stm (>= 0.4))
+  (qcheck-stm (>= 0.5))
   (ortac-runtime (= :version))))
 
 (package

ortac-qcheck-stm.opam

@@ -30,7 +30,7 @@ depends: [
   "ppxlib" {>= "0.26.0"}
   "mdx" {with-test & >= "2.3.0"}
   "gospel" {= "0.3.0"}
-  "qcheck-stm" {>= "0.4"}
+  "qcheck-stm" {>= "0.5"}
   "ortac-core" {= version}
   "ortac-runtime-qcheck-stm" {= version}
   "odoc" {with-doc}

ortac-runtime-qcheck-stm.opam

@@ -21,7 +21,7 @@ bug-reports: "https://github.com/ocaml-gospel/ortac/issues"
 depends: [
   "dune" {>= "3.8"}
   "ocaml" {>= "4.12.0"}
-  "qcheck-stm" {>= "0.4"}
+  "qcheck-stm" {>= "0.5"}
   "ortac-runtime" {= version}
   "odoc" {with-doc}
 ]

plugins/qcheck-stm/doc/example_limitations.mli

@@ -20,7 +20,8 @@ val of_list : 'a list -> 'a t
 val g : 'a t -> 'a * 'a * 'a * 'a * 'a * 'a * 'a * 'a * 'a * 'a -> bool
 (*@ b = g t x *)
 
-val for_all : ('a -> bool) -> 'a t -> bool
-(*@ b = for_all p t *)
+val large_second_order_arity :
+  ('a -> 'a -> 'a -> 'a -> 'a -> bool) -> 'a t -> bool
+(*@ b = large_second_order_arity p t *)
 
 (* $MDX part-end *)

plugins/qcheck-stm/doc/index.mld

@@ -392,7 +392,8 @@ val for_all : 'a t -> bool
 
 Finally, note that this tool is still fairly new and comes with limitations
 that should be lifted in the future. Fow now, we only support tuples with less
-than 10 elements and we only support first-order functions.
+than 10 elements and we only support first-order functions and second-order
+functions up to arity four.
 
 If we add the following declarations to our example file,
 
@@ -403,8 +404,9 @@ val of_list : 'a list -> 'a t
 val g : 'a t -> 'a * 'a * 'a * 'a * 'a * 'a * 'a * 'a * 'a * 'a -> bool
 (*@ b = g t x *)
 
-val for_all : ('a -> bool) -> 'a t -> bool
-(*@ b = for_all p t *)
+val large_second_order_arity :
+  ('a -> 'a -> 'a -> 'a -> 'a -> bool) -> 'a t -> bool
+(*@ b = large_second_order_arity p t *)
 ]}
 
 Ortac/QCheck-STM will generate the following warnings:
@@ -423,10 +425,11 @@ File "example_limitations.mli", line 20, characters 16-63:
 20 | val g : 'a t -> 'a * 'a * 'a * 'a * 'a * 'a * 'a * 'a * 'a * 'a -> bool
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 Warning: Skipping g: Can only test tuples with arity < 10.
-File "example_limitations.mli", line 23, characters 15-25:
-23 | val for_all : ('a -> bool) -> 'a t -> bool
-                    ^^^^^^^^^^
-Warning: Skipping for_all: functions are not supported yet as arguments.
+File "example_limitations.mli", line 24, characters 3-37:
+24 |   ('a -> 'a -> 'a -> 'a -> 'a -> bool) -> 'a t -> bool
+        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Warning: Skipping large_second_order_arity: Can only test function arguments
+         with arity < 5.
 ]}
 
 {2:returning-sut Functions returning a SUT}

plugins/qcheck-stm/src/ir.ml

@@ -44,6 +44,7 @@ type value = {
   sut_vars : Ident.t list;
       (* invariant: suts must be in the order in which they appear, so for
          example in [test (t1 : sut) (t2 : sut)] the list must be [t1; t2] *)
+  fun_vars : Ident.t list;
   args : (Ppxlib.core_type * Ident.t option) list;
       (* arguments of unit types can be nameless *)
   ret : Ident.t list;
@@ -67,12 +68,14 @@ let get_return_type value =
   in
   aux value.ty
 
-let value id ty inst sut_vars args ret ret_values next_states precond postcond =
+let value id ty inst sut_vars fun_vars args ret ret_values next_states precond
+    postcond =
   {
     id;
     ty;
     inst;
     sut_vars;
+    fun_vars;
     args;
     ret;
     ret_values;

plugins/qcheck-stm/src/ir_of_gospel.ml

@@ -10,22 +10,58 @@ let constant_test vd =
       (Constant_value (Fmt.str "%a" Ident.pp vd.vd_name), vd.vd_loc) |> error
   | _ -> ok ()
 
-let no_functional_arg_or_big_tuple vd =
+let no_third_order_fun_or_big_tuple vd =
   let open Reserr in
   let open Ppxlib in
   let rec contains_arrow ty =
     match ty.ptyp_desc with
-    | Ptyp_arrow (_, _, _) ->
-        error (Functional_argument vd.vd_name.id_str, ty.ptyp_loc)
-    | Ptyp_tuple xs when List.length xs > 9 ->
-        error (Tuple_arity vd.vd_name.id_str, ty.ptyp_loc)
-    | Ptyp_tuple xs | Ptyp_constr (_, xs) -> traverse_ contains_arrow xs
-    | _ -> ok ()
+    | Ptyp_arrow (_, _, _) -> true
+    | Ptyp_tuple xs | Ptyp_constr (_, xs) -> List.exists contains_arrow xs
+    | Ptyp_alias (t, _) -> contains_arrow t
+    | _ -> false
+  in
+  let rec contains_nested_arrow ty =
+    match ty.ptyp_desc with
+    | Ptyp_arrow (_, l, r) -> contains_arrow l || contains_nested_arrow r
+    | Ptyp_tuple xs | Ptyp_constr (_, xs) ->
+        List.exists contains_nested_arrow xs
+    | Ptyp_alias (t, _) -> contains_nested_arrow t
+    | _ -> false
+  in
+  let rec highest_arity_arrow ty =
+    match ty.ptyp_desc with
+    | Ptyp_arrow (_, _l, r) -> 1 + highest_arity_arrow r
+    | Ptyp_tuple xs | Ptyp_constr (_, xs) ->
+        List.fold_left max 0 (List.map highest_arity_arrow xs)
+    | Ptyp_alias (t, _) -> highest_arity_arrow t
+    | _ -> 0
+  in
+  let rec contains_big_tuple ty =
+    match ty.ptyp_desc with
+    | Ptyp_arrow (_, l, r) -> contains_big_tuple l || contains_big_tuple r
+    | Ptyp_tuple xs when List.length xs > 9 -> true
+    | Ptyp_tuple xs | Ptyp_constr (_, xs) -> List.exists contains_big_tuple xs
+    | Ptyp_alias (t, _) -> contains_big_tuple t
+    | _ -> false
   in
   let rec aux ty =
     match ty.ptyp_desc with
     | Ptyp_arrow (_, l, r) ->
-        let* _ = contains_arrow l in
+        let* _ =
+          if contains_big_tuple l then
+            error (Tuple_arity vd.vd_name.id_str, l.ptyp_loc)
+          else ok ()
+        in
+        let* _ =
+          if contains_nested_arrow l then
+            error (Third_order_function_argument vd.vd_name.id_str, l.ptyp_loc)
+          else ok ()
+        in
+        let* _ =
+          if highest_arity_arrow l > 4 then
+            error (Function_arity vd.vd_name.id_str, l.ptyp_loc)
+          else ok ()
+        in
         aux r
     | _ -> ok ()
   in
@@ -58,6 +94,8 @@ let unify case sut_ty ty =
           | _, Ptyp_any -> ok i
           | _, Ptyp_var a -> add_if_needed a x i
           | Ptyp_tuple xs, Ptyp_tuple ys -> aux i (xs, ys)
+          | Ptyp_arrow (_, l, r), Ptyp_arrow (_, l', r') ->
+              aux i ([ l; r ], [ l'; r' ])
           | Ptyp_constr (c, xs), Ptyp_constr (d, ys) when c.txt = d.txt ->
               aux i (xs, ys)
           | _ -> fail ()
@@ -128,20 +166,22 @@ let split_args config vd args =
   let open Ppxlib in
   let open Reserr in
   let module Ident = Identifier.Ident in
-  let rec aux suts acc ty args =
+  let rec aux suts funs acc ty args =
     match (ty.ptyp_desc, args) with
     | _, Lghost vs :: _ ->
         error (Ghost_values (vd.vd_name.id_str, `Arg), vs.vs_name.id_loc)
     | Ptyp_arrow (_, l, r), Lnone vs :: xs
     | Ptyp_arrow (_, l, r), Loptional vs :: xs
     | Ptyp_arrow (_, l, r), Lnamed vs :: xs ->
-        if Cfg.is_sut config l then aux (vs.vs_name :: suts) acc r xs
-        else aux suts ((l, Some vs.vs_name) :: acc) r xs
-    | Ptyp_arrow (_, l, r), Lunit :: xs -> aux suts ((l, None) :: acc) r xs
-    | _, [] -> ok (List.rev suts, List.rev acc)
+        if Cfg.is_sut config l then aux (vs.vs_name :: suts) funs acc r xs
+        else if is_a_function l then
+          aux suts (vs.vs_name :: funs) ((l, Some vs.vs_name) :: acc) r xs
+        else aux suts funs ((l, Some vs.vs_name) :: acc) r xs
+    | Ptyp_arrow (_, l, r), Lunit :: xs -> aux suts funs ((l, None) :: acc) r xs
+    | _, [] -> ok (List.rev suts, List.rev funs, List.rev acc)
     | _, _ -> failwith "shouldn't happen (too few parameters)"
   in
-  aux [] [] vd.vd_type args
+  aux [] [] [] vd.vd_type args
 
 let get_state_description_with_index is_t state spec =
   let open Tterm in
@@ -274,12 +314,12 @@ let postcond spec =
 
 let val_desc config state vd =
   let open Reserr in
-  let* () = constant_test vd and* () = no_functional_arg_or_big_tuple vd in
+  let* () = constant_test vd and* () = no_third_order_fun_or_big_tuple vd in
   let* inst = ty_var_substitution config vd
   and* spec =
     of_option ~default:(No_spec vd.vd_name.id_str, vd.vd_loc) vd.vd_spec
   in
-  let* suts, args = split_args config vd spec.sp_args
+  let* suts, fun_vars, args = split_args config vd spec.sp_args
   and* ret =
     let p = function
       | Lnone vs -> ok vs.vs_name
@@ -300,8 +340,8 @@ let val_desc config state vd =
     | None -> next_states
   in
   let postcond = postcond spec in
-  Ir.value vd.vd_name vd.vd_type inst suts args ret ret_values next_states
-    spec.sp_pre postcond
+  Ir.value vd.vd_name vd.vd_type inst suts fun_vars args ret ret_values
+    next_states spec.sp_pre postcond
   |> ok
 
 let sig_item config state s =

plugins/qcheck-stm/src/reserr.ml

@@ -13,6 +13,7 @@ type W.kind +=
   | Constant_value of string
   | Ensures_not_found_for_next_state of (string * string)
   | Ensures_not_found_for_ret_sut of (string * string list)
+  | Function_arity of string
   | Functional_argument of string
   | Ghost_values of (string * [ `Arg | `Ret ])
   | Impossible_init_state_generation of init_state_error
@@ -34,6 +35,7 @@ type W.kind +=
   | Sut_type_not_specified of string
   | Sut_type_not_supported of string
   | Syntax_error_in_config_module of string
+  | Third_order_function_argument of string
   | Tuple_arity of string
   | Type_not_supported of string
   | Type_not_supported_for_sut_parameter of string
@@ -42,11 +44,11 @@ type W.kind +=
 let level kind =
   match kind with
   | Constant_value _ | Ensures_not_found_for_next_state _
-  | Ensures_not_found_for_ret_sut _ | Functional_argument _ | Ghost_values _
-  | Impossible_term_substitution _ | Incompatible_type _
+  | Ensures_not_found_for_ret_sut _ | Function_arity _ | Functional_argument _
+  | Ghost_values _ | Impossible_term_substitution _ | Incompatible_type _
   | Incomplete_ret_val_computation _ | No_spec _ | Returning_nested_sut _
-  | Sut_as_type_inst _ | Sut_in_tuple _ | Tuple_arity _ | Type_not_supported _
-    ->
+  | Sut_as_type_inst _ | Sut_in_tuple _ | Third_order_function_argument _
+  | Tuple_arity _ | Type_not_supported _ ->
       W.Warning
   | Impossible_init_state_generation _ | Incompatible_sut _
   | Incomplete_configuration_module _ | No_configuration_file _
@@ -79,6 +81,9 @@ let pp_kind ppf kind =
         models text
         "where x is the returned SUT and expr can refer to other SUTs only \
          under an old operator"
+  | Function_arity fct ->
+      pf ppf "Skipping %s:@ %a" fct text
+        "Can only test function arguments with arity < 5"
   | Functional_argument f ->
       pf ppf "Skipping %s:@ %a" f text
         "functions are not supported yet as arguments"
@@ -120,6 +125,8 @@ let pp_kind ppf kind =
              next_state function"
       in
       pf ppf "Skipping clause:@ %a" text msg
+  | Third_order_function_argument name ->
+      pf ppf "Third-order functions and above are not supported in %s" name
   | Tuple_arity fct ->
       pf ppf "Skipping %s:@ %a" fct text "Can only test tuples with arity < 10"
   (* This following message is broad and used in seemingly different contexts

plugins/qcheck-stm/src/reserr.mli

@@ -13,6 +13,7 @@ type W.kind +=
   | Constant_value of string
   | Ensures_not_found_for_next_state of (string * string)
   | Ensures_not_found_for_ret_sut of (string * string list)
+  | Function_arity of string
   | Functional_argument of string
   | Ghost_values of (string * [ `Arg | `Ret ])
   | Impossible_init_state_generation of init_state_error
@@ -34,6 +35,7 @@ type W.kind +=
   | Sut_type_not_specified of string
   | Sut_type_not_supported of string
   | Syntax_error_in_config_module of string
+  | Third_order_function_argument of string
   | Tuple_arity of string
   | Type_not_supported of string
   | Type_not_supported_for_sut_parameter of string