Fix haddocks on infix constructors

CHANGELOG.md

@@ -1,3 +1,8 @@
+## Unreleased
+
+* Fix Haddock comments on infix constructors
+  [Issue 758](https://github.com/tweag/ormolu/issues/758)
+
 ## Ormolu 0.7.5.0
 
 * Switched to `ghc-lib-parser-9.10`, with the following new syntactic features/behaviors:

data/examples/declaration/data/infix-haddocks-out.hs

@@ -0,0 +1,51 @@
+{- https://github.com/tweag/ormolu/issues/758 -}
+
+data A
+  = -- | Docs for :#
+    A :# A
+
+data WithDocs
+  = forall left right.
+    (Show left) =>
+    -- | Docs for left arg
+    left
+      -- | Docs for op
+      :*:
+      -- | Docs for right arg
+      right
+
+data MixedDocs
+  = forall left right.
+    (Show left) =>
+    left -- ^ before
+      :*:
+      -- | after
+      right
+
+data DocPartial
+  = Left -- ^ left docs
+    -- on multiple
+    -- lines
+      :*:
+      Right
+  | -- | op
+    Left
+      :*:
+      Right
+  | Left
+      :*:
+      -- | right
+      Right
+  | -- | op
+    Left
+      :*:
+      Right
+  | -- | op
+    Left
+      :*:
+      Right
+
+data NoDocs
+  = Left
+      :*:
+      Right

data/examples/declaration/data/infix-haddocks.hs

@@ -0,0 +1,43 @@
+{- https://github.com/tweag/ormolu/issues/758 -}
+
+data A = A :# A -- ^ Docs for :#
+
+data WithDocs
+  = forall left right.
+    Show left =>
+    left -- ^ Docs for left arg
+    :*: -- ^ Docs for op
+    right -- ^ Docs for right arg
+
+data MixedDocs
+  -- | before
+  = forall left right.
+    Show left =>
+    left :*: right
+  -- ^ after
+
+data DocPartial
+  = Left -- ^ left docs
+         -- on multiple
+         -- lines
+      :*: Right
+  | Left
+      :*: -- ^ op
+      Right
+  | Left
+      :*:
+      -- | right
+      Right
+  | -- | op
+    Left
+      :*:
+      Right
+  | Left
+      :*:
+      Right
+    -- ^ op
+
+data NoDocs
+  = Left
+    :*:
+    Right

src/Ormolu/Printer/Meat/Declaration/Data.hs

@@ -13,7 +13,7 @@ where
 import Control.Monad
 import Data.List.NonEmpty (NonEmpty (..))
 import Data.List.NonEmpty qualified as NE
-import Data.Maybe (isJust, mapMaybe, maybeToList)
+import Data.Maybe (isJust, isNothing, mapMaybe, maybeToList)
 import Data.Void
 import GHC.Data.Strict qualified as Strict
 import GHC.Hs
@@ -125,99 +125,128 @@ p_dataDecl style name tyVars getTyVarLoc p_tyVar fixity HsDataDefn {..} = do
   unless (null dd_derivs) breakpoint
   inci $ sep newline (located' p_hsDerivingClause) dd_derivs
 
-p_conDecl ::
-  Bool ->
-  ConDecl GhcPs ->
-  R ()
-p_conDecl singleConstRec = \case
-  ConDeclGADT {..} -> do
-    mapM_ (p_hsDoc Pipe True) con_doc
-    let conDeclSpn =
-          fmap getLocA (NE.toList con_names)
-            <> [getLocA con_bndrs]
-            <> maybeToList (fmap getLocA con_mb_cxt)
-            <> conArgsSpans
-          where
-            conArgsSpans = case con_g_args of
-              PrefixConGADT NoExtField xs -> getLocA . hsScaledThing <$> xs
-              RecConGADT _ x -> [getLocA x]
-    switchLayout conDeclSpn $ do
-      let c :| cs = con_names
-      p_rdrName c
-      unless (null cs) . inci $ do
-        commaDel
-        sep commaDel p_rdrName cs
-      inci $ do
-        let conTy = case con_g_args of
-              PrefixConGADT NoExtField xs ->
-                let go (HsScaled a b) t = addCLocA t b (HsFunTy NoExtField a b t)
-                 in foldr go con_res_ty xs
-              RecConGADT _ r ->
-                addCLocA r con_res_ty $
-                  HsFunTy
-                    NoExtField
-                    (HsUnrestrictedArrow noAnn)
-                    (la2la $ HsRecTy noAnn <$> r)
-                    con_res_ty
-            qualTy = case con_mb_cxt of
-              Nothing -> conTy
-              Just qs ->
-                addCLocA qs conTy $
-                  HsQualTy NoExtField qs conTy
-            quantifiedTy :: LHsType GhcPs
-            quantifiedTy =
-              addCLocA con_bndrs qualTy $
-                hsOuterTyVarBndrsToHsType (unLoc con_bndrs) qualTy
-        space
-        txt "::"
-        if hasDocStrings (unLoc con_res_ty)
-          then newline
-          else breakpoint
-        located quantifiedTy p_hsType
-  ConDeclH98 {..} -> do
-    mapM_ (p_hsDoc Pipe True) con_doc
-    let conNameSpn = getLocA con_name
-        conNameWithContextSpn =
-          [ RealSrcSpan real Strict.Nothing
-            | EpaSpan (RealSrcSpan real _) <-
-                mapMaybe (matchAddEpAnn AnnForall) con_ext
-          ]
-            <> fmap getLocA con_ex_tvs
-            <> maybeToList (fmap getLocA con_mb_cxt)
-            <> [conNameSpn]
-        conDeclSpn = conNameSpn : conArgsSpans
-          where
-            conArgsSpans = case con_args of
-              PrefixCon [] xs -> getLocA . hsScaledThing <$> xs
-              PrefixCon (v : _) _ -> absurd v
-              RecCon l -> [getLocA l]
-              InfixCon x y -> getLocA . hsScaledThing <$> [x, y]
-    switchLayout conNameWithContextSpn $ do
-      when con_forall $ do
-        p_forallBndrs ForAllInvis p_hsTyVarBndr con_ex_tvs
-        breakpoint
-      forM_ con_mb_cxt p_lhsContext
-    switchLayout conDeclSpn $ case con_args of
-      PrefixCon [] xs -> do
+p_conDecl :: Bool -> ConDecl GhcPs -> R ()
+p_conDecl _ ConDeclGADT {..} = do
+  mapM_ (p_hsDoc Pipe True) con_doc
+  switchLayout conDeclSpn $ do
+    let c :| cs = con_names
+    p_rdrName c
+    unless (null cs) . inci $ do
+      commaDel
+      sep commaDel p_rdrName cs
+    inci $ do
+      let conTy = case con_g_args of
+            PrefixConGADT NoExtField xs ->
+              let go (HsScaled a b) t = addCLocA t b (HsFunTy NoExtField a b t)
+               in foldr go con_res_ty xs
+            RecConGADT _ r ->
+              addCLocA r con_res_ty $
+                HsFunTy
+                  NoExtField
+                  (HsUnrestrictedArrow noAnn)
+                  (la2la $ HsRecTy noAnn <$> r)
+                  con_res_ty
+          qualTy = case con_mb_cxt of
+            Nothing -> conTy
+            Just qs ->
+              addCLocA qs conTy $
+                HsQualTy NoExtField qs conTy
+          quantifiedTy :: LHsType GhcPs
+          quantifiedTy =
+            addCLocA con_bndrs qualTy $
+              hsOuterTyVarBndrsToHsType (unLoc con_bndrs) qualTy
+      space
+      txt "::"
+      if hasDocStrings (unLoc con_res_ty)
+        then newline
+        else breakpoint
+      located quantifiedTy p_hsType
+  where
+    conDeclSpn =
+      fmap getLocA (NE.toList con_names)
+        <> [getLocA con_bndrs]
+        <> maybeToList (fmap getLocA con_mb_cxt)
+        <> conArgsSpans
+    conArgsSpans = case con_g_args of
+      PrefixConGADT NoExtField xs -> getLocA . hsScaledThing <$> xs
+      RecConGADT _ x -> [getLocA x]
+p_conDecl singleConstRec ConDeclH98 {..} =
+  case con_args of
+    PrefixCon (_ :: [Void]) xs -> do
+      renderConDoc
+      renderContext
+      switchLayout conDeclSpn $ do
         p_rdrName con_name
         let args = hsScaledThing <$> xs
             argsHaveDocs = conArgsHaveHaddocks args
             delimiter = if argsHaveDocs then newline else breakpoint
         unless (null xs) delimiter
         inci . sitcc $
           sep delimiter (sitcc . located' p_hsType) args
-      PrefixCon (v : _) _ -> absurd v
-      RecCon l -> do
+    RecCon l -> do
+      renderConDoc
+      renderContext
+      switchLayout conDeclSpn $ do
         p_rdrName con_name
         breakpoint
         inciIf (not singleConstRec) (located l p_conDeclFields)
-      InfixCon (HsScaled _ x) (HsScaled _ y) -> do
-        located x p_hsType
-        breakpoint
+    InfixCon (HsScaled _ l) (HsScaled _ r) -> do
+      -- manually render these
+      let (lType, larg_doc) = splitDocTy l
+      let (rType, rarg_doc) = splitDocTy r
+
+      -- the constructor haddock can go on top of the entire constructor
+      -- only if neither argument has haddocks
+      let putConDocOnTop = isNothing larg_doc && isNothing rarg_doc
+
+      when putConDocOnTop renderConDoc
+      renderContext
+      switchLayout conDeclSpn $ do
+        -- the left arg haddock can use pipe only if the infix constructor has docs
+        if isJust con_doc
+          then do
+            mapM_ (p_hsDoc Pipe True) larg_doc
+            located lType p_hsType
+            breakpoint
+          else do
+            located lType p_hsType
+            case larg_doc of
+              Just doc -> space >> p_hsDoc Caret True doc
+              Nothing -> breakpoint
         inci $ do
+          unless putConDocOnTop renderConDoc
           p_rdrName con_name
-          space
-          located y p_hsType
+          case rarg_doc of
+            Just doc -> newline >> p_hsDoc Pipe True doc
+            Nothing -> breakpoint
+          located rType p_hsType
+  where
+    renderConDoc = mapM_ (p_hsDoc Pipe True) con_doc
+    renderContext =
+      switchLayout conNameWithContextSpn $ do
+        when con_forall $ do
+          p_forallBndrs ForAllInvis p_hsTyVarBndr con_ex_tvs
+          breakpoint
+        forM_ con_mb_cxt p_lhsContext
+
+    conNameWithContextSpn =
+      [ RealSrcSpan real Strict.Nothing
+        | EpaSpan (RealSrcSpan real _) <-
+            mapMaybe (matchAddEpAnn AnnForall) con_ext
+      ]
+        <> fmap getLocA con_ex_tvs
+        <> maybeToList (fmap getLocA con_mb_cxt)
+        <> [conNameSpn]
+    conDeclSpn = conNameSpn : conArgsSpans
+    conNameSpn = getLocA con_name
+    conArgsSpans = case con_args of
+      PrefixCon (_ :: [Void]) xs -> getLocA . hsScaledThing <$> xs
+      RecCon l -> [getLocA l]
+      InfixCon x y -> getLocA . hsScaledThing <$> [x, y]
+
+    splitDocTy = \case
+      L _ (HsDocTy _ ty doc) -> (ty, Just doc)
+      ty -> (ty, Nothing)
 
 p_lhsContext ::
   LHsContext GhcPs ->