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

5 files changed, 228 insertions, 1 deletions

diff --git a/libraries/template-haskell/Language/Haskell/TH.hs b/libraries/template-haskell/Language/Haskell/TH.hs
index 2da2bd61c6..83432c14e3 100644
--- a/libraries/template-haskell/Language/Haskell/TH.hs
+++ b/libraries/template-haskell/Language/Haskell/TH.hs
@@ -90,6 +90,9 @@ module Language.Haskell.TH(
         Syntax.Specificity(..),
         FamilyResultSig(..), Syntax.InjectivityAnn(..), PatSynType, BangType, VarBangType,
 
+    -- ** Documentation
+        putDoc, getDoc, DocLoc(..),
+
     -- * Library functions
     module Language.Haskell.TH.Lib,
 
diff --git a/libraries/template-haskell/Language/Haskell/TH/Lib.hs b/libraries/template-haskell/Language/Haskell/TH/Lib.hs
index 3e05081619..de90df2bfd 100644
--- a/libraries/template-haskell/Language/Haskell/TH/Lib.hs
+++ b/libraries/template-haskell/Language/Haskell/TH/Lib.hs
@@ -124,7 +124,11 @@ module Language.Haskell.TH.Lib (
     implicitParamBindD,
 
     -- ** Reify
-    thisModule
+    thisModule,
+
+    -- ** Documentation
+    withDecDoc, withDecsDoc, funD_doc, dataD_doc, newtypeD_doc, dataInstD_doc,
+    newtypeInstD_doc, patSynD_doc
 
    ) where
 
diff --git a/libraries/template-haskell/Language/Haskell/TH/Lib/Internal.hs b/libraries/template-haskell/Language/Haskell/TH/Lib/Internal.hs
index a41d0a47b3..706d4a8c6a 100644
--- a/libraries/template-haskell/Language/Haskell/TH/Lib/Internal.hs
+++ b/libraries/template-haskell/Language/Haskell/TH/Lib/Internal.hs
@@ -981,3 +981,171 @@ thisModule :: Q Module
 thisModule = do
   loc <- location
   pure $ Module (mkPkgName $ loc_package loc) (mkModName $ loc_module loc)
+
+--------------------------------------------------------------
+-- * Documentation combinators
+
+-- | Attaches Haddock documentation to the declaration provided. Unlike
+-- 'putDoc', the names do not need to be in scope when calling this function so
+-- it can be used for quoted declarations and anything else currently being
+-- spliced.
+-- Not all declarations can have documentation attached to them. For those that
+-- can't, 'withDecDoc' will return it unchanged without any side effects.
+withDecDoc :: String -> Q Dec -> Q Dec
+withDecDoc doc dec = do
+  dec' <- dec
+  case doc_loc dec' of
+    Just loc -> qAddModFinalizer $ qPutDoc loc doc
+    Nothing  -> pure ()
+  pure dec'
+  where
+    doc_loc (FunD n _)                                     = Just $ DeclDoc n
+    doc_loc (ValD (VarP n) _ _)                            = Just $ DeclDoc n
+    doc_loc (DataD _ n _ _ _ _)                            = Just $ DeclDoc n
+    doc_loc (NewtypeD _ n _ _ _ _)                         = Just $ DeclDoc n
+    doc_loc (TySynD n _ _)                                 = Just $ DeclDoc n
+    doc_loc (ClassD _ n _ _ _)                             = Just $ DeclDoc n
+    doc_loc (SigD n _)                                     = Just $ DeclDoc n
+    doc_loc (ForeignD (ImportF _ _ _ n _))                 = Just $ DeclDoc n
+    doc_loc (ForeignD (ExportF _ _ n _))                   = Just $ DeclDoc n
+    doc_loc (InfixD _ n)                                   = Just $ DeclDoc n
+    doc_loc (DataFamilyD n _ _)                            = Just $ DeclDoc n
+    doc_loc (OpenTypeFamilyD (TypeFamilyHead n _ _ _))     = Just $ DeclDoc n
+    doc_loc (ClosedTypeFamilyD (TypeFamilyHead n _ _ _) _) = Just $ DeclDoc n
+    doc_loc (PatSynD n _ _ _)                              = Just $ DeclDoc n
+    doc_loc (PatSynSigD n _)                               = Just $ DeclDoc n
+
+    -- For instances we just pass along the full type
+    doc_loc (InstanceD _ _ t _)           = Just $ InstDoc t
+    doc_loc (DataInstD _ _ t _ _ _)       = Just $ InstDoc t
+    doc_loc (NewtypeInstD _ _ t _ _ _)    = Just $ InstDoc t
+    doc_loc (TySynInstD (TySynEqn _ t _)) = Just $ InstDoc t
+
+    -- Declarations that can't have documentation attached to
+    -- ValDs that aren't a simple variable pattern
+    doc_loc (ValD _ _ _)             = Nothing
+    doc_loc (KiSigD _ _)             = Nothing
+    doc_loc (PragmaD _)              = Nothing
+    doc_loc (RoleAnnotD _ _)         = Nothing
+    doc_loc (StandaloneDerivD _ _ _) = Nothing
+    doc_loc (DefaultSigD _ _)        = Nothing
+    doc_loc (ImplicitParamBindD _ _) = Nothing
+
+-- | Variant of 'withDecDoc' that applies the same documentation to
+-- multiple declarations. Useful for documenting quoted declarations.
+withDecsDoc :: String -> Q [Dec] -> Q [Dec]
+withDecsDoc doc decs = decs >>= mapM (withDecDoc doc . pure)
+
+-- | Variant of 'funD' that attaches Haddock documentation.
+funD_doc :: Name -> [Q Clause]
+         -> Maybe String -- ^ Documentation to attach to function
+         -> [Maybe String] -- ^ Documentation to attach to arguments
+         -> Q Dec
+funD_doc nm cs mfun_doc arg_docs = do
+  qAddModFinalizer $ sequence_
+    [putDoc (ArgDoc nm i) s | (i, Just s) <- zip [0..] arg_docs]
+  let dec = funD nm cs
+  case mfun_doc of
+    Just fun_doc -> withDecDoc fun_doc dec
+    Nothing -> funD nm cs
+
+-- | Variant of 'dataD' that attaches Haddock documentation.
+dataD_doc :: Q Cxt -> Name -> [Q (TyVarBndr ())] -> Maybe (Q Kind)
+          -> [(Q Con, Maybe String, [Maybe String])]
+          -- ^ List of constructors, documentation for the constructor, and
+          -- documentation for the arguments
+          -> [Q DerivClause]
+          -> Maybe String
+          -- ^ Documentation to attach to the data declaration
+          -> Q Dec
+dataD_doc ctxt tc tvs ksig cons_with_docs derivs mdoc = do
+  qAddModFinalizer $ mapM_ docCons cons_with_docs
+  let dec = dataD ctxt tc tvs ksig (map (\(con, _, _) -> con) cons_with_docs) derivs
+  maybe dec (flip withDecDoc dec) mdoc
+
+-- | Variant of 'newtypeD' that attaches Haddock documentation.
+newtypeD_doc :: Q Cxt -> Name -> [Q (TyVarBndr ())] -> Maybe (Q Kind)
+             -> (Q Con, Maybe String, [Maybe String])
+             -- ^ The constructor, documentation for the constructor, and
+             -- documentation for the arguments
+             -> [Q DerivClause]
+             -> Maybe String
+             -- ^ Documentation to attach to the newtype declaration
+             -> Q Dec
+newtypeD_doc ctxt tc tvs ksig con_with_docs@(con, _, _) derivs mdoc = do
+  qAddModFinalizer $ docCons con_with_docs
+  let dec = newtypeD ctxt tc tvs ksig con derivs
+  maybe dec (flip withDecDoc dec) mdoc
+
+-- | Variant of 'dataInstD' that attaches Haddock documentation.
+dataInstD_doc :: Q Cxt -> (Maybe [Q (TyVarBndr ())]) -> Q Type -> Maybe (Q Kind)
+              -> [(Q Con, Maybe String, [Maybe String])]
+              -- ^ List of constructors, documentation for the constructor, and
+              -- documentation for the arguments
+              -> [Q DerivClause]
+              -> Maybe String
+              -- ^ Documentation to attach to the instance declaration
+              -> Q Dec
+dataInstD_doc ctxt mb_bndrs ty ksig cons_with_docs derivs mdoc = do
+  qAddModFinalizer $ mapM_ docCons cons_with_docs
+  let dec = dataInstD ctxt mb_bndrs ty ksig (map (\(con, _, _) -> con) cons_with_docs)
+              derivs
+  maybe dec (flip withDecDoc dec) mdoc
+
+-- | Variant of 'newtypeInstD' that attaches Haddock documentation.
+newtypeInstD_doc :: Q Cxt -> (Maybe [Q (TyVarBndr ())]) -> Q Type
+                 -> Maybe (Q Kind)
+                 -> (Q Con, Maybe String, [Maybe String])
+                 -- ^ The constructor, documentation for the constructor, and
+                 -- documentation for the arguments
+                 -> [Q DerivClause]
+                 -> Maybe String
+                 -- ^ Documentation to attach to the instance declaration
+                 -> Q Dec
+newtypeInstD_doc ctxt mb_bndrs ty ksig con_with_docs@(con, _, _) derivs mdoc = do
+  qAddModFinalizer $ docCons con_with_docs
+  let dec = newtypeInstD ctxt mb_bndrs ty ksig con derivs
+  maybe dec (flip withDecDoc dec) mdoc
+
+-- | Variant of 'patSynD' that attaches Haddock documentation.
+patSynD_doc :: Name -> Q PatSynArgs -> Q PatSynDir -> Q Pat
+            -> Maybe String   -- ^ Documentation to attach to the pattern synonym
+            -> [Maybe String] -- ^ Documentation to attach to the pattern arguments
+            -> Q Dec
+patSynD_doc name args dir pat mdoc arg_docs = do
+  qAddModFinalizer $ sequence_
+    [putDoc (ArgDoc name i) s | (i, Just s) <- zip [0..] arg_docs]
+  let dec = patSynD name args dir pat
+  maybe dec (flip withDecDoc dec) mdoc
+
+-- | Document a data/newtype constructor with its arguments.
+docCons :: (Q Con, Maybe String, [Maybe String]) -> Q ()
+docCons (c, md, arg_docs) = do
+  c' <- c
+  -- Attach docs to the constructors
+  sequence_ [ putDoc (DeclDoc nm) d | Just d <- [md], nm <- get_cons_names c' ]
+  -- Attach docs to the arguments
+  case c' of
+    -- Record selector documentation isn't stored in the argument map,
+    -- but in the declaration map instead
+    RecC _ var_bang_types ->
+      sequence_ [ putDoc (DeclDoc nm) arg_doc
+                  | (Just arg_doc, (nm, _, _)) <- zip arg_docs var_bang_types
+                ]
+    _ ->
+      sequence_ [ putDoc (ArgDoc nm i) arg_doc
+                    | nm <- get_cons_names c'
+                    , (i, Just arg_doc) <- zip [0..] arg_docs
+                ]
+  where
+    get_cons_names :: Con -> [Name]
+    get_cons_names (NormalC n _) = [n]
+    get_cons_names (RecC n _) = [n]
+    get_cons_names (InfixC _ n _) = [n]
+    get_cons_names (ForallC _ _ cons) = get_cons_names cons
+    -- GadtC can have multiple names, e.g
+    -- > data Bar a where
+    -- >   MkBar1, MkBar2 :: a -> Bar a
+    -- Will have one GadtC with [MkBar1, MkBar2] as names
+    get_cons_names (GadtC ns _ _) = ns
+    get_cons_names (RecGadtC ns _ _) = ns
diff --git a/libraries/template-haskell/Language/Haskell/TH/Syntax.hs b/libraries/template-haskell/Language/Haskell/TH/Syntax.hs
index 3cb5a44ee8..d3c5a5eb45 100644
--- a/libraries/template-haskell/Language/Haskell/TH/Syntax.hs
+++ b/libraries/template-haskell/Language/Haskell/TH/Syntax.hs
@@ -123,6 +123,9 @@ class (MonadIO m, MonadFail m) => Quasi m where
   qIsExtEnabled :: Extension -> m Bool
   qExtsEnabled :: m [Extension]
 
+  qPutDoc :: DocLoc -> String -> m ()
+  qGetDoc :: DocLoc -> m (Maybe String)
+
 -----------------------------------------------------
 --      The IO instance of Quasi
 --
@@ -161,6 +164,8 @@ instance Quasi IO where
   qPutQ _               = badIO "putQ"
   qIsExtEnabled _       = badIO "isExtEnabled"
   qExtsEnabled          = badIO "extsEnabled"
+  qPutDoc _ _           = badIO "putDoc"
+  qGetDoc _             = badIO "getDoc"
 
 instance Quote IO where
   newName = newNameIO
@@ -745,6 +750,32 @@ isExtEnabled ext = Q (qIsExtEnabled ext)
 extsEnabled :: Q [Extension]
 extsEnabled = Q qExtsEnabled
 
+-- | Add Haddock documentation to the specified location. This will overwrite
+-- any documentation at the location if it already exists. This will reify the
+-- specified name, so it must be in scope when you call it. If you want to add
+-- documentation to something that you are currently splicing, you can use
+-- 'addModFinalizer' e.g.
+--
+-- > do
+-- >   let nm = mkName "x"
+-- >   addModFinalizer $ putDoc (DeclDoc nm) "Hello"
+-- >   [d| $(varP nm) = 42 |]
+--
+-- The helper functions 'withDecDoc' and 'withDecsDoc' will do this for you, as
+-- will the 'funD_doc' and other @_doc@ combinators.
+-- You most likely want to have the @-haddock@ flag turned on when using this.
+-- Adding documentation to anything outside of the current module will cause an
+-- error.
+putDoc :: DocLoc -> String -> Q ()
+putDoc t s = Q (qPutDoc t s)
+
+-- | Retreives the Haddock documentation at the specified location, if one
+-- exists.
+-- It can be used to read documentation on things defined outside of the current
+-- module, provided that those modules were compiled with the @-haddock@ flag.
+getDoc :: DocLoc -> Q (Maybe String)
+getDoc n = Q (qGetDoc n)
+
 instance MonadIO Q where
   liftIO = runIO
 
@@ -772,6 +803,8 @@ instance Quasi Q where
   qPutQ               = putQ
   qIsExtEnabled       = isExtEnabled
   qExtsEnabled        = extsEnabled
+  qPutDoc             = putDoc
+  qGetDoc             = getDoc
 
 
 ----------------------------------------------------
@@ -2625,6 +2658,17 @@ constructors):
                     (PromotedConsT  `AppT` IO `AppT` PromotedNilT)
 -}
 
