Share argument locations and clarify callback layout rules

Signed-off-by: Christoph Knittel <ck@cca.io>

Author: cknitt

compiler/syntax/src/res_comments_table.ml

@@ -986,13 +986,8 @@ and walk_expression expr t comments =
     attach t.trailing call_expr.Parsetree.pexp_loc after_expr;
     walk_list
       (arguments
-      |> List.map (fun (lbl, expr) ->
-          let loc =
-            match lbl with
-            | Asttypes.Labelled {loc} | Optional {loc} ->
-              {loc with loc_end = expr.Parsetree.pexp_loc.loc_end}
-            | _ -> expr.pexp_loc
-          in
+      |> List.map (fun ((_, expr) as argument) ->
+          let loc = Parsetree_viewer.argument_loc argument in
           ExprArgument {expr; loc}))
       t rest
   in

compiler/syntax/src/res_parsetree_viewer.ml

@@ -1,5 +1,11 @@
 open Parsetree
 
+let argument_loc (lbl, (arg : Parsetree.expression)) =
+  match lbl with
+  | Asttypes.Labelled {loc} | Optional {loc} ->
+    {loc with loc_end = arg.pexp_loc.loc_end}
+  | Nolabel -> arg.pexp_loc
+
 let arrow_type ct =
   match ct with
   | {ptyp_desc = Ptyp_arrow {params; ret}; ptyp_attributes = attrs} ->

compiler/syntax/src/res_parsetree_viewer.mli

@@ -1,3 +1,7 @@
+(* Full argument span used for comment attachment and printing. For labeled
+ * and optional arguments it starts at the label and ends at the expression. *)
+val argument_loc : Asttypes.arg_label * Parsetree.expression -> Location.t
+
 (* Restructures a nested tree of arrow types into its args & returnType
  * The parsetree contains: a => b => c => d, for printing purposes
  * we restructure the tree into (a, b, c) and its returnType d *)

compiler/syntax/src/res_printer.ml

@@ -93,18 +93,29 @@ let has_trailing_comments tbl loc =
   | None -> false
   | _ -> true
 
-(* Match the full argument location used by comment attachment. *)
-let argument_loc (lbl, (arg : Parsetree.expression)) =
-  match lbl with
-  | Asttypes.Labelled {loc} | Optional {loc} ->
-    {loc with loc_end = arg.pexp_loc.loc_end}
-  | Nolabel -> arg.pexp_loc
-
 let has_leading_comments tbl loc =
   match Hashtbl.find_opt tbl.Comment_table.leading loc with
   | None -> false
   | _ -> true
 
+(* Compact layouts print expression comments but not comments on the full
+ * labeled argument. For unlabeled arguments, leading expression comments
+ * already work in the compact layout; forcing a different layout can make
+ * comments exposed by removing parameter parentheses unstable. *)
+let argument_requires_regular_layout cmt_tbl ((lbl, arg) as argument) =
+  let loc = Parsetree_viewer.argument_loc argument in
+  let has_leading_label_comments =
+    match lbl with
+    | Asttypes.Nolabel -> false
+    | Labelled _ | Optional _ ->
+      has_leading_comments cmt_tbl loc
+      || has_leading_comments cmt_tbl arg.Parsetree.pexp_loc
+  in
+  (* Trailing comments can escape compact layouts when the body breaks. *)
+  has_leading_label_comments
+  || has_trailing_comments cmt_tbl loc
+  || has_trailing_comments cmt_tbl arg.pexp_loc
+
 let print_multiline_comment_content txt =
   (* Turns
    *         |* first line
@@ -4533,35 +4544,19 @@ and print_pexp_apply ~state expr cmt_tbl =
       | Braced braces -> print_braces doc call_expr braces
       | Nothing -> doc
     in
-    (* Use the regular layout for comments attached to arguments. Compact
-     * callback layouts can detach trailing comments when the body breaks and
-     * skip leading comments attached to the full labeled argument. *)
-    let args_have_comments =
-      List.exists
-        (fun (lbl, (arg : Parsetree.expression)) ->
-          let loc = argument_loc (lbl, arg) in
-          let has_leading_label_comments =
-            match lbl with
-            | Asttypes.Nolabel -> false
-            | Labelled _ | Optional _ ->
-              has_leading_comments cmt_tbl loc
-              || has_leading_comments cmt_tbl arg.pexp_loc
-          in
-          has_leading_label_comments
-          || has_trailing_comments cmt_tbl loc
-          || has_trailing_comments cmt_tbl arg.pexp_loc)
-        args
+    let requires_regular_layout =
+      List.exists (argument_requires_regular_layout cmt_tbl) args
     in
     let args_doc, maybe_break_parent =
       if
-        (not args_have_comments)
+        (not requires_regular_layout)
         && Parsetree_viewer.requires_special_callback_printing_first_arg args
       then
         ( print_arguments_with_callback_in_first_position ~state ~partial args
             cmt_tbl,
           Doc.nil )
       else if
-        (not args_have_comments)
+        (not requires_regular_layout)
         && Parsetree_viewer.requires_special_callback_printing_last_arg args
       then
         let args_doc =
@@ -5137,7 +5132,8 @@ and print_arguments ~state ~partial
       List.exists
         (fun ((_, arg) as argument) ->
           Parsetree_viewer.is_fun_expr arg
-          && (has_any_trailing_line_comment cmt_tbl (argument_loc argument)
+          && (has_any_trailing_line_comment cmt_tbl
+                (Parsetree_viewer.argument_loc argument)
              || has_any_trailing_line_comment cmt_tbl arg.pexp_loc))
         args
     in

tests/syntax_tests/res_test.ml

@@ -33,39 +33,38 @@ let () =
     Filename.concat data_dir "printer/comments/callbackTrailing.res"
   in
   let source = IO.read_file ~filename in
-  let format ~width source =
+  let parse source =
     let result =
-      Res_driver.parse_implementation_from_source ~for_printer:true
-        ~display_filename:filename ~source
+      Res_driver.parse_implementation_from_source ~display_filename:filename
+        ~source
     in
     assert (not result.invalid);
-    let printed =
-      Res_printer.print_implementation ~width result.parsetree
-        ~comments:result.comments
-    in
-    let reparsed =
-      Res_driver.parse_implementation_from_source ~for_printer:true
-        ~display_filename:filename ~source:printed
-    in
-    assert (not reparsed.invalid);
-    let comment_texts comments =
-      List.map (fun comment -> String.trim (Res_comment.txt comment)) comments
-    in
-    if comment_texts result.comments <> comment_texts reparsed.comments then
-      failwith
-        (Printf.sprintf
-           "Callback formatting changed comments at width %d.\n\
-            Source:\n\
-            %s\n\
-            Printed:\n\
-            %s"
-           width source printed);
-    printed
+    result
+  in
+  let format ~width result =
+    Res_printer.print_implementation ~width result.Res_driver.parsetree
+      ~comments:result.comments
+  in
+  let comment_texts result =
+    List.map
+      (fun comment -> String.trim (Res_comment.txt comment))
+      result.Res_driver.comments
   in
   List.iter
     (fun width ->
-      let printed = format ~width source in
-      let reprinted = format ~width printed in
+      let original = parse source in
+      let printed = format ~width original in
+      let reparsed = parse printed in
+      if comment_texts original <> comment_texts reparsed then
+        failwith
+          (Printf.sprintf
+             "Callback formatting changed comments at width %d.\n\
+              Source:\n\
+              %s\n\
+              Printed:\n\
+              %s"
+             width source printed);
+      let reprinted = format ~width reparsed in
       if printed <> reprinted then
         failwith
           (Printf.sprintf