Fix sail fmt's handling of comments (#543)

* FIx formatter

* fix typo and testcase

src/lib/chunk_ast.ml

@@ -95,7 +95,7 @@ let comment_type_delimiters = function Lexer.Comment_line -> ("//", "") | Lexer.
 type infix_chunk = Infix_prefix of string | Infix_op of string | Infix_chunks of chunks
 
 and chunk =
-  | Comment of Lexer.comment_type * int * int * string
+  | Comment of Lexer.comment_type * int * int * string * bool
   | Spacer of bool * int
   | Function of {
       id : id;
@@ -152,9 +152,9 @@ let add_chunk q chunk = Queue.add chunk q
 
 [@@@coverage off]
 let rec prerr_chunk indent = function
-  | Comment (comment_type, n, col, contents) ->
+  | Comment (comment_type, n, col, contents, trailing) ->
       let s, e = comment_type_delimiters comment_type in
-      Printf.eprintf "%sComment: blank=%d col=%d %s%s%s\n" indent n col s contents e
+      Printf.eprintf "%sComment: blank=%d col=%d trailing=%b %s%s%s\n" indent n col trailing s contents e
   | Spacer (line, w) -> Printf.eprintf "%sSpacer:%b %d\n" indent line w
   | Atom str -> Printf.eprintf "%sAtom:%s\n" indent str
   | String_literal str -> Printf.eprintf "%sString_literal:%s\n" indent str
@@ -409,7 +409,9 @@ let rec pop_header_comments comments chunks l lnum =
       match Reporting.simp_loc l with
       | Some (s, _) when e.pos_cnum < s.pos_cnum && comment_s.pos_lnum = lnum ->
           let _ = Stack.pop comments in
-          Queue.add (Comment (comment_type, 0, comment_s.pos_cnum - comment_s.pos_bol, contents)) chunks;
+          Queue.add
+            (Comment (comment_type, 0, comment_s.pos_cnum - comment_s.pos_bol, contents, e.pos_lnum == lnum))
+            chunks;
           Queue.add (Spacer (true, 1)) chunks;
           pop_header_comments comments chunks l (lnum + 1)
       | _ -> ()
@@ -423,16 +425,36 @@ let chunk_header_comments comments chunks = function
 let rec pop_comments ?(spacer = true) comments chunks l =
   match Stack.top_opt comments with
   | None -> ()
-  | Some (Lexer.Comment (comment_type, comment_s, e, contents)) -> begin
+  | Some (Lexer.Comment (comment_type, comment_s, comment_e, contents)) -> begin
       match Reporting.simp_loc l with
-      | Some (s, _) when e.pos_cnum <= s.pos_cnum ->
+      | Some (s, e) when comment_e.pos_cnum <= s.pos_cnum ->
           let _ = Stack.pop comments in
-          Queue.add (Comment (comment_type, 0, comment_s.pos_cnum - comment_s.pos_bol, contents)) chunks;
-          if spacer && e.pos_lnum < s.pos_lnum then Queue.add (Spacer (true, 1)) chunks;
+          Queue.add
+            (Comment
+               (comment_type, 0, comment_s.pos_cnum - comment_s.pos_bol, contents, comment_s.pos_lnum == e.pos_lnum)
+            )
+            chunks;
+          if spacer && comment_e.pos_lnum < s.pos_lnum then Queue.add (Spacer (true, 1)) chunks;
           pop_comments comments chunks l
       | _ -> ()
     end
 
+let rec pop_comments_until_loc_end comments chunks l =
+  match Stack.top_opt comments with
+  | None -> ()
+  | Some (Lexer.Comment (comment_type, comment_s, comment_e, contents)) -> begin
+      match Reporting.simp_loc l with
+      | Some (_, e) when comment_s.pos_cnum < e.pos_cnum ->
+          let _ = Stack.pop comments in
+          Queue.add
+            (Comment
+               (comment_type, 0, comment_s.pos_cnum - comment_s.pos_bol, contents, comment_s.pos_lnum == e.pos_lnum)
+            )
+            chunks;
+          pop_comments_until_loc_end comments chunks l
+      | _ -> ()
+    end
+
 let rec discard_comments comments (pos : Lexing.position) =
   match Stack.top_opt comments with
   | None -> ()
@@ -449,7 +471,7 @@ let pop_trailing_comment ?space:(n = 0) comments chunks line_num =
       match Stack.top_opt comments with
       | Some (Lexer.Comment (comment_type, s, _, contents)) when s.pos_lnum = lnum ->
           let _ = Stack.pop comments in
-          Queue.add (Comment (comment_type, n, s.pos_cnum - s.pos_bol, contents)) chunks;
+          Queue.add (Comment (comment_type, n, s.pos_cnum - s.pos_bol, contents, true)) chunks;
           begin
             match comment_type with Lexer.Comment_line -> true | _ -> false
           end
@@ -869,11 +891,16 @@ let rec chunk_exp comments chunks (E_aux (aux, l)) =
                   Queue.add (Block_binder (Var_binder, lexp_chunks, exp_chunks)) chunks
             end;
 
-            (* TODO: Do we need to do something different for multiple trailing comments at end of a block? *)
             let next_line_num = Option.bind next (fun bexp -> block_exp_locs bexp |> fst |> starting_line_num) in
             if have_linebreak (ending_line_num e_l) next_line_num || Option.is_none next then
               ignore (pop_trailing_comment comments chunks (ending_line_num e_l));
-
+            begin
+              match next with
+              | Some next ->
+                  let next_s_l, _ = block_exp_locs block_exp in
+                  pop_comments comments chunks next_s_l
+              | _ -> pop_comments_until_loc_end comments chunks l
+            end;
             (chunks, have_blank_linebreak (ending_line_num e_l) next_line_num)
           )
           false block_exps
