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

10 files changed, 323 insertions, 49 deletions

diff --git a/compiler/GHC/Hs/Doc.hs b/compiler/GHC/Hs/Doc.hs
index 207b65f2fd..425cc03bf0 100644
--- a/compiler/GHC/Hs/Doc.hs
+++ b/compiler/GHC/Hs/Doc.hs
@@ -19,6 +19,8 @@ module GHC.Hs.Doc
 
   , ArgDocMap(..)
   , emptyArgDocMap
+
+  , ExtractedTHDocs(..)
   ) where
 
 #include "HsVersions.h"
@@ -35,6 +37,8 @@ import Data.ByteString (ByteString)
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Char8 as C8
 import Data.Data
+import Data.IntMap (IntMap)
+import qualified Data.IntMap as IntMap
 import Data.Map (Map)
 import qualified Data.Map as Map
 import Data.Maybe
@@ -126,21 +130,34 @@ emptyDeclDocMap :: DeclDocMap
 emptyDeclDocMap = DeclDocMap Map.empty
 
 -- | Docs for arguments. E.g. function arguments, method arguments.
-newtype ArgDocMap = ArgDocMap (Map Name (Map Int HsDocString))
+newtype ArgDocMap = ArgDocMap (Map Name (IntMap HsDocString))
 
 instance Binary ArgDocMap where
-  put_ bh (ArgDocMap m) = put_ bh (Map.toList (Map.toAscList <$> m))
+  put_ bh (ArgDocMap m) = put_ bh (Map.toList (IntMap.toAscList <$> m))
   -- We can't rely on a deterministic ordering of the `Name`s here.
   -- See the comments on `Name`'s `Ord` instance for context.
-  get bh = ArgDocMap . fmap Map.fromDistinctAscList . Map.fromList <$> get bh
+  get bh = ArgDocMap . fmap IntMap.fromDistinctAscList . Map.fromList <$> get bh
 
 instance Outputable ArgDocMap where
   ppr (ArgDocMap m) = vcat (map pprPair (Map.toAscList m))
     where
       pprPair (name, int_map) =
         ppr name Outputable.<> colon $$ nest 2 (pprIntMap int_map)
-      pprIntMap im = vcat (map pprIPair (Map.toAscList im))
+      pprIntMap im = vcat (map pprIPair (IntMap.toAscList im))
       pprIPair (i, doc) = ppr i Outputable.<> colon $$ nest 2 (ppr doc)
 
 emptyArgDocMap :: ArgDocMap
 emptyArgDocMap = ArgDocMap Map.empty
+
+-- | Maps of docs that were added via Template Haskell's @putDoc@.
+data ExtractedTHDocs =
+  ExtractedTHDocs
+    { ethd_mod_header :: Maybe HsDocString
+      -- ^ The added module header documentation, if it exists.
+    , ethd_decl_docs  :: DeclDocMap
+      -- ^ The documentation added to declarations.
+    , ethd_arg_docs   :: ArgDocMap
+      -- ^ The documentation added to function arguments.
+    , ethd_inst_docs  :: DeclDocMap
+      -- ^ The documentation added to class and family instances.
+    }
diff --git a/compiler/GHC/HsToCore.hs b/compiler/GHC/HsToCore.hs
index bf15fd2e10..fafcdb6533 100644
--- a/compiler/GHC/HsToCore.hs
+++ b/compiler/GHC/HsToCore.hs
@@ -214,7 +214,7 @@ deSugar hsc_env
 
         ; foreign_files <- readIORef th_foreign_files_var
 
