Separate monomorphisation from top-level type checking

3 files changed, 13 insertions, 10 deletions

diff --git a/src/process_file.ml b/src/process_file.ml
index 5eb9f1ee..0601bfab 100644
--- a/src/process_file.ml
+++ b/src/process_file.ml
@@ -92,23 +92,19 @@ let opt_new_typecheck = ref false
 let opt_just_check = ref false
 let opt_ddump_tc_ast = ref false
 let opt_dno_cast = ref false
-let opt_mono_split = ref ([]:((string * int) * string) list)
 
 let check_ast (defs : unit Ast.defs) : Type_check.tannot Ast.defs * Type_check.Env.t =
   let ienv = if !opt_dno_cast then Type_check.Env.no_casts Type_check.initial_env else Type_check.initial_env in
   let ast, env = Type_check.check ienv defs in
-  let ast = match !opt_mono_split with
-    | [] -> ast
-    | l ->
-       let ast = Monomorphise.split_defs l ast in
-       let ienv = Type_check.Env.no_casts Type_check.initial_env in
-       let ast, _ = Type_check.check ienv ast in
-       ast
-  in
   let () = if !opt_ddump_tc_ast then Pretty_print.pp_defs stdout ast else () in
   let () = if !opt_just_check then exit 0 else () in
   (ast, env)
 
+let monomorphise_ast locs ast =
+  let ast = Monomorphise.split_defs locs ast in
+  let ienv = Type_check.Env.no_casts Type_check.initial_env in
+  Type_check.check ienv ast
+
 let rewrite_ast (defs: Type_check.tannot Ast.defs) = Rewriter.rewrite_defs defs
 let rewrite_ast_lem (defs: Type_check.tannot Ast.defs) = Rewriter.rewrite_defs_lem defs
 let rewrite_ast_ocaml (defs: Type_check.tannot Ast.defs) = Rewriter.rewrite_defs_ocaml defs
diff --git a/src/process_file.mli b/src/process_file.mli
index f91fa064..b15523bb 100644
--- a/src/process_file.mli
+++ b/src/process_file.mli
@@ -43,6 +43,7 @@
 val parse_file : string -> Parse_ast.defs
 val convert_ast : Parse_ast.defs -> unit Ast.defs
 val check_ast: unit Ast.defs -> Type_check.tannot Ast.defs * Type_check.Env.t
+val monomorphise_ast : ((string * int) * string) list -> Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs * Type_check.Env.t
 val rewrite_ast: Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs
 val rewrite_ast_lem : Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs
 val rewrite_ast_ocaml : Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs
@@ -53,7 +54,6 @@ val opt_new_typecheck : bool ref
 val opt_just_check : bool ref
 val opt_ddump_tc_ast : bool ref
 val opt_dno_cast : bool ref
-val opt_mono_split : ((string * int) * string) list ref
 
 type out_type =
   | Lem_ast_out
diff --git a/src/sail.ml b/src/sail.ml
index 163a0af1..bb2fdaa5 100644
--- a/src/sail.ml
+++ b/src/sail.ml
@@ -53,6 +53,7 @@ let opt_print_ocaml = ref false
 let opt_libs_lem = ref ([]:string list)
 let opt_libs_ocaml = ref ([]:string list)
 let opt_file_arguments = ref ([]:string list)
+let opt_mono_split = ref ([]:((string * int) * string) list)
 let options = Arg.align ([
   ( "-o",
     Arg.String (fun f -> opt_file_out := Some f),
@@ -131,6 +132,12 @@ let main() =
     let ast = convert_ast ast in
     let (ast, type_envs) = check_ast ast in
 
+    let (ast, type_envs) =
+      match !opt_mono_split with
+      | [] -> ast, type_envs
+      | locs -> monomorphise_ast locs ast
+    in
+
     let ast = rewrite_ast ast in
     let out_name = match !opt_file_out with
       | None -> fst (List.hd parsed)