EPA: Add comments to EpaDeltawip/az/epadelta-comments

The EpaDelta variant of EpaLocation cannot be sorted by location.

So we capture any comments that need to be printed between the prior
output and this location, when creating an EpaDelta offset in ghc-exactprint.

And make the EpaLocation fields strict.

7 files changed, 51 insertions, 45 deletions

diff --git a/compiler/GHC/Hs/Dump.hs b/compiler/GHC/Hs/Dump.hs
index 247e8099da..e059cda6b9 100644
--- a/compiler/GHC/Hs/Dump.hs
+++ b/compiler/GHC/Hs/Dump.hs
@@ -139,7 +139,9 @@ showAstData bs ba a0 = blankLine $$ showAstData' a0
 
             epaAnchor :: EpaLocation -> SDoc
             epaAnchor (EpaSpan r)  = parens $ text "EpaSpan" <+> realSrcSpan r
-            epaAnchor (EpaDelta d) = parens $ text "EpaDelta" <+> deltaPos d
+            epaAnchor (EpaDelta d cs) = case ba of
+              NoBlankEpAnnotations -> parens $ text "EpaDelta" <+> deltaPos d <+> showAstData' cs
+              BlankEpAnnotations -> parens $ text "EpaDelta" <+> deltaPos d <+> text "blanked"
 
             deltaPos :: DeltaPos -> SDoc
             deltaPos (SameLine c) = parens $ text "SameLine" <+> ppr c
