From 70cd7728995845ee1a6efea14a6ef79fd71d4a4a Mon Sep 17 00:00:00 2001
From: David Vulakh <dvulakh@janestreet.com>
Date: Mon, 2 Dec 2024 18:48:37 -0500
Subject: [PATCH] fix [type%ext nonrec]

Signed-off-by: David Vulakh <dvulakh@janestreet.com>
---
 lib/Fmt_ast.ml                                | 25 +++++++++++--------
 .../passing/tests/extensions-indent.ml.js-ref |  3 +++
 test/passing/tests/extensions-indent.ml.ref   |  6 +++++
 test/passing/tests/extensions.ml.js-ref       |  3 +++
 test/passing/tests/extensions.ml.ref          |  6 +++++
 5 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/lib/Fmt_ast.ml b/lib/Fmt_ast.ml
index 14ed2af016..6fe894f00b 100644
--- a/lib/Fmt_ast.ml
+++ b/lib/Fmt_ast.ml
@@ -3722,8 +3722,7 @@ and fmt_class_params c ctx params =
     (hvbox 0
        (wrap_fits_breaks c.conf "[" "]" (list_fl params fmt_param) $ fmt "@ ") )
 
-and fmt_type_declaration c ?ext ?(pre = "") ?name ?(eq = "=") {ast= decl; _}
-    =
+and fmt_type_declaration c ?(pre = noop) ?name ?(eq = "=") {ast= decl; _} =
   protect c (Td decl)
   @@
   let decl =
@@ -3762,9 +3761,7 @@ and fmt_type_declaration c ?ext ?(pre = "") ?name ?(eq = "=") {ast= decl; _}
       ( hvbox_if
           (Option.is_some ptype_jkind)
           c.conf.fmt_opts.type_decl_indent.v
-          ( str pre
-          $ fmt_extension_suffix c ext
-          $ str " "
+          ( pre $ str " "
           $ hvbox_if
               (not (List.is_empty ptype_params))
               0
@@ -4584,13 +4581,16 @@ and fmt_module_statement c ~attributes ?keyword mod_expr =
 
 and fmt_with_constraint c ctx ~pre = function
   | Pwith_type (lid, td) ->
-      fmt_type_declaration ~pre:(pre ^ " type") c ~name:lid (sub_td ~ctx td)
+      fmt_type_declaration
+        ~pre:(str (pre ^ " type"))
+        c ~name:lid (sub_td ~ctx td)
   | Pwith_module (m1, m2) ->
       str pre $ str " module " $ fmt_longident_loc c m1 $ str " = "
       $ fmt_longident_loc c m2
   | Pwith_typesubst (lid, td) ->
-      fmt_type_declaration ~pre:(pre ^ " type") c ~eq:":=" ~name:lid
-        (sub_td ~ctx td)
+      fmt_type_declaration
+        ~pre:(str (pre ^ " type"))
+        c ~eq:":=" ~name:lid (sub_td ~ctx td)
   | Pwith_modsubst (m1, m2) ->
       str pre $ str " module " $ fmt_longident_loc c m1 $ str " := "
       $ fmt_longident_loc c m2
@@ -4840,10 +4840,13 @@ and fmt_type c ?ext ?eq rec_flag decls ctx =
   let fmt_decl c ctx ~prev ~next:_ decl =
     let first = Option.is_none prev in
     let pre =
-      if first then if is_rec then "type" else "type nonrec" else "and"
+      if first then
+        str "type"
+        $ fmt_extension_suffix c ext
+        $ fmt_if (not is_rec) " nonrec"
+      else str "and"
     in
-    let ext = if first then ext else None in
-    fmt_type_declaration c ~pre ?eq ?ext (sub_td ~ctx decl)
+    fmt_type_declaration c ~pre ?eq (sub_td ~ctx decl)
   in
   let ast x = Td x in
   fmt_item_list c ctx update_config ast fmt_decl decls
diff --git a/test/passing/tests/extensions-indent.ml.js-ref b/test/passing/tests/extensions-indent.ml.js-ref
index ac0206f445..b4538a0a76 100644
--- a/test/passing/tests/extensions-indent.ml.js-ref
+++ b/test/passing/tests/extensions-indent.ml.js-ref
@@ -194,6 +194,9 @@ let%lwt f = function
 ;;
 
 type%any_extension t = < a : 'a >
+[%%any_extension type t = < a : 'a > ]
+type%any_extension nonrec t = < a : 'a >
+[%%any_extension type nonrec t = < a : 'a > ]
 
 let value =
   f
diff --git a/test/passing/tests/extensions-indent.ml.ref b/test/passing/tests/extensions-indent.ml.ref
index 7b730b8c8a..85a8de16ab 100644
--- a/test/passing/tests/extensions-indent.ml.ref
+++ b/test/passing/tests/extensions-indent.ml.ref
@@ -144,6 +144,12 @@ let%lwt f = function _ -> ()
 
 type%any_extension t = < a: 'a >
 
+[%%any_extension type t = < a: 'a > ]
+
+type%any_extension nonrec t = < a: 'a >
+
+[%%any_extension type nonrec t = < a: 'a > ]
+
 let value = f [%any_extension function 0 -> false | _ -> true]
 
 let value = [%any_extension fun x -> y] x
diff --git a/test/passing/tests/extensions.ml.js-ref b/test/passing/tests/extensions.ml.js-ref
index ac0206f445..b4538a0a76 100644
--- a/test/passing/tests/extensions.ml.js-ref
+++ b/test/passing/tests/extensions.ml.js-ref
@@ -194,6 +194,9 @@ let%lwt f = function
 ;;
 
 type%any_extension t = < a : 'a >
+[%%any_extension type t = < a : 'a > ]
+type%any_extension nonrec t = < a : 'a >
+[%%any_extension type nonrec t = < a : 'a > ]
 
 let value =
   f
diff --git a/test/passing/tests/extensions.ml.ref b/test/passing/tests/extensions.ml.ref
index 3dbe10d019..8687fbac29 100644
--- a/test/passing/tests/extensions.ml.ref
+++ b/test/passing/tests/extensions.ml.ref
@@ -144,6 +144,12 @@ let%lwt f = function _ -> ()
 
 type%any_extension t = < a: 'a >
 
+[%%any_extension type t = < a: 'a > ]
+
+type%any_extension nonrec t = < a: 'a >
+
+[%%any_extension type nonrec t = < a: 'a > ]
+
 let value = f [%any_extension function 0 -> false | _ -> true]
 
 let value = [%any_extension fun x -> y] x