EPA: correctly capture comments between 'where' and binds

In the following

    foo = x
      where -- do stuff
            doStuff = do stuff

The "-- do stuff" comment is captured in the HsValBinds.

Closes #20297

8 files changed, 45 insertions, 32 deletions

diff --git a/compiler/GHC/Hs/Expr.hs b/compiler/GHC/Hs/Expr.hs
index d0c5dbef0c..06cc4cd946 100644
--- a/compiler/GHC/Hs/Expr.hs
+++ b/compiler/GHC/Hs/Expr.hs
@@ -1274,7 +1274,12 @@ matchGroupArity (MG { mg_alts = alts })
 hsLMatchPats :: LMatch (GhcPass id) body -> [LPat (GhcPass id)]
 hsLMatchPats (L _ (Match { m_pats = pats })) = pats
 
-type instance XCGRHSs (GhcPass _) _ = NoExtField
+-- We keep the type checker happy by providing EpAnnComments.  They
+-- can only be used if they follow a `where` keyword with no binds,
+-- but in that case the comment is attached to the following parsed
+-- item. So this can never be used in practice.
+type instance XCGRHSs (GhcPass _) _ = EpAnnComments
+
 type instance XXGRHSs (GhcPass _) _ = NoExtCon
 
 data GrhsAnn
diff --git a/compiler/GHC/Hs/Utils.hs b/compiler/GHC/Hs/Utils.hs
index ac73720456..42669f4c2c 100644
--- a/compiler/GHC/Hs/Utils.hs
+++ b/compiler/GHC/Hs/Utils.hs
@@ -191,7 +191,7 @@ unguardedGRHSs :: Anno (GRHS (GhcPass p) (LocatedA (body (GhcPass p))))
                => SrcSpan -> LocatedA (body (GhcPass p)) -> EpAnn GrhsAnn
                -> GRHSs (GhcPass p) (LocatedA (body (GhcPass p)))
 unguardedGRHSs loc rhs an
-  = GRHSs noExtField (unguardedRHS an loc rhs) emptyLocalBinds
+  = GRHSs emptyComments (unguardedRHS an loc rhs) emptyLocalBinds
 
 unguardedRHS :: Anno (GRHS (GhcPass p) (LocatedA (body (GhcPass p))))
                      ~ SrcSpan
@@ -922,7 +922,7 @@ mkMatch ctxt pats expr binds
   = noLocA (Match { m_ext   = noAnn
                   , m_ctxt  = ctxt
                   , m_pats  = map mkParPat pats
-                  , m_grhss = GRHSs noExtField (unguardedRHS noAnn noSrcSpan expr) binds })
+                  , m_grhss = GRHSs emptyComments (unguardedRHS noAnn noSrcSpan expr) binds })
 
 {-
 ************************************************************************
diff --git a/compiler/GHC/HsToCore/Expr.hs b/compiler/GHC/HsToCore/Expr.hs
index 0241f611ed..1208061735 100644
--- a/compiler/GHC/HsToCore/Expr.hs
+++ b/compiler/GHC/HsToCore/Expr.hs
@@ -389,7 +389,7 @@ dsExpr (HsMultiIf res_ty alts)
   = mkErrorExpr
 
   | otherwise
-  = do { let grhss = GRHSs noExtField alts emptyLocalBinds
+  = do { let grhss = GRHSs emptyComments  alts emptyLocalBinds
        ; rhss_nablas  <- pmcGRHSs IfAlt grhss
        ; match_result <- dsGRHSs IfAlt grhss res_ty rhss_nablas
        ; error_expr   <- mkErrorExpr
diff --git a/compiler/GHC/Parser.y b/compiler/GHC/Parser.y
index 732a03f7d5..6d0a276ab7 100644
--- a/compiler/GHC/Parser.y
+++ b/compiler/GHC/Parser.y
@@ -1807,10 +1807,12 @@ binds   ::  { Located (HsLocalBinds GhcPs) }
                                              $ HsIPBinds (EpAnn (glR $1) (AnnList (Just $ glR $2) Nothing Nothing [] []) cs) (IPBinds noExtField (reverse $ unLoc $2)))) }
 
 
