ApiAnnotations : preserve parens in GADTs

A cleanup in 7f418acf61e accidentally discarded some parens in
ConDeclGADT.

Make sure these stay in the AST in a usable format.

Also ensure the AnnLolly does not get lost in a GADT.

5 files changed, 65 insertions, 53 deletions

diff --git a/compiler/GHC/Hs/Type.hs b/compiler/GHC/Hs/Type.hs
index 818fe75475..f1f688fdea 100644
--- a/compiler/GHC/Hs/Type.hs
+++ b/compiler/GHC/Hs/Type.hs
@@ -106,6 +106,7 @@ import GHC.Utils.Misc ( count )
 
 import Data.Data hiding ( Fixity, Prefix, Infix )
 import Data.Maybe
+import GHC.Parser.Annotation
 
 {-
 ************************************************************************
@@ -1325,17 +1326,20 @@ mkHsAppKindTy ext ty k
 -- splitHsFunType decomposes a type (t1 -> t2 ... -> tn)
 -- Breaks up any parens in the result type:
 --      splitHsFunType (a -> (b -> c)) = ([a,b], c)
+-- It returns API Annotations for any parens removed
 splitHsFunType ::
      LHsType (GhcPass p)
-  -> ([HsScaled (GhcPass p) (LHsType (GhcPass p))], LHsType (GhcPass p))
-splitHsFunType (L _ (HsParTy _ ty))
-  = splitHsFunType ty
+  -> ([HsScaled (GhcPass p) (LHsType (GhcPass p))], LHsType (GhcPass p), [AddAnn])
+splitHsFunType ty = go ty []
+  where
+    go (L l (HsParTy _ ty)) anns
+      = go ty (anns ++ mkParensApiAnn l)
 
-splitHsFunType (L _ (HsFunTy _ mult x y))
-  | (args, res) <- splitHsFunType y
-  = (HsScaled mult x:args, res)
+    go (L _ (HsFunTy _ mult x y)) anns
+      | (args, res, anns') <- go y anns
+      = (HsScaled mult x:args, res, anns')
 
-splitHsFunType other = ([], other)
+    go other anns = ([], other, anns)
 
 -- | Retrieve the name of the \"head\" of a nested type application.
 -- This is somewhat like @GHC.Tc.Gen.HsType.splitHsAppTys@, but a little more
diff --git a/compiler/GHC/Parser.y b/compiler/GHC/Parser.y
index 16eff90656..b8a777a8fe 100644
--- a/compiler/GHC/Parser.y
+++ b/compiler/GHC/Parser.y
@@ -2055,12 +2055,14 @@ type :: { LHsType GhcPs }
                                        >> ams (sLL $1 $> $ HsFunTy noExtField HsUnrestrictedArrow $1 $3)
                                               [mu AnnRarrow $2] }
 
-        | btype mult '->' ctype        {% hintLinear (getLoc $2) >>
-                                          ams (sLL $1 $> $ HsFunTy noExtField (unLoc $2) $1 $4)
+        | btype mult '->' ctype        {% hintLinear (getLoc $2)
+                                       >> ams $1 [mu AnnRarrow $3] -- See note [GADT decl discards annotations]
+                                       >> ams (sLL $1 $> $ HsFunTy noExtField (unLoc $2) $1 $4)
                                               [mu AnnRarrow $3] }
 
-        | btype '->.' ctype            {% hintLinear (getLoc $2) >>
-                                          ams (sLL $1 $> $ HsFunTy noExtField HsLinearArrow $1 $3)
+        | btype '->.' ctype            {% hintLinear (getLoc $2)
+                                       >> ams $1 [mu AnnLollyU $2] -- See note [GADT decl discards annotations]
+                                       >> ams (sLL $1 $> $ HsFunTy noExtField HsLinearArrow $1 $3)
                                               [mu AnnLollyU $2] }
 
 mult :: { Located (HsArrow GhcPs) }
@@ -2285,9 +2287,9 @@ gadt_constr :: { LConDecl GhcPs }
     -- see Note [Difference in parsing GADT and data constructors]
     -- Returns a list because of:   C,D :: ty
         : optSemi con_list '::' sigtype
-                {% do { decl <- mkGadtDecl (unLoc $2) $4
+                {% do { (decl, anns) <- mkGadtDecl (unLoc $2) $4
                       ; ams (sLL $2 $> decl)
-                            [mu AnnDcolon $3] } }
+                            (mu AnnDcolon $3:anns) } }
 
 {- Note [Difference in parsing GADT and data constructors]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/compiler/GHC/Parser/Annotation.hs b/compiler/GHC/Parser/Annotation.hs
index 6560d5e735..0cbf44296f 100644
--- a/compiler/GHC/Parser/Annotation.hs
+++ b/compiler/GHC/Parser/Annotation.hs
@@ -6,6 +6,7 @@ module GHC.Parser.Annotation (
   ApiAnns(..),
   ApiAnnKey,
   AnnKeywordId(..),
+  AddAnn(..),mkParensApiAnn,
   AnnotationComment(..),
   IsUnicodeSyntax(..),
   unicodeAnn,
@@ -148,6 +149,44 @@ data ApiAnns =
 type ApiAnnKey = (RealSrcSpan,AnnKeywordId)
 
 
+-- ---------------------------------------------------------------------
+
+-- | Encapsulated call to addAnnotation, requiring only the SrcSpan of
+--   the AST construct the annotation belongs to; together with the
+--   AnnKeywordId, this is the key of the annotation map.
+--
+--   This type is useful for places in the parser where it is not yet
+--   known what SrcSpan an annotation should be added to.  The most
+--   common situation is when we are parsing a list: the annotations
+--   need to be associated with the AST element that *contains* the
+--   list, not the list itself.  'AddAnn' lets us defer adding the
+--   annotations until we finish parsing the list and are now parsing
+--   the enclosing element; we then apply the 'AddAnn' to associate
+--   the annotations.  Another common situation is where a common fragment of
+--   the AST has been factored out but there is no separate AST node for
+--   this fragment (this occurs in class and data declarations). In this
+--   case, the annotation belongs to the parent data declaration.
+--
+--   The usual way an 'AddAnn' is created is using the 'mj' ("make jump")
+--   function, and then it can be discharged using the 'ams' function.
+data AddAnn = AddAnn AnnKeywordId SrcSpan
+
+-- |Given a 'SrcSpan' that surrounds a 'HsPar' or 'HsParTy', generate
+-- 'AddAnn' values for the opening and closing bordering on the start
+-- and end of the span
+mkParensApiAnn :: SrcSpan -> [AddAnn]
+mkParensApiAnn (UnhelpfulSpan _)  = []
+mkParensApiAnn (RealSrcSpan ss _) = [AddAnn AnnOpenP lo,AddAnn AnnCloseP lc]
+  where
+    f = srcSpanFile ss
+    sl = srcSpanStartLine ss
+    sc = srcSpanStartCol ss
+    el = srcSpanEndLine ss
+    ec = srcSpanEndCol ss
+    lo = RealSrcSpan (mkRealSrcSpan (realSrcSpanStart ss)        (mkRealSrcLoc f sl (sc+1))) Nothing
+    lc = RealSrcSpan (mkRealSrcSpan (mkRealSrcLoc f el (ec - 1)) (realSrcSpanEnd ss))        Nothing
+
+-- ---------------------------------------------------------------------
 -- | Retrieve a list of annotation 'SrcSpan's based on the 'SrcSpan'
 -- of the annotated AST element, and the known type of the annotation.
 getAnnotation :: ApiAnns -> RealSrcSpan -> AnnKeywordId -> [RealSrcSpan]
diff --git a/compiler/GHC/Parser/Lexer.x b/compiler/GHC/Parser/Lexer.x
index b31e445740..f1b6e4efc6 100644
--- a/compiler/GHC/Parser/Lexer.x
+++ b/compiler/GHC/Parser/Lexer.x
@@ -64,7 +64,6 @@ module GHC.Parser.Lexer (
    ExtBits(..),
    xtest, xunset, xset,
    lexTokenStream,
-   AddAnn(..),mkParensApiAnn,
    addAnnsAt,
    commentToAnnotation,
    HdkComment(..),
@@ -3299,45 +3298,12 @@ clean_pragma prag = canon_ws (map toLower (unprefix prag))
 %************************************************************************
 -}
 
--- | Encapsulated call to addAnnotation, requiring only the SrcSpan of
---   the AST construct the annotation belongs to; together with the
---   AnnKeywordId, this is the key of the annotation map.
---
---   This type is useful for places in the parser where it is not yet
---   known what SrcSpan an annotation should be added to.  The most
---   common situation is when we are parsing a list: the annotations
---   need to be associated with the AST element that *contains* the
---   list, not the list itself.  'AddAnn' lets us defer adding the
---   annotations until we finish parsing the list and are now parsing
---   the enclosing element; we then apply the 'AddAnn' to associate
---   the annotations.  Another common situation is where a common fragment of
---   the AST has been factored out but there is no separate AST node for
---   this fragment (this occurs in class and data declarations). In this
---   case, the annotation belongs to the parent data declaration.
---
---   The usual way an 'AddAnn' is created is using the 'mj' ("make jump")
---   function, and then it can be discharged using the 'ams' function.
-data AddAnn = AddAnn AnnKeywordId SrcSpan
 
 addAnnotationOnly :: RealSrcSpan -> AnnKeywordId -> RealSrcSpan -> P ()
 addAnnotationOnly l a v = P $ \s -> POk s {
   annotations = ((l,a), [v]) : annotations s
   } ()
 
--- |Given a 'SrcSpan' that surrounds a 'HsPar' or 'HsParTy', generate
--- 'AddAnn' values for the opening and closing bordering on the start
--- and end of the span
-mkParensApiAnn :: SrcSpan -> [AddAnn]
-mkParensApiAnn (UnhelpfulSpan _)  = []
-mkParensApiAnn (RealSrcSpan ss _) = [AddAnn AnnOpenP lo,AddAnn AnnCloseP lc]
-  where
-    f = srcSpanFile ss
-    sl = srcSpanStartLine ss
-    sc = srcSpanStartCol ss
-    el = srcSpanEndLine ss
-    ec = srcSpanEndCol ss
-    lo = RealSrcSpan (mkRealSrcSpan (realSrcSpanStart ss)        (mkRealSrcLoc f sl (sc+1))) Nothing
-    lc = RealSrcSpan (mkRealSrcSpan (mkRealSrcLoc f el (ec - 1)) (realSrcSpanEnd ss))        Nothing
 
 queueComment :: RealLocated Token -> P()
 queueComment c = P $ \s -> POk s {
diff --git a/compiler/GHC/Parser/PostProcess.hs b/compiler/GHC/Parser/PostProcess.hs
index 2fd38dbdba..4294046d1f 100644
--- a/compiler/GHC/Parser/PostProcess.hs
+++ b/compiler/GHC/Parser/PostProcess.hs
@@ -633,16 +633,16 @@ mkConDeclH98 name mb_forall mb_cxt args
 --   Note [GADT abstract syntax] in "GHC.Hs.Decls" for more details.
 mkGadtDecl :: [Located RdrName]
            -> LHsType GhcPs
-           -> P (ConDecl GhcPs)
+           -> P (ConDecl GhcPs, [AddAnn])
 mkGadtDecl names ty = do
-  let (args, res_ty)
+  let (args, res_ty, anns)
         | L _ (HsFunTy _ _w (L loc (HsRecTy _ rf)) res_ty) <- body_ty
-        = (RecCon (L loc rf), res_ty)
+        = (RecCon (L loc rf), res_ty, [])
         | otherwise
-        = let (arg_types, res_type) = splitHsFunType body_ty
-          in (PrefixCon arg_types, res_type)
+        = let (arg_types, res_type, anns) = splitHsFunType body_ty
+          in (PrefixCon arg_types, res_type, anns)
 
-  pure $ ConDeclGADT { con_g_ext  = noExtField
+  pure ( ConDeclGADT { con_g_ext  = noExtField
                      , con_names  = names
                      , con_forall = L (getLoc ty) $ isJust mtvs
                      , con_qvars  = fromMaybe [] mtvs
@@ -650,6 +650,7 @@ mkGadtDecl names ty = do
                      , con_args   = args
                      , con_res_ty = res_ty
                      , con_doc    = Nothing }
+       , anns )
   where
     (mtvs, mcxt, body_ty) = splitLHsGadtTy ty