template-haskell: Add putDoc, getDoc, withDecDoc and friends

This adds two new methods to the Quasi class, putDoc and getDoc. They
allow Haddock documentation to be added to declarations, module headers,
function arguments and class/type family instances, as well as looked
up.

It works by building up a map of names to attach pieces of
documentation to, which are then added in the extractDocs function in
GHC.HsToCore.Docs. However because these template haskell names need to
be resolved to GHC names at the time they are added, putDoc cannot
directly add documentation to declarations that are currently being
spliced. To remedy this, withDecDoc/withDecsDoc wraps the operation with
addModFinalizer, and provides a more ergonomic interface for doing so.
Similarly, the funD_doc, dataD_doc etc. combinators provide a more
ergonomic interface for documenting functions and their arguments
simultaneously.

This also changes ArgDocMap to use an IntMap rather than an Map Int, for
efficiency.

Part of the work towards #5467

1 files changed, 104 insertions, 26 deletions

diff --git a/compiler/GHC/HsToCore/Docs.hs b/compiler/GHC/HsToCore/Docs.hs
index 56f089a756..fa278b7983 100644
--- a/compiler/GHC/HsToCore/Docs.hs
+++ b/compiler/GHC/HsToCore/Docs.hs
@@ -27,15 +27,22 @@ import GHC.Types.SrcLoc
 import GHC.Tc.Types
 
 import Control.Applicative
+import Control.Monad.IO.Class
 import Data.Bifunctor (first)
+import Data.IntMap (IntMap)
+import qualified Data.IntMap as IM
 import Data.Map (Map)
 import qualified Data.Map as M
 import Data.Maybe
 import Data.Semigroup
+import GHC.IORef (readIORef)
 
 -- | Extract docs from renamer output.
-extractDocs :: TcGblEnv
-            -> (Maybe HsDocString, DeclDocMap, ArgDocMap)
+-- This is monadic since we need to be able to read documentation added from
+-- Template Haskell's @putDoc@, which is stored in 'tcg_th_docs'.
+extractDocs :: MonadIO m
+            => TcGblEnv
+            -> m (Maybe HsDocString, DeclDocMap, ArgDocMap)
             -- ^
             -- 1. Module header
             -- 2. Docs on top level declarations
@@ -45,8 +52,20 @@ extractDocs TcGblEnv { tcg_semantic_mod = mod
                      , tcg_insts = insts
                      , tcg_fam_insts = fam_insts
                      , tcg_doc_hdr = mb_doc_hdr
-                     } =
-    (unLoc <$> mb_doc_hdr, DeclDocMap doc_map, ArgDocMap arg_map)
+                     , tcg_th_docs = th_docs_var
+                     } = do
+    th_docs <- liftIO $ readIORef th_docs_var
+    let doc_hdr = th_doc_hdr <|> (unLoc <$> mb_doc_hdr)
+        ExtractedTHDocs
+          th_doc_hdr
+          (DeclDocMap th_doc_map)
+          (ArgDocMap th_arg_map)
+          (DeclDocMap th_inst_map) = extractTHDocs th_docs
+    return
+      ( doc_hdr
+      , DeclDocMap (th_doc_map <> th_inst_map <> doc_map)
+      , ArgDocMap (th_arg_map `unionArgMaps` arg_map)
+      )
   where
     (doc_map, arg_map) = maybe (M.empty, M.empty)
                                (mkMaps local_insts)
@@ -59,10 +78,10 @@ extractDocs TcGblEnv { tcg_semantic_mod = mod
 -- For each declaration, find its names, its subordinates, and its doc strings.
 mkMaps :: [Name]
        -> [(LHsDecl GhcRn, [HsDocString])]
-       -> (Map Name (HsDocString), Map Name (Map Int (HsDocString)))
+       -> (Map Name (HsDocString), Map Name (IntMap HsDocString))
 mkMaps instances decls =
     ( f' (map (nubByName fst) decls')
-    , f  (filterMapping (not . M.null) args)
+    , f  (filterMapping (not . IM.null) args)
     )
   where
     (decls', args) = unzip (map mappings decls)
@@ -78,7 +97,7 @@ mkMaps instances decls =
 
     mappings :: (LHsDecl GhcRn, [HsDocString])
              -> ( [(Name, HsDocString)]
-                , [(Name, Map Int (HsDocString))]
+                , [(Name, IntMap HsDocString)]
                 )
     mappings (L (RealSrcSpan l _) decl, docStrs) =
            (dm, am)
@@ -86,7 +105,7 @@ mkMaps instances decls =
         doc = concatDocs docStrs
         args = declTypeDocs decl
 
-        subs :: [(Name, [(HsDocString)], Map Int (HsDocString))]
+        subs :: [(Name, [HsDocString], IntMap HsDocString)]
         subs = subordinates instanceMap decl
 
         (subDocs, subArgs) =
@@ -162,13 +181,13 @@ getInstLoc = \case
 -- family of a type class.
 subordinates :: Map RealSrcSpan Name
              -> HsDecl GhcRn
-             -> [(Name, [(HsDocString)], Map Int (HsDocString))]
+             -> [(Name, [HsDocString], IntMap HsDocString)]
 subordinates instMap decl = case decl of
   InstD _ (ClsInstD _ d) -> do
     DataFamInstDecl { dfid_eqn =
       FamEqn { feqn_tycon = L l _
              , feqn_rhs   = defn }} <- unLoc <$> cid_datafam_insts d
-    [ (n, [], M.empty) | Just n <- [lookupSrcSpan l instMap] ] ++ dataSubs defn
+    [ (n, [], IM.empty) | Just n <- [lookupSrcSpan l instMap] ] ++ dataSubs defn
 
   InstD _ (DataFamInstD _ (DataFamInstDecl d))
     -> dataSubs (feqn_rhs d)
@@ -181,7 +200,7 @@ subordinates instMap decl = case decl of
                    , name <- getMainDeclBinder d, not (isValD d)
                    ]
     dataSubs :: HsDataDefn GhcRn
