Include comments in AST type

4 files changed, 13 insertions, 5 deletions

diff --git a/src/ast_defs.ml b/src/ast_defs.ml
index f6b83c2c..0826fb42 100644
--- a/src/ast_defs.ml
+++ b/src/ast_defs.ml
@@ -51,9 +51,11 @@
 open Ast
 
 type 'a ast = {
-    defs : 'a def list
+    defs : 'a def list;
+    comments : (string * Lexer.comment list) list
   }
 
 let empty_ast = {
-    defs = []
+    defs = [];
+    comments = []
   }
diff --git a/src/ast_util.ml b/src/ast_util.ml
index 40cb7d98..f1df1ff5 100644
--- a/src/ast_util.ml
+++ b/src/ast_util.ml
@@ -1546,7 +1546,7 @@ let split_defs f defs =
   | Some (pre_defs, def, post_defs) ->
      Some (List.rev pre_defs, def, post_defs)
 
-let append_ast ast1 ast2 = { defs = ast1.defs @ ast2.defs }
+let append_ast ast1 ast2 = { defs = ast1.defs @ ast2.defs; comments = ast1.comments @ ast2.comments }
 let append_ast_defs ast defs = { ast with defs = ast.defs @ defs }
 let concat_ast asts = List.fold_right append_ast asts empty_ast
 
diff --git a/src/initial_check.ml b/src/initial_check.ml
index a958dced..1b21e2be 100644
--- a/src/initial_check.ml
+++ b/src/initial_check.ml
@@ -847,7 +847,7 @@ let to_ast ctx (P.Defs files) =
         let defs', ctx = to_ast_defs ctx file in (defs @ wrap_file (fst file) defs', ctx)
       ) ([], ctx) files
   in
-  { defs = defs }, ctx
+  { defs = defs; comments = [] }, ctx
 
 let initial_ctx = {
     type_constructors =
diff --git a/src/process_file.ml b/src/process_file.ml
index b121c87c..d67fba76 100644
--- a/src/process_file.ml
+++ b/src/process_file.ml
@@ -275,8 +275,14 @@ let load_files ?check:(check=false) options type_envs files =
   if !opt_memo_z3 then Constraint.load_digests () else ();
 
   let t = Profile.start () in
-  let ast = Parse_ast.Defs (List.map (fun f -> (f, parse_file f |> snd |> preprocess options)) files) in
+
+  let parsed_files = List.map (fun f -> (f, parse_file f)) files in
+
+  let comments = List.map (fun (f, (comments, _)) -> (f, comments)) parsed_files in
+  let ast = Parse_ast.Defs (List.map (fun (f, (_, file_ast)) -> (f, preprocess options file_ast)) parsed_files) in
   let ast = Initial_check.process_ast ~generate:(not check) ast in
+  let ast = { ast with comments = comments } in
+  
   let () = if !opt_ddump_initial_ast then Pretty_print_sail.pp_ast stdout ast else () in
 
   begin match !opt_reformat with