@@ -915,6 +942,7 @@ let rec chunk_exp comments chunks (E_aux (aux, l)) =
       let i_chunks = rec_chunk_exp i in
       pop_comments ~spacer:false comments i_chunks keywords.then_loc;
       let t_chunks = rec_chunk_exp t in
+      ignore (pop_trailing_comment comments t_chunks (ending_line_num keywords.then_loc));
       (match keywords.else_loc with Some l -> pop_comments comments t_chunks l | None -> ());
       let e_chunks = rec_chunk_exp e in
       Queue.add (If_then_else (if_format, i_chunks, t_chunks, e_chunks)) chunks
@@ -1323,4 +1351,14 @@ let chunk_defs source comments defs =
   let chunks = Queue.create () in
   chunk_header_comments comments chunks defs;
   let _ = List.fold_left (fun last_span def -> chunk_def source last_span comments chunks def) (None, Some 0) defs in
+
+  (* pop remaining comments *)
+  if not (Stack.is_empty comments) then Queue.add (Spacer (true, 1)) chunks;
+  Stack.iter
+    (fun c ->
+      match c with
+      | Lexer.Comment (comment_type, comment_s, e, contents) ->
+          Queue.add (Comment (comment_type, 0, comment_s.pos_cnum - comment_s.pos_bol, contents, false)) chunks
+    )
+    comments;
   chunks

src/lib/chunk_ast.mli

@@ -87,7 +87,7 @@ val comment_type_delimiters : Lexer.comment_type -> string * string
 type infix_chunk = Infix_prefix of string | Infix_op of string | Infix_chunks of chunks
 
 and chunk =
