Properly handle comments at the end of imports

source/library/Language/Haskell/Brittany/Internal/S2_SplitModule.hs

@@ -13,6 +13,7 @@ module Language.Haskell.Brittany.Internal.S2_SplitModule
 import           Language.Haskell.Brittany.Internal.Prelude
 
 import qualified Data.Generics                 as SYB
+import qualified Data.List.Extra
 import qualified Data.Map                      as Map
 import qualified GHC
 import           GHC                            ( AddEpAnn(AddEpAnn)
@@ -91,16 +92,42 @@ extractDeclMap modul =
 
 splitModuleDecls :: GHC.ParsedSource -> FinalList ModuleElement ExactPrint.Pos
 splitModuleDecls lmod = do
-  let L moduleSpan modl          = lmod
-      HsModule _ _layout _name _exports _imports decls _ _ = modl
-      (hsModAnn', finalComments) = case GHC.hsmodAnn modl of
-        EpAnn a modAnns (EpaCommentsBalanced prior post) ->
-          (EpAnn a modAnns (EpaCommentsBalanced prior []), post)
-        _ -> (GHC.hsmodAnn modl, [])
-      moduleWithoutComments = L moduleSpan modl { GHC.hsmodAnn   = hsModAnn'
-                                                , GHC.hsmodDecls = []
-                                                }
-  lastSpan <- do
+  let
+    L moduleSpan modl          = lmod
+    HsModule _ _layout _name _exports imports decls _ _ = modl
+    (hsModAnn', finalComments) = case GHC.hsmodAnn modl of
+      EpAnn a modAnns (EpaCommentsBalanced prior post) ->
+        (EpAnn a modAnns (EpaCommentsBalanced prior []), post)
+      _ -> (GHC.hsmodAnn modl, [])
+    (newImports, commsAfterImports) = case Data.List.Extra.unsnoc imports of
+      Just (allButLast, L (SrcSpanAnn epAnn s@(RealSrcSpan span _)) lastImp) ->
+        case epAnn of
+          EpAnnNotUsed -> (imports, [])
+          EpAnn anch anns (EpaComments cs) ->
+            let
+              (keepImports, moveImports) =
+                partition
+                  (\(L cAnch _) ->
+                    GHC.srcSpanEndLine (anchor cAnch) <= GHC.srcSpanEndLine span
+                  )
+                  cs
+              newLastImport =
+                L (SrcSpanAnn (EpAnn anch anns (EpaComments keepImports)) s)
+                  lastImp
+            in
+              ( allButLast ++ [newLastImport]
+              , List.sortOn (\(L l _) -> l) moveImports
+              )
+          EpAnn anch anns (EpaCommentsBalanced cs1 cs2) ->
+            let newLastImport =
+                  L (SrcSpanAnn (EpAnn anch anns (EpaComments cs1)) s) lastImp
+            in  (allButLast ++ [newLastImport], List.sortOn (\(L l _) -> l) cs2)
+      _ -> ([], [])
+    moduleWithoutComments = L moduleSpan modl { GHC.hsmodAnn     = hsModAnn'
+                                              , GHC.hsmodDecls   = []
+                                              , GHC.hsmodImports = newImports
+                                              }
+  spanAfterImports <- do
     finalYield $ MEExactModuleHead moduleWithoutComments
     pure
       $  maybe (0, 1) (ExactPrint.ss2posEnd)
@@ -111,17 +138,14 @@ splitModuleDecls lmod = do
            EpAnnNotUsed -> error "unexpected EpAnnNotUsed"
          ]
       ++ [ s | AddEpAnn _ (EpaSpan s) <- GHC.am_main $ GHC.anns hsModAnn' ]
-      ++ [ GHC.anchor a
-         | L da _ <- GHC.hsmodImports modl
-         , L a  _ <- case GHC.ann da of
-           EpAnn _ _ (EpaComments l          ) -> l
-           EpAnn _ _ (EpaCommentsBalanced _ l) -> l
-           EpAnnNotUsed                        -> []
+      ++ [ span
+         | L (SrcSpanAnn _ (RealSrcSpan span _)) _ <- GHC.hsmodImports modl
          ]
       ++ [ span
          | L (SrcSpanAnn _ (GHC.RealSrcSpan span _)) _ <- GHC.hsmodImports modl
          ]
-  spanAfterDecls <- enrichDecls lastSpan decls
+  spanBeforeDecls <- enrichComms spanAfterImports commsAfterImports
+  spanAfterDecls  <- enrichDecls spanBeforeDecls decls
   enrichComms spanAfterDecls finalComments
 
 splitModuleStart