Introduce Strict.Maybe, Strict.Pair (#19156)

This patch fixes a space leak related to the use of
Maybe in RealSrcSpan by introducing a strict variant
of Maybe.

In addition to that, it also introduces a strict pair
and uses the newly introduced strict data types in a few
other places (e.g. the lexer/parser state) to reduce
allocations.

Includes a regression test.

5 files changed, 54 insertions, 45 deletions

diff --git a/compiler/GHC/Parser/Annotation.hs b/compiler/GHC/Parser/Annotation.hs
index 20ac4bde62..abece56898 100644
--- a/compiler/GHC/Parser/Annotation.hs
+++ b/compiler/GHC/Parser/Annotation.hs
@@ -95,6 +95,7 @@ import GHC.Types.SrcLoc
 import GHC.Utils.Binary
 import GHC.Utils.Outputable hiding ( (<>) )
 import GHC.Utils.Panic
+import qualified GHC.Data.Strict as Strict
 
 {-
 Note [exact print annotations]
@@ -928,7 +929,7 @@ widenSpan :: SrcSpan -> [AddEpAnn] -> SrcSpan
 widenSpan s as = foldl combineSrcSpans s (go as)
   where
     go [] = []
-    go (AddEpAnn _ (EpaSpan s):rest) = RealSrcSpan s Nothing : go rest
+    go (AddEpAnn _ (EpaSpan s):rest) = RealSrcSpan s Strict.Nothing : go rest
     go (AddEpAnn _ (EpaDelta _):rest) = go rest
 
 -- | The annotations need to all come after the anchor.  Make sure
diff --git a/compiler/GHC/Parser/Header.hs b/compiler/GHC/Parser/Header.hs
index 7a60830d34..52a98d86dc 100644
--- a/compiler/GHC/Parser/Header.hs
+++ b/compiler/GHC/Parser/Header.hs
@@ -55,6 +55,7 @@ import GHC.Data.StringBuffer
 import GHC.Data.Maybe
 import GHC.Data.Bag         (Bag, isEmptyBag )
 import GHC.Data.FastString
+import qualified GHC.Data.Strict as Strict
 
 import Control.Monad
 import System.IO
@@ -347,7 +348,7 @@ toArgs starting_loc orig_str
   advance_src_loc_many = foldl' advanceSrcLoc
 
   locate :: RealSrcLoc -> RealSrcLoc -> a -> Located a
-  locate begin end x = L (RealSrcSpan (mkRealSrcSpan begin end) Nothing) x
+  locate begin end x = L (RealSrcSpan (mkRealSrcSpan begin end) Strict.Nothing) x
 
   toArgs' :: RealSrcLoc -> String -> Either String [Located String]
   -- Remove outer quotes:
diff --git a/compiler/GHC/Parser/Lexer.x b/compiler/GHC/Parser/Lexer.x
index dc13d44493..f9494afa6a 100644
--- a/compiler/GHC/Parser/Lexer.x
+++ b/compiler/GHC/Parser/Lexer.x
@@ -78,9 +78,11 @@ module GHC.Parser.Lexer (
   ) where
 
 import GHC.Prelude
+import qualified GHC.Data.Strict as Strict
 
 -- base
 import Control.Monad
+import Control.Applicative
 import Data.Char
 import Data.List (stripPrefix, isInfixOf, partition)
 import Data.Maybe
@@ -1581,7 +1583,7 @@ varid span buf len =
     Just (ITcase, _) -> do
       lastTk <- getLastTk
       keyword <- case lastTk of
-        Just (L _ ITlam) -> do
+        Strict.Just (L _ ITlam) -> do
           lambdaCase <- getBit LambdaCaseBit
           unless lambdaCase $ do
             pState <- getPState
@@ -2256,7 +2258,7 @@ warnTab srcspan _buf _len = do
 
 warnThen :: WarningFlag -> (SrcSpan -> PsWarning) -> Action -> Action
 warnThen flag warning action srcspan buf len = do
-    addWarning flag (warning (RealSrcSpan (psRealSpan srcspan) Nothing))
+    addWarning flag (warning (RealSrcSpan (psRealSpan srcspan) Strict.Nothing))
     action srcspan buf len
 
 -- -----------------------------------------------------------------------------
@@ -2324,9 +2326,10 @@ data PState = PState {
         options    :: ParserOpts,
         warnings   :: Bag PsWarning,
         errors     :: Bag PsError,
-        tab_first  :: Maybe RealSrcSpan, -- pos of first tab warning in the file
+        tab_first  :: Strict.Maybe RealSrcSpan,
+                                         -- pos of first tab warning in the file
         tab_count  :: !Word,             -- number of tab warnings in the file
-        last_tk    :: Maybe (PsLocated Token), -- last non-comment token
+        last_tk    :: Strict.Maybe (PsLocated Token), -- last non-comment token
         prev_loc   :: PsSpan,      -- pos of previous token, including comments,
         prev_loc2  :: PsSpan,      -- pos of two back token, including comments,
                                    -- see Note [PsSpan in Comments]
@@ -2359,8 +2362,8 @@ 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 [exact print annotations] in GHC.Parser.Annotation
-        eof_pos :: Maybe (RealSrcSpan, RealSrcSpan), -- pos, gap to prior token
-        header_comments :: Maybe [LEpaComment],
+        eof_pos :: Strict.Maybe (Strict.Pair RealSrcSpan RealSrcSpan), -- pos, gap to prior token
+        header_comments :: Strict.Maybe [LEpaComment],
         comment_q :: [LEpaComment],
 
         -- Haddock comments accumulated in ascending order of their location
@@ -2418,7 +2421,7 @@ failMsgP f = do
 
 failLocMsgP :: RealSrcLoc -> RealSrcLoc -> (SrcSpan -> PsError) -> P a
 failLocMsgP loc1 loc2 f =
-  addFatalError (f (RealSrcSpan (mkRealSrcSpan loc1 loc2) Nothing))
+  addFatalError (f (RealSrcSpan (mkRealSrcSpan loc1 loc2) Strict.Nothing))
 
 getPState :: P PState
 getPState = P $ \s -> POk s s
@@ -2448,7 +2451,7 @@ addSrcFile :: FastString -> P ()
 addSrcFile f = P $ \s -> POk s{ srcfiles = f : srcfiles s } ()
 
 setEofPos :: RealSrcSpan -> RealSrcSpan -> P ()
-setEofPos span gap = P $ \s -> POk s{ eof_pos = Just (span, gap) } ()
+setEofPos span gap = P $ \s -> POk s{ eof_pos = Strict.Just (span `Strict.And` gap) } ()
 
 setLastToken :: PsSpan -> Int -> P ()
 setLastToken loc len = P $ \s -> POk s {
@@ -2457,7 +2460,7 @@ setLastToken loc len = P $ \s -> POk s {
   } ()
 
 setLastTk :: PsLocated Token -> P ()
-setLastTk tk@(L l _) = P $ \s -> POk s { last_tk = Just tk
+setLastTk tk@(L l _) = P $ \s -> POk s { last_tk = Strict.Just tk
                                        , prev_loc = l
                                        , prev_loc2 = prev_loc s} ()
 
@@ -2465,7 +2468,7 @@ setLastComment :: PsLocated Token -> P ()
 setLastComment (L l _) = P $ \s -> POk s { prev_loc = l
                                          , prev_loc2 = prev_loc s} ()
 
-getLastTk :: P (Maybe (PsLocated Token))
+getLastTk :: P (Strict.Maybe (PsLocated Token))
 getLastTk = P $ \s@(PState { last_tk = last_tk }) -> POk s last_tk
 
 -- see Note [PsSpan in Comments]
@@ -2844,9 +2847,9 @@ initParserState options buf loc =
       options       = options,
       errors        = emptyBag,
       warnings      = emptyBag,
-      tab_first     = Nothing,
+      tab_first     = Strict.Nothing,
       tab_count     = 0,
-      last_tk       = Nothing,
+      last_tk       = Strict.Nothing,
       prev_loc      = mkPsSpan init_loc init_loc,
       prev_loc2     = mkPsSpan init_loc init_loc,
       last_loc      = mkPsSpan init_loc init_loc,
@@ -2861,8 +2864,8 @@ initParserState options buf loc =
       alr_context = [],
       alr_expecting_ocurly = Nothing,
       alr_justClosedExplicitLetBlock = False,
-      eof_pos = Nothing,
-      header_comments = Nothing,
+      eof_pos = Strict.Nothing,
+      header_comments = Strict.Nothing,
       comment_q = [],
       hdk_comments = nilOL
     }
@@ -2944,7 +2947,7 @@ instance MonadP P where
       POk s {
          header_comments = header_comments',
          comment_q = comment_q'
-       } (EpaCommentsBalanced (fromMaybe [] header_comments') (reverse newAnns))
+       } (EpaCommentsBalanced (Strict.fromMaybe [] header_comments') (reverse newAnns))
 
 getCommentsFor :: (MonadP m) => SrcSpan -> m EpAnnComments
 getCommentsFor (RealSrcSpan l _) = allocateCommentsP l
@@ -2958,13 +2961,13 @@ getFinalCommentsFor :: (MonadP m) => SrcSpan -> m EpAnnComments
 getFinalCommentsFor (RealSrcSpan l _) = allocateFinalCommentsP l
 getFinalCommentsFor _ = return emptyComments
 
-getEofPos :: P (Maybe (RealSrcSpan, RealSrcSpan))
+getEofPos :: P (Strict.Maybe (Strict.Pair RealSrcSpan RealSrcSpan))
 getEofPos = P $ \s@(PState { eof_pos = pos }) -> POk s pos
 
 addTabWarning :: RealSrcSpan -> P ()
 addTabWarning srcspan
  = P $ \s@PState{tab_first=tf, tab_count=tc, options=o} ->
-       let tf' = if isJust tf then tf else Just srcspan
+       let tf' = tf <|> Strict.Just srcspan
            tc' = tc + 1
            s' = if warnopt Opt_WarnTabs o
                 then s{tab_first = tf', tab_count = tc'}
@@ -2984,8 +2987,9 @@ getMessages p =
       -- we add the tabulation warning on the fly because
       -- we count the number of occurrences of tab characters
       ws' = case tab_first p of
-               Nothing -> ws
-               Just tf -> PsWarnTab (RealSrcSpan tf Nothing) (tab_count p)
+        Strict.Nothing -> ws
+        Strict.Just tf ->
+          PsWarnTab (RealSrcSpan tf Strict.Nothing) (tab_count p)
                            `consBag` ws
   in (ws', errors p)
 
@@ -3482,8 +3486,8 @@ allocateComments ss comment_q =
 allocatePriorComments
   :: RealSrcSpan
   -> [LEpaComment]
-  -> Maybe [LEpaComment]
-  -> (Maybe [LEpaComment], [LEpaComment], [LEpaComment])
+  -> Strict.Maybe [LEpaComment]
+  -> (Strict.Maybe [LEpaComment], [LEpaComment], [LEpaComment])
 allocatePriorComments ss comment_q mheader_comments =
   let
     cmp (L l _) = anchor l <= ss
@@ -3492,14 +3496,14 @@ allocatePriorComments ss comment_q mheader_comments =
     comment_q'= after
   in
     case mheader_comments of
-      Nothing -> (Just newAnns, comment_q', [])
-      Just _ -> (mheader_comments, comment_q', newAnns)
+      Strict.Nothing -> (Strict.Just newAnns, comment_q', [])
+      Strict.Just _ -> (mheader_comments, comment_q', newAnns)
 
 allocateFinalComments
   :: RealSrcSpan
   -> [LEpaComment]
-  -> Maybe [LEpaComment]
-  -> (Maybe [LEpaComment], [LEpaComment], [LEpaComment])
+  -> Strict.Maybe [LEpaComment]
+  -> (Strict.Maybe [LEpaComment], [LEpaComment], [LEpaComment])
 allocateFinalComments ss comment_q mheader_comments =
   let
     cmp (L l _) = anchor l <= ss
@@ -3508,8 +3512,8 @@ allocateFinalComments ss comment_q mheader_comments =
     comment_q'= before
   in
     case mheader_comments of
-      Nothing -> (Just newAnns,    [], comment_q')
-      Just _ -> (mheader_comments, [], comment_q' ++ newAnns)
+      Strict.Nothing -> (Strict.Just newAnns, [], comment_q')
+      Strict.Just _ -> (mheader_comments, [], comment_q' ++ newAnns)
 
 commentToAnnotation :: RealLocated Token -> LEpaComment
 commentToAnnotation (L l (ITdocCommentNext s ll))  = mkLEpaComment l ll (EpaDocCommentNext s)
diff --git a/compiler/GHC/Parser/PostProcess.hs b/compiler/GHC/Parser/PostProcess.hs
index d6b36b9d51..8cd5105e42 100644
--- a/compiler/GHC/Parser/PostProcess.hs
+++ b/compiler/GHC/Parser/PostProcess.hs
@@ -147,6 +147,7 @@ import qualified Data.Semigroup as Semi
 import GHC.Utils.Panic
 import GHC.Utils.Panic.Plain
 import qualified GHC.LanguageExtensions as LangExt
+import qualified GHC.Data.Strict as Strict
 
 import Control.Monad
 import Text.ParserCombinators.ReadP as ReadP
@@ -963,13 +964,13 @@ checkTyClHdr is_cls ty
         lr = combineRealSrcSpans (realSrcSpan l) (anchor as)
         -- lr = widenAnchorR as (realSrcSpan l)
         an = (EpAnn (Anchor lr UnchangedAnchor) (NameAnn NameParens o (EpaSpan $ realSrcSpan l) c []) cs)
-      in SrcSpanAnn an (RealSrcSpan lr Nothing)
+      in SrcSpanAnn an (RealSrcSpan lr Strict.Nothing)
     newAnns _ EpAnnNotUsed = panic "missing AnnParen"
     newAnns (SrcSpanAnn (EpAnn ap (AnnListItem ta) csp) l) (EpAnn as (AnnParen _ o c) cs) =
       let
         lr = combineRealSrcSpans (anchor ap) (anchor as)
         an = (EpAnn (Anchor lr UnchangedAnchor) (NameAnn NameParens o (EpaSpan $ realSrcSpan l) c ta) (csp Semi.<> cs))
-      in SrcSpanAnn an (RealSrcSpan lr Nothing)
+      in SrcSpanAnn an (RealSrcSpan lr Strict.Nothing)
 
 -- | Yield a parse error if we have a function applied directly to a do block
 -- etc. and BlockArguments is not enabled.
@@ -1053,18 +1054,18 @@ checkImportDecl mPre mPost = do
   -- 'ImportQualifiedPost' is not in effect.
   whenJust mPost $ \post ->
     when (not importQualifiedPostEnabled) $
-      failOpNotEnabledImportQualifiedPost (RealSrcSpan (epaLocationRealSrcSpan post) Nothing)
+      failOpNotEnabledImportQualifiedPost (RealSrcSpan (epaLocationRealSrcSpan post) Strict.Nothing)
 
   -- Error if 'qualified' occurs in both pre and postpositive
   -- positions.
   whenJust mPost $ \post ->
     when (isJust mPre) $
-      failOpImportQualifiedTwice (RealSrcSpan (epaLocationRealSrcSpan post) Nothing)
+      failOpImportQualifiedTwice (RealSrcSpan (epaLocationRealSrcSpan post) Strict.Nothing)
 
   -- Warn if 'qualified' found in prepositive position and
   -- 'Opt_WarnPrepositiveQualifiedModule' is enabled.
   whenJust mPre $ \pre ->
-    warnPrepositiveQualifiedModule (RealSrcSpan (epaLocationRealSrcSpan pre) Nothing)
+    warnPrepositiveQualifiedModule (RealSrcSpan (epaLocationRealSrcSpan pre) Strict.Nothing)
 
 -- -------------------------------------------------------------------------
 -- Checking Patterns.
@@ -2727,7 +2728,7 @@ data PV_Accum =
   PV_Accum
     { pv_warnings        :: Bag PsWarning
     , pv_errors          :: Bag PsError
-    , pv_header_comments :: Maybe [LEpaComment]
+    , pv_header_comments :: Strict.Maybe [LEpaComment]
     , pv_comment_q       :: [LEpaComment]
     }
 
@@ -2831,7 +2832,7 @@ instance MonadP PV where
       PV_Ok s {
          pv_header_comments = header_comments',
          pv_comment_q = comment_q'
-       } (EpaCommentsBalanced (fromMaybe [] header_comments') (reverse newAnns))
+       } (EpaCommentsBalanced (Strict.fromMaybe [] header_comments') (reverse newAnns))
 
 {- Note [Parser-Validator Hint]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/compiler/GHC/Parser/PostProcess/Haddock.hs b/compiler/GHC/Parser/PostProcess/Haddock.hs
index 189ddce29c..301e902f8b 100644
--- a/compiler/GHC/Parser/PostProcess/Haddock.hs
+++ b/compiler/GHC/Parser/PostProcess/Haddock.hs
@@ -73,6 +73,7 @@ import qualified Data.Monoid
 import GHC.Parser.Lexer
 import GHC.Parser.Errors
 import GHC.Utils.Misc (mergeListsBy, filterOut, mapLastM, (<&&>))
+import qualified GHC.Data.Strict as Strict
 
 {- Note [Adding Haddock comments to the syntax tree]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -1023,7 +1024,8 @@ instance HasHaddock (LocatedA (HsType GhcPs)) where
 --  which it is used.
 data HdkA a =
   HdkA
-    !(Maybe BufSpan) -- Just b  <=> BufSpan occupied by the processed AST element.
+    !(Strict.Maybe BufSpan)
+                     -- Just b  <=> BufSpan occupied by the processed AST element.
                      --             The surrounding computations will not look inside.
                      --
                      -- Nothing <=> No BufSpan (e.g. when the HdkA is constructed by 'pure' or 'liftHdkA').
@@ -1056,9 +1058,9 @@ instance Applicative HdkA where
                                 -- These delim1/delim2 are key to how HdkA operates.
     where
       -- Delimit the LHS by the location information from the RHS
-      delim1 = inLocRange (locRangeTo (fmap @Maybe bufSpanStart l2))
+      delim1 = inLocRange (locRangeTo (fmap @Strict.Maybe bufSpanStart l2))
       -- Delimit the RHS by the location information from the LHS
-      delim2 = inLocRange (locRangeFrom (fmap @Maybe bufSpanEnd l1))
+      delim2 = inLocRange (locRangeFrom (fmap @Strict.Maybe bufSpanEnd l1))
 
   pure a =
     -- Return a value without performing any stateful computation, and without
@@ -1377,14 +1379,14 @@ instance Monoid LocRange where
   mempty = LocRange mempty mempty mempty
 
 -- The location range from the specified position to the end of the file.
-locRangeFrom :: Maybe BufPos -> LocRange
-locRangeFrom (Just l) = mempty { loc_range_from = StartLoc l }
-locRangeFrom Nothing = mempty
+locRangeFrom :: Strict.Maybe BufPos -> LocRange
+locRangeFrom (Strict.Just l) = mempty { loc_range_from = StartLoc l }
+locRangeFrom Strict.Nothing = mempty
 
 -- The location range from the start of the file to the specified position.
-locRangeTo :: Maybe BufPos -> LocRange
-locRangeTo (Just l) = mempty { loc_range_to = EndLoc l }
-locRangeTo Nothing = mempty
+locRangeTo :: Strict.Maybe BufPos -> LocRange
+locRangeTo (Strict.Just l) = mempty { loc_range_to = EndLoc l }
+locRangeTo Strict.Nothing = mempty
 
 -- Represents a predicate on BufPos:
 --