Parser API annotations: RealSrcLoc

During parsing, GHC collects lexical information about AST nodes and
stores it in a map. It is needed to faithfully restore original source
code, e.g. compare these expressions:

	a =  b
	a  = b

The position of the equality sign is not recorded in the AST, so it must
be stored elsewhere.

This system is described in Note [Api annotations].

Before this patch, the mapping was represented by:

	Map (SrcSpan, AnnKeywordId) SrcSpan

After this patch, the mapping is represented by:

	Map (RealSrcSpan, AnnKeywordId) RealSrcSpan

The motivation behind this change is to avoid using the Ord SrcSpan
instance (required by Map here), as it interferes with #17632 (see the
discussion there).

SrcSpan is isomorphic to  Either String RealSrcSpan,  but we shouldn't
use those strings as Map keys. Those strings are intended as hints to
the user, e.g. "<interactive>" or "<compiler-generated code>", so they
are not a valid way to identify nodes in the source code.

7 files changed, 109 insertions, 79 deletions

diff --git a/compiler/backpack/DriverBkp.hs b/compiler/backpack/DriverBkp.hs
index ee4d9fb5e4..4246a04dd3 100644
--- a/compiler/backpack/DriverBkp.hs
+++ b/compiler/backpack/DriverBkp.hs
@@ -23,6 +23,7 @@ import GhcPrelude
 -- In a separate module because it hooks into the parser.
 import BkpSyn
 
+import ApiAnnotation
 import GHC hiding (Failed, Succeeded)
 import Packages
 import Parser
@@ -702,7 +703,7 @@ summariseRequirement pn mod_name = do
                         hsmodHaddockModHeader = Nothing
                     }),
                 hpm_src_files = [],
-                hpm_annotations = (Map.empty, Map.empty)
+                hpm_annotations = ApiAnns Map.empty Nothing Map.empty []
             }),
         ms_hspp_file = "", -- none, it came inline
         ms_hspp_opts = dflags,
@@ -812,7 +813,7 @@ hsModuleToModSummary pn hsc_src modname
             ms_parsed_mod = Just (HsParsedModule {
                     hpm_module = hsmod,
                     hpm_src_files = [], -- TODO if we preprocessed it
-                    hpm_annotations = (Map.empty, Map.empty) -- BOGUS
+                    hpm_annotations = ApiAnns Map.empty Nothing Map.empty [] -- BOGUS
                 }),
             ms_hs_date = time,
             ms_obj_date = Nothing, -- TODO do this, but problem: hi_timestamp is BOGUS
diff --git a/compiler/basicTypes/SrcLoc.hs b/compiler/basicTypes/SrcLoc.hs
index 4ddbf2b683..199888ced6 100644
--- a/compiler/basicTypes/SrcLoc.hs
+++ b/compiler/basicTypes/SrcLoc.hs
@@ -83,7 +83,7 @@ module SrcLoc (
         -- ** Combining and comparing Located values
         eqLocated, cmpLocated, combineLocs, addCLoc,
         leftmost_smallest, leftmost_largest, rightmost,
-        spans, isSubspanOf, sortLocated,
+        spans, isSubspanOf, isRealSubspanOf, sortLocated,
 
         liftL
     ) where
@@ -180,7 +180,7 @@ advanceSrcLoc (SrcLoc f l c) _    = SrcLoc f  l (c + 1)
 ************************************************************************
 -}
 
-sortLocated :: [Located a] -> [Located a]
+sortLocated :: Ord l => [GenLocated l a] -> [GenLocated l a]
 sortLocated things = sortBy (comparing getLoc) things
 
 instance Outputable RealSrcLoc where
@@ -596,10 +596,17 @@ spans (RealSrcSpan span) (l,c) = realSrcSpanStart span <= loc && loc <= realSrcS
 isSubspanOf :: SrcSpan -- ^ The span that may be enclosed by the other
             -> SrcSpan -- ^ The span it may be enclosed by
             -> Bool
