Lean: split up files along import lines (#1173)

Author: javra

src/sail_lean_backend/pretty_print_lean.ml

@@ -1116,25 +1116,40 @@ let doc_val ctx pat exp =
   let base_pp = doc_exp false ctx exp in
   nest 2 (group (string "def" ^^ space ^^ idpp ^^ typpp ^^ space ^^ coloneq ^/^ base_pp))
 
-let rec doc_defs_rec ctx defs types docdefs =
+let should_print_function_def def =
+  match def with
+  | DEF_aux (DEF_fundef fdef, dannot) -> not (Env.is_extern (id_of_fundef fdef) dannot.env "lean")
+  | DEF_aux (DEF_let (LB_aux (LB_val (pat, exp), _)), _) -> true
+  | _ -> false
+
+let rec doc_defs_rec ctx defs types (former_funcs : document list) (docdefs : document) =
   match defs with
-  | [] -> (types, docdefs)
+  | [] -> (types, former_funcs @ [docdefs])
   | DEF_aux (DEF_fundef fdef, dannot) :: defs' ->
       let env = dannot.env in
       let pp_f =
         if Env.is_extern (id_of_fundef fdef) env "lean" then docdefs
         else docdefs ^^ group (doc_fundef ctx fdef) ^/^ hardline
       in
-      doc_defs_rec ctx defs' types pp_f
+      doc_defs_rec ctx defs' types former_funcs pp_f
   | DEF_aux (DEF_type tdef, _) :: defs' when List.mem (string_of_id (id_of_type_def tdef)) !opt_extern_types ->
-      doc_defs_rec ctx defs' types docdefs
+      doc_defs_rec ctx defs' types former_funcs docdefs
   | DEF_aux (DEF_type tdef, _) :: defs' ->
-      doc_defs_rec ctx defs' (types ^^ group (doc_typdef ctx tdef) ^/^ hardline) docdefs
+      doc_defs_rec ctx defs' (types ^^ group (doc_typdef ctx tdef) ^/^ hardline) former_funcs docdefs
   | DEF_aux (DEF_let (LB_aux (LB_val (pat, exp), _)), _) :: defs' ->
-      doc_defs_rec ctx defs' types (docdefs ^^ group (doc_val ctx pat exp) ^/^ hardline)
-  | _ :: defs' -> doc_defs_rec ctx defs' types docdefs
-
-let doc_defs ctx defs = doc_defs_rec ctx defs empty empty
+      doc_defs_rec ctx defs' types former_funcs (docdefs ^^ group (doc_val ctx pat exp) ^/^ hardline)
+  | DEF_aux (DEF_pragma ("include_start", Pragma_line (file, _)), _) :: defs'
+  | DEF_aux (DEF_pragma ("file_start", Pragma_line (file, _)), _) :: defs'
+  | DEF_aux (DEF_pragma ("include_end", Pragma_line (file, _)), _) :: defs'
+  | DEF_aux (DEF_pragma ("file_end", Pragma_line (file, _)), _) :: defs'
+    when Filename.check_suffix file ".sail" ->
+      if docdefs = empty then doc_defs_rec ctx defs' types former_funcs docdefs
+      else doc_defs_rec ctx defs' types (former_funcs @ [docdefs]) empty
+  | d :: defs' ->
+      if should_print_function_def d then failwith "this case of doc_defs_rec should be unreachable"
+      else doc_defs_rec ctx defs' types former_funcs docdefs
+
+let doc_defs ctx defs = doc_defs_rec ctx defs empty [] empty
 
 (* Remove all imports for now, they will be printed in other files. Probably just for testing. *)
 let rec remove_imports (defs : (Libsail.Type_check.tannot, Libsail.Type_check.env) def list) depth =
@@ -1258,8 +1273,40 @@ let populate_fun_args defs =
   in
   List.fold_left (fun args d -> add_args args d) Bindings.empty defs
 
-let pp_ast_lean (env : Type_check.env) effect_info ({ defs; _ } as ast : Libsail.Type_check.typed_ast) types_file
-    funcs_file =
+let rec collect_import_files_aux defs file_stack last_namespace ret =
+  match defs with
+  | [] -> ret
+  | DEF_aux (DEF_pragma ("include_start", Pragma_line (file, _)), _) :: ds
+  | DEF_aux (DEF_pragma ("file_start", Pragma_line (file, _)), _) :: ds
+    when Filename.check_suffix file ".sail" ->
+      collect_import_files_aux ds (file :: file_stack) last_namespace ret
+  | DEF_aux (DEF_pragma ("include_end", Pragma_line (file, _)), _) :: ds
+  | DEF_aux (DEF_pragma ("file_end", Pragma_line (file, _)), _) :: ds
+    when Filename.check_suffix file ".sail" -> (
+      match file_stack with
+      | f :: fs -> collect_import_files_aux ds fs last_namespace ret
+      | _ -> failwith "should not be reachable"
+    )
+  | d :: ds -> (
+      match file_stack with
+      | f :: _ ->
+          if should_print_function_def d && not (last_namespace = Some f) then
+            collect_import_files_aux ds file_stack (Some f) (ret @ [f])
+          else collect_import_files_aux ds file_stack last_namespace ret
+      | _ -> failwith "should not be reachable"
+    )
+
+let collect_import_files defs base =
+  let res = collect_import_files_aux defs [base] None [] in
+  if res = [] then [base] else res
+
+let rec take n xs = match (n, xs) with 0, _ -> [] | n, x :: xs -> x :: take (n - 1) xs | n, xs -> xs
+
+let rec last xs =
+  match xs with [] -> failwith "cannot take last element of empty list" | [x] -> x | x :: xs -> last xs
+
+let pp_ast_lean (env : Type_check.env) effect_info ({ defs; _ } as ast : Libsail.Type_check.typed_ast) out_name_camel
+    types_file imp_funcs_files funcs_file noncomputable =
   let regs = State.find_registers defs in
   let fun_args = populate_fun_args defs in
   let global = { effect_info; fun_args } in
@@ -1271,12 +1318,29 @@ let pp_ast_lean (env : Type_check.env) effect_info ({ defs; _ } as ast : Libsail
   in
   let monad = doc_monad_abbrev defs has_registers in
   let instantiations = doc_instantiations ctx env in
-  let types, fundefs = doc_defs ctx defs in
-  let fundefs = string "namespace Functions\n\n" ^^ fundefs ^^ string "end Functions\n" in
+  let types, all_fundefss = doc_defs ctx defs in
+  let imp_fundefss, main_fundefs =
+    if imp_funcs_files = [] then ([], concat all_fundefss)
+    else (
+      let imp_fundefss = take (List.length all_fundefss - 1) all_fundefss in
+      let main_fundefs = last all_fundefss in
+      (imp_fundefss, main_fundefs)
+    )
+  in
+  let main_fundefs = main_fundefs ^^ string ("end " ^ out_name_camel ^ ".Functions") ^^ hardline in
   let main_function =
-    if !the_main_function_has_been_seen then main_function_stub effect_info has_registers else empty
+    if !the_main_function_has_been_seen then (
+      let stub = main_function_stub effect_info has_registers in
+      [string ("open " ^ out_name_camel ^ ".Functions\n\n") ^^ stub]
+    )
+    else []
   in
   let opens = IdSet.fold (fun id doc -> string "open " ^^ doc_id_ctor id ^^ hardline ^^ doc) !opens empty in
   print types_file (types ^^ register_refs ^^ monad ^^ instantiations);
-  print funcs_file (opens ^^ hardline ^^ fundefs ^^ string "open Functions\n\n" ^^ main_function);
+  let _ =
+    List.map2
+      (fun file defs -> print file (separate hardline (remove_empties [opens; defs])))
+      imp_funcs_files imp_fundefss
+  in
+  print funcs_file (separate hardline (remove_empties ([opens; main_fundefs] @ main_function)));
   !the_main_function_has_been_seen

src/sail_lean_backend/sail_plugin_lean.ml

@@ -77,6 +77,8 @@ let opt_lean_import_files : string list ref = ref []
 
 let opt_lean_noncomputable : bool ref = ref false
 
+let opt_single_file : bool ref = ref false
+
 let lean_version : string = "lean4:nightly-2025-03-17"
 
 let lean_options =
@@ -89,6 +91,10 @@ let lean_options =
       Arg.Unit (fun () -> opt_lean_force_output := true),
       "removes the content of the output directory if it is non-empty"
     );
+    ( Flag.create ~prefix:["lean"] "single_file",
+      Arg.Unit (fun () -> opt_single_file := true),
+      "puts the entire output in a single .lean file"
+    );
     ( Flag.create ~prefix:["lean"] "noncomputable",
       Arg.Unit (fun () -> opt_lean_noncomputable := true),
       "add a 'noncomputable section' at the beginning of the output"
@@ -183,6 +189,7 @@ type lean_context = {
   sail_dir : string;
   types_file : out_channel;
   funcs_file : out_channel;
+  import_files : out_channel list;
   lakefile : out_channel;
 }
 
@@ -205,7 +212,24 @@ let path_to_static_libarary sail_dir str = Filename.quote (sail_dir ^ "/src/sail
 let copy_from_static_library sail_dir lean_sail_dir str =
   Unix.system ("cp " ^ path_to_static_libarary sail_dir str ^ " " ^ Filename.quote lean_sail_dir)
 
-let start_lean_output (out_name : string) default_sail_dir =
+let print_function_file_prelude file out_name_camel (prev_file : string option) =
+  let _ =
+    match prev_file with
+    | None ->
+        output_string file ("import " ^ out_name_camel ^ ".Sail.Sail\n");
+        output_string file ("import " ^ out_name_camel ^ ".Sail.BitVec\n");
+        output_string file ("import " ^ out_name_camel ^ ".Sail.IntRange\n");
+        output_string file ("import " ^ out_name_camel ^ ".Defs\n");
+        List.iter
+          (fun filename -> output_string file ("import " ^ out_name_camel ^ "." ^ file_to_module filename ^ "\n"))
+          !opt_lean_import_files
+    | Some n -> output_string file ("import " ^ out_name_camel ^ "." ^ n ^ "\n")
+  in
+  output_string file ("\n" ^ file_prelude);
+  if !opt_lean_noncomputable then output_string file "noncomputable section\n\n";
+  output_string file ("namespace " ^ out_name_camel ^ ".Functions\n\n")
+
+let start_lean_output (out_name : string) (import_names : string list) default_sail_dir =
   let base_dir = match !opt_lean_output_dir with Some dir -> dir | None -> "." in
   let project_dir = Filename.concat base_dir out_name in
   if !opt_lean_force_output && Sys.file_exists project_dir && Sys.is_directory project_dir then (
@@ -222,6 +246,7 @@ let start_lean_output (out_name : string) default_sail_dir =
   close_out lean_toolchain;
   let sail_dir = Reporting.get_sail_dir default_sail_dir in
   let out_name_camel = Libsail.Util.to_upper_camel_case out_name in
+  let import_names_camel = List.map Libsail.Util.to_upper_camel_case import_names in
   let lean_src_dir = Filename.concat project_dir out_name_camel in
   if not (Sys.file_exists lean_src_dir) then Unix.mkdir lean_src_dir 0o775;
   let lean_sail_dir = lean_src_dir ^ "/Sail/" in
@@ -244,17 +269,21 @@ let start_lean_output (out_name : string) default_sail_dir =
   output_string types_file "open PreSail\n\n";
   output_string types_file file_prelude;
   let funcs_file = open_out (Filename.concat project_dir (out_name_camel ^ ".lean")) in
-  output_string funcs_file ("import " ^ out_name_camel ^ ".Sail.Sail\n");
-  output_string funcs_file ("import " ^ out_name_camel ^ ".Sail.BitVec\n");
-  output_string funcs_file ("import " ^ out_name_camel ^ ".Sail.IntRange\n");
-  output_string funcs_file ("import " ^ out_name_camel ^ ".Defs\n");
-  List.iter
-    (fun filename -> output_string funcs_file ("import " ^ out_name_camel ^ "." ^ file_to_module filename ^ "\n"))
-    !opt_lean_import_files;
-  output_string funcs_file ("\n" ^ file_prelude);
-  if !opt_lean_noncomputable then output_string funcs_file "noncomputable section\n\n";
   let lakefile = open_out (Filename.concat project_dir "lakefile.toml") in
-  { out_name; out_name_camel; sail_dir; types_file; funcs_file; lakefile }
+  let import_files =
+    List.map (fun name -> open_out (Filename.concat lean_src_dir (name ^ ".lean"))) import_names_camel
+  in
+  let last_import_name =
+    List.fold_left
+      (fun prev_name (out, n) ->
+        print_function_file_prelude out out_name_camel prev_name;
+        Some n
+      )
+      None
+      (List.combine import_files import_names_camel)
+  in
+  print_function_file_prelude funcs_file out_name_camel last_import_name;
+  { out_name; out_name_camel; sail_dir; types_file; funcs_file; import_files; lakefile }
 
 let close_context ctx =
   close_out ctx.types_file;
@@ -265,27 +294,48 @@ let create_lake_project (ctx : lean_context) executable =
   (* Change the base directory if the option '--lean-output-dir' is set *)
   output_string ctx.lakefile
     ("name = \"" ^ ctx.out_name ^ "\"\ndefaultTargets = [\"" ^ ctx.out_name_camel
-   ^ "\"]\n\
-      moreLeanArgs = [\"--tstack=400000\"]\n\
-      moreGlobalServerArgs = [\"--tstack=400000\"]\n\n\
-      [[lean_lib]]\n\
-      name = \"" ^ ctx.out_name_camel ^ "\""
+   ^ "\"]\nmoreLeanArgs = [\"--tstack=400000\"]\n\n[[lean_lib]]\nname = \"" ^ ctx.out_name_camel ^ "\""
     );
   if executable then (
     output_string ctx.lakefile "\n\n[[lean_exe]]\n";
     output_string ctx.lakefile "name = \"run\"\n";
     output_string ctx.lakefile ("root = \"" ^ ctx.out_name_camel ^ "\"\n")
   )
 
-let output (out_name : string) env effect_info ast default_sail_dir =
-  let ctx = start_lean_output out_name default_sail_dir in
-  let executable = Pretty_print_lean.pp_ast_lean env effect_info ast ctx.types_file ctx.funcs_file in
+let rec dedup_files (files : string list) (acc : string list) =
+  match files with
+  | [] -> acc
+  | f :: fs -> (
+      match List.fold_left (fun n y -> if y = f then n + 1 else n) 0 acc with
+      | 0 -> dedup_files fs (acc @ [f])
+      | n -> dedup_files fs (acc @ [f ^ Int.to_string (n - 1)])
+    )
+
+let output (out_name : string) env effect_info ({ defs; _ } as ast : Libsail.Type_check.typed_ast) default_sail_dir
+    single_file noncomputable =
+  let imports =
+    if single_file then []
+    else (
+      (* Collect all non-empty slices between include pragmas in the file *)
+      let imports = Pretty_print_lean.collect_import_files defs (out_name ^ ".sail") in
+      let imports = List.map file_to_module imports in
+      (* Discard the last import file, as we will use the main file instead *)
+      let imports = Pretty_print_lean.take (List.length imports - 1) imports in
+      dedup_files imports []
+    )
+  in
+  let ctx = start_lean_output out_name imports default_sail_dir in
+  let out_name_camel = Libsail.Util.to_upper_camel_case out_name in
+  let executable =
+    Pretty_print_lean.pp_ast_lean env effect_info ast out_name_camel ctx.types_file ctx.import_files ctx.funcs_file
+      noncomputable
+  in
   create_lake_project ctx executable
 (* Uncomment for debug output of the Sail code after the rewrite passes *)
 (* Pretty_print_sail.output_ast stdout (Type_check.strip_ast ast) *)
 
 let lean_target out_name { default_sail_dir; ctx; ast; effect_info; env; _ } =
   let out_name = match out_name with Some f -> f | None -> "out" in
-  output out_name env effect_info ast default_sail_dir
+  output out_name env effect_info ast default_sail_dir !opt_single_file !opt_lean_noncomputable
 
 let _ = Target.register ~name:"lean" ~options:lean_options ~rewrites:lean_rewrites ~asserts_termination:true lean_target

test/lean/SailTinyArm.expected.lean

@@ -474,6 +474,8 @@ set_option match.ignoreUnusedAlts true
 
 open Sail
 
+namespace Out.Functions
+
 open option
 open ast
 open arm_acc_type
@@ -504,8 +506,6 @@ open CacheOp
 open Barrier
 open AccessType
 
-namespace Functions
-
 /-- Type quantifiers: k_ex6347# : Bool, k_ex6346# : Bool -/
 def neq_bool (x : Bool) (y : Bool) : Bool :=
   (Bool.not (BEq.beq x y))
@@ -2160,6 +2160,4 @@ def initialize_registers (_ : Unit) : SailM Unit := do
 def sail_model_init (x_0 : Unit) : SailM Unit := do
   (initialize_registers ())
 
-end Functions
-open Functions
-
+end Out.Functions

test/lean/append.expected.lean

@@ -41,9 +41,9 @@ set_option match.ignoreUnusedAlts true
 
 open Sail
 
-open option
+namespace Out.Functions
 
-namespace Functions
+open option
 
 /-- Type quantifiers: k_ex699# : Bool, k_ex698# : Bool -/
 def neq_bool (x : Bool) (y : Bool) : Bool :=
@@ -135,6 +135,4 @@ def initialize_registers (_ : Unit) : Unit :=
 def sail_model_init (x_0 : Unit) : Unit :=
   (initialize_registers ())
 
-end Functions
-open Functions
-
+end Out.Functions

test/lean/atom_bool.expected.lean

@@ -33,8 +33,7 @@ set_option match.ignoreUnusedAlts true
 
 open Sail
 
-
-namespace Functions
+namespace Out.Functions
 
 def foo (_ : Unit) : Bool :=
   true
@@ -45,6 +44,4 @@ def initialize_registers (_ : Unit) : Unit :=
 def sail_model_init (x_0 : Unit) : Unit :=
   (initialize_registers ())
 
-end Functions
-open Functions
-
+end Out.Functions

test/lean/bitfield.expected.lean

@@ -50,11 +50,11 @@ set_option match.ignoreUnusedAlts true
 
 open Sail
 
+namespace Out.Functions
+
 open option
 open Register
 
-namespace Functions
-
 /-- Type quantifiers: k_ex700# : Bool, k_ex699# : Bool -/
 def neq_bool (x : Bool) (y : Bool) : Bool :=
   (Bool.not (BEq.beq x y))
@@ -198,6 +198,4 @@ def initialize_registers (_ : Unit) : SailM Unit := do
 def sail_model_init (x_0 : Unit) : SailM Unit := do
   (initialize_registers ())
 
-end Functions
-open Functions
-
+end Out.Functions

test/lean/bitvec_operation.expected.lean

@@ -41,9 +41,9 @@ set_option match.ignoreUnusedAlts true
 
 open Sail
 
-open option
+namespace Out.Functions
 
-namespace Functions
+open option
 
 /-- Type quantifiers: k_ex820# : Bool, k_ex819# : Bool -/
 def neq_bool (x : Bool) (y : Bool) : Bool :=
@@ -185,6 +185,4 @@ def initialize_registers (_ : Unit) : Unit :=
 def sail_model_init (x_0 : Unit) : Unit :=
   (initialize_registers ())
 
-end Functions
-open Functions
-
+end Out.Functions