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

11 files changed, 393 insertions, 0 deletions

diff --git a/testsuite/tests/showIface/DocsInHiFileTH.hs b/testsuite/tests/showIface/DocsInHiFileTH.hs
new file mode 100644
index 0000000000..73b46c8876
--- /dev/null
+++ b/testsuite/tests/showIface/DocsInHiFileTH.hs
@@ -0,0 +1,218 @@
+{-# LANGUAGE TemplateHaskell, FlexibleInstances, TypeFamilies, DataKinds #-}
+{-# LANGUAGE MultiParamTypeClasses, StandaloneKindSignatures, PolyKinds #-}
+{-# LANGUAGE PatternSynonyms #-}
+
+-- |This is the module header
+module DocInHiFilesTH where
+
+import Language.Haskell.TH
+import Language.Haskell.TH.Syntax
+import DocsInHiFileTHExternal
+
+f :: Int
+f = 42
+
+$(putDoc (DeclDoc 'f) "The meaning of life" >> pure [])
+
+-- |A data type
+data Foo =
+    -- |A constructor
+    Foo
+
+do
+  Just "A data type" <- getDoc (DeclDoc ''Foo)
+  Just "A constructor" <- getDoc (DeclDoc 'Foo)
+  putDoc (DeclDoc ''Foo) "A new data type"
+  putDoc (DeclDoc 'Foo) "A new constructor"
+  Just "A new data type" <- getDoc (DeclDoc ''Foo)
+  Just "A new constructor" <- getDoc (DeclDoc 'Foo)
+  pure []
+
+-- |Some documentation
+g :: String
+g = "Hello world"
+
+do
+  Just "Some documentation" <- getDoc (DeclDoc 'g)
+  pure []
+
+-- Testing module headers
+
+do
+  Just "This is the module header" <- getDoc ModuleDoc
+  putDoc ModuleDoc "This is the new module header"
+  Just "This is the new module header" <- getDoc ModuleDoc
+  pure []
+
+-- Testing argument documentation
+
+h :: Int -- ^Your favourite number
+  -> Bool -- ^Your favourite element in the Boolean algebra
+  -> String -- ^A return value
+h _ _ = "Hello world"
+
+do
+  Just "Your favourite number" <- getDoc (ArgDoc 'h 0)
+  Just "Your favourite element in the Boolean algebra" <- getDoc (ArgDoc 'h 1)
+  Just "A return value" <- getDoc (ArgDoc 'h 2)
+  Nothing <- getDoc (ArgDoc 'h 3)
+  putDoc (ArgDoc 'h 1) "Your least favourite Boolean"
+  Just "Your least favourite Boolean" <- getDoc (ArgDoc 'h 1)
+  pure []
+
+
+-- Testing classes and instances
+
+-- |A fancy class
+class C a where
+
+-- |A fancy instance
+instance C Int where
+instance C String where
+
+class D a where
+-- |Another fancy instance
+instance D a where
+
+-- |A type family
+type family E a
+
+-- |A type family instance
+type instance E Bool = Int
+
+i :: E Bool
+i = 42
+
+do
+  Just "A fancy class" <- getDoc (DeclDoc ''C)
+  Just "A fancy instance" <- getDoc . InstDoc =<< [t| C Int |]
+  Just "Another fancy instance" <- getDoc (InstDoc (AppT (ConT ''D) (VarT (mkName "a"))))
+  Just "Another fancy instance" <- getDoc (InstDoc (AppT (ConT ''D) (VarT (mkName "b"))))
+  Nothing <- getDoc . InstDoc =<< [t| C String |]
+
+  putDoc (DeclDoc ''C) "A new class"
+  putDoc (InstDoc (AppT (ConT ''C) (ConT ''Int))) "A new instance"
+  putDoc (InstDoc (AppT (ConT ''C) (ConT ''String))) "Another new instance"
+  putDoc (InstDoc (AppT (ConT ''D) (VarT (mkName "a")))) "Another new instance"
+  Just "A new class" <- getDoc (DeclDoc ''C)
+  Just "A new instance" <- getDoc . InstDoc =<< [t| C Int |]
+  Just "Another new instance" <- getDoc . InstDoc =<< [t| C String |]
+  Just "Another new instance" <- getDoc (InstDoc (AppT (ConT ''D) (VarT (mkName "a"))))
+
+  Just "A type family" <- getDoc (DeclDoc ''E)
+  -- Doesn't work just yet. See T18241
+  -- https://gitlab.haskell.org/ghc/ghc/issues/18241
+  Just "A type family instance" <- getDoc . InstDoc =<< [t| E Bool |]
+
+  pure []
+
+-- Testing documentation from external modules
+do
+  Just "This is an external function" <- getDoc (DeclDoc 'externalFunc)
+  Just "Some integer" <- getDoc (ArgDoc 'externalFunc 0)
+
+  Just "This is an external class" <- getDoc (DeclDoc ''ExternalClass)
+  Just "This is an external instance" <-
+    getDoc . InstDoc =<< [t| ExternalClass Int |]
+
+  pure []
+
+data family WD11 a
+type family WD13 a
+
+wd8 = ()
+
+class F
+
+-- Testing combinators
+
+withDecsDoc "1" [d| wd1 x = () |]
+withDecsDoc "2" [d| wd2 = () |]
+withDecsDoc "3" [d| data WD3 = WD3 |]
+withDecsDoc "4" [d| newtype WD4 = WD4 () |]
+withDecsDoc "5" [d| type WD5 = () |]
+withDecsDoc "6" [d| class WD6 a where |]
+withDecsDoc "7" [d| instance C Foo where |]
+do
+  d <- withDecDoc "8" $ sigD 'wd8 [t| () |]
+  pure [d]
+--  this gives 'Illegal variable name: ‘WD9’' when splicing
+--  withDoc "9"  [sigD ''WD9 [t| Type -> Type |]]
+withDecsDoc "10" [d| data family WD10 a|]
+withDecsDoc "11" [d| data instance WD11 Foo = WD11Foo |]
+withDecsDoc "12" [d| type family WD12 a |]
+withDecsDoc "13" [d| type instance WD13 Foo = Int |]
+
+--  testing nullary classes here
+withDecsDoc "14" [d| instance F |]
+
+withDecsDoc "15" [d| foreign import ccall "math.h sin" sin :: Double -> Double |]
+-- this gives 'Foreign export not (yet) handled by Template Haskell'
+-- withDecsDoc "16" [d| foreign export ccall "addInt" (+) :: Int -> Int -> Int |]
+
+wd17 = 42
+
+do
+  d <- withDecDoc "17" (sigD 'wd17 [t| Int |])
+  pure [d]
+
+do
+  let nm = mkName "wd18"
+  d' <- withDecDoc "18" $ sigD nm [t| Int |]
+  d <-  withDecDoc "19" $ valD (varP nm) (normalB [| 42 |]) []
+  pure [d, d']
+
+-- Doing this to test that wd20 is documented as "20" and not "2020"
+withDecsDoc "20" [d|
+  wd20 :: Int
+  wd20 = 42
+  |]
+
+do
+  let defBang = bang noSourceUnpackedness noSourceStrictness
+  patSynVarName <- newName "a"
+  sequenceA
+    [ funD_doc (mkName "qux") [clause [ [p| a |], [p| b |] ] (normalB [e| () |]) []]
+        (Just "This is qux") [Just "Arg uno", Just "Arg dos"]
+
+    , dataD_doc (cxt []) (mkName "Quux") [] Nothing
+        [ ( normalC (mkName "Quux1") [bangType defBang (reifyType ''Int)]
+          , Just "This is Quux1",  [Just "I am an integer"])
+        , ( normalC (mkName "Quux2")
+              [ bangType defBang (reifyType ''String)
+              , bangType defBang (reifyType ''Bool)
+              ]
+          , Just "This is Quux2", map Just ["I am a string", "I am a bool"])
+        ] [] (Just "This is Quux")
+
+    , dataD_doc (cxt []) (mkName "Quuz") [] Nothing
+        [ ( recC (mkName "Quuz") [varBangType (mkName "quuz1_a") (bangType defBang (reifyType ''String))]
+        , Just "This is a record constructor", [Just "This is the record constructor's argument"])
+        ] [] (Just "This is a record type")
+
+    , newtypeD_doc (cxt []) (mkName "Corge") [] Nothing
+        ( recC (mkName ("Corge")) [varBangType (mkName "runCorge") (bangType defBang [t| Int |])]
+        , Just "This is a newtype record constructor", [Just "This is the newtype record constructor's argument"]
+        ) [] (Just "This is a record newtype")
+
+    , dataInstD_doc (cxt []) Nothing [t| WD11 Int |] Nothing
+        [ ( normalC (mkName "WD11Int") [bangType defBang [t| Int |]]
+        , Just "This is a data instance constructor", [Just "This is a data instance constructor argument"])
+        ] [] (Just "This is a data instance")
+
+    , newtypeInstD_doc (cxt []) Nothing [t| WD11 Bool |] Nothing
+        (normalC (mkName "WD11Bool") [bangType defBang [t| Bool |]]
+        , Just "This is a newtype instance constructor", [Just "This is a newtype instance constructor argument"])
+        [] (Just "This is a newtype instance")
+
+    , patSynD_doc (mkName "Tup2") (prefixPatSyn [patSynVarName]) unidir
+        [p| ($(varP patSynVarName), $(varP patSynVarName)) |]
+        (Just "Matches a tuple of (a, a)") [Just "The thing to match twice"]
+
+    , withDecDoc "My cool class" $ do
+        tyVar <- newName "a"
+        classD (cxt []) (mkName "Pretty") [plainTV tyVar] []
+          [ withDecDoc "Prettily prints the object" $
+              sigD (mkName "prettyPrint") [t| $(varT tyVar) -> String |]
+          ]
+    ]
diff --git a/testsuite/tests/showIface/DocsInHiFileTH.stdout b/testsuite/tests/showIface/DocsInHiFileTH.stdout
new file mode 100644
index 0000000000..6951b9a1e5
--- /dev/null
+++ b/testsuite/tests/showIface/DocsInHiFileTH.stdout
@@ -0,0 +1,118 @@
+module header:
+  Just "This is the new module header"
+declaration docs:
+  Tup2:
+    "Matches a tuple of (a, a)"
+  f:
+    "The meaning of life"
+  g:
+    "Some documentation"
+  qux:
+    "This is qux"
+  sin:
+    "15"
+  wd1:
+    "1"
+  wd17:
+    "17"
+  wd18:
+    "18"
+  wd2:
+    "2"
+  wd20:
+    "20"
+  wd8:
+    "8"
+  C:
+    "A new class"
+  Corge:
+    "This is a newtype record constructor"
+  runCorge:
+    "This is the newtype record constructor's argument"
+  E:
+    "A type family"
+  Foo:
+    "A new data type"
+  Foo:
+    "A new constructor"
+  Pretty:
+    "My cool class"
+  prettyPrint:
+    "Prettily prints the object"
+  Quux:
+    "This is Quux"
+  Quux1:
+    "This is Quux1"
+  Quux2:
+    "This is Quux2"
+  Quuz:
+    "This is a record constructor"
+  quuz1_a:
+    "This is the record constructor's argument"
+  WD10:
+    "10"
+  WD11Bool:
+    "This is a newtype instance constructor"
+  WD11Int:
+    "This is a data instance constructor"
+  WD12:
+    "12"
+  WD3:
+    "3"
+  WD4:
+    "4"
+  WD5:
+    "5"
+  WD6:
+    "6"
+  $fCTYPEFoo:
+    "7"
+  $fCTYPEInt:
+    "A new instance"
+  $fCTYPE[]:
+    "Another new instance"
+  $fDka:
+    "Another new instance"
+  $fF:
+    "14"
+  D:R:EBool:
+    "A type family instance"
+  D:R:WD11Bool0:
+    "This is a newtype instance"
+  D:R:WD11Foo0:
+    "11"
+  D:R:WD11Int0:
+    "This is a data instance"
+  D:R:WD13Foo:
+    "13"
+arg docs:
+  Tup2:
+    0:
+      "The thing to match twice"
+  h:
+    0:
+      "Your favourite number"
+    1:
+      "Your least favourite Boolean"
+    2:
+      "A return value"
+  qux:
+    0:
+      "Arg uno"
+    1:
+      "Arg dos"
+  Quux1:
+    0:
+      "I am an integer"
+  Quux2:
+    0:
+      "I am a string"
+    1:
+      "I am a bool"
+  WD11Bool:
+    0:
+      "This is a newtype instance constructor argument"
+  WD11Int:
+    0:
+      "This is a data instance constructor argument"
+extensible fields:
diff --git a/testsuite/tests/showIface/DocsInHiFileTHExternal.hs b/testsuite/tests/showIface/DocsInHiFileTHExternal.hs
new file mode 100644
index 0000000000..9a1d46b05e
--- /dev/null
+++ b/testsuite/tests/showIface/DocsInHiFileTHExternal.hs
@@ -0,0 +1,12 @@
+module DocsInHiFileTHExternal where
+
+-- |This is an external function
+externalFunc :: Int -- ^Some integer
+             -> Int -- ^Another integer
+externalFunc = const 42
+
+-- |This is an external class
+class ExternalClass a where
+
+-- |This is an external instance
+instance ExternalClass Int where
diff --git a/testsuite/tests/showIface/Makefile b/testsuite/tests/showIface/Makefile
index 7eafdfc9d2..c45f38684e 100644
--- a/testsuite/tests/showIface/Makefile
+++ b/testsuite/tests/showIface/Makefile
@@ -13,3 +13,7 @@ DocsInHiFile0:
 DocsInHiFile1:
 	'$(TEST_HC)' $(TEST_HC_OPTS) -c -haddock DocsInHiFile.hs
 	'$(TEST_HC)' $(TEST_HC_OPTS) --show-iface DocsInHiFile.hi | grep -A 100 'module header:'
+
+DocsInHiFileTH:
+	'$(TEST_HC)' $(TEST_HC_OPTS) -c -haddock DocsInHiFileTHExternal.hs DocsInHiFileTH.hs
+	'$(TEST_HC)' $(TEST_HC_OPTS) --show-iface DocsInHiFileTH.hi | grep -A 200 'module header:'
diff --git a/testsuite/tests/showIface/all.T b/testsuite/tests/showIface/all.T
index e2ec264431..a5e5f5f085 100644
--- a/testsuite/tests/showIface/all.T
+++ b/testsuite/tests/showIface/all.T
@@ -6,3 +6,6 @@ test('DocsInHiFile1',
      extra_files(['DocsInHiFile.hs']),
      makefile_test, ['DocsInHiFile1'])
 test('T17871', [extra_files(['T17871a.hs'])], multimod_compile, ['T17871', '-v0'])
+test('DocsInHiFileTH',
+     extra_files(['DocsInHiFileTHExternal.hs', 'DocsInHiFileTH.hs']),
+     makefile_test, ['DocsInHiFileTH'])
diff --git a/testsuite/tests/showIface/should_fail/THPutDocExternal.hs b/testsuite/tests/showIface/should_fail/THPutDocExternal.hs
new file mode 100644
index 0000000000..f9a180af4c
--- /dev/null
+++ b/testsuite/tests/showIface/should_fail/THPutDocExternal.hs
@@ -0,0 +1,8 @@
+{-# LANGUAGE TemplateHaskell #-}
+
+module THPutDocExternal where
+
+import Language.Haskell.TH
+import THPutDocExternalA
+
+putDoc (DeclDoc 'f) "Hello world" >> pure []
diff --git a/testsuite/tests/showIface/should_fail/THPutDocExternal.stderr b/testsuite/tests/showIface/should_fail/THPutDocExternal.stderr
new file mode 100644
index 0000000000..3063fe9350
--- /dev/null
+++ b/testsuite/tests/showIface/should_fail/THPutDocExternal.stderr
@@ -0,0 +1,2 @@
+THPutDocExternal.hs:8:1:
+    Can't add documentation to THPutDocExternalA.f as it isn't inside the current module
diff --git a/testsuite/tests/showIface/should_fail/THPutDocExternalA.hs b/testsuite/tests/showIface/should_fail/THPutDocExternalA.hs
new file mode 100644
index 0000000000..694266bbe9
--- /dev/null
+++ b/testsuite/tests/showIface/should_fail/THPutDocExternalA.hs
@@ -0,0 +1,4 @@
+module THPutDocExternalA where
+
+f :: Int
+f = 42
diff --git a/testsuite/tests/showIface/should_fail/THPutDocNonExistent.hs b/testsuite/tests/showIface/should_fail/THPutDocNonExistent.hs
new file mode 100644
index 0000000000..d0b1d7a162
--- /dev/null
+++ b/testsuite/tests/showIface/should_fail/THPutDocNonExistent.hs
@@ -0,0 +1,13 @@
+{-# LANGUAGE TemplateHaskell #-}
+
+module THPutDocNonExistent where
+
+import Language.Haskell.TH
+
+class A a where
+data B
+
+do
+  t <- [t| A B |]
+  putDoc (InstDoc t) "a"
+  pure []
diff --git a/testsuite/tests/showIface/should_fail/THPutDocNonExistent.stderr b/testsuite/tests/showIface/should_fail/THPutDocNonExistent.stderr
new file mode 100644
index 0000000000..ce3a64a1d9
--- /dev/null
+++ b/testsuite/tests/showIface/should_fail/THPutDocNonExistent.stderr
@@ -0,0 +1,2 @@
+THPutDocNonExistent.hs:10:1:
+    Couldn't find any instances of THPutDocNonExistent.A THPutDocNonExistent.B to add documentation to
diff --git a/testsuite/tests/showIface/should_fail/all.T b/testsuite/tests/showIface/should_fail/all.T
new file mode 100644
index 0000000000..0dd8106b81
--- /dev/null
+++ b/testsuite/tests/showIface/should_fail/all.T
@@ -0,0 +1,9 @@
+test('THPutDocExternal',
+     normal,
+     multimod_compile_fail,
+     ['THPutDocExternal', '-no-hs-main -haddock -c -v0'])
+
+test('THPutDocNonExistent',
+     normal,
+     multimod_compile_fail,
+     ['THPutDocNonExistent', '-no-hs-main -haddock -c -v0'])