-isSubspanOf src parent
-    | srcSpanFileName_maybe parent /= srcSpanFileName_maybe src = False
-    | otherwise = srcSpanStart parent <= srcSpanStart src &&
-                  srcSpanEnd parent   >= srcSpanEnd src
+isSubspanOf (RealSrcSpan src) (RealSrcSpan parent) = isRealSubspanOf src parent
+isSubspanOf _ _ = False
+
+-- | Determines whether a span is enclosed by another one
+isRealSubspanOf :: RealSrcSpan -- ^ The span that may be enclosed by the other
+                -> RealSrcSpan -- ^ The span it may be enclosed by
+                -> Bool
+isRealSubspanOf src parent
+    | srcSpanFile parent /= srcSpanFile src = False
+    | otherwise = realSrcSpanStart parent <= realSrcSpanStart src &&
+                  realSrcSpanEnd parent   >= realSrcSpanEnd src
 
 liftL :: Monad m => (a -> m b) -> GenLocated l a -> m (GenLocated l b)
 liftL f (L loc a) = do
diff --git a/compiler/main/GHC.hs b/compiler/main/GHC.hs
index 1bbf4a4929..b15803eed1 100644
--- a/compiler/main/GHC.hs
+++ b/compiler/main/GHC.hs
@@ -275,7 +275,7 @@ module GHC (
         parser,
 
         -- * API Annotations
-        ApiAnns,AnnKeywordId(..),AnnotationComment(..),
+        ApiAnns(..),AnnKeywordId(..),AnnotationComment(..),
         getAnnotation, getAndRemoveAnnotation,
         getAnnotationComments, getAndRemoveAnnotationComments,
         unicodeAnn,
diff --git a/compiler/main/HscMain.hs b/compiler/main/HscMain.hs
index 4ca05e9657..879d8a05ec 100644
--- a/compiler/main/HscMain.hs
+++ b/compiler/main/HscMain.hs
@@ -102,6 +102,7 @@ import Panic
 import ConLike
 import Control.Concurrent
 
+import ApiAnnotation
 import Module
 import Packages
 import RdrName
@@ -392,13 +393,16 @@ hscParse' mod_summary
             -- filter them out:
             srcs2 <- liftIO $ filterM doesFileExist srcs1
 
-            let res = HsParsedModule {
+            let api_anns = ApiAnns {
+                      apiAnnItems = M.fromListWith (++) $ annotations pst,
+                      apiAnnEofPos = eof_pos pst,
+                      apiAnnComments = M.fromList (annotations_comments pst),
+                      apiAnnRogueComments = comment_q pst
+                   }
+                res = HsParsedModule {
                       hpm_module    = rdr_module,
                       hpm_src_files = srcs2,
-                      hpm_annotations
-                              = (M.fromListWith (++) $ annotations pst,
-                                 M.fromList $ ((noSrcSpan,comment_q pst)
-                                                 :(annotations_comments pst)))
+                      hpm_annotations = api_anns
                    }
 
             -- apply parse transformation of plugins
diff --git a/compiler/parser/ApiAnnotation.hs b/compiler/parser/ApiAnnotation.hs
index ca88716f34..52905902b6 100644
--- a/compiler/parser/ApiAnnotation.hs
+++ b/compiler/parser/ApiAnnotation.hs
@@ -3,7 +3,7 @@
 module ApiAnnotation (
   getAnnotation, getAndRemoveAnnotation,
   getAnnotationComments,getAndRemoveAnnotationComments,
-  ApiAnns,
+  ApiAnns(..),
   ApiAnnKey,
   AnnKeywordId(..),
   AnnotationComment(..),
@@ -41,8 +41,13 @@ pm_annotations field of the ParsedModule type.
 
 The full ApiAnns type is
 
-> type ApiAnns = ( Map.Map ApiAnnKey [SrcSpan]                  -- non-comments
->                , Map.Map SrcSpan [Located AnnotationComment]) -- comments
+> data ApiAnns =
+>  ApiAnns
+>    { apiAnnItems :: Map.Map ApiAnnKey [RealSrcSpan],
+>      apiAnnEofPos :: Maybe RealSrcSpan,
+>      apiAnnComments :: Map.Map RealSrcSpan [RealLocated AnnotationComment],
+>      apiAnnRogueComments :: [RealLocated AnnotationComment]
+>    }
 
 NON-COMMENT ELEMENTS
 
@@ -52,13 +57,13 @@ can show up multiple times before the next AST element), each of which
 needs to be associated with its location in the original source code.
 
 Consequently, the structure that records non-comment elements is logically
-a two level map, from the SrcSpan of the AST element containing it, to
+a two level map, from the RealSrcSpan of the AST element containing it, to
 a map from keywords ('AnnKeyWord') to all locations of the keyword directly
 in the AST element:
 
-> type ApiAnnKey = (SrcSpan,AnnKeywordId)
+> type ApiAnnKey = (RealSrcSpan,AnnKeywordId)
 >
-> Map.Map ApiAnnKey [SrcSpan]
+> Map.Map ApiAnnKey [RealSrcSpan]
 
 So
 
@@ -87,16 +92,16 @@ Every comment is associated with a *located* AnnotationComment.
 We associate comments with the lowest (most specific) AST element
 enclosing them:
 
-> Map.Map SrcSpan [Located AnnotationComment]
+> Map.Map RealSrcSpan [RealLocated AnnotationComment]
 
 PARSER STATE
 
 There are three fields in PState (the parser state) which play a role
 with annotations.
 
->  annotations :: [(ApiAnnKey,[SrcSpan])],
->  comment_q :: [Located AnnotationComment],
->  annotations_comments :: [(SrcSpan,[Located AnnotationComment])]
+>  annotations :: [(ApiAnnKey,[RealSrcSpan])],
+>  comment_q :: [RealLocated AnnotationComment],
+>  annotations_comments :: [(RealSrcSpan,[RealLocated AnnotationComment])]
 
 The 'annotations' and 'annotations_comments' fields are simple: they simply
 accumulate annotations that will end up in 'ApiAnns' at the end
@@ -105,21 +110,21 @@ accumulate annotations that will end up in 'ApiAnns' at the end
 The 'comment_q' field captures comments as they are seen in the token stream,
 so that when they are ready to be allocated via the parser they are
 available (at the time we lex a comment, we don't know what the enclosing
-AST node of it is, so we can't associate it with a SrcSpan in
+AST node of it is, so we can't associate it with a RealSrcSpan in
 annotations_comments).
 
 PARSER EMISSION OF ANNOTATIONS
 
 The parser interacts with the lexer using the function
 
-> addAnnotation :: SrcSpan -> AnnKeywordId -> SrcSpan -> P ()
+> addAnnotation :: RealSrcSpan -> AnnKeywordId -> RealSrcSpan -> P ()
 
-which takes the AST element SrcSpan, the annotation keyword and the
-target SrcSpan.
+which takes the AST element RealSrcSpan, the annotation keyword and the
+target RealSrcSpan.
 
 This adds the annotation to the `annotations` field of `PState` and
 transfers any comments in `comment_q` WHICH ARE ENCLOSED by
-the SrcSpan of this element to the `annotations_comments`
+the RealSrcSpan of this element to the `annotations_comments`
 field.  (Comments which are outside of this annotation are deferred
 until later. 'allocateComments' in 'Lexer' is responsible for
 making sure we only attach comments that actually fit in the 'SrcSpan'.)
@@ -131,49 +136,59 @@ https://gitlab.haskell.org/ghc/ghc/wikis/api-annotations
 -- ---------------------------------------------------------------------
 
 -- If you update this, update the Note [Api annotations] above
-type ApiAnns = ( Map.Map ApiAnnKey [SrcSpan]
-               , Map.Map SrcSpan [Located AnnotationComment])
+data ApiAnns =
+  ApiAnns
+    { apiAnnItems :: Map.Map ApiAnnKey [RealSrcSpan],
+      apiAnnEofPos :: Maybe RealSrcSpan,
+      apiAnnComments :: Map.Map RealSrcSpan [RealLocated AnnotationComment],
+      apiAnnRogueComments :: [RealLocated AnnotationComment]
+    }
 
 -- If you update this, update the Note [Api annotations] above
-type ApiAnnKey = (SrcSpan,AnnKeywordId)
+type ApiAnnKey = (RealSrcSpan,AnnKeywordId)
 
 
 -- | 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 -> SrcSpan -> AnnKeywordId -> [SrcSpan]
-getAnnotation (anns,_) span ann
-   = case Map.lookup (span,ann) anns of
-       Nothing -> []
-       Just ss -> ss
+getAnnotation :: ApiAnns -> RealSrcSpan -> AnnKeywordId -> [RealSrcSpan]
+getAnnotation anns span ann =
+  case Map.lookup ann_key ann_items of
+    Nothing -> []
+    Just ss -> ss
+  where ann_items = apiAnnItems anns
+        ann_key = (span,ann)
 
 -- | Retrieve a list of annotation 'SrcSpan's based on the 'SrcSpan'
 -- of the annotated AST element, and the known type of the annotation.
 -- The list is removed from the annotations.
-getAndRemoveAnnotation :: ApiAnns -> SrcSpan -> AnnKeywordId
-                       -> ([SrcSpan],ApiAnns)
-getAndRemoveAnnotation (anns,cs) span ann
-   = case Map.lookup (span,ann) anns of
-       Nothing -> ([],(anns,cs))
-       Just ss -> (ss,(Map.delete (span,ann) anns,cs))
+getAndRemoveAnnotation :: ApiAnns -> RealSrcSpan -> AnnKeywordId
+                       -> ([RealSrcSpan],ApiAnns)
+getAndRemoveAnnotation anns span ann =
+  case Map.lookup ann_key ann_items of
+    Nothing -> ([],anns)
+    Just ss -> (ss,anns{ apiAnnItems = Map.delete ann_key ann_items })
+  where ann_items = apiAnnItems anns
+        ann_key = (span,ann)
 
 -- |Retrieve the comments allocated to the current 'SrcSpan'
 --
 --  Note: A given 'SrcSpan' may appear in multiple AST elements,
 --  beware of duplicates
-getAnnotationComments :: ApiAnns -> SrcSpan -> [Located AnnotationComment]
-getAnnotationComments (_,anns) span =
-  case Map.lookup span anns of
+getAnnotationComments :: ApiAnns -> RealSrcSpan -> [RealLocated AnnotationComment]
+getAnnotationComments anns span =
+  case Map.lookup span (apiAnnComments anns) of
     Just cs -> cs
     Nothing -> []
 
 -- |Retrieve the comments allocated to the current 'SrcSpan', and
 -- remove them from the annotations
-getAndRemoveAnnotationComments :: ApiAnns -> SrcSpan
-                               -> ([Located AnnotationComment],ApiAnns)
-getAndRemoveAnnotationComments (anns,canns) span =
-  case Map.lookup span canns of
-    Just cs -> (cs,(anns,Map.delete span canns))
-    Nothing -> ([],(anns,canns))
+getAndRemoveAnnotationComments :: ApiAnns -> RealSrcSpan
+                               -> ([RealLocated AnnotationComment],ApiAnns)
+getAndRemoveAnnotationComments anns span =
+  case Map.lookup span ann_comments of
+    Just cs -> (cs, anns{ apiAnnComments = Map.delete span ann_comments })
+    Nothing -> ([], anns)
+  where ann_comments = apiAnnComments anns
 
 -- --------------------------------------------------------------------
 
@@ -296,7 +311,6 @@ data AnnKeywordId
     | AnnLarrowtailU -- ^ '-<<', unicode variant
     | AnnRarrowtail -- ^ '>>-'
     | AnnRarrowtailU -- ^ '>>-', unicode variant
-    | AnnEofPos
     deriving (Eq, Ord, Data, Show)
 
 instance Outputable AnnKeywordId where
diff --git a/compiler/parser/Lexer.x b/compiler/parser/Lexer.x
index bee441362f..7b280086ad 100644
--- a/compiler/parser/Lexer.x
+++ b/compiler/parser/Lexer.x
@@ -2122,9 +2122,10 @@ data PState = PState {
         -- locations of 'noise' tokens in the source, so that users of
         -- the GHC API can do source to source conversions.
         -- See note [Api annotations] in ApiAnnotation.hs
-        annotations :: [(ApiAnnKey,[SrcSpan])],
-        comment_q :: [Located AnnotationComment],
-        annotations_comments :: [(SrcSpan,[Located AnnotationComment])]
+        annotations :: [(ApiAnnKey,[RealSrcSpan])],
+        eof_pos :: Maybe RealSrcSpan,
+        comment_q :: [RealLocated AnnotationComment],
+        annotations_comments :: [(RealSrcSpan,[RealLocated AnnotationComment])]
      }
         -- last_loc and last_len are used when generating error messages,
         -- and in pushCurrentContext only.  Sigh, if only Happy passed the
@@ -2196,6 +2197,9 @@ getRealSrcLoc = P $ \s@(PState{ loc=loc }) -> POk s loc
 addSrcFile :: FastString -> P ()
 addSrcFile f = P $ \s -> POk s{ srcfiles = f : srcfiles s } ()
 
+setEofPos :: RealSrcSpan -> P ()
+setEofPos span = P $ \s -> POk s{ eof_pos = Just span } ()
+
 setLastToken :: RealSrcSpan -> Int -> P ()
 setLastToken loc len = P $ \s -> POk s {
   last_loc=loc,
@@ -2591,6 +2595,7 @@ mkPStatePure options buf loc =
       alr_expecting_ocurly = Nothing,
       alr_justClosedExplicitLetBlock = False,
       annotations = [],
+      eof_pos = Nothing,
       comment_q = [],
       annotations_comments = []
     }
@@ -2670,9 +2675,10 @@ instance MonadP P where
     addError span msg >> P PFailed
   getBit ext = P $ \s -> let b =  ext `xtest` pExtsBitmap (options s)
                          in b `seq` POk s b
-  addAnnotation l a v = do
+  addAnnotation (RealSrcSpan l) a (RealSrcSpan v) = do
     addAnnotationOnly l a v
     allocateCommentsP l
+  addAnnotation _ _ _ = return ()
 
 addAnnsAt :: MonadP m => SrcSpan -> [AddAnn] -> m ()
 addAnnsAt l = mapM_ (\(AddAnn a v) -> addAnnotation l a v)
@@ -2809,16 +2815,12 @@ lexer queueComments cont = do
   (L span tok) <- lexTokenFun
   --trace ("token: " ++ show tok) $ do
 
-  case tok of
-    ITeof -> addAnnotationOnly noSrcSpan AnnEofPos (RealSrcSpan span)
-    _ -> return ()
-
   if (queueComments && isDocComment tok)
-    then queueComment (L (RealSrcSpan span) tok)
+    then queueComment (L span tok)
     else return ()
 
   if (queueComments && isComment tok)
-    then queueComment (L (RealSrcSpan span) tok) >> lexer queueComments cont
+    then queueComment (L span tok) >> lexer queueComments cont
     else cont (L (RealSrcSpan span) tok)
 
 -- Use this instead of 'lexer' in Parser.y to dump the tokens for debugging.
@@ -3055,6 +3057,7 @@ lexToken = do
   case alexScanUser exts inp sc of
     AlexEOF -> do
         let span = mkRealSrcSpan loc1 loc1
+        setEofPos span
         setLastToken span 0
         return (L span ITeof)
     AlexError (AI loc2 buf) ->
@@ -3203,7 +3206,7 @@ clean_pragma prag = canon_ws (map toLower (unprefix prag))
 --   function, and then it can be discharged using the 'ams' function.
 data AddAnn = AddAnn AnnKeywordId SrcSpan
 
-addAnnotationOnly :: SrcSpan -> AnnKeywordId -> SrcSpan -> P ()
+addAnnotationOnly :: RealSrcSpan -> AnnKeywordId -> RealSrcSpan -> P ()
 addAnnotationOnly l a v = P $ \s -> POk s {
   annotations = ((l,a), [v]) : annotations s
   } ()
@@ -3213,24 +3216,24 @@ addAnnotationOnly l a v = P $ \s -> POk s {
 -- and end of the span
 mkParensApiAnn :: SrcSpan -> [AddAnn]
 mkParensApiAnn (UnhelpfulSpan _)  = []
-mkParensApiAnn s@(RealSrcSpan ss) = [AddAnn AnnOpenP lo,AddAnn AnnCloseP lc]
+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 = mkSrcSpan (srcSpanStart s)         (mkSrcLoc f sl (sc+1))
-    lc = mkSrcSpan (mkSrcLoc f el (ec - 1)) (srcSpanEnd s)
+    lo = RealSrcSpan (mkRealSrcSpan (realSrcSpanStart ss)        (mkRealSrcLoc f sl (sc+1)))
+    lc = RealSrcSpan (mkRealSrcSpan (mkRealSrcLoc f el (ec - 1)) (realSrcSpanEnd ss))
 
-queueComment :: Located Token -> P()
+queueComment :: RealLocated Token -> P()
 queueComment c = P $ \s -> POk s {
   comment_q = commentToAnnotation c : comment_q s
   } ()
 
 -- | Go through the @comment_q@ in @PState@ and remove all comments
 -- that belong within the given span
-allocateCommentsP :: SrcSpan -> P ()
+allocateCommentsP :: RealSrcSpan -> P ()
 allocateCommentsP ss = P $ \s ->
   let (comment_q', newAnns) = allocateComments ss (comment_q s) in
     POk s {
@@ -3239,13 +3242,13 @@ allocateCommentsP ss = P $ \s ->
      } ()
 
 allocateComments
-  :: SrcSpan
-  -> [Located AnnotationComment]
-  -> ([Located AnnotationComment], [(SrcSpan,[Located AnnotationComment])])
+  :: RealSrcSpan
+  -> [RealLocated AnnotationComment]
+  -> ([RealLocated AnnotationComment], [(RealSrcSpan,[RealLocated AnnotationComment])])
 allocateComments ss comment_q =
   let
-    (before,rest)  = break (\(L l _) -> isSubspanOf l ss) comment_q
-    (middle,after) = break (\(L l _) -> not (isSubspanOf l ss)) rest
+    (before,rest)  = break (\(L l _) -> isRealSubspanOf l ss) comment_q
+    (middle,after) = break (\(L l _) -> not (isRealSubspanOf l ss)) rest
     comment_q' = before ++ after
     newAnns = if null middle then []
                              else [(ss,middle)]
@@ -3253,7 +3256,7 @@ allocateComments ss comment_q =
     (comment_q', newAnns)
 
 
-commentToAnnotation :: Located Token -> Located AnnotationComment
+commentToAnnotation :: RealLocated Token -> RealLocated AnnotationComment
 commentToAnnotation (L l (ITdocCommentNext s))  = L l (AnnDocCommentNext s)
 commentToAnnotation (L l (ITdocCommentPrev s))  = L l (AnnDocCommentPrev s)
 commentToAnnotation (L l (ITdocCommentNamed s)) = L l (AnnDocCommentNamed s)
diff --git a/compiler/parser/RdrHsSyn.hs b/compiler/parser/RdrHsSyn.hs
index ef64ce25e6..e8229a9443 100644
--- a/compiler/parser/RdrHsSyn.hs
+++ b/compiler/parser/RdrHsSyn.hs
@@ -2851,9 +2851,9 @@ data PV_Context =
 data PV_Accum =
   PV_Accum
     { pv_messages :: DynFlags -> Messages
-    , pv_annotations :: [(ApiAnnKey,[SrcSpan])]
-    , pv_comment_q :: [Located AnnotationComment]
-    , pv_annotations_comments :: [(SrcSpan,[Located AnnotationComment])]
+    , pv_annotations :: [(ApiAnnKey,[RealSrcSpan])]
+    , pv_comment_q :: [RealLocated AnnotationComment]
+    , pv_annotations_comments :: [(RealSrcSpan,[RealLocated AnnotationComment])]
     }
 
 data PV_Result a = PV_Ok PV_Accum a | PV_Failed PV_Accum
@@ -2918,7 +2918,7 @@ instance MonadP PV where
     PV $ \ctx acc ->
       let b = ext `xtest` pExtsBitmap (pv_options ctx) in
       PV_Ok acc $! b
-  addAnnotation l a v =
+  addAnnotation (RealSrcSpan l) a (RealSrcSpan v) =
     PV $ \_ acc ->
       let
         (comment_q', new_ann_comments) = allocateComments l (pv_comment_q acc)
@@ -2930,6 +2930,7 @@ instance MonadP PV where
           , pv_annotations_comments = annotations_comments' }
       in
         PV_Ok acc' ()
+  addAnnotation _ _ _ = return ()
 
 {- Note [Parser-Validator]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~