-wherebinds :: { Maybe (Located (HsLocalBinds GhcPs)) }
+wherebinds :: { Maybe (Located (HsLocalBinds GhcPs, Maybe EpAnnComments )) }
                                                 -- May have implicit parameters
                                                 -- No type declarations
-        : 'where' binds                 { Just (sLL $1 $> (annBinds (mj AnnWhere $1) (unLoc $2))) }
+        : 'where' binds                 {% do { r <- acs (\cs ->
+                                                (sLL $1 $> (annBinds (mj AnnWhere $1) cs (unLoc $2))))
+                                              ; return $ Just r} }
         | {- empty -}                   { Nothing }
 
 -----------------------------------------------------------------------------
@@ -2520,12 +2522,14 @@ decl    :: { LHsDecl GhcPs }
 
 rhs     :: { Located (GRHSs GhcPs (LHsExpr GhcPs)) }
         : '=' exp wherebinds    {% runPV (unECP $2) >>= \ $2 ->
-                                  do { let loc = (comb3 $1 (reLoc $2) (adaptWhereBinds $3))
+                                  do { let L l (bs, csw) = adaptWhereBinds $3
+                                     ; let loc = (comb3 $1 (reLoc $2) (L l bs))
                                      ; acs (\cs ->
-                                       sL loc (GRHSs NoExtField (unguardedRHS (EpAnn (anc $ rs loc) (GrhsAnn Nothing (mj AnnEqual $1)) cs) loc $2)
-                                                      (unLoc $ (adaptWhereBinds $3)))) } }
-        | gdrhs wherebinds      { sL (comb2 $1 (adaptWhereBinds $>))
-                                    (GRHSs noExtField (reverse (unLoc $1)) (unLoc $ (adaptWhereBinds $2))) }
+                                       sL loc (GRHSs csw (unguardedRHS (EpAnn (anc $ rs loc) (GrhsAnn Nothing (mj AnnEqual $1)) cs) loc $2)
+                                                      bs)) } }
+        | gdrhs wherebinds      {% do { let {L l (bs, csw) = adaptWhereBinds $2}
+                                      ; acs (\cs -> sL (comb2 $1 (L l bs))
+                                                (GRHSs (cs Semi.<> csw) (reverse (unLoc $1)) bs)) }}
 
 gdrhs :: { Located [LGRHS GhcPs (LHsExpr GhcPs)] }
         : gdrhs gdrh            { sLL $1 $> ($2 : unLoc $1) }
@@ -3244,7 +3248,8 @@ alt     :: { forall b. DisambECP b => PV (LMatch GhcPs (LocatedA b)) }
 
 alt_rhs :: { forall b. DisambECP b => PV (Located (GRHSs GhcPs (LocatedA b))) }
         : ralt wherebinds           { $1 >>= \alt ->
-                                      return $ sLL alt (adaptWhereBinds $>) (GRHSs noExtField (unLoc alt) (unLoc $ adaptWhereBinds $2)) }
+                                      do { let {L l (bs, csw) = adaptWhereBinds $2}
+                                         ; acs (\cs -> sLL alt (L l bs) (GRHSs (cs Semi.<> csw) (unLoc alt) bs)) }}
 
 ralt :: { forall b. DisambECP b => PV (Located [LGRHS GhcPs (LocatedA b)]) }
         : '->' exp            { unECP $2 >>= \ $2 ->
@@ -4429,8 +4434,9 @@ addTrailingDarrowC (L (SrcSpanAnn (EpAnn lr (AnnContext _ o c) csc) l) a) lt cs
 -- We need a location for the where binds, when computing the SrcSpan
 -- for the AST element using them.  Where there is a span, we return
 -- it, else noLoc, which is ignored in the comb2 call.
-adaptWhereBinds :: Maybe (Located (HsLocalBinds GhcPs)) -> Located (HsLocalBinds GhcPs)
-adaptWhereBinds Nothing = noLoc (EmptyLocalBinds noExtField)
-adaptWhereBinds (Just b) = b
+adaptWhereBinds :: Maybe (Located (HsLocalBinds GhcPs, Maybe EpAnnComments))
+                ->        Located (HsLocalBinds GhcPs,       EpAnnComments)
+adaptWhereBinds Nothing = noLoc (EmptyLocalBinds noExtField, emptyComments)
+adaptWhereBinds (Just (L l (b, mc))) = L l (b, maybe emptyComments id mc)
 
 }
