diff --git a/source/library/Language/Haskell/Brittany/Internal/S2_SplitModule.hs b/source/library/Language/Haskell/Brittany/Internal/S2_SplitModule.hs index bc93371..defba49 100644 --- a/source/library/Language/Haskell/Brittany/Internal/S2_SplitModule.hs +++ b/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