Add module-prefix cli option (WIP)

n-osborne · Oct 17, 2024 · 50d6e9f · 50d6e9f
1 parent 8008c59
commit 50d6e9f
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 7 deletions.
diff --git a/plugins/qcheck-stm/src/ortac_qcheck_stm.ml b/plugins/qcheck-stm/src/ortac_qcheck_stm.ml
@@ -4,14 +4,14 @@ module Ir_of_gospel = Ir_of_gospel
 module Reserr = Reserr
 module Stm_of_ir = Stm_of_ir
 
-let main path config output quiet () =
+let main path config output module_prefix quiet () =
   let open Reserr in
   let fmt = Registration.get_out_formatter output in
   let pp = pp quiet Ppxlib_ast.Pprintast.structure fmt in
   pp
     (let* sigs, config = Config.init path config in
      let* ir = Ir_of_gospel.run sigs config in
-     Stm_of_ir.stm config ir)
+     Stm_of_ir.stm module_prefix config ir)
 
 open Cmdliner
 
@@ -22,6 +22,14 @@ end = struct
     Cmd.info "qcheck-stm"
       ~doc:"Generate QCheck-stm test file according to Gospel specifications."
 
+  let module_prefix =
+    let parse s = Ok (Some s) and docv = "MODULE_PREFIX" in
+    Arg.(
+      value
+      & opt (conv ~docv (parse, Fmt.(option string))) None
+      & info [ "p"; "module-prefix" ] ~absent:"" ~docv
+          ~doc:"Prefix module corresponding to FILE with MODULE_PREFIX.")
+
   let config =
     Arg.(
       required
@@ -34,7 +42,14 @@ end = struct
 
   let term =
     let open Registration in
-    Term.(const main $ ocaml_file $ config $ output_file $ quiet $ setup_log)
+    Term.(
+      const main
+      $ ocaml_file
+      $ config
+      $ output_file
+      $ module_prefix
+      $ quiet
+      $ setup_log)
 
   let cmd = Cmd.v info term
 end

diff --git a/plugins/qcheck-stm/src/stm_of_ir.ml b/plugins/qcheck-stm/src/stm_of_ir.ml
@@ -1542,7 +1542,7 @@ let sut_stm_ty_defs config ir =
       let sut = (SUT, fun _ -> "<sut>")]
   else []
 
-let stm config ir =
+let stm module_prefix config ir =
   let open Reserr in
   let* ghost_types = ghost_types config ir.ghost_types in
   let* config, ghost_functions = ghost_functions config ir.ghost_functions in
@@ -1564,12 +1564,14 @@ let stm config ir =
     in
     Option.value config.cleanup ~default
   in
-  let open_mod m = pstr_open Ast_helper.(Opn.mk (Mod.ident (lident m))) in
+  let open_mod i =
+    pstr_open (open_infos ~expr:(pmod_ident i) ~override:Fresh)
+  in
   let sut_defs = sut_defs ir in
   let state_defs = state_defs ir in
   let spec_expr =
     pmod_structure
-      ((open_mod "STM" :: qcheck config)
+      ((open_mod (lident "STM") :: qcheck config)
       @ util config
       @ Option.value config.ty_mod ~default:[]
       @ integer_ty_ext
@@ -1615,9 +1617,21 @@ let stm config ir =
   in
   let sut_mod = sut_module config in
   let* model_mod = model_module config ir in
+  let opened_mod =
+    match module_prefix with
+    | None -> Lident module_name
+    | Some s -> (
+        match String.split_on_char '.' s with
+        | m :: ms ->
+            let pref =
+              List.fold_left (fun acc x -> Ldot (acc, x)) (Lident m) ms
+            in
+            Ldot (pref, module_name)
+        | _ -> assert false)
+  in
   ok
     (warn
-     :: open_mod module_name
+     :: open_mod (noloc opened_mod)
      :: [%stri module Ortac_runtime = Ortac_runtime_qcheck_stm]
      :: ghost_types
     @ ghost_functions