Fix comments getting dropped by letop-punning=always

WardBrian · WardBrian · commit 68d651a12cd0 · 2025-11-03T10:47:32.000-05:00
diff --git a/lib/Extended_ast.ml b/lib/Extended_ast.ml
@@ -241,8 +241,8 @@ module Parse = struct
               b.pbop_is_pun
               ||
               match (b.pbop_pat.ppat_desc, b.pbop_exp.pexp_desc) with
-              | Ppat_var {txt= v; _}, Pexp_ident {txt= Lident e; _} ->
-                  String.equal v e
+              | Ppat_var {txt; _}, Pexp_ident {txt= Lident e; _} ->
+                  String.equal txt e
               | _ -> false )
         in
         {b with pbop_loc= {b.pbop_loc with loc_start; loc_end}; pbop_is_pun}
@@ -252,6 +252,8 @@ module Parse = struct
     let value_bindings (m : Ast_mapper.mapper) vbs =
       let punning is_extension vb =
         let is_extension =
+          (* [and] nodes don't have extensions, so we need to track if the
+             earlier [let] did *)
           is_extension || Option.is_some vb.pvb_attributes.attrs_extension
         in
         let pvb_is_pun =
@@ -264,10 +266,10 @@ module Parse = struct
               vb.pvb_is_pun
               ||
               match (vb.pvb_pat.ppat_desc, vb.pvb_body) with
-              | ( Ppat_var {txt= v; _}
+              | ( Ppat_var {txt; _}
                 , Pfunction_body {pexp_desc= Pexp_ident {txt= Lident e; _}; _}
                 ) ->
-                  String.equal v e
+                  String.equal txt e
               | _ -> false )
         in
         (is_extension, {vb with pvb_is_pun})
diff --git a/lib/Fmt_ast.ml b/lib/Fmt_ast.ml
@@ -2572,14 +2572,16 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
              $ fmt_atrs ) )
   | Pexp_let (lbs, body, loc_in) ->
       let bindings =
-        Sugar.Let_binding.of_let_bindings ~ctx lbs.pvbs_bindings
+        Sugar.Let_binding.of_let_bindings ~ctx ~cmts:c.cmts lbs.pvbs_bindings
       in
       let fmt_expr = fmt_expression c (sub_exp ~ctx body) in
       pro
       $ fmt_let_bindings c ~ctx0:ctx ~parens ~fmt_atrs ~fmt_expr ~has_attr
           ~loc_in lbs.pvbs_rec bindings body
   | Pexp_letop {let_; ands; body; loc_in} ->
-      let bd = Sugar.Let_binding.of_binding_ops (let_ :: ands) in
+      let bd =
+        Sugar.Let_binding.of_binding_ops ~cmts:c.cmts (let_ :: ands)
+      in
       let fmt_expr = fmt_expression c (sub_exp ~ctx body) in
       pro
       $ fmt_let_bindings c ~ctx0:ctx ~parens ~fmt_atrs ~fmt_expr ~has_attr
@@ -3268,7 +3270,7 @@ and fmt_class_expr c ({ast= exp; ctx= ctx0} as xexp) =
         | _ -> c.conf.fmt_opts.indent_after_in.v
       in
       let bindings =
-        Sugar.Let_binding.of_let_bindings ~ctx lbs.pvbs_bindings
+        Sugar.Let_binding.of_let_bindings ~ctx ~cmts:c.cmts lbs.pvbs_bindings
       in
       let fmt_expr = fmt_class_expr c (sub_cl ~ctx body) in
       let has_attr = not (List.is_empty pcl_attributes) in
@@ -4696,7 +4698,9 @@ and fmt_structure_item c ~last:last_item ~semisemi {ctx= parent_ctx; ast= si}
       let fmt_item c ctx ~prev ~next b =
         let first = Option.is_none prev in
         let last = Option.is_none next in
