letop-punning for extension nodes

Author: WardBrian

CHANGES.md

@@ -14,6 +14,7 @@ profile. This started with version 0.26.0.
   `let+ x in ...` when `letop-punning=always`. With `letop-punning=never`, it
   becomes `let+ x = x in ...`. The default is `preserve`, which will
   only use punning when it exists in the source.
+  This also applies to `let%ext` bindings (#<PR_NUMBER>, @WardBrian).
 
 ### Fixed
 

doc/manpage_ocamlformat.mld

@@ -272,14 +272,14 @@ OPTIONS (CODE FORMATTING STYLE)
            a single line. The default value is compact.
 
        --letop-punning={preserve|always|never}
-           Name punning in bindings using extended let operators. preserve
-           uses let-punning only when it exists in the source; the code "let*
-           foo and* z = z in ..." will be left unchanged. always uses
-           let-punning whenever possible; the code "let* foo and* z = z in
-           ..." will be rewritten to "let* foo and* z in ...". never never
-           uses let-punning; the code "let* foo and* z = z in ..." will be
-           rewritten to "let* foo = foo and* z = z in ...". The default value
-           is preserve.
+           Name punning in bindings using extended let operators and let%ext
+           bindings. preserve uses let-punning only when it exists in the
+           source; the code "let* foo and* z = z in ..." will be left
+           unchanged. always uses let-punning whenever possible; the code
+           "let* foo and* z = z in ..." will be rewritten to "let* foo and* z
+           in ...". never never uses let-punning; the code "let* foo and* z =
+           z in ..." will be rewritten to "let* foo = foo and* z = z in ...".
+           The default value is preserve.
 
        --line-endings={lf|crlf}
            Line endings used. lf uses Unix line endings. crlf uses Windows

lib/Conf.ml

@@ -998,7 +998,10 @@ module Formatting = struct
       (fun conf -> conf.fmt_opts.let_module)
 
   let letop_punning =
-    let doc = "Name punning in bindings using extended let operators." in
+    let doc =
+      "Name punning in bindings using extended let operators and \
+       $(i,let%ext) bindings."
+    in
     let names = ["letop-punning"] in
     let all =
       [ Decl.Value.make ~name:"preserve" `Preserve

lib/Extended_ast.ml

@@ -238,15 +238,47 @@ module Parse = struct
           | None -> b.pbop_is_pun
           | Some false -> false
           | Some true -> (
-            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
-            | _ -> false )
+              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
+              | _ -> false )
         in
         {b with pbop_loc= {b.pbop_loc with loc_start; loc_end}; pbop_is_pun}
       in
       Ast_mapper.default_mapper.binding_op m b'
     in
+    let value_bindings (m : Ast_mapper.mapper) vbs =
+      let punning is_extension vb =
+        let is_extension =
+          is_extension || Option.is_some vb.pvb_attributes.attrs_extension
+        in
+        let pvb_is_pun =
+          is_extension
+          &&
+          match prefer_let_puns with
+          | None -> vb.pvb_is_pun
+          | Some false -> false
+          | Some true -> (
+              vb.pvb_is_pun
+              ||
+              match (vb.pvb_pat.ppat_desc, vb.pvb_body) with
+              | ( Ppat_var {txt= v; _}
+                , Pfunction_body {pexp_desc= Pexp_ident {txt= Lident e; _}; _}
+                ) ->
+                  String.equal v e
+              | _ -> false )
+        in
+        (is_extension, {vb with pvb_is_pun})
+      in
+      let vbs' =
+        { vbs with
+          pvbs_bindings=
+            snd @@ List.fold_map ~init:false ~f:punning vbs.pvbs_bindings }
+      in
+      Ast_mapper.default_mapper.value_bindings m vbs'
+    in
     let pat m = function
       | {ppat_desc= Ppat_cons (_ :: _ :: _ :: _ as l); _} as p
         when match List.last_exn l with
@@ -315,7 +347,7 @@ module Parse = struct
           {p with pexp_desc= Pexp_tuple l}
       | e -> Ast_mapper.default_mapper.expr m e
     in
-    Ast_mapper.{default_mapper with expr; pat; binding_op}
+    Ast_mapper.{default_mapper with expr; pat; binding_op; value_bindings}
 
   let ast (type a) (fg : a t) ~ocaml_version ~preserve_beginend
       ~prefer_let_puns ~input_name str : a =

test/failing/tests/letop_broken.ml.broken-ref

@@ -2,7 +2,13 @@ ocamlformat: Cannot process "tests/letop_broken.ml".
   Please report this bug at https://github.com/ocaml-ppx/ocamlformat/issues.
   BUG: comment changed.
 File "tests/letop_broken.ml", line 13, characters 33-40:
+Error: comment (*  3  *) dropped.
+  BUG: comment changed.
+File "tests/letop_broken.ml", line 16, characters 36-43:
 Error: comment (*  3  *) dropped.
   BUG: comment changed.
 File "tests/letop_broken.ml", line 13, characters 43-50:
+Error: comment (*  4  *) dropped.
+  BUG: comment changed.
+File "tests/letop_broken.ml", line 16, characters 46-53:
 Error: comment (*  4  *) dropped.

test/passing/refs.ahrefs/let_punning-denied.ml.ref

@@ -15,7 +15,7 @@ let p =
   x, y, z
 
 let q =
-  let%foo x
+  let%foo x = x
   and y = y
-  and z in
+  and z = z in
   x, y, z

test/passing/refs.ahrefs/let_punning-preferred.ml.ref

@@ -16,6 +16,6 @@ let p =
 
 let q =
   let%foo x
-  and y = y
+  and y
   and z in
   x, y, z

test/passing/refs.default/let_punning-denied.ml.ref

@@ -10,5 +10,5 @@ let p =
   (x, y, z)
 
 let q =
-  let%foo x and y = y and z in
+  let%foo x = x and y = y and z = z in
   (x, y, z)

test/passing/refs.default/let_punning-preferred.ml.ref

@@ -10,5 +10,5 @@ let p =
   (x, y, z)
 
 let q =
-  let%foo x and y = y and z in
+  let%foo x and y and z in
   (x, y, z)

test/passing/refs.janestreet/let_punning-denied.ml.ref

@@ -13,8 +13,8 @@ let p =
 ;;
 
 let q =
-  let%foo x
+  let%foo x = x
   and y = y
-  and z in
+  and z = z in
   x, y, z
 ;;