EPA: Add comments to EpaDelta

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.

(cherry picked from commit 085f46fc37003172c963ecc219c701ec4d3af953)

4 files changed, 26 insertions, 25 deletions

diff --git a/utils/check-exact/ExactPrint.hs b/utils/check-exact/ExactPrint.hs
index 0c5dcff989..3965c1ca35 100644
--- a/utils/check-exact/ExactPrint.hs
+++ b/utils/check-exact/ExactPrint.hs
@@ -300,7 +300,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
@@ -433,7 +433,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 fd9e5eca84..e13fdf832e 100644
--- a/utils/check-exact/Main.hs
+++ b/utils/check-exact/Main.hs
@@ -441,7 +441,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 f78d9c073b..22b1adbcfc 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
@@ -955,22 +955,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)
@@ -1122,8 +1122,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
@@ -1399,7 +1399,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
@@ -1414,7 +1414,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
@@ -1431,7 +1431,7 @@ newWhereAnnotation ww = do
   let anc  = Anchor (rs newSpan) (MovedAnchor (DifferentLine 1 2))
   let anc2 = Anchor (rs newSpan) (MovedAnchor (DifferentLine 1 4))
   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 94328355ef..ad5eb9690e 100644
--- a/utils/check-exact/Utils.hs
+++ b/utils/check-exact/Utils.hs
@@ -260,11 +260,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