-  | Comment of Lexer.comment_type * int * int * string
+  | Comment of Lexer.comment_type * int * int * string * bool
   | Spacer of bool * int
   | Function of {
       id : Parse_ast.id;

src/lib/format_sail.ml

@@ -77,7 +77,7 @@ let rec map_last f = function
       let x = f false x in
       x :: map_last f xs
 
-let line_comment_opt = function Comment (Lexer.Comment_line, _, _, contents) -> Some contents | _ -> None
+let line_comment_opt = function Comment (Lexer.Comment_line, _, _, contents, _trailing) -> Some contents | _ -> None
 
 (* Remove additional (> 1) trailing newlines at the end of a string *)
 let discard_extra_trailing_newlines s =
@@ -522,7 +522,7 @@ module Make (Config : CONFIG) = struct
         surround indent 1 (char '{')
           (doc_chunks opts exp ^^ space ^^ string "with" ^^ break 1 ^^ separate_map (break 1) (doc_chunks opts) fexps)
           (char '}')
-    | Comment (comment_type, n, col, contents) -> begin
+    | Comment (comment_type, n, col, contents, _) -> begin
         match comment_type with
         | Lexer.Comment_line -> blank n ^^ string "//" ^^ string contents ^^ require_hardline
         | Lexer.Comment_block -> (
@@ -696,23 +696,31 @@ module Make (Config : CONFIG) = struct
      In this case a hardline must be inserted by the block formatter, so
      we return and additional boolean to indicate this. *)
   and doc_block_exp_chunks opts no_semi chunks =
-    let chunks = Queue.to_seq chunks |> List.of_seq in
     let requires_hardline = ref false in
     let terminator = if no_semi then empty else char ';' in
-    let doc =
-      concat_map_last
-        (fun last chunk ->
-          if last then (
-            match line_comment_opt chunk with
-            | Some contents ->
-                requires_hardline := true;
-                terminator ^^ space ^^ string "//" ^^ string contents ^^ require_hardline
-            | None -> doc_chunk opts chunk ^^ terminator
-          )
-          else doc_chunk opts chunk
-        )
-        chunks
+    let rec splice_into_doc chunks doc_acc =
+      match Queue.peek_opt chunks with
+      | Some chunk ->
+          let _ = Queue.pop chunks in
+          let doc_acc = ref (doc_acc ^^ doc_chunk opts chunk) in
+          let doc_acc =
+            match (chunk, Queue.peek_opt chunks) with
+            | Comment _, _ -> !doc_acc
+            | Spacer _, _ -> !doc_acc
+            | _, Some (Comment (_, _, _, _, trailing)) ->
+                doc_acc := !doc_acc ^^ terminator;
+                (* if current is not a Comment or Spacer, and next is not trailing, then insert a hardline *)
+                if not trailing then doc_acc := !doc_acc ^^ hardline;
+                doc_acc := !doc_acc ^^ doc_chunk opts (Queue.pop chunks);
+                if Queue.peek_opt chunks = None then requires_hardline := true;
+                !doc_acc
+            | _, None -> !doc_acc ^^ terminator
+            | _, _ -> !doc_acc
+          in
+          splice_into_doc chunks doc_acc
+      | None -> doc_acc
     in
+    let doc = splice_into_doc chunks empty in
     (group doc, !requires_hardline)
 
   and doc_chunks ?(ungroup_tuple = false) opts chunks =

test/format/comments.sail

@@ -0,0 +1,68 @@
+function a () -> int = {
+    *R = baz; // comment
+       // comment
+     1// comment
+    // comment
+   // comment
+                                /* com-
+                                ment */
+ }
+
+        // comment
+
+function b () -> int = {
+    let a = {
+        1;
+        // comment
+    };
+
+    let b = {
+        1;// comment
+    };
+
+    let c = {
+        1;// comment
+        // comment
+    };
+
+    let d = {
+        1;// comment
+        // comment
+    };// comment
+
+    let f = {
+        let g1 = {
+            1// comment
+        };
+
+        let g2 = {
+            1
+            // comment
+        };
+
+        let g3 = {
+            1// comment
+          // comment
+        }// comment
+    };
+
+    1
+    // comment
+}
+
+let c=1
+// comment
+let c2=1// comment
+
+
+let  d   ={
+        1
+    // comment
+     }// comment
+
+// comment
+
+// comment
+
+
+// comment

test/format/default/comments.expect

@@ -0,0 +1,65 @@
+function a () -> int = {
+    *R = baz; // comment
+
+    // comment
+    1 // comment
+    // comment
+    // comment
+    /* com-
+    ment */
+}
+
+// comment
+function b () -> int = {
+    let a = {
+        1
+        // comment
+    };
+
+    let b = {
+        1 // comment
+    };
+
+    let c = {
+        1 // comment
+        // comment
+    };
+
+    let d = {
+        1 // comment
+        // comment
+    }; // comment
+
+    let f = {
+        let g1 = {
+            1 // comment
+        };
+
+        let g2 = {
+            1
+            // comment
+        };
+
+        let g3 = {
+            1 // comment
+            // comment
+        } // comment
+    };
+
+    1
+    // comment
+}
+
+let c = 1
+
+// comment
+let c2 =
+    1 // comment
+let d =
+    {
+        1
+        // comment
+    } // comment
+// comment
+// comment
+// comment