diff --git a/compiler/GHC/Parser/Annotation.hs b/compiler/GHC/Parser/Annotation.hs
index b414e70be5..1692d394b5 100644
--- a/compiler/GHC/Parser/Annotation.hs
+++ b/compiler/GHC/Parser/Annotation.hs
@@ -394,22 +394,25 @@ instance Outputable EpaComment where
 -- The usual way an 'AddEpAnn' is created is using the 'mj' ("make
 -- jump") function, and then it can be inserted into the appropriate
 -- annotation.
-data AddEpAnn = AddEpAnn AnnKeywordId EpaLocation deriving (Data,Show,Eq,Ord)
-
--- | The anchor for an @'AnnKeywordId'@. The Parser inserts the @'EpaSpan'@
--- variant, giving the exact location of the original item in the
--- parsed source.  This can be replaced by the @'EpaDelta'@ version, to
--- provide a position for the item relative to the end of the previous
--- item in the source.  This is useful when editing an AST prior to
--- exact printing the changed one.
-data EpaLocation = EpaSpan RealSrcSpan
-                 | EpaDelta DeltaPos
-               deriving (Data,Show,Eq,Ord)
+data AddEpAnn = AddEpAnn AnnKeywordId EpaLocation deriving (Data,Eq,Ord)
+
+-- | The anchor for an @'AnnKeywordId'@. The Parser inserts the
+-- @'EpaSpan'@ variant, giving the exact location of the original item
+-- in the parsed source.  This can be replaced by the @'EpaDelta'@
+-- version, to provide a position for the item relative to the end of
+-- the previous item in the source.  This is useful when editing an
+-- AST prior to exact printing the changed one. The list of comments
+-- in the @'EpaDelta'@ variant captures any comments between the prior
+-- output and the thing being marked here, since we cannot otherwise
+-- sort the relative order.
+data EpaLocation = EpaSpan !RealSrcSpan
+                 | EpaDelta !DeltaPos ![LEpaComment]
+               deriving (Data,Eq,Ord)
 
 -- | Tokens embedded in the AST have an EpaLocation, unless they come from
 -- generated code (e.g. by TH).
 data TokenLocation = NoTokenLoc | TokenLoc !EpaLocation
-               deriving (Data,Show,Eq,Ord)
+               deriving (Data,Eq,Ord)
 
 -- | Spacing between output items when exact printing.  It captures
 -- the spacing from the current print position on the page to the
@@ -444,7 +447,7 @@ getDeltaLine (DifferentLine r _) = r
 -- partial function is safe.
 epaLocationRealSrcSpan :: EpaLocation -> RealSrcSpan
 epaLocationRealSrcSpan (EpaSpan r) = r
-epaLocationRealSrcSpan (EpaDelta _) = panic "epaLocationRealSrcSpan"
+epaLocationRealSrcSpan (EpaDelta _ _) = panic "epaLocationRealSrcSpan"
 
 epaLocationFromSrcAnn :: SrcAnn ann -> EpaLocation
 epaLocationFromSrcAnn (SrcSpanAnn EpAnnNotUsed l) = EpaSpan (realSrcSpan l)
@@ -452,7 +455,7 @@ epaLocationFromSrcAnn (SrcSpanAnn (EpAnn anc _ _) _) = EpaSpan (anchor anc)
 
 instance Outputable EpaLocation where
   ppr (EpaSpan r) = text "EpaSpan" <+> ppr r
-  ppr (EpaDelta d) = text "EpaDelta" <+> ppr d
+  ppr (EpaDelta d cs) = text "EpaDelta" <+> ppr d <+> ppr cs
 
 instance Outputable AddEpAnn where
   ppr (AddEpAnn kw ss) = text "AddEpAnn" <+> ppr kw <+> ppr ss
@@ -631,7 +634,7 @@ data TrailingAnn
   = AddSemiAnn EpaLocation    -- ^ Trailing ';'
   | AddCommaAnn EpaLocation   -- ^ Trailing ','
   | AddVbarAnn EpaLocation    -- ^ Trailing '|'
-  deriving (Data,Show,Eq, Ord)
+  deriving (Data, Eq, Ord)
 
 instance Outputable TrailingAnn where
   ppr (AddSemiAnn ss)    = text "AddSemiAnn"    <+> ppr ss
@@ -934,7 +937,7 @@ widenSpan s as = foldl combineSrcSpans s (go as)
   where
     go [] = []
     go (AddEpAnn _ (EpaSpan s):rest) = RealSrcSpan s Strict.Nothing : go rest
-    go (AddEpAnn _ (EpaDelta _):rest) = go rest
+    go (AddEpAnn _ (EpaDelta _ _):rest) = go rest
 
 -- | The annotations need to all come after the anchor.  Make sure
 -- this is the case.
@@ -943,7 +946,7 @@ widenRealSpan s as = foldl combineRealSrcSpans s (go as)
   where
     go [] = []
     go (AddEpAnn _ (EpaSpan s):rest) = s : go rest
-    go (AddEpAnn _ (EpaDelta _):rest) =     go rest
+    go (AddEpAnn _ (EpaDelta _ _):rest) =     go rest
 
 widenAnchor :: Anchor -> [AddEpAnn] -> Anchor
 widenAnchor (Anchor s op) as = Anchor (widenRealSpan s as) op
diff --git a/compiler/GHC/Parser/PostProcess.hs b/compiler/GHC/Parser/PostProcess.hs
index 198a14ec72..ae19e7b7b3 100644
--- a/compiler/GHC/Parser/PostProcess.hs
+++ b/compiler/GHC/Parser/PostProcess.hs
@@ -456,7 +456,7 @@ add_where an@(AddEpAnn _ (EpaSpan rs)) (EpAnn a (AnnList anc o c r t) cs) cs2
 add_where an@(AddEpAnn _ (EpaSpan rs)) EpAnnNotUsed cs
   = EpAnn (Anchor rs UnchangedAnchor)
            (AnnList (Just $ Anchor rs UnchangedAnchor) Nothing Nothing [an] []) cs
-add_where (AddEpAnn _ (EpaDelta _)) _ _ = panic "add_where"
+add_where (AddEpAnn _ (EpaDelta _ _)) _ _ = panic "add_where"
  -- EpaDelta should only be used for transformations
 
 valid_anchor :: RealSrcSpan -> Bool
@@ -3011,7 +3011,7 @@ token_location_widenR NoTokenLoc _ = NoTokenLoc
 token_location_widenR tl (UnhelpfulSpan _) = tl
 token_location_widenR (TokenLoc (EpaSpan r1)) (RealSrcSpan r2 _) =
                       (TokenLoc (EpaSpan (combineRealSrcSpans r1 r2)))
-token_location_widenR (TokenLoc (EpaDelta _)) _ =
+token_location_widenR (TokenLoc (EpaDelta _ _)) _ =
   -- Never happens because the parser does not produce EpaDelta.
   panic "token_location_widenR: EpaDelta"
 
diff --git a/utils/check-exact/ExactPrint.hs b/utils/check-exact/ExactPrint.hs
index 583adc682f..9de262547e 100644
--- a/utils/check-exact/ExactPrint.hs
+++ b/utils/check-exact/ExactPrint.hs
@@ -327,7 +327,7 @@ annotationsToComments ans kws = do
     doOne :: AnnKeywordId -> EPP [Comment]
     doOne kw = do
       let sps =getSpans kw ans
-      return $ map (mkKWComment kw ) sps
+      return $ concatMap (mkKWComment kw ) sps
     -- TODO:AZ make sure these are sorted/merged properly when the invariant for
     -- allocateComments is re-established.
   newComments <- mapM doOne kws
@@ -461,7 +461,8 @@ printStringAtAnn (EpAnn _ a _) f str = printStringAtAA (f a) str
 
 printStringAtAA :: EpaLocation -> String -> EPP ()
 printStringAtAA (EpaSpan r) s = printStringAtKw' r s
-printStringAtAA (EpaDelta d) s = do
+printStringAtAA (EpaDelta d cs) s = do
+  mapM_ (printOneComment . tokComment) cs
   pe <- getPriorEndD
   p1 <- getPosP
   printStringAtLsDelta d s
diff --git a/utils/check-exact/Main.hs b/utils/check-exact/Main.hs
index 73c99345f2..71af87b760 100644
--- a/utils/check-exact/Main.hs
+++ b/utils/check-exact/Main.hs
@@ -442,7 +442,7 @@ changeLetIn1 _libdir parsed
              (L (SrcSpanAnn _ le) e) = expr
              a = (SrcSpanAnn (EpAnn (Anchor (realSrcSpan le) (MovedAnchor (SameLine 1))) mempty emptyComments) le)
              expr' = L a e
-         in (HsLet (EpAnn anc (AnnsLet l (EpaDelta (DifferentLine 1 0))) cs)
+         in (HsLet (EpAnn anc (AnnsLet l (EpaDelta (DifferentLine 1 0) [])) cs)
                 (HsValBinds x (ValBinds xv bagDecls' sigs)) expr')
 
     replace x = x
diff --git a/utils/check-exact/Transform.hs b/utils/check-exact/Transform.hs
index 867133f195..11c986b644 100644
--- a/utils/check-exact/Transform.hs
+++ b/utils/check-exact/Transform.hs
@@ -294,8 +294,8 @@ captureTypeSigSpacing (L l (SigD x (TypeSig (EpAnn anc (AnnSig dc rs') cs) ns (H
       L (SrcSpanAnn (EpAnn anc' _ _) _) _ -> anchor anc' -- TODO MovedAnchor?
     -- DP (line, col) = ss2delta (ss2pos $ anchor $ getLoc lc) r
     dc' = case dca of
-      EpaSpan r -> AddEpAnn kw (EpaDelta $ ss2delta (ss2posEnd rd) r)
-      EpaDelta _ -> AddEpAnn kw dca
+      EpaSpan r -> AddEpAnn kw (EpaDelta (ss2delta (ss2posEnd rd) r) [])
+      EpaDelta _ _ -> AddEpAnn kw dca
 
     -- ---------------------------------
 
@@ -305,7 +305,7 @@ captureTypeSigSpacing (L l (SigD x (TypeSig (EpAnn anc (AnnSig dc rs') cs) ns (H
         -> let
              op = case dca of
                EpaSpan r -> MovedAnchor (ss2delta (ss2posEnd r) (realSrcSpan ll))
-               EpaDelta _ -> MovedAnchor (SameLine 1)
+               EpaDelta _ _ -> MovedAnchor (SameLine 1)
            in (L (SrcSpanAnn (EpAnn (Anchor (realSrcSpan ll) op) mempty emptyComments) ll) b)
       (L (SrcSpanAnn (EpAnn (Anchor r op) a c) ll) b)
         -> let
@@ -313,7 +313,7 @@ captureTypeSigSpacing (L l (SigD x (TypeSig (EpAnn anc (AnnSig dc rs') cs) ns (H
                 MovedAnchor _ -> op
                 _ -> case dca of
                   EpaSpan dcr -> MovedAnchor (ss2delta (ss2posEnd dcr) r)
-                  EpaDelta _ -> MovedAnchor (SameLine 1)
+                  EpaDelta _ _ -> MovedAnchor (SameLine 1)
            in (L (SrcSpanAnn (EpAnn (Anchor r op') a c) ll) b)
 
 captureTypeSigSpacing s = s
@@ -500,13 +500,13 @@ setEntryDP _ast _dp anns = anns
 -- ---------------------------------------------------------------------
 
 addEpaLocationDelta :: LayoutStartCol -> RealSrcSpan -> EpaLocation -> EpaLocation
-addEpaLocationDelta _off _anc (EpaDelta d) = EpaDelta d
+addEpaLocationDelta _off _anc (EpaDelta d cs) = EpaDelta d cs
 addEpaLocationDelta  off  anc (EpaSpan r)
-  = EpaDelta (adjustDeltaForOffset 0 off (ss2deltaEnd anc r))
+  = EpaDelta (adjustDeltaForOffset 0 off (ss2deltaEnd anc r)) []
 
 -- Set the entry DP for an element coming after an existing keyword annotation
 setEntryDPFromAnchor :: LayoutStartCol -> EpaLocation -> LocatedA t -> LocatedA t
-setEntryDPFromAnchor _off (EpaDelta _) (L la a) = L la a
+setEntryDPFromAnchor _off (EpaDelta _ _) (L la a) = L la a
 setEntryDPFromAnchor  off (EpaSpan anc) ll@(L la _) = setEntryDP' ll dp'
   where
     r = case la of
@@ -1017,22 +1017,22 @@ noAnnSrcSpanDPn :: (Monoid ann) => SrcSpan -> Int -> SrcSpanAnn' (EpAnn ann)
 noAnnSrcSpanDPn l s = noAnnSrcSpanDP l (SameLine s)
 
 d0 :: EpaLocation
-d0 = EpaDelta $ SameLine 0
+d0 = EpaDelta (SameLine 0) []
 
 d1 :: EpaLocation
-d1 = EpaDelta $ SameLine 1
+d1 = EpaDelta (SameLine 1) []
 
 dn :: Int -> EpaLocation
-dn n = EpaDelta $ SameLine n
+dn n = EpaDelta (SameLine n) []
 
 m0 :: AnchorOperation
-m0 = MovedAnchor $ SameLine 0
+m0 = MovedAnchor (SameLine 0)
 
 m1 :: AnchorOperation
-m1 = MovedAnchor $ SameLine 1
+m1 = MovedAnchor (SameLine 1)
 
 mn :: Int -> AnchorOperation
-mn n = MovedAnchor $ SameLine n
+mn n = MovedAnchor (SameLine n)
 
 addComma :: SrcSpanAnnA -> SrcSpanAnnA
 addComma (SrcSpanAnn EpAnnNotUsed l)
@@ -1186,8 +1186,8 @@ instance HasDecls (LocatedA (HsExpr GhcPs)) where
                 let
                   off = case l of
                           (EpaSpan r) -> LayoutStartCol $ snd $ ss2pos r
-                          (EpaDelta (SameLine _)) -> LayoutStartCol 0
-                          (EpaDelta (DifferentLine _ c)) -> LayoutStartCol c
+                          (EpaDelta (SameLine _) _) -> LayoutStartCol 0
+                          (EpaDelta (DifferentLine _ c) _) -> LayoutStartCol c
                   ex'' = setEntryDPFromAnchor off i ex
                   newDecls'' = case newDecls of
                     [] -> newDecls
@@ -1463,7 +1463,7 @@ oldWhereAnnotation :: (Monad m)
 oldWhereAnnotation EpAnnNotUsed ww _oldSpan = do
   newSpan <- uniqueSrcSpanT
   let w = case ww of
-        WithWhere -> [AddEpAnn AnnWhere (EpaDelta (SameLine 0))]
+        WithWhere -> [AddEpAnn AnnWhere (EpaDelta (SameLine 0) [])]
         WithoutWhere -> []
   let anc2' = Anchor (rs newSpan) (MovedAnchor (SameLine 1))
   (anc, anc2) <- do
@@ -1478,7 +1478,7 @@ oldWhereAnnotation (EpAnn anc an cs) ww _oldSpan = do
   -- TODO: when we set DP (0,0) for the HsValBinds EpEpaLocation, change the AnnList anchor to have the correct DP too
   let (AnnList ancl o c _r t) = an
   let w = case ww of
-        WithWhere -> [AddEpAnn AnnWhere (EpaDelta (SameLine 0))]
+        WithWhere -> [AddEpAnn AnnWhere (EpaDelta (SameLine 0) [])]
         WithoutWhere -> []
   (anc', ancl') <- do
         case ww of
@@ -1495,7 +1495,7 @@ newWhereAnnotation ww = do
   let anc  = Anchor (rs newSpan) (MovedAnchor (DifferentLine 1 3))
   let anc2 = Anchor (rs newSpan) (MovedAnchor (DifferentLine 1 5))
   let w = case ww of
-        WithWhere -> [AddEpAnn AnnWhere (EpaDelta (SameLine 0))]
+        WithWhere -> [AddEpAnn AnnWhere (EpaDelta (SameLine 0) [])]
         WithoutWhere -> []
   let an = EpAnn anc
                   (AnnList (Just anc2) Nothing Nothing w [])
diff --git a/utils/check-exact/Utils.hs b/utils/check-exact/Utils.hs
index 74a861e773..a9b7640107 100644
--- a/utils/check-exact/Utils.hs
+++ b/utils/check-exact/Utils.hs
@@ -263,11 +263,11 @@ normaliseCommentText ('\r':xs) = normaliseCommentText xs
 normaliseCommentText (x:xs) = x:normaliseCommentText xs
 
 -- | Makes a comment which originates from a specific keyword.
-mkKWComment :: AnnKeywordId -> EpaLocation -> Comment
+mkKWComment :: AnnKeywordId -> EpaLocation -> [Comment]
 mkKWComment kw (EpaSpan ss)
-  = Comment (keywordToString $ G kw) (Anchor ss UnchangedAnchor) (Just kw)
-mkKWComment kw (EpaDelta dp)
-  = Comment (keywordToString $ G kw) (Anchor placeholderRealSpan (MovedAnchor dp)) (Just kw)
+  = [Comment (keywordToString $ G kw) (Anchor ss UnchangedAnchor) (Just kw)]
+mkKWComment kw (EpaDelta dp cs)
+  = (map tokComment cs) ++ [Comment (keywordToString $ G kw) (Anchor placeholderRealSpan (MovedAnchor dp)) (Just kw)]
 
 comment2dp :: (Comment,  DeltaPos) -> (KeywordId, DeltaPos)
 comment2dp = first AnnComment