diff --git a/compiler/GHC/Parser/PostProcess.hs b/compiler/GHC/Parser/PostProcess.hs
index 957e0f28a5..688464dd9d 100644
--- a/compiler/GHC/Parser/PostProcess.hs
+++ b/compiler/GHC/Parser/PostProcess.hs
@@ -437,21 +437,23 @@ fromSpecTyVarBndr bndr = case bndr of
                                      PsErrInferredTypeVarNotAllowed
 
 -- | Add the annotation for a 'where' keyword to existing @HsLocalBinds@
-annBinds :: AddEpAnn -> HsLocalBinds GhcPs -> HsLocalBinds GhcPs
-annBinds a (HsValBinds an bs)  = (HsValBinds (add_where a an) bs)
-annBinds a (HsIPBinds an bs)   = (HsIPBinds (add_where a an) bs)
-annBinds _ (EmptyLocalBinds x) = (EmptyLocalBinds x)
-
-add_where :: AddEpAnn -> EpAnn AnnList -> EpAnn AnnList
-add_where an@(AddEpAnn _ (EpaSpan rs)) (EpAnn a (AnnList anc o c r t) cs)
+annBinds :: AddEpAnn -> EpAnnComments -> HsLocalBinds GhcPs
+  -> (HsLocalBinds GhcPs, Maybe EpAnnComments)
+annBinds a cs (HsValBinds an bs)  = (HsValBinds (add_where a an cs) bs, Nothing)
+annBinds a cs (HsIPBinds an bs)   = (HsIPBinds (add_where a an cs) bs, Nothing)
+annBinds _ cs  (EmptyLocalBinds x) = (EmptyLocalBinds x, Just cs)
+
+add_where :: AddEpAnn -> EpAnn AnnList -> EpAnnComments -> EpAnn AnnList
+add_where an@(AddEpAnn _ (EpaSpan rs)) (EpAnn a (AnnList anc o c r t) cs) cs2
   | valid_anchor (anchor a)
-  = EpAnn (widenAnchor a [an]) (AnnList anc o c (an:r) t) cs
+  = EpAnn (widenAnchor a [an]) (AnnList anc o c (an:r) t) (cs Semi.<> cs2)
   | otherwise
-  = EpAnn (patch_anchor rs a) (AnnList (fmap (patch_anchor rs) anc) o c (an:r) t) cs
-add_where an@(AddEpAnn _ (EpaSpan rs)) EpAnnNotUsed
+  = EpAnn (patch_anchor rs a)
+          (AnnList (fmap (patch_anchor rs) anc) o c (an:r) t) (cs Semi.<> cs2)
+add_where an@(AddEpAnn _ (EpaSpan rs)) EpAnnNotUsed cs
   = EpAnn (Anchor rs UnchangedAnchor)
-           (AnnList (Just $ Anchor rs UnchangedAnchor) Nothing Nothing [an] []) emptyComments
-add_where (AddEpAnn _ (EpaDelta _)) _ = panic "add_where"
+           (AnnList (Just $ Anchor rs UnchangedAnchor) Nothing Nothing [an] []) cs
+add_where (AddEpAnn _ (EpaDelta _)) _ _ = panic "add_where"
  -- EpaDelta should only be used for transformations
 
 valid_anchor :: RealSrcSpan -> Bool
