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, 164 insertions, 5 deletions

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