+-- | A location at which to attach Haddock documentation.
+-- Note that adding documentation to a 'Name' defined oustide of the current
+-- module will cause an error.
+data DocLoc
+  = ModuleDoc         -- ^ At the current module's header.
+  | DeclDoc Name      -- ^ At a declaration, not necessarily top level.
+  | ArgDoc Name Int   -- ^ At a specific argument of a function, indexed by its
+                      -- position.
+  | InstDoc Type      -- ^ At a class or family instance.
+  deriving ( Show, Eq, Ord, Data, Generic )
+
 -----------------------------------------------------
 --              Internal helper functions
 -----------------------------------------------------
diff --git a/libraries/template-haskell/changelog.md b/libraries/template-haskell/changelog.md
index 6d6e06b8ce..0a570a89ee 100644
--- a/libraries/template-haskell/changelog.md
+++ b/libraries/template-haskell/changelog.md
@@ -1,5 +1,13 @@
 # Changelog for [`template-haskell` package](http://hackage.haskell.org/package/template-haskell)
 
+## 2.18.0.0
+  * Add `putDoc` and `getDoc` which allow Haddock documentation to be attached
+    to module headers, declarations, function arguments and instances, as well
+    as queried. These are quite low level operations, so for convenience there
+    are several combinators that can be used with `Dec`s directly, including
+    `withDecDoc`/`withDecsDoc` as well as `_doc` counterparts to many of the
+    `Dec` helper functions.
+
 ## 2.17.0.0
   * Typed Quotations now return a value of type `Code m a` (GHC Proposal #195).
     The main motiviation is to make writing instances easier and make it easier to