Improve haddock comments

CODEOWNERS

@@ -16,7 +16,7 @@
 /plugins/hls-floskell-plugin @Ailrun
 /plugins/hls-fourmolu-plugin @georgefst
 /plugins/hls-gadt-plugin @July541
-/plugins/hls-haddock-comments-plugin @berberman
+/plugins/hls-haddock-comments-plugin @berberman @kokobd
 /plugins/hls-hlint-plugin @jneira @eddiemundo
 /plugins/hls-module-name-plugin
 /plugins/hls-ormolu-plugin @georgefst

plugins/hls-haddock-comments-plugin/hls-haddock-comments-plugin.cabal

@@ -26,7 +26,10 @@ library
     buildable: False
   else
     buildable: True
-  exposed-modules:    Ide.Plugin.HaddockComments
+  exposed-modules:
+    Ide.Plugin.HaddockComments
+    Ide.Plugin.HaddockComments.Data
+    Ide.Plugin.HaddockComments.Prelude
   hs-source-dirs:     src
   ghc-options:
     -Wall -Wno-name-shadowing -Wredundant-constraints
@@ -43,6 +46,8 @@ library
     , lsp-types
     , text
     , unordered-containers
+    , transformers
+    , mtl
 
   default-language:   Haskell2010
   default-extensions:

plugins/hls-haddock-comments-plugin/src/Ide/Plugin/HaddockComments.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE CPP                       #-}
+{-# LANGUAGE DuplicateRecordFields     #-}
 {-# LANGUAGE ExistentialQuantification #-}
 {-# LANGUAGE FlexibleContexts          #-}
 {-# LANGUAGE NamedFieldPuns            #-}
@@ -8,8 +9,9 @@
 
 module Ide.Plugin.HaddockComments (descriptor) where
 
-import           Control.Monad                         (join)
+import           Control.Monad                         (join, when)
 import           Control.Monad.IO.Class
+import           Control.Monad.Trans.Class             (lift)
 import qualified Data.HashMap.Strict                   as HashMap
 import qualified Data.Map                              as Map
 import qualified Data.Text                             as T
@@ -19,6 +21,8 @@ import           Development.IDE.GHC.Compat.ExactPrint
 import           Development.IDE.GHC.ExactPrint        (GetAnnotatedParsedSource (..))
 import qualified Development.IDE.GHC.ExactPrint        as E
 import           Development.IDE.Plugin.CodeAction
+import           Ide.Plugin.HaddockComments.Data       (genForDataDecl)
+import           Ide.Plugin.HaddockComments.Prelude
 import           Ide.Types
 import           Language.Haskell.GHC.ExactPrint
 import           Language.Haskell.GHC.ExactPrint.Types hiding (GhcPs)
@@ -40,47 +44,49 @@ codeActionProvider ideState _pId (CodeActionParams _ _ (TextDocumentIdentifier u
     (join -> pm) <- liftIO $ runAction "HaddockComments.GetAnnotatedParsedSource" ideState $ use GetAnnotatedParsedSource `traverse` nfp
     let locDecls = hsmodDecls . unLoc . astA <$> pm
         anns = annsA <$> pm
-        edits = [runGenComments gen locDecls anns range | noErr, gen <- genList]
+        edits = [gen locDecls anns range | noErr, gen <- genList]
     return $ Right $ List [InR $ toAction title uri edit | (Just (title, edit)) <- edits]
 
-genList :: [GenComments]
+genList :: [Maybe [LHsDecl GhcPs] -> Maybe Anns -> Range -> Maybe (T.Text, TextEdit)]
 genList =
-  [ genForSig,
-    genForRecord
+  [ runGenCommentsSimple genForSig,
+    runGenComments genForDataDecl
   ]
 
 -----------------------------------------------------------------------------
 
--- | Defines how to generate haddock comments by tweaking annotations of AST
-data GenComments = forall a.
-  GenComments
-  { title         :: T.Text,
-    fromDecl      :: HsDecl GhcPs -> Maybe a,
-    collectKeys   :: a -> [AnnKey],
-    isFresh       :: Annotation -> Bool,
-    updateAnn     :: Annotation -> Annotation,
-    updateDeclAnn :: Annotation -> Annotation
-  }
-
 runGenComments :: GenComments -> Maybe [LHsDecl GhcPs] -> Maybe Anns -> Range -> Maybe (T.Text, TextEdit)
-runGenComments GenComments {..} mLocDecls mAnns range
+runGenComments GenComments{..} mLocDecls mAnns range
   | Just locDecls <- mLocDecls,
     Just anns <- mAnns,
-    [(locDecl, src, x)] <- [(locDecl, l, x) | locDecl@(L l (fromDecl -> Just x)) <- locDecls, range `isIntersectWith` l],
-    annKeys <- collectKeys x,
-    not $ null annKeys,
-    and $ maybe False isFresh . flip Map.lookup anns <$> annKeys,
-    declKey <- mkAnnKey locDecl,
-    anns' <- Map.adjust updateDeclAnn declKey $ foldr (Map.adjust updateAnn) anns annKeys,
+    [(locDecl, src)] <- [(locDecl, l) | locDecl@(L l _) <- locDecls, range `isIntersectWith` l],
     Just range' <- toRange src,
-    result <- T.strip . T.pack $ exactPrint locDecl anns' =
-    Just (title, TextEdit range' result)
+    Just (_, (anns', _), _) <- runTransformT anns (updateAnns locDecl),
+    result <- T.strip . T.pack $ exactPrint locDecl anns'
+    = Just (title, TextEdit range' result)
   | otherwise = Nothing
 
+runGenCommentsSimple :: GenCommentsSimple -> Maybe [LHsDecl GhcPs] -> Maybe Anns -> Range -> Maybe (T.Text, TextEdit)
+runGenCommentsSimple GenCommentsSimple {..} = runGenComments GenComments {
+        title = title,
+        updateAnns = updateAnns
+    }
+  where
+    updateAnns :: LHsDecl GhcPs -> TransformT Maybe ()
+    updateAnns locDecl@(L _ decl) = do
+        x <- lift $ fromDecl decl
+        let annKeys = collectKeys x
+        anns <- getAnnsT
+        when (null annKeys || not (and $ maybe False isFresh . flip Map.lookup anns <$> annKeys)) $
+            lift Nothing
+        let declKey = mkAnnKey locDecl
+            anns' = Map.adjust updateDeclAnn declKey $ foldr (Map.adjust updateAnn) anns annKeys
+        putAnnsT anns'
+
 -----------------------------------------------------------------------------
 
-genForSig :: GenComments
-genForSig = GenComments {..}
+genForSig :: GenCommentsSimple
+genForSig = GenCommentsSimple {..}
   where
     title = "Generate signature comments"
 
@@ -102,30 +108,6 @@ genForSig = GenComments {..}
 #endif
     dp = [(AnnComment comment, DP (0, 1)), (G AnnRarrow, DP (1, 2))]
 
-genForRecord :: GenComments
-genForRecord = GenComments {..}
-  where
-    title = "Generate fields comments"
-
-    fromDecl (TyClD _ DataDecl {tcdDataDefn = HsDataDefn {dd_cons = cons}}) =
-      Just [x | (L _ ConDeclH98 {con_args = x}) <- cons]
-    fromDecl _ = Nothing
-
-    updateAnn x = x {annEntryDelta = DP (1, 2), annPriorComments = [(comment, DP (1, 2))]}
-    updateDeclAnn = cleanPriorComments
-
-    isFresh Ann {annPriorComments} = null annPriorComments
-
-    collectKeys = keyFromCon
-
-#if MIN_VERSION_ghc(9,2,0)
-    comment = mkComment "-- | " (spanAsAnchor noSrcSpan)
-#elif MIN_VERSION_ghc(9,0,0)
-    comment = mkComment "-- | " badRealSrcSpan
-#else
-    comment = mkComment "-- | " noSrcSpan
-#endif
-
 -----------------------------------------------------------------------------
 
 toAction :: T.Text -> Uri -> TextEdit -> CodeAction
@@ -176,7 +158,4 @@ keyFromTyVar dep (L _ (HsParTy _ x)) = keyFromTyVar (succ dep) x
 keyFromTyVar dep (L _ (HsBangTy _ _ x)) = keyFromTyVar dep x
 keyFromTyVar _ _ = []
 
-keyFromCon :: [HsConDeclDetails GhcPs] -> [AnnKey]
-keyFromCon cons = mconcat [mkAnnKey <$> xs | (RecCon (L _ xs)) <- cons]
-
 -----------------------------------------------------------------------------

plugins/hls-haddock-comments-plugin/src/Ide/Plugin/HaddockComments/Data.hs

@@ -0,0 +1,168 @@
+{-# LANGUAGE CPP                   #-}
+{-# LANGUAGE DuplicateRecordFields #-}
+{-# LANGUAGE LambdaCase            #-}
+{-# LANGUAGE OverloadedStrings     #-}
+
+module Ide.Plugin.HaddockComments.Data
+    ( genForDataDecl
+    ) where
+
+import           Control.Monad                         (unless, when)
+import           Control.Monad.Trans.Class             (lift)
+import           Data.Data                             (Data)
+import           Data.Foldable                         (for_)
+import           Data.List                             (isPrefixOf)
+import qualified Data.Map.Strict                       as Map
+import           Data.Maybe                            (fromMaybe, isJust)
+import           Development.IDE                       (realSpan)
+import           Development.IDE.GHC.Compat
+import           Development.IDE.GHC.ExactPrint
+import           Ide.Plugin.HaddockComments.Prelude
+import           Language.Haskell.GHC.ExactPrint
+import           Language.Haskell.GHC.ExactPrint.Types hiding (GhcPs)
+import           Language.Haskell.GHC.ExactPrint.Utils (mkComment)
+
+genForDataDecl :: GenComments
+genForDataDecl =
+    GenComments {
+        title = "Generate haddock comments",
+        updateAnns = updateDataAnns
+    }
+
+updateDataAnns :: LHsDecl GhcPs -> TransformT Maybe ()
+updateDataAnns decl@(L declLoc (TyClD _ DataDecl {tcdDataDefn = HsDataDefn { dd_cons = cons }})) = do
+    -- skip if all constructors and fields already have a haddock comment
+    getAnnsT >>= (\anns -> unless (missingSomeHaddock anns cons) (lift Nothing))
+
+    -- visit each constructor and field
+    addHaddockCommentsToList True declLoc (G AnnVbar) cons
+    for_ cons $ \case
+        L conLoc ConDeclH98 { con_args = RecCon (L _ fields) } -> addHaddockCommentsToList False conLoc (G AnnComma) fields
+        _ -> pure ()
+    modifyAnnsT $ Map.adjust (\ann -> ann {annPriorComments = []}) (mkAnnKey decl)
+updateDataAnns _ = lift Nothing
+
+-- | Add haddock comments to a list of nodes
+addHaddockCommentsToList
+    :: (Data a, Monad m)
+    => Bool -- ^ If true, for each node, use previous node in the list as the anchor. Otherwise, use the outer node
+    -> SrcSpan -- ^ The outer node
+    -> KeywordId -- ^ The seperator between adjacent nodes
+    -> [Located a] -- ^ The list of nodes. Haddock comments will be added to each of them
+    -> TransformT m ()
+addHaddockCommentsToList usePrevNodeAsAnchor outerLoc seperator nodes =
+    -- If you want to understand this function, please first read this page carefully:
+    --     https://hackage.haskell.org/package/ghc-exactprint-0.6.4/docs/Language-Haskell-GHC-ExactPrint-Delta.html
+    -- The important part is that for DP(r,c), if r is zero, c is the offset start from the end of the previous node.
+    -- However, if r is greater than zero, c is the offset start from the 'anchor'.
+    -- Generally speaking, the 'anchor' is the node that "enclose" the current node. But it's not always the case.
+    --     Sometimes 'anchor' is just the previous node. It depends on the the syntactic structure.
+    --     For constructors, the anchor is the previous node (if there is any).
+    --     For record fields, the anchor is always the constructor they belong to.
+    for_ (zip nodes (Nothing: fmap Just nodes)) $ \(node, prevNode) -> do
+        addHaddockCommentToCurrentNode <- fmap (not . fromMaybe True . flip hasHaddock node) getAnnsT
+        -- We don't add new haddock comments to nodes with existing ones.
+        when addHaddockCommentToCurrentNode $ do
+            -- 'sameLineAsPrev' is a flag to determine the inline case, for example:
+            --     data T = A { a :: Int, b :: String } | B { b :: Double }
+            -- Note that it's a 'Maybe (Located a)', containing the previous node if the current node
+            -- and the previous node are on the same line.
+            --
+            -- For the multiline case (which is the most common), we keep the original indentation of each constructor
+            -- and field.
+            --
+            -- For the inline case, we use the first construcotr/field as the base, and align all following items
+            -- to them.
+            let sameLineAsPrev = prevNode >>= (
+                    \prevNode' -> if notSeperatedByLineEnding prevNode' node
+                        then pure prevNode'
+                        else Nothing
+                    )
+            -- For the inline case, we need to move the seperator to the next line.
+            -- For constructors, it's vertical bar; for fields, it's comma.
+            -- The seperator is passed in as function argument.
+            when (isJust sameLineAsPrev) $ modifyAnnsT $ \anns ->
+                let newSepCol :: Annotation -> Int
+                    newSepCol ann =
+                        if usePrevNodeAsAnchor then 0 else deltaColumn (annEntryDelta ann)
+                    updateSepAnn :: Annotation -> Annotation
+                    updateSepAnn ann = ann {annsDP =
+                        Map.toList . Map.adjust (const $ DP (1, newSepCol ann)) seperator . Map.fromList $ annsDP ann}
+                 in flip (maybe anns) prevNode $ \prevNode' -> Map.adjust updateSepAnn (mkAnnKey prevNode') anns
+            -- Calculate the real column of the anchor
+            let anchorCol = maybe 0 srcSpanStartCol . realSpan . maybe outerLoc getLoc $
+                    if usePrevNodeAsAnchor then prevNode else Nothing
+            -- 'dpCol' is what we will use for the current node's entry delta's column
+            dpCol <- flip fmap getAnnsT $ \anns ->
+                case sameLineAsPrev of
+                    Just prevNode' ->
+                        -- If the previous node is the anchor, using 0 as column will make current code align with
+                        -- the previous one.
+                        -- Otherwise, use the column of entry delta of the previous node.
+                        -- The map lookup should not fail. '2' is used as a fallback value to make sure the syntax
+                        -- is correct after the changes.
+                        if usePrevNodeAsAnchor then 0 else maybe 2 (deltaColumn . annEntryDelta)
+                            $ anns Map.!? mkAnnKey prevNode'
+                    -- We subtract the real column to get dp column.
+                    Nothing -> (maybe 2 srcSpanStartCol . realSpan $ getLoc node) - anchorCol
+            -- Modify the current node
+            modifyAnnsT $
+                let updateCurrent :: Annotation -> Annotation
+                    updateCurrent ann = ann {
+                            -- If there exist non-haddock comments, we simply inherit the first one's delta pos,
+                            -- and move them two lines below, to seperate them from our newly added haddock comments
+                            -- Otherwise, inherit the node's entry delta pos.
+                            annPriorComments = case annPriorComments ann of
+                                (c, dp) : rem -> (emptyPriorHaddockComment, dp) : (c, DP (2,0)) : rem
+                                _ -> [(emptyPriorHaddockComment, annEntryDelta ann)],
+                            annEntryDelta = DP (1, dpCol)
+                        }
+                 in Map.adjust updateCurrent (mkAnnKey node)
+
+-- | Determine if a list of constructor declarations is missing some haddock comments.
+missingSomeHaddock :: Anns -> [LConDecl GhcPs] -> Bool
+missingSomeHaddock anns = any $ \lcon@(L _ conDecl) -> case conDecl of
+    ConDeclH98 { con_args = RecCon (L _ fields) } ->
+        elem (Just False) $ hasHaddock anns lcon : fmap (hasHaddock anns) fields
+    _                                       -> False -- GADT is not supported yet
+
+-- | Returns 'True' if the end of the first node and the start of the second node are on the same line.
+notSeperatedByLineEnding :: Located a
+                         -> Located a
+                         -> Bool
+notSeperatedByLineEnding (L (RealSrcSpan x _) _) (L (RealSrcSpan y _) _) =
+    srcLocLine (realSrcSpanEnd x) == srcLocLine (realSrcSpanStart y)
+notSeperatedByLineEnding _ _ = False
+
+-- | Empty haddock, suitable for being added to 'annPriorComments'
+emptyPriorHaddockComment :: Comment
+emptyPriorHaddockComment = mkComment "-- |"
+#if MIN_VERSION_ghc(9,0,0)
+    badRealSrcSpan
+#else
+    noSrcSpan
+#endif
+
+-- | Determines the given node has haddock comments attached to it.
+hasHaddock :: Data a => Anns -> Located a -> Maybe Bool
+hasHaddock anns node = fmap annHasHaddock (anns Map.!? key)
+  where
+    key = mkAnnKey node
+    annHasHaddock ann =
+        any (matchCommentPrefix priorCommentPrefix . fst) (annPriorComments ann)
+        || any (matchCommentPrefix followingCommentPrefix . fst) (annFollowingComments ann)
+        || any (keywordIdIsHaddockComment . fst) (annsDP ann)
+
+-- | Checks if the given 'KeywordId' is a comment, and specifically, a haddock comment.
+keywordIdIsHaddockComment :: KeywordId -> Bool
+keywordIdIsHaddockComment (AnnComment comment) = any (`isPrefixOf` commentContents comment) (priorCommentPrefix ++ followingCommentPrefix)
+keywordIdIsHaddockComment _ = False
+
+priorCommentPrefix :: [String]
+priorCommentPrefix = ["-- |", "{-|", "{- |"]
+
+followingCommentPrefix :: [String]
+followingCommentPrefix = ["-- ^", "{-^", "{- ^"]
+
+matchCommentPrefix :: [String] -> Comment -> Bool
+matchCommentPrefix prefix comment = any (`isPrefixOf` commentContents comment) prefix

plugins/hls-haddock-comments-plugin/src/Ide/Plugin/HaddockComments/Prelude.hs

@@ -0,0 +1,29 @@
+{-# LANGUAGE DuplicateRecordFields     #-}
+{-# LANGUAGE ExistentialQuantification #-}
+
+module Ide.Plugin.HaddockComments.Prelude where
+import qualified Data.Text                       as T
+import           Development.IDE.GHC.Compat
+import           Development.IDE.GHC.ExactPrint
+import           Language.Haskell.GHC.ExactPrint (AnnKey, Annotation)
+
+-- | A more generic comments generator
+data GenComments = GenComments
+  { title      :: T.Text,
+    -- | Use 'Maybe' monad to exit early. 'Nothing' means a code action for haddock comments
+    -- in the given context is not possible.
+    updateAnns :: LHsDecl GhcPs -> TransformT Maybe ()
+  }
+
+-- | Defines how to generate haddock comments by tweaking annotations of AST
+--
+-- This is left here for compatibility reason, so that we don't break the existing code.
+data GenCommentsSimple = forall a.
+  GenCommentsSimple
+  { title         :: T.Text,
+    fromDecl      :: HsDecl GhcPs -> Maybe a,
+    collectKeys   :: a -> [AnnKey],
+    isFresh       :: Annotation -> Bool,
+    updateAnn     :: Annotation -> Annotation,
+    updateDeclAnn :: Annotation -> Annotation
+  }

plugins/hls-haddock-comments-plugin/test/Main.hs

@@ -30,9 +30,11 @@ tests =
       goldenWithHaddockComments "MultivariateFunction" Signature 4 8,
       goldenWithHaddockComments "QualFunction" Signature 2 10,
       goldenWithHaddockComments "Record" Record 7 2,
+      goldenWithHaddockComments "Record2" Record 3 6,
+      goldenWithHaddockComments "InlineRecord" Record 3 20,
       expectedNothing "ConstFunction" Signature 2 2,
       expectedNothing "StaleFunction" Signature 3 3,
-      expectedNothing "StaleRecord" Record 3 12
+      expectedNothing "StaleRecord" Record 4 9
     ]
 
 goldenWithHaddockComments :: FilePath -> GenCommentsType -> UInt -> UInt -> TestTree
@@ -54,7 +56,7 @@ data GenCommentsType = Signature | Record
 
 toTitle :: GenCommentsType -> Text
 toTitle Signature = "Generate signature comments"
-toTitle Record    = "Generate fields comments"
+toTitle Record    = "Generate haddock comments"
 
 caTitle :: (Command |? CodeAction) -> Maybe Text
 caTitle (InR CodeAction {_title}) = Just _title