diff --git a/compiler/GHC/Rename/Bind.hs b/compiler/GHC/Rename/Bind.hs
index dd52c85351..a4c3ab9865 100644
--- a/compiler/GHC/Rename/Bind.hs
+++ b/compiler/GHC/Rename/Bind.hs
@@ -1244,7 +1244,7 @@ rnGRHSs :: AnnoBody body
 rnGRHSs ctxt rnBody (GRHSs _ grhss binds)
   = rnLocalBindsAndThen binds   $ \ binds' _ -> do
     (grhss', fvGRHSs) <- mapFvRn (rnGRHS ctxt rnBody) grhss
-    return (GRHSs noExtField grhss' binds', fvGRHSs)
+    return (GRHSs emptyComments grhss' binds', fvGRHSs)
 
 rnGRHS :: AnnoBody body
        => HsMatchContext GhcRn
diff --git a/compiler/GHC/Tc/Gen/Match.hs b/compiler/GHC/Tc/Gen/Match.hs
index 9b8b68aad6..d1bca5d663 100644
--- a/compiler/GHC/Tc/Gen/Match.hs
+++ b/compiler/GHC/Tc/Gen/Match.hs
@@ -277,7 +277,7 @@ tcGRHSs ctxt (GRHSs _ grhss binds) res_ty
                mapM (tcCollectingUsage . wrapLocM (tcGRHS ctxt res_ty)) grhss
         ; let (usages, grhss') = unzip ugrhss
         ; tcEmitBindingUsage $ supUEs usages
-        ; return (GRHSs noExtField grhss' binds') }
+        ; return (GRHSs emptyComments grhss' binds') }
 
 -------------
 tcGRHS :: TcMatchCtxt body -> ExpRhoType -> GRHS GhcRn (LocatedA (body GhcRn))
diff --git a/compiler/GHC/ThToHs.hs b/compiler/GHC/ThToHs.hs
index 96a27a528c..eb92fe1240 100644
--- a/compiler/GHC/ThToHs.hs
+++ b/compiler/GHC/ThToHs.hs
@@ -190,7 +190,7 @@ cvtDec (TH.ValD pat body ds)
         ; ds' <- cvtLocalDecs (text "a where clause") ds
         ; returnJustLA $ Hs.ValD noExtField $
           PatBind { pat_lhs = pat'
-                  , pat_rhs = GRHSs noExtField body' ds'
+                  , pat_rhs = GRHSs emptyComments body' ds'
                   , pat_ext = noAnn
                   , pat_ticks = ([],[]) } }
 
@@ -911,7 +911,7 @@ cvtClause ctxt (Clause ps body wheres)
         ; let pps = map (parenthesizePat appPrec) ps'
         ; g'  <- cvtGuard body
         ; ds' <- cvtLocalDecs (text "a where clause") wheres
-        ; returnLA $ Hs.Match noAnn ctxt pps (GRHSs noExtField g' ds') }
+        ; returnLA $ Hs.Match noAnn ctxt pps (GRHSs emptyComments g' ds') }
 
 cvtImplicitParamBind :: String -> TH.Exp -> CvtM (LIPBind GhcPs)
 cvtImplicitParamBind n e = do
@@ -1223,7 +1223,7 @@ cvtMatch ctxt (TH.Match p body decs)
                      _                -> p'
         ; g' <- cvtGuard body
         ; decs' <- cvtLocalDecs (text "a where clause") decs
-        ; returnLA $ Hs.Match noAnn ctxt [lp] (GRHSs noExtField g' decs') }
+        ; returnLA $ Hs.Match noAnn ctxt [lp] (GRHSs emptyComments g' decs') }
 
 cvtGuard :: TH.Body -> CvtM [LGRHS GhcPs (LHsExpr GhcPs)]
 cvtGuard (GuardedB pairs) = mapM cvtpair pairs