-        let b = Sugar.Let_binding.of_let_binding ~ctx ~first b in
+        let b =
+          Sugar.Let_binding.of_let_binding ~ctx ~first ~cmts:c.cmts b
+        in
         let epi =
           match c.conf.fmt_opts.let_binding_spacing.v with
           | `Compact -> None
diff --git a/lib/Sugar.ml b/lib/Sugar.ml
@@ -123,7 +123,15 @@ module Let_binding = struct
     ; lb_attrs: ext_attrs
     ; lb_loc: Location.t }
 
-  let of_let_binding ~ctx ~first vb =
+  let of_let_binding ~ctx ~first ~cmts vb =
+    if vb.pvb_is_pun then
+      Cmts.relocate cmts
+        ~src:
+          ( match vb.pvb_body with
+          | Pfunction_body e -> e.pexp_loc
+          | Pfunction_cases (_, l, _) ->
+              (* NB: should be impossible for puns *) l )
+        ~before:vb.pvb_pat.ppat_loc ~after:vb.pvb_pat.ppat_loc ;
     { lb_op= Location.{txt= (if first then "let" else "and"); loc= none}
     ; lb_pat= sub_pat ~ctx vb.pvb_pat
     ; lb_args= vb.pvb_args
@@ -133,11 +141,14 @@ module Let_binding = struct
     ; lb_attrs= vb.pvb_attributes
     ; lb_loc= vb.pvb_loc }
 
-  let of_let_bindings ~ctx =
-    List.mapi ~f:(fun i -> of_let_binding ~ctx ~first:(i = 0))
+  let of_let_bindings ~ctx ~cmts =
+    List.mapi ~f:(fun i -> of_let_binding ~ctx ~first:(i = 0) ~cmts)
 
-  let of_binding_ops bos =
+  let of_binding_ops ~cmts bos =
     List.map bos ~f:(fun bo ->
+        if bo.pbop_is_pun then
+          Cmts.relocate cmts ~src:bo.pbop_exp.pexp_loc
+            ~before:bo.pbop_pat.ppat_loc ~after:bo.pbop_pat.ppat_loc ;
         let ctx = Bo bo in
         { lb_op= bo.pbop_op
         ; lb_pat= sub_pat ~ctx bo.pbop_pat
diff --git a/lib/Sugar.mli b/lib/Sugar.mli
@@ -47,9 +47,11 @@ module Let_binding : sig
     ; lb_attrs: ext_attrs
     ; lb_loc: Location.t }
 
-  val of_let_binding : ctx:Ast.t -> first:bool -> value_binding -> t
+  val of_let_binding :
+    ctx:Ast.t -> first:bool -> cmts:Cmts.t -> value_binding -> t
 
-  val of_let_bindings : ctx:Ast.t -> value_binding list -> t list
+  val of_let_bindings :
+    ctx:Ast.t -> cmts:Cmts.t -> value_binding list -> t list
 
-  val of_binding_ops : binding_op list -> t list
+  val of_binding_ops : cmts:Cmts.t -> binding_op list -> t list
 end
diff --git a/test/failing/dune.inc b/test/failing/dune.inc
@@ -181,19 +181,6 @@
  (package ocamlformat)
  (action (diff tests/js_test.ml.broken-ref js_test.ml.output)))
 
-(rule
- (deps tests/.ocamlformat )
- (package ocamlformat)
- (action
-   (with-outputs-to letop_broken.ml.output
-     (with-accepted-exit-codes 1
-       (run %{bin:ocamlformat} --ocaml-version=4.14 --letop-punning=always %{dep:tests/letop_broken.ml})))))
-
-(rule
- (alias runtest)
- (package ocamlformat)
- (action (diff tests/letop_broken.ml.broken-ref letop_broken.ml.output)))
-
 (rule
  (deps tests/.ocamlformat )
  (package ocamlformat)
diff --git a/test/failing/tests/letop_broken.ml b/test/failing/tests/letop_broken.ml
diff --git a/test/failing/tests/letop_broken.ml.broken-ref b/test/failing/tests/letop_broken.ml.broken-ref
diff --git a/test/failing/tests/letop_broken.ml.opts b/test/failing/tests/letop_broken.ml.opts
diff --git a/test/passing/refs.ahrefs/let_punning-denied.ml.ref b/test/passing/refs.ahrefs/let_punning-denied.ml.ref
@@ -19,3 +19,19 @@ let q =
   and y = y
   and z = z in
   x, y, z
+
+let r =
+  let* (* 1 *) x (* 2 *) = (* 3 *) x (* 4 *)
+  and* (* 5 *) y =
+    y
+    (* 6 *)
+  in
+  x, y
+
+let s =
+  let%foo (* 1 *) x (* 2 *) = (* 3 *) x (* 4 *)
+  and (* 5 *) y =
+    y
+    (* 6 *)
+  in
+  x, y
diff --git a/test/passing/refs.ahrefs/let_punning-preferred.ml.ref b/test/passing/refs.ahrefs/let_punning-preferred.ml.ref
@@ -19,3 +19,25 @@ let q =
   and y
   and z in
   x, y, z
+
+let r =
+  let*
+    (* 1 *)
+    (* 3 *) x
+  (* 2 *)
+  (* 4 *)
+  and* (* 5 *) y
+    (* 6 *)
+  in
+  x, y
+
+let s =
+  let%foo
+    (* 1 *)
+    (* 3 *) x
+  (* 2 *)
+  (* 4 *)
+  and (* 5 *) y
+    (* 6 *)
+  in
+  x, y
diff --git a/test/passing/refs.ahrefs/let_punning.ml.ref b/test/passing/refs.ahrefs/let_punning.ml.ref
@@ -19,3 +19,17 @@ let q =
   and y = y
   and z in
   x, y, z
+
+let r =
+  let* (* 1 *) x (* 2 *) = (* 3 *) x (* 4 *)
+  and* (* 5 *) y
+    (* 6 *)
+  in
+  x, y
+
+let s =
+  let%foo (* 1 *) x (* 2 *) = (* 3 *) x (* 4 *)
+  and (* 5 *) y
+    (* 6 *)
+  in
+  x, y
diff --git a/test/passing/refs.default/let_punning-denied.ml.ref b/test/passing/refs.default/let_punning-denied.ml.ref
@@ -12,3 +12,19 @@ let p =
 let q =
   let%foo x = x and y = y and z = z in
   (x, y, z)
+
+let r =
+  let* (* 1 *) x (* 2 *) = (* 3 *) x (* 4 *)
+  and* (* 5 *) y =
+    y
+    (* 6 *)
+  in
+  (x, y)
+
+let s =
+  let%foo (* 1 *) x (* 2 *) = (* 3 *) x (* 4 *)
+  and (* 5 *) y =
+    y
+    (* 6 *)
+  in
+  (x, y)
diff --git a/test/passing/refs.default/let_punning-preferred.ml.ref b/test/passing/refs.default/let_punning-preferred.ml.ref
@@ -12,3 +12,25 @@ let p =
 let q =
   let%foo x and y and z in
   (x, y, z)
+
+let r =
+  let*
+      (* 1 *)
+      (* 3 *) x
+  (* 2 *)
+  (* 4 *)
+  and* (* 5 *) y
+    (* 6 *)
+  in
+  (x, y)
+
+let s =
+  let%foo
+      (* 1 *)
+      (* 3 *) x
+  (* 2 *)
+  (* 4 *)
+  and (* 5 *) y
+    (* 6 *)
+  in
+  (x, y)
diff --git a/test/passing/refs.default/let_punning.ml.ref b/test/passing/refs.default/let_punning.ml.ref
@@ -12,3 +12,19 @@ let p =
 let q =
   let%foo x = x and y = y and z = z in
   (x, y, z)
+
+let r =
+  let* (* 1 *) x (* 2 *) = (* 3 *) x (* 4 *)
+  and* (* 5 *) y =
+    y
+    (* 6 *)
+  in
+  (x, y)
+
+let s =
+  let%foo (* 1 *) x (* 2 *) = (* 3 *) x (* 4 *)
+  and (* 5 *) y =
+    y
+    (* 6 *)
+  in
+  (x, y)
diff --git a/test/passing/refs.janestreet/let_punning-denied.ml.ref b/test/passing/refs.janestreet/let_punning-denied.ml.ref
@@ -18,3 +18,21 @@ let q =
   and z = z in
   x, y, z
 ;;
+
+let r =
+  let* (* 1 *) x (* 2 *) = (* 3 *) x (* 4 *)
+  and* (* 5 *) y =
+    y
+    (* 6 *)
+  in
+  x, y
+;;
+
+let s =
+  let%foo (* 1 *) x (* 2 *) = (* 3 *) x (* 4 *)
+  and (* 5 *) y =
+    y
+    (* 6 *)
+  in
+  x, y
+;;
diff --git a/test/passing/refs.janestreet/let_punning-preferred.ml.ref b/test/passing/refs.janestreet/let_punning-preferred.ml.ref
@@ -18,3 +18,29 @@ let q =
   and z in
   x, y, z
 ;;
+
+let r =
+  let*
+      (* 1 *)
+      (* 3 *)
+        x
+  (* 2 *)
+  (* 4 *)
+  and* (* 5 *) y
+    (* 6 *)
+  in
+  x, y
+;;
+
+let s =
+  let%foo
+      (* 1 *)
+      (* 3 *)
+        x
+  (* 2 *)
+  (* 4 *)
+  and (* 5 *) y
+    (* 6 *)
+  in
+  x, y
+;;
diff --git a/test/passing/refs.janestreet/let_punning.ml.ref b/test/passing/refs.janestreet/let_punning.ml.ref
@@ -18,3 +18,21 @@ let q =
   and z = z in
   x, y, z
 ;;
+
+let r =
+  let* (* 1 *) x (* 2 *) = (* 3 *) x (* 4 *)
+  and* (* 5 *) y =
+    y
+    (* 6 *)
+  in
+  x, y
+;;
+
+let s =
+  let%foo (* 1 *) x (* 2 *) = (* 3 *) x (* 4 *)
+  and (* 5 *) y =
+    y
+    (* 6 *)
+  in
+  x, y
+;;
diff --git a/test/passing/refs.ocamlformat/let_punning-denied.ml.ref b/test/passing/refs.ocamlformat/let_punning-denied.ml.ref
diff --git a/test/passing/refs.ocamlformat/let_punning-preferred.ml.ref b/test/passing/refs.ocamlformat/let_punning-preferred.ml.ref
diff --git a/test/passing/refs.ocamlformat/let_punning.ml.ref b/test/passing/refs.ocamlformat/let_punning.ml.ref
diff --git a/test/passing/tests/let_punning.ml b/test/passing/tests/let_punning.ml
