Add reformat option to Sail

Reformats input source code using the pretty printer preserving the file
structure. Probably not useable as an ocamlformat or rustfmt like tool,
but good enough to take generated code without formatting and make it
readable.

4 files changed, 98 insertions, 9 deletions

diff --git a/src/pretty_print_sail.ml b/src/pretty_print_sail.ml
index 10a15245..85684887 100644
--- a/src/pretty_print_sail.ml
+++ b/src/pretty_print_sail.ml
@@ -53,6 +53,7 @@ open Ast_util
 open PPrint
 
 let opt_use_heuristics = ref false
+let opt_insert_braces = ref false
 
 module Big_int = Nat_big_num
 
@@ -397,6 +398,12 @@ let rec doc_exp (E_aux (e_aux, _) as exp) =
   | E_app_infix _ -> doc_infix 0 exp
   | E_tuple exps -> parens (separate_map (comma ^^ space) doc_exp exps)
 
+  | E_if (if_exp, then_exp, else_exp) when !opt_insert_braces ->
+     separate space [string "if"; doc_exp if_exp; string "then"] ^^ space
+     ^^ doc_exp_as_block then_exp
+     ^^ space ^^ string "else" ^^ space
+     ^^ doc_exp_as_block else_exp
+    
   (* Various rules to try to format if blocks nicely based on content *)
   | E_if (if_exp, then_exp, else_exp) when if_block_then then_exp && if_block_else else_exp ->
      (separate space [string "if"; doc_exp if_exp; string "then"] ^//^ doc_exp then_exp)
@@ -428,7 +435,8 @@ let rec doc_exp (E_aux (e_aux, _) as exp) =
   | E_case (exp, pexps) ->
      separate space [string "match"; doc_exp exp; doc_pexps pexps]
   | E_let (LB_aux (LB_val (pat, binding), _), exp) ->
-     separate space [string "let"; doc_pat pat; equals; doc_exp binding; string "in"; doc_exp exp]
+     separate space [string "let"; doc_pat pat; equals; doc_exp binding; string "in"]
+     ^//^ doc_exp exp
   | E_internal_plet (pat, exp1, exp2) ->
      let le =
        prefix 2 1
@@ -519,6 +527,8 @@ and doc_block = function
   | [] -> string "()"
   | [E_aux (E_let (LB_aux (LB_val (pat, binding), _), E_aux (E_block exps, _)), _)] ->
      separate space [string "let"; doc_pat pat; equals; doc_exp binding] ^^ semi ^^ hardline ^^ doc_block exps
+  | [E_aux (E_let (LB_aux (LB_val (pat, binding), _), exp), _)] ->
+     separate space [string "let"; doc_pat pat; equals; doc_exp binding] ^^ semi ^^ hardline ^^ doc_block [exp]
   | [E_aux (E_var (lexp, binding, E_aux (E_block exps, _)), _)] ->
      separate space [string "var"; doc_lexp lexp; equals; doc_exp binding] ^^ semi ^^ hardline ^^ doc_block exps
   | [exp] -> doc_exp exp
@@ -548,13 +558,19 @@ and doc_letbind (LB_aux (lb_aux, _)) =
   match lb_aux with
   | LB_val (pat, exp) ->
      separate space [doc_pat pat; equals; doc_exp exp]
-
+and doc_exp_as_block (E_aux (aux, _) as exp) =
+  match aux with
+  | E_block _ | E_lit _ -> doc_exp exp
+  | _ when !opt_insert_braces ->
+     group (lbrace ^^ nest 4 (hardline ^^ doc_block [exp]) ^^ hardline ^^ rbrace)
+  | _ -> doc_exp exp
+ 
 let doc_funcl (FCL_aux (FCL_Funcl (id, Pat_aux (pexp,_)), _)) =
   match pexp with
   | Pat_exp (pat,exp) ->
-     group (separate space [doc_id id; doc_pat pat; equals; doc_exp exp])
+     group (separate space [doc_id id; doc_pat pat; equals; doc_exp_as_block exp])
   | Pat_when (pat,wh,exp) ->
-     group (separate space [doc_id id; parens (separate space [doc_pat pat; string "if"; doc_exp wh]); string "="; doc_exp exp])
+     group (separate space [doc_id id; parens (separate space [doc_pat pat; string "if"; doc_exp wh]); string "="; doc_exp_as_block exp])
 
 let doc_default (DT_aux (DT_order ord, _)) = separate space [string "default"; string "Order"; doc_ord ord]
 
@@ -750,6 +766,58 @@ let rec doc_def def = group (match def with
 let doc_defs (Defs(defs)) =
   separate_map hardline doc_def (List.filter doc_filter defs)
 
+let reformat dir (Defs defs) =
+  let file_stack = ref [] in
+
+  let pop () = match !file_stack with
+    | [] -> Reporting.unreachable Parse_ast.Unknown __POS__ "Unbalanced file structure"
+    | Some chan :: chans ->
+       close_out chan;
+       file_stack := chans
+    | None :: chans ->
+       file_stack := chans
+  in
+
+  let push = function
+    | Some path ->
+       file_stack := Some (open_out path) :: !file_stack
+    | None ->
+       file_stack := None :: !file_stack
+  in
+
+  let adjust_path path = Filename.concat (Filename.concat (Filename.dirname path) dir) (Filename.basename path) in
+
+  let output_def def = match !file_stack with
+    | Some chan :: _ ->
+       ToChannel.pretty 1. 120 chan (hardline ^^ doc_def def)
+    | None :: _ -> ()
+    | [] -> Reporting.unreachable Parse_ast.Unknown __POS__ "No file for definition"
+  in
+
+  let output_include path = match !file_stack with
+    | Some chan :: _ ->
+       if Filename.is_relative path then
+         Printf.fprintf chan "$include \"%s\"\n" path
+       else
+         Printf.fprintf chan "$include <%s>\n" (Filename.basename path)
+    | None :: _ -> ()
+    | [] -> Reporting.unreachable Parse_ast.Unknown __POS__ "No file for include"
+  in
+          
+  let format_def = function
+    | DEF_pragma ("file_start", path, _) -> push (Some (adjust_path path))
+    | DEF_pragma ("file_end", _, _) -> pop ()
+    | DEF_pragma ("include_start", path, _) ->
+       output_include path;
+       if Filename.is_relative path then push (Some (adjust_path path)) else push None
+    | DEF_pragma ("include_end", path, _) -> pop ()
+    | def -> output_def def
+  in
+
+  opt_insert_braces := true;
+  List.iter format_def defs;
+  opt_insert_braces := false
+
 let pp_defs f d = ToChannel.pretty 1. 80 f (doc_defs d)
 
 let pretty_sail f doc = ToChannel.pretty 1. 120 f doc
diff --git a/src/process_file.ml b/src/process_file.ml
index d1325e00..a0245c8f 100644
--- a/src/process_file.ml
+++ b/src/process_file.ml
@@ -161,11 +161,13 @@ let all_pragmas =
       "include_end"
     ]
 
-let wrap_include l file defs =
-  [Parse_ast.DEF_pragma ("include_start", file, l)]
-  @ defs
-  @ [Parse_ast.DEF_pragma ("include_end", file, l)]
-  
+let wrap_include l file = function
+  | [] -> []
+  | defs ->
+     [Parse_ast.DEF_pragma ("include_start", file, l)]
+     @ defs
+     @ [Parse_ast.DEF_pragma ("include_end", file, l)]
+ 
 let rec preprocess opts = function
   | [] -> []
   | Parse_ast.DEF_pragma ("define", symbol, _) :: defs ->
@@ -255,6 +257,8 @@ let rec preprocess opts = function
   | def :: defs -> def :: preprocess opts defs
 
 let opt_just_check = ref false
+let opt_reformat = ref None
+let opt_ddump_initial_ast = ref false
 let opt_ddump_tc_ast = ref false
 let opt_ddump_rewrite_ast = ref None
 let opt_dno_cast = ref false
@@ -272,6 +276,15 @@ let load_files ?check:(check=false) options type_envs files =
   let t = Profile.start () in
   let ast = Parse_ast.Defs (List.map (fun f -> (f, parse_file f |> snd |> preprocess options)) files) in
   let ast = Initial_check.process_ast ~generate:(not check) ast in
+  let () = if !opt_ddump_initial_ast then Pretty_print_sail.pp_defs stdout ast else () in
+
+  begin match !opt_reformat with
+  | Some dir ->
+     Pretty_print_sail.reformat dir ast;
+     exit 0
+  | None -> ()
+  end;
+  
   (* The separate loop measures declarations would be awkward to type check, so
      move them into the definitions beforehand. *)
   let ast = Rewrites.move_loop_measures ast in
diff --git a/src/process_file.mli b/src/process_file.mli
index 3b23b1bd..ea00dd29 100644
--- a/src/process_file.mli
+++ b/src/process_file.mli
@@ -65,6 +65,8 @@ val rewrite_ast_check : Type_check.Env.t -> Type_check.tannot Ast.defs -> Type_c
 val opt_file_out : string option ref
 val opt_memo_z3 : bool ref
 val opt_just_check : bool ref
+val opt_reformat : string option ref
+val opt_ddump_initial_ast : bool ref
 val opt_ddump_tc_ast : bool ref
 val opt_ddump_rewrite_ast : ((string * int) option) ref
 val opt_dno_cast : bool ref
diff --git a/src/sail.ml b/src/sail.ml
index b537705c..31fc7717 100644
--- a/src/sail.ml
+++ b/src/sail.ml
@@ -383,6 +383,12 @@ let options = Arg.align ([
   ( "-output_sail",
     set_target "sail",
     " print Sail code after type checking and initial rewriting");
+  ( "-reformat",
+    Arg.String (fun dir -> opt_reformat := Some dir),
+    "<dir> reformat the input Sail code placing output into a separate directory");
+  ( "-ddump_initial_ast",
+    Arg.Set opt_ddump_initial_ast,
+    " (debug) dump the initial ast to stdout");
   ( "-ddump_tc_ast",
     Arg.Set opt_ddump_tc_ast,
     " (debug) dump the typechecked ast to stdout");