-             -> [(Name, [HsDocString], Map Int (HsDocString))]
+             -> [(Name, [HsDocString], IntMap HsDocString)]
     dataSubs dd = constrs ++ fields ++ derivs
       where
         cons = map unLoc $ (dd_cons dd)
@@ -189,11 +208,11 @@ subordinates instMap decl = case decl of
                     , maybeToList $ fmap unLoc $ con_doc c
                     , conArgDocs c)
                   | c <- cons, cname <- getConNames c ]
-        fields  = [ (extFieldOcc n, maybeToList $ fmap unLoc doc, M.empty)
+        fields  = [ (extFieldOcc n, maybeToList $ fmap unLoc doc, IM.empty)
                   | Just flds <- map getRecConArgs_maybe cons
                   , (L _ (ConDeclField _ ns _ doc)) <- (unLoc flds)
                   , (L _ n) <- ns ]
-        derivs  = [ (instName, [unLoc doc], M.empty)
+        derivs  = [ (instName, [unLoc doc], IM.empty)
                   | (l, doc) <- concatMap (extract_deriv_clause_tys .
                                            deriv_clause_tys . unLoc) $
                                 unLoc $ dd_derivs dd
@@ -213,26 +232,26 @@ subordinates instMap decl = case decl of
             _               -> Nothing
 
 -- | Extract constructor argument docs from inside constructor decls.
-conArgDocs :: ConDecl GhcRn -> Map Int HsDocString
+conArgDocs :: ConDecl GhcRn -> IntMap HsDocString
 conArgDocs (ConDeclH98{con_args = args}) =
   h98ConArgDocs args
 conArgDocs (ConDeclGADT{con_g_args = args, con_res_ty = res_ty}) =
   gadtConArgDocs args (unLoc res_ty)
 
-h98ConArgDocs :: HsConDeclH98Details GhcRn -> Map Int HsDocString
+h98ConArgDocs :: HsConDeclH98Details GhcRn -> IntMap HsDocString
 h98ConArgDocs con_args = case con_args of
   PrefixCon _ args   -> con_arg_docs 0 $ map (unLoc . hsScaledThing) args
   InfixCon arg1 arg2 -> con_arg_docs 0 [ unLoc (hsScaledThing arg1)
                                        , unLoc (hsScaledThing arg2) ]
-  RecCon _           -> M.empty
+  RecCon _           -> IM.empty
 
-gadtConArgDocs :: HsConDeclGADTDetails GhcRn -> HsType GhcRn -> Map Int HsDocString
+gadtConArgDocs :: HsConDeclGADTDetails GhcRn -> HsType GhcRn -> IntMap HsDocString
 gadtConArgDocs con_args res_ty = case con_args of
   PrefixConGADT args -> con_arg_docs 0 $ map (unLoc . hsScaledThing) args ++ [res_ty]
   RecConGADT _       -> con_arg_docs 1 [res_ty]
 
-con_arg_docs :: Int -> [HsType GhcRn] -> Map Int HsDocString
-con_arg_docs n = M.fromList . catMaybes . zipWith f [n..]
+con_arg_docs :: Int -> [HsType GhcRn] -> IntMap HsDocString
+con_arg_docs n = IM.fromList . catMaybes . zipWith f [n..]
   where
     f n (HsDocTy _ _ lds) = Just (n, unLoc lds)
     f n (HsBangTy _ _ (L _ (HsDocTy _ _ lds))) = Just (n, unLoc lds)
@@ -254,14 +273,14 @@ classDecls class_ = filterDecls . collectDocs . sortLocated $ decls
     ats   = mkDecls tcdATs (TyClD noExtField . FamDecl noExtField) class_
 
 -- | Extract function argument docs from inside top-level decls.
-declTypeDocs :: HsDecl GhcRn -> Map Int (HsDocString)
+declTypeDocs :: HsDecl GhcRn -> IntMap (HsDocString)
 declTypeDocs = \case
   SigD  _ (TypeSig _ _ ty)          -> sigTypeDocs (unLoc (dropWildCards ty))
   SigD  _ (ClassOpSig _ _ _ ty)     -> sigTypeDocs (unLoc ty)
   SigD  _ (PatSynSig _ _ ty)        -> sigTypeDocs (unLoc ty)
   ForD  _ (ForeignImport _ _ ty _)  -> sigTypeDocs (unLoc ty)
   TyClD _ (SynDecl { tcdRhs = ty }) -> typeDocs (unLoc ty)
-  _                                 -> M.empty
+  _                                 -> IM.empty
 
 nubByName :: (a -> Name) -> [a] -> [a]
 nubByName f ns = go emptyNameSet ns
@@ -275,19 +294,19 @@ nubByName f ns = go emptyNameSet ns
         y = f x
 
 -- | Extract function argument docs from inside types.
-typeDocs :: HsType GhcRn -> Map Int (HsDocString)
+typeDocs :: HsType GhcRn -> IntMap HsDocString
 typeDocs = go 0
   where
     go n = \case
       HsForAllTy { hst_body = ty }          -> go n (unLoc ty)
       HsQualTy   { hst_body = ty }          -> go n (unLoc ty)
-      HsFunTy _ _ (unLoc->HsDocTy _ _ x) ty -> M.insert n (unLoc x) $ go (n+1) (unLoc ty)
+      HsFunTy _ _ (unLoc->HsDocTy _ _ x) ty -> IM.insert n (unLoc x) $ go (n+1) (unLoc ty)
       HsFunTy _ _ _ ty                      -> go (n+1) (unLoc ty)
-      HsDocTy _ _ doc                       -> M.singleton n (unLoc doc)
-      _                                     -> M.empty
+      HsDocTy _ _ doc                       -> IM.singleton n (unLoc doc)
+      _                                     -> IM.empty
 
 -- | Extract function argument docs from inside types.
-sigTypeDocs :: HsSigType GhcRn -> Map Int HsDocString
+sigTypeDocs :: HsSigType GhcRn -> IntMap HsDocString
 sigTypeDocs (HsSig{sig_body = body}) = typeDocs (unLoc body)
 
 -- | The top-level declarations of a module that we care about,
@@ -372,3 +391,62 @@ mkDecls :: (struct -> [Located decl])
         -> struct
         -> [Located hsDecl]
 mkDecls field con = map (mapLoc con) . field
+
+-- | Extracts out individual maps of documentation added via Template Haskell's
+-- @putDoc@.
+extractTHDocs :: THDocs
+              -> ExtractedTHDocs
+extractTHDocs docs =
+  -- Split up docs into separate maps for each 'DocLoc' type
+  ExtractedTHDocs
+    docHeader
+    (DeclDocMap (searchDocs decl))
+    (ArgDocMap (searchDocs args))
+    (DeclDocMap (searchDocs insts))
+  where
+    docHeader :: Maybe HsDocString
+    docHeader
+      | ((_, s):_) <- filter isModDoc (M.toList docs) = Just (mkHsDocString s)
+      | otherwise = Nothing
+
+    isModDoc (ModuleDoc, _) = True
+    isModDoc _ = False
+
+    -- Folds over the docs, applying 'f' as the accumulating function.
+    -- We use different accumulating functions to sift out the specific types of
+    -- documentation
+    searchDocs :: Monoid a => (a -> (DocLoc, String) -> a) -> a
+    searchDocs f = foldl' f mempty $ M.toList docs
+
+    -- Pick out the declaration docs
+    decl acc ((DeclDoc name), s) = M.insert name (mkHsDocString s) acc
+    decl acc _ = acc
+
+    -- Pick out the instance docs
+    insts acc ((InstDoc name), s) = M.insert name (mkHsDocString s) acc
+    insts acc _ = acc
+
+    -- Pick out the argument docs
+    args :: Map Name (IntMap HsDocString)
+         -> (DocLoc, String)
+         -> Map Name (IntMap HsDocString)
+    args acc ((ArgDoc name i), s) =
+      -- Insert the doc for the arg into the argument map for the function. This
+      -- means we have to search to see if an map already exists for the
+      -- function, and insert the new argument if it exists, or create a new map
+      let ds = mkHsDocString s
+       in M.insertWith (\_ m -> IM.insert i ds m) name (IM.singleton i ds) acc
+    args acc _ = acc
+
+-- | Unions together two 'ArgDocMaps' (or ArgMaps in haddock-api), such that two
+-- maps with values for the same key merge the inner map as well.
+-- Left biased so @unionArgMaps a b@ prefers @a@ over @b@.
+unionArgMaps :: Map Name (IntMap b)
+             -> Map Name (IntMap b)
+             -> Map Name (IntMap b)
+unionArgMaps a b = M.foldlWithKey go b a
+  where
+    go acc n newArgMap
+      | Just oldArgMap <- M.lookup n acc =
+          M.insert n (newArgMap `IM.union` oldArgMap) acc
+      | otherwise = M.insert n newArgMap acc