Collect cleanup function in configuration module

The `cleanup` function is optional.
n-osborne · Apr 17, 2024 · 596c61c · 596c61c
1 parent 662cdde
commit 596c61c
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 5 deletions.
diff --git a/plugins/qcheck-stm/src/config.ml b/plugins/qcheck-stm/src/config.ml
@@ -9,6 +9,7 @@ type config_under_construction = {
   gen_mod' : Ppxlib.structure option;
   pp_mod' : Ppxlib.structure option;
   ty_mod' : Ppxlib.structure option;
+  cleanup' : Ppxlib.structure_item option;
 }
 
 let config_under_construction =
@@ -19,6 +20,7 @@ let config_under_construction =
     gen_mod' = None;
     pp_mod' = None;
     ty_mod' = None;
+    cleanup' = None;
   }
 
 type t = {
@@ -29,6 +31,7 @@ type t = {
   gen_mod : Ppxlib.structure option;
   pp_mod : Ppxlib.structure option;
   ty_mod : Ppxlib.structure option;
+  cleanup : Ppxlib.structure_item option;
 }
 
 let mk_config context cfg_uc =
@@ -47,8 +50,19 @@ let mk_config context cfg_uc =
   let init_sut_txt = Option.get cfg_uc.init_sut_txt'
   and gen_mod = cfg_uc.gen_mod'
   and pp_mod = cfg_uc.pp_mod'
-  and ty_mod = cfg_uc.ty_mod' in
-  ok { context; sut_core_type; init_sut; init_sut_txt; gen_mod; pp_mod; ty_mod }
+  and ty_mod = cfg_uc.ty_mod'
+  and cleanup = cfg_uc.cleanup' in
+  ok
+    {
+      context;
+      sut_core_type;
+      init_sut;
+      init_sut_txt;
+      gen_mod;
+      pp_mod;
+      ty_mod;
+      cleanup;
+    }
 
 let get_sut_type_name config =
   let open Ppxlib in
@@ -111,6 +125,11 @@ let value_bindings cfg_uc =
           Some (Fmt.str "%a" Pprintast.expression vb.pvb_expr)
         in
         ok { cfg_uc with init_sut'; init_sut_txt' }
+    | Ppat_var s when String.equal "cleanup" s.txt ->
+        let cleanup' =
+          Option.some @@ Ortac_core.Builder.pstr_value Nonrecursive [ vb ]
+        in
+        ok { cfg_uc with cleanup' }
     | _ -> ok cfg_uc
   in
   fold_left aux cfg_uc

diff --git a/plugins/qcheck-stm/src/stm_of_ir.ml b/plugins/qcheck-stm/src/stm_of_ir.ml
@@ -957,9 +957,12 @@ let stm config ir =
   let* init_state = init_state config ir in
   let* check_init_state = check_init_state config ir in
   let cleanup =
-    let pat = pvar "cleanup" in
-    let expr = efun [ (Nolabel, ppat_any) ] eunit in
-    pstr_value Nonrecursive [ value_binding ~pat ~expr ]
+    let default =
+      let pat = pvar "cleanup" in
+      let expr = efun [ (Nolabel, ppat_any) ] eunit in
+      pstr_value Nonrecursive [ value_binding ~pat ~expr ]
+    in
+    Option.value config.cleanup ~default
   in
   let init_sut =
     let pat = pvar "init_sut" in