From 62e1219812a8d68fde9859b11ed3ebd954ba779c Mon Sep 17 00:00:00 2001
From: Alan Zimmerman <alan.zimm@gmail.com>
Date: Mon, 1 Apr 2024 19:22:55 +0100
Subject: [PATCH] Reinstate fixModuleHeaderComments.

Now zero failure, running all bar the makeDelta tests
---
 .../Haskell/GHC/ExactPrint/Parsers.hs         | 48 +++++++++++++++++--
 tests/Test.hs                                 |  4 +-
 2 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/src/Language/Haskell/GHC/ExactPrint/Parsers.hs b/src/Language/Haskell/GHC/ExactPrint/Parsers.hs
index bd0a5765..418721b9 100644
--- a/src/Language/Haskell/GHC/ExactPrint/Parsers.hs
+++ b/src/Language/Haskell/GHC/ExactPrint/Parsers.hs
@@ -260,7 +260,7 @@ parseModuleEpAnnsWithCppInternal cppOptions dflags file = do
       GHC.PFailed pst
         -> Left (GHC.GhcPsMessage <$> GHC.getPsErrorMessages pst)
       GHC.POk _ pmod
-        -> Right $ (injectedComments, dflags', fixModuleTrailingComments pmod)
+        -> Right $ (injectedComments, dflags', fixModuleComments pmod)
 
 -- | Internal function. Exposed if you want to muck with DynFlags
 -- before parsing. Or after parsing.
@@ -269,8 +269,10 @@ postParseTransform
   -> Either a (GHC.ParsedSource)
 postParseTransform parseRes = fmap mkAnns parseRes
   where
-    -- TODO:AZ perhaps inject the comments into the parsedsource here already
-    mkAnns (_cs, _, m) = fixModuleTrailingComments m
+    mkAnns (_cs, _, m) = fixModuleComments m
+
+fixModuleComments :: GHC.ParsedSource -> GHC.ParsedSource
+fixModuleComments p = fixModuleHeaderComments $ fixModuleTrailingComments p
 
 fixModuleTrailingComments :: GHC.ParsedSource -> GHC.ParsedSource
 fixModuleTrailingComments (GHC.L l p) = GHC.L l p'
@@ -293,6 +295,46 @@ fixModuleTrailingComments (GHC.L l p) = GHC.L l p'
             in cs''
           _ -> cs
 
+-- Deal with https://gitlab.haskell.org/ghc/ghc/-/issues/23984
+-- The Lexer works bottom-up, so does not have module declaration info
+-- when the first top decl processed
+fixModuleHeaderComments :: GHC.ParsedSource -> GHC.ParsedSource
+fixModuleHeaderComments (GHC.L l p) = GHC.L l p'
+  where
+    moveComments :: GHC.EpaLocation -> GHC.LHsDecl GHC.GhcPs -> GHC.EpAnnComments
+                 -> (GHC.LHsDecl GHC.GhcPs, GHC.EpAnnComments)
+    moveComments GHC.EpaDelta{} dd cs = (dd,cs)
+    moveComments (GHC.EpaSpan (GHC.RealSrcSpan r _)) (GHC.L (GHC.EpAnn anc an csd) a) cs = (dd,css)
+      where
+        -- Move any comments on the decl that occur prior to the location
+        pc = GHC.priorComments csd
+        fc = GHC.getFollowingComments csd
+        bf (GHC.L anch _) = GHC.anchor anch > r
+        (move,keep) = break bf pc
+        csd' = GHC.EpaCommentsBalanced keep fc
+
+        dd = GHC.L (GHC.EpAnn anc an csd') a
+        css = cs <> GHC.EpaComments move
+
+    (ds',an') = rebalance (GHC.hsmodDecls p, GHC.hsmodAnn $ GHC.hsmodExt p)
+    p' = p { GHC.hsmodExt = (GHC.hsmodExt p){ GHC.hsmodAnn = an' },
+             GHC.hsmodDecls = ds'
+           }
+
+    rebalance :: ([GHC.LHsDecl GHC.GhcPs], GHC.EpAnn GHC.AnnsModule)
+              -> ([GHC.LHsDecl GHC.GhcPs], GHC.EpAnn GHC.AnnsModule)
+    rebalance (ds, GHC.EpAnn a an cs) = (ds1, GHC.EpAnn a an cs')
+      where
+        (ds1,cs') = case break (\(GHC.AddEpAnn k _) -> k == GHC.AnnWhere) (GHC.am_main an) of
+                     (_, (GHC.AddEpAnn _ whereLoc:_)) ->
+                           case GHC.hsmodDecls p of
+                               (d:ds0) -> (d':ds0, cs0)
+                                   where (d',cs0) = moveComments whereLoc d cs
+                               ds0 -> (ds0,cs)
+                     _ -> (ds,cs)
+
+
+
 -- | Internal function. Initializes DynFlags value for parsing.
 --
 -- Passes "-hide-all-packages" to the GHC API to prevent parsing of
diff --git a/tests/Test.hs b/tests/Test.hs
index 50fda61e..cf1341fd 100644
--- a/tests/Test.hs
+++ b/tests/Test.hs
@@ -213,10 +213,10 @@ tt' = do
     -- mkParserTest libdir "ghc94" "Haddock.hs"
 
     -- mkParserTest libdir "ghc94" "Haddock2.hs"
-    mkParserTestBC libdir "ghc94" "Haddock2.hs"
+    -- mkParserTestBC libdir "ghc94" "Haddock2.hs"
 
     -- mkParserTest libdir "ghc98" "ExportWarnings_aux.hs"
-    -- mkParserTest libdir "ghc98" "IndentedModule2.hs"
+    mkParserTest libdir "ghc98" "IndentedModule2.hs"
 
     -- mkParserTest libdir "ghc92" "TopLevelSemis.hs"
     -- mkParserTestBC libdir "ghc92" "TopLevelSemis.hs"