Attach res.doc to polyvariant (#8155)

* Attach res.doc to polyvariant

* Refactor

* Update CHANGELOG

* Handle inherited poly types

Author: shulhi

CHANGELOG.md

@@ -28,6 +28,7 @@
 - Fix rewatch swallowing parse warnings (%todo). https://github.com/rescript-lang/rescript/pull/8135
 - Rewatch: log errors and warnings to `stderr`. https://github.com/rescript-lang/rescript/pull/8147 https://github.com/rescript-lang/rescript/pull/8148
 - Rewatch: warn about deprecated package specs `es6`/`es6-global`. https://github.com/rescript-lang/rescript/pull/8146
+- Attach res.doc to polyvariant. https://github.com/rescript-lang/rescript/pull/8155
 
 #### :memo: Documentation
 

compiler/syntax/src/res_core.ml

@@ -5265,24 +5265,25 @@ and parse_constr_decl_args p =
   in
   (constr_args, res)
 
+(* Helper to check if current token is a bar or doc comment followed by a bar *)
+and is_bar_or_doc_comment_then_bar p =
+  Parser.lookahead p (fun state ->
+      match state.Parser.token with
+      | DocComment _ -> (
+        Parser.next state;
+        match state.token with
+        | Bar -> true
+        | _ -> false)
+      | Bar -> true
+      | _ -> false)
+
 (* constr-decl ::=
  *  | constr-name
  *  | attrs constr-name
  *  | constr-name const-args
  *  | attrs constr-name const-args *)
 and parse_type_constructor_declaration_with_bar p =
-  let is_constructor_with_bar p =
-    Parser.lookahead p (fun state ->
-        match state.Parser.token with
-        | DocComment _ -> (
-          Parser.next state;
-          match state.token with
-          | Bar -> true
-          | _ -> false)
-        | Bar -> true
-        | _ -> false)
-  in
-  if is_constructor_with_bar p then (
+  if is_bar_or_doc_comment_then_bar p then (
     let doc_comment_attrs =
       match p.Parser.token with
       | DocComment (loc, s) ->
@@ -5905,29 +5906,84 @@ and parse_tag_spec_full p =
 
 and parse_tag_specs p =
   match p.Parser.token with
-  | Bar ->
-    Parser.next p;
-    let row_field = parse_tag_spec p in
-    row_field :: parse_tag_specs p
+  | (Bar | DocComment _) when is_bar_or_doc_comment_then_bar p ->
+    let doc_comment_attrs =
+      match p.Parser.token with
+      | DocComment (loc, s) ->
+        Parser.next p;
+        [doc_comment_to_attribute loc s]
+      | _ -> []
+    in
+    Parser.expect Bar p;
+    let tag = parse_tag_spec p in
+    let tag_with_doc =
+      match tag with
+      | Parsetree.Rtag (name, attrs, contains_constant, types) ->
+        Parsetree.Rtag
+          (name, doc_comment_attrs @ attrs, contains_constant, types)
+      | Rinherit typ ->
+        Rinherit
+          {typ with ptyp_attributes = doc_comment_attrs @ typ.ptyp_attributes}
+    in
+    tag_with_doc :: parse_tag_specs p
   | _ -> []
 
 and parse_tag_spec p =
-  let attrs = parse_attributes p in
+  let doc_comment_attrs =
+    match p.Parser.token with
+    | DocComment (loc, s) ->
+      Parser.next p;
+      [doc_comment_to_attribute loc s]
+    | _ -> []
+  in
+  let attrs = doc_comment_attrs @ parse_attributes p in
   match p.Parser.token with
   | Hash -> parse_polymorphic_variant_type_spec_hash ~attrs ~full:false p
   | _ ->
     let typ = parse_typ_expr ~attrs p in
     Parsetree.Rinherit typ
 
 and parse_tag_spec_first p =
-  let attrs = parse_attributes p in
   match p.Parser.token with
-  | Bar ->
-    Parser.next p;
-    [parse_tag_spec p]
-  | Hash -> [parse_polymorphic_variant_type_spec_hash ~attrs ~full:false p]
+  | (Bar | DocComment _) when is_bar_or_doc_comment_then_bar p ->
+    let doc_comment_attrs =
+      match p.Parser.token with
+      | DocComment (loc, s) ->
+        Parser.next p;
+        [doc_comment_to_attribute loc s]
+      | _ -> []
+    in
+    Parser.expect Bar p;
+    let tag = parse_tag_spec p in
+    (match tag with
+    | Parsetree.Rtag (name, attrs, contains_constant, types) ->
+      Parsetree.Rtag (name, doc_comment_attrs @ attrs, contains_constant, types)
+    | Rinherit typ ->
+      Rinherit
+        {typ with ptyp_attributes = doc_comment_attrs @ typ.ptyp_attributes})
+    :: parse_tag_specs p
+  | DocComment _ | Hash | At -> (
+    let doc_comment_attrs =
+      match p.Parser.token with
+      | DocComment (loc, s) ->
+        Parser.next p;
+        [doc_comment_to_attribute loc s]
+      | _ -> []
+    in
+    let attrs = doc_comment_attrs @ parse_attributes p in
+    match p.Parser.token with
+    | Hash -> [parse_polymorphic_variant_type_spec_hash ~attrs ~full:false p]
+    | _ -> (
+      let typ = parse_typ_expr ~attrs p in
+      match p.token with
+      | Rbracket ->
+        (* example: [ListStyleType.t] *)
+        [Parsetree.Rinherit typ]
+      | _ ->
+        Parser.expect Bar p;
+        [Parsetree.Rinherit typ; parse_tag_spec p]))
   | _ -> (
-    let typ = parse_typ_expr ~attrs p in
+    let typ = parse_typ_expr p in
     match p.token with
     | Rbracket ->
       (* example: [ListStyleType.t] *)

compiler/syntax/src/res_printer.ml

@@ -1893,18 +1893,44 @@ and print_typ_expr ?inline_record_definitions ~(state : State.t)
         typ_expr.ptyp_loc.Location.loc_start.pos_lnum
         < typ_expr.ptyp_loc.loc_end.pos_lnum
       in
-      let print_row_field = function
+      let print_row_field i = function
         | Parsetree.Rtag ({txt; loc}, attrs, true, []) ->
-          let doc =
+          let comment_attrs, attrs =
+            ParsetreeViewer.partition_doc_comment_attributes attrs
+          in
+          let comment_doc =
+            match comment_attrs with
+            | [] -> Doc.nil
+            | comment_attrs ->
+              print_doc_comments ~sep:Doc.hard_line ~state cmt_tbl comment_attrs
+          in
+          let bar =
+            if i > 0 || comment_attrs <> [] then Doc.text "| "
+            else Doc.if_breaks (Doc.text "| ") Doc.nil
+          in
+          let tag_doc =
             Doc.group
               (Doc.concat
                  [
                    print_attributes ~state attrs cmt_tbl;
                    Doc.concat [Doc.text "#"; print_poly_var_ident txt];
                  ])
           in
-          print_comments doc cmt_tbl loc
-        | Rtag ({txt}, attrs, truth, types) ->
+          Doc.concat [comment_doc; bar; print_comments tag_doc cmt_tbl loc]
+        | Rtag ({txt; loc}, attrs, truth, types) ->
+          let comment_attrs, attrs =
+            ParsetreeViewer.partition_doc_comment_attributes attrs
+          in
+          let comment_doc =
+            match comment_attrs with
+            | [] -> Doc.nil
+            | comment_attrs ->
+              print_doc_comments ~sep:Doc.hard_line ~state cmt_tbl comment_attrs
+          in
+          let bar =
+            if i > 0 || comment_attrs <> [] then Doc.text "| "
+            else Doc.if_breaks (Doc.text "| ") Doc.nil
+          in
           let do_type t =
             match t.Parsetree.ptyp_desc with
             | Ptyp_tuple _ -> print_typ_expr ~state t cmt_tbl
@@ -1919,21 +1945,25 @@ and print_typ_expr ?inline_record_definitions ~(state : State.t)
           let cases =
             if truth then Doc.concat [Doc.line; Doc.text "& "; cases] else cases
           in
-          Doc.group
-            (Doc.concat
-               [
-                 print_attributes ~state attrs cmt_tbl;
-                 Doc.concat [Doc.text "#"; print_poly_var_ident txt];
-                 cases;
-               ])
-        | Rinherit core_type -> print_typ_expr ~state core_type cmt_tbl
-      in
-      let docs = List.map print_row_field row_fields in
-      let cases = Doc.join ~sep:(Doc.concat [Doc.line; Doc.text "| "]) docs in
-      let cases =
-        if docs = [] then cases
-        else Doc.concat [Doc.if_breaks (Doc.text "| ") Doc.nil; cases]
+          let tag_doc =
+            Doc.group
+              (Doc.concat
+                 [
+                   print_attributes ~state attrs cmt_tbl;
+                   Doc.concat [Doc.text "#"; print_poly_var_ident txt];
+                   cases;
+                 ])
+          in
+          Doc.concat [comment_doc; bar; print_comments tag_doc cmt_tbl loc]
+        | Rinherit core_type ->
+          let bar =
+            if i > 0 then Doc.text "| "
+            else Doc.if_breaks (Doc.text "| ") Doc.nil
+          in
+          Doc.concat [bar; print_typ_expr ~state core_type cmt_tbl]
       in
+      let docs = List.mapi print_row_field row_fields in
+      let cases = Doc.join ~sep:Doc.line docs in
       let opening_symbol =
         if closed_flag = Open then Doc.concat [Doc.greater_than; Doc.line]
         else if labels_opt = None then Doc.soft_line

tests/syntax_tests/data/parsing/grammar/typexpr/expected/polyVariant.res.txt

@@ -12,4 +12,11 @@ module type Conjunctive  =
 type nonrec t = [ s]
 type nonrec t = [ ListStyleType.t]
 type nonrec number = [ `1  | `42  | `4244 ]
-type nonrec complexNumbericPolyVar = [ `1 of string  | `2 of (int * string) ]
+type nonrec complexNumbericPolyVar = [ `1 of string  | `2 of (int * string) ]
+type nonrec withDocComments =
+  [ `Foo [@res.doc " First variant "] | `Bar [@res.doc " Second variant "]
+  | `Baz of (int * string) [@res.doc " Third variant with args "]]
+type nonrec singleDocComment = [ `Only [@res.doc " Single variant "]]
+type nonrec mixedDocComments =
+  [ `NoComment  | `WithComment [@res.doc " With comment "]
+  | `AnotherNoComment ]

tests/syntax_tests/data/parsing/grammar/typexpr/polyVariant.res

@@ -26,3 +26,25 @@ type complexNumbericPolyVar = [
   | #1(string)
   | #2(int, string)
 ]
+
+// Doc comments on polymorphic variants
+type withDocComments = [
+  /** First variant */
+  | #Foo
+  /** Second variant */
+  | #Bar
+  /** Third variant with args */
+  | #Baz(int, string)
+]
+
+type singleDocComment = [
+  /** Single variant */
+  #Only
+]
+
+type mixedDocComments = [
+  | #NoComment
+  /** With comment */
+  | #WithComment
+  | #AnotherNoComment
+]

tests/syntax_tests/data/printer/typexpr/expected/variant.res.txt

@@ -123,3 +123,13 @@ type foo = [
   | // before baz
   #Baz
 ] // after baz
+
+// doc comments on inherited polyvariant types
+type rgb = [#Red | #Green | #Blue]
+type color = [
+  | /** Primary colors from RGB */ rgb
+  /** A warm color */
+  | #Orange
+  /** Another warm color */
+  | #Yellow
+]

tests/syntax_tests/data/printer/typexpr/variant.res

@@ -117,3 +117,14 @@ type foo = [
   // before baz
   | #Baz // after baz
 ]
+
+// doc comments on inherited polyvariant types
+type rgb = [#Red | #Green | #Blue]
+type color = [
+  /** Primary colors from RGB */
+  | rgb
+  /** A warm color */
+  | #Orange
+  /** Another warm color */
+  | #Yellow
+]