From 3d7deae104539af932b0452a9d923ca25b361535 Mon Sep 17 00:00:00 2001
From: Jun Furuse <jun.furuse@dailambda.jp>
Date: Fri, 1 Mar 2024 08:55:50 +0900
Subject: [PATCH] Apply [@gen_wf] conversions to preambles

---
 lib/error_monad.ml  |  4 ++++
 lib/error_monad.mli |  1 +
 lib/gen_mlw.ml      | 28 ++++++++++++----------------
 3 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/lib/error_monad.ml b/lib/error_monad.ml
index cd78477..0ec8bc1 100644
--- a/lib/error_monad.ml
+++ b/lib/error_monad.ml
@@ -91,4 +91,8 @@ module List = struct
 
   let iter_e (f : 'a -> unit iresult) (l : 'a list) : unit iresult =
     fold_left_e (fun () x -> f x) () l
+
+  let concat_map_e f l =
+    let* xs = map_e f l in
+    return @@ concat xs
 end
diff --git a/lib/error_monad.mli b/lib/error_monad.mli
index 217cbb4..ab566fd 100644
--- a/lib/error_monad.mli
+++ b/lib/error_monad.mli
@@ -43,4 +43,5 @@ module List : sig
 
   val map_e : ('a -> 'b iresult) -> 'a list -> 'b list iresult
   val iter_e : ('a -> unit iresult) -> 'a list -> unit iresult
+  val concat_map_e : ('a -> 'b list iresult) -> 'a list -> 'b list iresult
 end
diff --git a/lib/gen_mlw.ml b/lib/gen_mlw.ml
index e37f2ef..b018922 100644
--- a/lib/gen_mlw.ml
+++ b/lib/gen_mlw.ml
@@ -369,6 +369,16 @@ module Is_type_wf = struct
       let* decls = List.map_e gen_decl type_decls in
       return [ Dlogic decls ]
     else return []
+
+  let add_wfs (decls : decl list) : decl list iresult =
+    (* Adds predicate [is_type_wf] for type decls with [@gen_wf] in the preambles *)
+    List.concat_map_e
+      (function
+        | Dtype dts as d ->
+            let* xs = gen_decls dts in
+            return ([ d ] @ xs)
+        | d -> return [ d ])
+      decls
 end
 
 let sort_wf (s : Sort.t) (p : expr) : term =
@@ -929,22 +939,7 @@ let convert_contract (epp : Sort.t list StringMap.t StringMap.t)
     |> Option.to_iresult ~none:(error_of_fmt "")
   in
   let* param_wf = gen_param_wf ep in
-
-  let other_decls =
-    List.concat_map
-      (fun d ->
-        match d with
-        | Dtype dts -> (
-            [ d ]
-            @
-            (* Adds is_type_wf if [@gen_wf] is attached *)
-            match Is_type_wf.gen_decls dts with
-            | Error _e -> assert false
-            | Ok xs -> xs)
-        | _ -> [ d ])
-      c.c_other_decls
-  in
-
+  let* other_decls = Is_type_wf.add_wfs c.c_other_decls in
   return
   (* scope Contract .. *)
   @@ Dscope
@@ -1079,6 +1074,7 @@ let convert_mlw (tzw : Tzw.t) =
     | Decls ds -> return (ds @ tzw.tzw_preambles)
     | _ -> error_with "invalid prembles: preambles must be list of declarations"
   in
+  let* preambles = Is_type_wf.add_wfs preambles in
   let module G = Generator (struct
     let desc = { d_contracts; d_whyml = [] }
   end) in