-        ; let (doc_hdr, decl_docs, arg_docs) = extractDocs tcg_env
+        ; (doc_hdr, decl_docs, arg_docs) <- extractDocs tcg_env
 
         ; let mod_guts = ModGuts {
                 mg_module       = mod,
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
diff --git a/compiler/GHC/Iface/Make.hs b/compiler/GHC/Iface/Make.hs
index 53f0032f28..323f69f0d3 100644
--- a/compiler/GHC/Iface/Make.hs
+++ b/compiler/GHC/Iface/Make.hs
@@ -212,7 +212,7 @@ mkIfaceTc hsc_env safe_mode mod_details
           usages <- mkUsageInfo hsc_env this_mod (imp_mods imports) used_names
                       dep_files merged pluginModules
 
-          let (doc_hdr', doc_map, arg_map) = extractDocs tc_result
+          (doc_hdr', doc_map, arg_map) <- extractDocs tc_result
 
           let partial_iface = mkIface_ hsc_env
                    this_mod hsc_src
diff --git a/compiler/GHC/Runtime/Eval.hs b/compiler/GHC/Runtime/Eval.hs
index c2626ce6b3..fc2f8b8ab3 100644
--- a/compiler/GHC/Runtime/Eval.hs
+++ b/compiler/GHC/Runtime/Eval.hs
@@ -117,9 +117,9 @@ import GHC.Unit.Home.ModInfo
 import System.Directory
 import Data.Dynamic
 import Data.Either
+import Data.IntMap (IntMap)
 import qualified Data.IntMap as IntMap
 import Data.List (find,intercalate)
-import Data.Map (Map)
 import qualified Data.Map as Map
 import Control.Monad
 import Control.Monad.Catch as MC
@@ -879,7 +879,7 @@ parseName str = withSession $ \hsc_env -> liftIO $
 
 getDocs :: GhcMonad m
         => Name
-        -> m (Either GetDocsFailure (Maybe HsDocString, Map Int HsDocString))
+        -> m (Either GetDocsFailure (Maybe HsDocString, IntMap HsDocString))
            -- TODO: What about docs for constructors etc.?
 getDocs name =
   withSession $ \hsc_env -> do
@@ -896,7 +896,7 @@ getDocs name =
              if isNothing mb_doc_hdr && Map.null dmap && Map.null amap
                then pure (Left (NoDocsInIface mod compiled))
                else pure (Right ( Map.lookup name dmap
-                                , Map.findWithDefault Map.empty name amap))
+                                , Map.findWithDefault mempty name amap))
   where
     compiled =
       -- TODO: Find a more direct indicator.
diff --git a/compiler/GHC/Tc/Errors/Hole.hs b/compiler/GHC/Tc/Errors/Hole.hs
index 109a4416bc..b89f5c8a6c 100644
--- a/compiler/GHC/Tc/Errors/Hole.hs
+++ b/compiler/GHC/Tc/Errors/Hole.hs
@@ -419,7 +419,7 @@ addDocs :: [HoleFit] -> TcM [HoleFit]
 addDocs fits =
   do { showDocs <- goptM Opt_ShowDocsOfHoleFits
      ; if showDocs
-       then do { (_, DeclDocMap lclDocs, _) <- extractDocs <$> getGblEnv
+       then do { (_, DeclDocMap lclDocs, _) <- getGblEnv >>= extractDocs
                ; mapM (upd lclDocs) fits }
        else return fits }
   where
diff --git a/compiler/GHC/Tc/Gen/Splice.hs b/compiler/GHC/Tc/Gen/Splice.hs
index 7ae4ccb0f6..89ba997d8a 100644
--- a/compiler/GHC/Tc/Gen/Splice.hs
+++ b/compiler/GHC/Tc/Gen/Splice.hs
@@ -68,6 +68,7 @@ import GHC.Builtin.Names
 import GHC.Builtin.Types
 
 import GHC.ThToHs
+import GHC.HsToCore.Docs
 import GHC.HsToCore.Expr
 import GHC.HsToCore.Monad
 import GHC.IfaceToCore
@@ -147,6 +148,7 @@ import Data.Maybe
 import qualified Data.ByteString as B
 import qualified Data.ByteString.Lazy as LB
 import Data.Dynamic  ( fromDynamic, toDyn )
+import qualified Data.IntMap as IntMap
 import qualified Data.Map as Map
 import Data.Typeable ( typeOf, Typeable, TypeRep, typeRep )
 import Data.Data (Data)
@@ -1220,6 +1222,148 @@ instance TH.Quasi TcM where
   qExtsEnabled =
     EnumSet.toList . extensionFlags . hsc_dflags <$> getTopEnv
 
+  qPutDoc doc_loc s = do
+    th_doc_var <- tcg_th_docs <$> getGblEnv
+    resolved_doc_loc <- resolve_loc doc_loc
+    is_local <- checkLocalName resolved_doc_loc
+    unless is_local $ failWithTc $ text
+      "Can't add documentation to" <+> ppr_loc doc_loc <+>
+      text "as it isn't inside the current module"
+    updTcRef th_doc_var (Map.insert resolved_doc_loc s)
+    where
+      resolve_loc (TH.DeclDoc n) = DeclDoc <$> lookupThName n
+      resolve_loc (TH.ArgDoc n i) = ArgDoc <$> lookupThName n <*> pure i
+      resolve_loc (TH.InstDoc t) = InstDoc <$> fmap getName (lookupThInstName t)
+      resolve_loc TH.ModuleDoc = pure ModuleDoc
+
+      ppr_loc (TH.DeclDoc n) = ppr_th n
+      ppr_loc (TH.ArgDoc n _) = ppr_th n
+      ppr_loc (TH.InstDoc t) = ppr_th t
+      ppr_loc TH.ModuleDoc = text "the module header"
+
+      -- It doesn't make sense to add documentation to something not inside
+      -- the current module. So check for it!
+      checkLocalName (DeclDoc n) = nameIsLocalOrFrom <$> getModule <*> pure n
+      checkLocalName (ArgDoc n _) = nameIsLocalOrFrom <$> getModule <*> pure n
+      checkLocalName (InstDoc n) = nameIsLocalOrFrom <$> getModule <*> pure n
+      checkLocalName ModuleDoc = pure True
+
+
+  qGetDoc (TH.DeclDoc n) = lookupThName n >>= lookupDeclDoc
+  qGetDoc (TH.InstDoc t) = lookupThInstName t >>= lookupDeclDoc
+  qGetDoc (TH.ArgDoc n i) = lookupThName n >>= lookupArgDoc i
+  qGetDoc TH.ModuleDoc = do
+    (moduleDoc, _, _) <- getGblEnv >>= extractDocs
+    return (fmap unpackHDS moduleDoc)
+
+-- | Looks up documentation for a declaration in first the current module,
+-- otherwise tries to find it in another module via 'hscGetModuleInterface'.
+lookupDeclDoc :: Name -> TcM (Maybe String)
+lookupDeclDoc nm = do
+  (_, DeclDocMap declDocs, _) <- getGblEnv >>= extractDocs
+  fam_insts <- tcg_fam_insts <$> getGblEnv
+  traceTc "lookupDeclDoc" (ppr nm <+> ppr declDocs <+> ppr fam_insts)
+  case Map.lookup nm declDocs of
+    Just doc -> pure $ Just (unpackHDS doc)
+    Nothing -> do
+      -- Wasn't in the current module. Try searching other external ones!
+      mIface <- getExternalModIface nm
+      case mIface of
+        Nothing -> pure Nothing
+        Just ModIface { mi_decl_docs = DeclDocMap dmap } ->
+          pure $ unpackHDS <$> Map.lookup nm dmap
+
+-- | Like 'lookupDeclDoc', looks up documentation for a function argument. If
+-- it can't find any documentation for a function in this module, it tries to
+-- find it in another module.
+lookupArgDoc :: Int -> Name -> TcM (Maybe String)
+lookupArgDoc i nm = do
+  (_, _, ArgDocMap argDocs) <- getGblEnv >>= extractDocs
+  case Map.lookup nm argDocs of
+    Just m -> pure $ unpackHDS <$> IntMap.lookup i m
+    Nothing -> do
+      mIface <- getExternalModIface nm
+      case mIface of
+        Nothing -> pure Nothing
+        Just ModIface { mi_arg_docs = ArgDocMap amap } ->
+          pure $ unpackHDS <$> (Map.lookup nm amap >>= IntMap.lookup i)
+
+-- | Returns the module a Name belongs to, if it is isn't local.
+getExternalModIface :: Name -> TcM (Maybe ModIface)
+getExternalModIface nm = do
+  isLocal <- nameIsLocalOrFrom <$> getModule <*> pure nm
+  if isLocal
+    then pure Nothing
+    else case nameModule_maybe nm of
+          Nothing -> pure Nothing
+          Just modNm -> do
+            hsc_env <- getTopEnv
+            iface <- liftIO $ hscGetModuleInterface hsc_env modNm
+            pure (Just iface)
+
+-- | Find the GHC name of the first instance that matches the TH type
+lookupThInstName :: TH.Type -> TcM Name
+lookupThInstName th_type = do
+  cls_name <- inst_cls_name th_type
+  insts <- reifyInstances' cls_name (inst_arg_types th_type)
+  case insts of   -- This expands any type synonyms
+    Left  (_, (inst:_)) -> return $ getName inst
+    Left  (_, [])       -> noMatches
+    Right (_, (inst:_)) -> return $ getName inst
+    Right (_, [])       -> noMatches
+  where
+    noMatches = failWithTc $
+      text "Couldn't find any instances of"
+        <+> ppr_th th_type
+        <+> text "to add documentation to"
+
+    -- | Get the name of the class for the instance we are documenting
+    -- > inst_cls_name (Monad Maybe) == Monad
+    -- > inst_cls_name C = C
+    inst_cls_name :: TH.Type -> TcM TH.Name
+    inst_cls_name (TH.AppT t _)           = inst_cls_name t
+    inst_cls_name (TH.SigT n _)           = inst_cls_name n
+    inst_cls_name (TH.VarT n)             = pure n
+    inst_cls_name (TH.ConT n)             = pure n
+    inst_cls_name (TH.PromotedT n)        = pure n
+    inst_cls_name (TH.InfixT _ n _)       = pure n
+    inst_cls_name (TH.UInfixT _ n _)      = pure n
+    inst_cls_name (TH.ParensT t)          = inst_cls_name t
+
+    inst_cls_name (TH.ForallT _ _ _)      = inst_cls_name_err
+    inst_cls_name (TH.ForallVisT _ _)     = inst_cls_name_err
+    inst_cls_name (TH.AppKindT _ _)       = inst_cls_name_err
+    inst_cls_name (TH.TupleT _)           = inst_cls_name_err
+    inst_cls_name (TH.UnboxedTupleT _)    = inst_cls_name_err
+    inst_cls_name (TH.UnboxedSumT _)      = inst_cls_name_err
+    inst_cls_name TH.ArrowT               = inst_cls_name_err
+    inst_cls_name TH.MulArrowT            = inst_cls_name_err
+    inst_cls_name TH.EqualityT            = inst_cls_name_err
+    inst_cls_name TH.ListT                = inst_cls_name_err
+    inst_cls_name (TH.PromotedTupleT _)   = inst_cls_name_err
+    inst_cls_name TH.PromotedNilT         = inst_cls_name_err
+    inst_cls_name TH.PromotedConsT        = inst_cls_name_err
+    inst_cls_name TH.StarT                = inst_cls_name_err
+    inst_cls_name TH.ConstraintT          = inst_cls_name_err
+    inst_cls_name (TH.LitT _)             = inst_cls_name_err
+    inst_cls_name TH.WildCardT            = inst_cls_name_err
+    inst_cls_name (TH.ImplicitParamT _ _) = inst_cls_name_err
+
+    inst_cls_name_err = failWithTc $
+      text "Couldn't work out what instance"
+        <+> ppr_th th_type
+        <+> text "is supposed to be"
+
+    -- | Basically does the opposite of 'mkThAppTs'
+    -- > inst_arg_types (Monad Maybe) == [Maybe]
+    -- > inst_arg_types C == []
+    inst_arg_types :: TH.Type -> [TH.Type]
+    inst_arg_types (TH.AppT _ args) =
+      let go (TH.AppT t ts) = t:go ts
+          go t = [t]
+        in go args
+    inst_arg_types _ = []
+
 -- | Adds a mod finalizer reference to the local environment.
 addModFinalizerRef :: ForeignRef (TH.Q ()) -> TcM ()
 addModFinalizerRef finRef = do
@@ -1411,6 +1555,8 @@ handleTHMessage msg = case msg of
   AddForeignFilePath lang str -> wrapTHResult $ TH.qAddForeignFilePath lang str
   IsExtEnabled ext -> wrapTHResult $ TH.qIsExtEnabled ext
   ExtsEnabled -> wrapTHResult $ TH.qExtsEnabled
+  PutDoc l s -> wrapTHResult $ TH.qPutDoc l s
+  GetDoc l -> wrapTHResult $ TH.qGetDoc l
   FailIfErrs -> wrapTHResult failIfErrsM
   _ -> panic ("handleTHMessage: unexpected message " ++ show msg)
 
@@ -1434,6 +1580,19 @@ getAnnotationsByTypeRep th_name tyrep
 
 reifyInstances :: TH.Name -> [TH.Type] -> TcM [TH.Dec]
 reifyInstances th_nm th_tys
+  = do { insts <- reifyInstances' th_nm th_tys
+       ; case insts of
+           Left (cls, cls_insts) ->
+             reifyClassInstances cls cls_insts
+           Right (tc, fam_insts) ->
+             reifyFamilyInstances tc fam_insts }
+
+reifyInstances' :: TH.Name
+                -> [TH.Type]
+                -> TcM (Either (Class, [ClsInst]) (TyCon, [FamInst]))
+                -- ^ Returns 'Left' in the case that the instances were found to
+                -- be class instances, or 'Right' if they are family instances.
+reifyInstances' th_nm th_tys
    = addErrCtxt (text "In the argument of reifyInstances:"
                  <+> ppr_th th_nm <+> sep (map ppr_th th_tys)) $
      do { loc <- getSrcSpanM
@@ -1467,19 +1626,19 @@ reifyInstances th_nm th_tys
                 -- In particular, the type might have kind
                 -- variables inside it (#7477)
 
-        ; traceTc "reifyInstances" (ppr ty $$ ppr (tcTypeKind ty))
+        ; traceTc "reifyInstances'" (ppr ty $$ ppr (tcTypeKind ty))
         ; case splitTyConApp_maybe ty of   -- This expands any type synonyms
             Just (tc, tys)                 -- See #7910
                | Just cls <- tyConClass_maybe tc
                -> do { inst_envs <- tcGetInstEnvs
                      ; let (matches, unifies, _) = lookupInstEnv False inst_envs cls tys
-                     ; traceTc "reifyInstances1" (ppr matches)
-                     ; reifyClassInstances cls (map fst matches ++ unifies) }
+                     ; traceTc "reifyInstances'1" (ppr matches)
+                     ; return $ Left (cls, map fst matches ++ unifies) }
                | isOpenFamilyTyCon tc
                -> do { inst_envs <- tcGetFamInstEnvs
                      ; let matches = lookupFamInstEnv inst_envs tc tys
-                     ; traceTc "reifyInstances2" (ppr matches)
-                     ; reifyFamilyInstances tc (map fim_instance matches) }
+                     ; traceTc "reifyInstances'2" (ppr matches)
+                     ; return $ Right (tc, map fim_instance matches) }
             _  -> bale_out (hang (text "reifyInstances:" <+> quotes (ppr ty))
                                2 (text "is not a class constraint or type family application")) }
   where
diff --git a/compiler/GHC/Tc/Module.hs b/compiler/GHC/Tc/Module.hs
index 9e9e82bca4..81cf5ea408 100644
--- a/compiler/GHC/Tc/Module.hs
+++ b/compiler/GHC/Tc/Module.hs
@@ -283,6 +283,14 @@ tcRnModuleTcRnM hsc_env mod_sum
           tcg_env <- {-# SCC "tcRnImports" #-}
                      tcRnImports hsc_env all_imports
 
+       ;  -- Don't need to rename the Haddock documentation,
+          -- it's not parsed by GHC anymore.
+          -- Make sure to do this before 'tcRnSrcDecls', because we need the
+          -- module header when we're splicing TH, since it can be accessed via
+          -- 'getDoc'.
+          tcg_env <- return (tcg_env
+                              { tcg_doc_hdr = maybe_doc_hdr })
+
         ; -- If the whole module is warned about or deprecated
           -- (via mod_deprec) record that in tcg_warns. If we do thereby add
           -- a WarnAll, it will override any subsequent deprecations added to tcg_warns
@@ -320,13 +328,8 @@ tcRnModuleTcRnM hsc_env mod_sum
                         -- because the latter might add new bindings for
                         -- boot_dfuns, which may be mentioned in imported
                         -- unfoldings.
-
-                        -- Don't need to rename the Haddock documentation,
-                        -- it's not parsed by GHC anymore.
-                        tcg_env <- return (tcg_env
-                                           { tcg_doc_hdr = maybe_doc_hdr })
-                      ; -- Report unused names
-                        -- Do this /after/ type inference, so that when reporting
+                        -- Report unused names
+                        -- Do this /after/ typeinference, so that when reporting
                         -- a function with no type signature we can give the
                         -- inferred type
                         reportUnusedNames tcg_env hsc_src
diff --git a/compiler/GHC/Tc/Types.hs b/compiler/GHC/Tc/Types.hs
index 0003a93169..2c9be13dff 100644
--- a/compiler/GHC/Tc/Types.hs
+++ b/compiler/GHC/Tc/Types.hs
@@ -55,7 +55,7 @@ module GHC.Tc.Types(
         ThStage(..), SpliceType(..), PendingStuff(..),
         topStage, topAnnStage, topSpliceStage,
         ThLevel, impLevel, outerLevel, thLevel,
-        ForeignSrcLang(..),
+        ForeignSrcLang(..), THDocs, DocLoc(..),
 
         -- Arrows
         ArrowCtxt(..),
@@ -522,6 +522,9 @@ data TcGblEnv
         tcg_th_remote_state :: TcRef (Maybe (ForeignRef (IORef QState))),
         -- ^ Template Haskell state
 
+        tcg_th_docs   :: TcRef THDocs,
+        -- ^ Docs added in Template Haskell via @putDoc@.
+
         tcg_ev_binds  :: Bag EvBind,        -- Top-level evidence bindings
 
         -- Things defined in this module, or (in GHCi)
@@ -1738,3 +1741,15 @@ lintGblEnv logger dflags tcg_env =
   liftIO $ lintAxioms logger dflags (text "TcGblEnv axioms") axioms
   where
     axioms = typeEnvCoAxioms (tcg_type_env tcg_env)
+
+-- | This is a mirror of Template Haskell's DocLoc, but the TH names are
+-- resolved to GHC names.
+data DocLoc = DeclDoc Name
+            | ArgDoc Name Int
+            | InstDoc Name
+            | ModuleDoc
+  deriving (Eq, Ord)
+
+-- | The current collection of docs that Template Haskell has built up via
+-- putDoc.
+type THDocs = Map DocLoc String
diff --git a/compiler/GHC/Tc/Utils/Monad.hs b/compiler/GHC/Tc/Utils/Monad.hs
index a3c087c4da..873c9b9fd2 100644
--- a/compiler/GHC/Tc/Utils/Monad.hs
+++ b/compiler/GHC/Tc/Utils/Monad.hs
@@ -257,6 +257,7 @@ initTc hsc_env hsc_src keep_rn_syntax mod loc do_this
         th_coreplugins_var <- newIORef [] ;
         th_state_var         <- newIORef Map.empty ;
         th_remote_state_var  <- newIORef Nothing ;
+        th_docs_var          <- newIORef Map.empty ;
         let {
              -- bangs to avoid leaking the env (#19356)
              !dflags = hsc_dflags hsc_env ;
@@ -284,6 +285,7 @@ initTc hsc_env hsc_src keep_rn_syntax mod loc do_this
                 tcg_th_coreplugins = th_coreplugins_var,
                 tcg_th_state         = th_state_var,
                 tcg_th_remote_state  = th_remote_state_var,
+                tcg_th_docs          = th_docs_var,
 
                 tcg_mod            = mod,
                 tcg_semantic_mod   = homeModuleInstantiation home_unit mod,