{-# LANGUAGE NoImplicitPrelude #-} {-# OPTIONS_GHC -fno-warn-unused-imports #-} module Language.Haskell.Brittany.Internal.SplitExactModule ( extractDeclMap , splitModuleDecls , splitModuleStart , getDeclBindingNames ) where 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 Data.Text as Text import qualified GHC import GHC ( AddEpAnn(AddEpAnn) , Anchor(Anchor) , EpAnn(EpAnn, EpAnnNotUsed) , EpAnnComments ( EpaComments , EpaCommentsBalanced ) , EpaComment(EpaComment) , EpaCommentTok ( EpaBlockComment , EpaDocCommentNamed , EpaDocCommentNext , EpaDocCommentPrev , EpaDocOptions , EpaDocSection , EpaEofComment , EpaLineComment ) , EpaLocation(EpaSpan) , GenLocated(L) , HsModule(HsModule) , LEpaComment , LHsDecl , LImportDecl , SrcSpan ( RealSrcSpan , UnhelpfulSpan ) , SrcSpanAnn'(SrcSpanAnn) , anchor , getLoc , ideclName , moduleName , moduleNameString , srcLocCol , srcLocLine , unLoc ) import GHC.Types.Name ( getOccString ) import GHC.Types.Name.Occurrence ( occNameString ) import GHC.Types.Name.Reader ( RdrName ( Exact , Orig , Qual , Unqual ) ) import GHC.Types.SourceText ( SourceText(NoSourceText) ) import qualified GHC.OldList as List import GHC.Parser.Annotation ( DeltaPos ( DifferentLine , SameLine ) , EpaCommentTok(EpaEofComment) ) import GHC.Types.SrcLoc ( realSrcSpanEnd , realSrcSpanStart ) import qualified Language.Haskell.GHC.ExactPrint.Types as ExactPrint import qualified Language.Haskell.GHC.ExactPrint.Utils as ExactPrint import Safe ( maximumMay ) import qualified Control.Monad.Trans.Writer.Strict as WriterS import Language.Haskell.Brittany.Internal.Components.BriDoc import Language.Haskell.Brittany.Internal.ToBriDoc.Module import Language.Haskell.Brittany.Internal.Types extractDeclMap :: GHC.ParsedSource -> Map GHC.RealSrcSpan [String] extractDeclMap modul = Map.fromList [ ( case span of GHC.RealSrcSpan s _ -> s GHC.UnhelpfulSpan _ -> error "unexpected UnhelpfulSpan" , getDeclBindingNames decl ) | decl <- decls , let (L (GHC.SrcSpanAnn _ span) _) = decl ] where (L _ (GHC.HsModule _ _ _ _ _ decls _ _)) = 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, []) (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 $ maximumMay $ [ GHC.anchor a | L a _ <- GHC.priorComments $ case hsModAnn' of EpAnn _ _ cs -> cs EpAnnNotUsed -> error "unexpected EpAnnNotUsed" ] ++ [ s | AddEpAnn _ (EpaSpan s) <- GHC.am_main $ GHC.anns hsModAnn' ] ++ [ span | L (SrcSpanAnn _ (RealSrcSpan span _)) _ <- GHC.hsmodImports modl ] ++ [ span | L (SrcSpanAnn _ (GHC.RealSrcSpan span _)) _ <- GHC.hsmodImports modl ] spanBeforeDecls <- enrichComms spanAfterImports commsAfterImports spanAfterDecls <- enrichDecls spanBeforeDecls decls enrichComms spanAfterDecls finalComments splitModuleStart :: GHC.ParsedSource -> Maybe GHC.RealSrcLoc -> FinalList ModuleElement ExactPrint.Pos splitModuleStart modul posWhere = do finalYield $ MEPrettyModuleHead modul let locBeforeImports = maximumMay $ [ realSrcSpanEnd $ anchor a | L a _ <- case GHC.hsmodAnn $ unLoc modul of EpAnn _ _ (EpaComments cs ) -> cs EpAnn _ _ (EpaCommentsBalanced cs1 cs2) -> cs1 ++ cs2 EpAnnNotUsed -> error "unexpected EpAnnNotUsed" ] ++ [ pos | Just pos <- [posWhere] ] let (importLines, lastSpan) = finalToList $ transformToImportLine ( maybe 0 srcLocLine locBeforeImports , maybe 1 srcLocCol locBeforeImports ) (GHC.hsmodImports $ unLoc modul) let commentedImports = groupifyImportLines importLines sortCommentedImports commentedImports `forM_` \case EmptyLines n -> finalYield $ MEWhitespace $ DifferentLine n 1 SamelineComment{} -> error "brittany internal error: splitModuleStart SamelineComment" NewlineComment comm -> finalYield $ MEComment comm ImportStatement record -> do forM_ (commentsBefore record) $ finalYield . MEComment finalYield $ MEImportDecl (importStatement record) (commentsSameline record) forM_ (commentsAfter record) $ finalYield . MEComment pure $ lastSpan enrichComms :: ExactPrint.Pos -> [LEpaComment] -> FinalList ModuleElement ExactPrint.Pos enrichComms lastSpanEnd = \case [] -> pure lastSpanEnd (L (Anchor span _) (EpaComment EpaEofComment _) : commRest) -> do finalYield $ MEWhitespace $ case ExactPrint.ss2delta lastSpanEnd span of -- TODO92 move this (l-1) bit into utility function SameLine i -> SameLine i DifferentLine l c -> DifferentLine (l - 1) c enrichComms (ExactPrint.ss2posEnd span) commRest lcomm@(L (Anchor span _) _) : commRest -> do case ExactPrint.ss2delta lastSpanEnd span of SameLine i -> do finalYield $ MEComment (i, lcomm) DifferentLine l c -> do finalYield $ MEWhitespace $ DifferentLine (l - 1) c finalYield $ MEComment (0, lcomm) enrichComms (ExactPrint.ss2posEnd span) commRest enrichDecls :: ExactPrint.Pos -> [LHsDecl GhcPs] -> FinalList ModuleElement ExactPrint.Pos enrichDecls lastSpanEnd = \case [] -> finalPure $ lastSpanEnd ldecl@(L (SrcSpanAnn dAnn (GHC.RealSrcSpan span _)) _) : declRest -> case dAnn of EpAnn _dAnchor _items (EpaComments _dComments) -> do let commentExtract :: [LEpaComment] -> WriterS.Writer [LEpaComment] [LEpaComment] commentExtract comms = do let (innerComments, outerComments) = partition (\(L (Anchor anch _) _) -> ( realSrcSpanStart anch < realSrcSpanEnd span && realSrcSpanEnd anch > realSrcSpanStart span ) ) comms WriterS.tell outerComments pure innerComments (ldecl', extractedComments) = WriterS.runWriter $ SYB.everywhereM (SYB.mkM commentExtract) ldecl case ExactPrint.ss2delta lastSpanEnd span of SameLine{} -> pure () DifferentLine n _ -> finalYield $ MEWhitespace $ DifferentLine (n - 1) 1 let fixedSpanEnd = ExactPrint.ss2posEnd span let (afterComms, span2) = finalToList $ enrichComms fixedSpanEnd (List.sortOn (\(L l _) -> l) extractedComments) let (immediate, later) = List.span (\case MEComment{} -> True _ -> False ) afterComms finalYield $ MEDecl ldecl' [ (ind, GHC.ac_tok comm) | MEComment (ind, L _ comm) <- immediate ] later `forM_` finalYield enrichDecls span2 declRest EpAnn _anchor _items (EpaCommentsBalanced{}) -> error "EpaCommentsBalanced" EpAnnNotUsed -> error "EpAnnNotUsed" (L (SrcSpanAnn _ann (GHC.UnhelpfulSpan{})) _decl : _declRest) -> error "UnhelpfulSpan" -- module head pretty-printing data ImportLine = EmptyLines Int | SamelineComment (Int, LEpaComment) | NewlineComment (Int, LEpaComment) -- indentation and comment | ImportStatement ImportStatementRecord instance Show ImportLine where show = \case EmptyLines n -> "EmptyLines " ++ show n SamelineComment{} -> "SamelineComment" NewlineComment{} -> "NewlineComment" ImportStatement r -> "ImportStatement " ++ show (length $ commentsBefore r) ++ " " ++ show (length $ commentsAfter r) data ImportStatementRecord = ImportStatementRecord { commentsBefore :: [(Int, LEpaComment)] , importStatement :: LImportDecl GhcPs , commentsSameline :: [(Int, EpaCommentTok)] , commentsAfter :: [(Int, LEpaComment)] } instance Show ImportStatementRecord where show r = "ImportStatement " ++ show (length $ commentsBefore r) ++ " " ++ show (length $ commentsAfter r) transformToImportLine :: ExactPrint.Pos -> [LImportDecl GhcPs] -> FinalList ImportLine ExactPrint.Pos transformToImportLine startPos is = let flattenComms :: [LEpaComment] -> ExactPrint.Pos -> FinalList ImportLine ExactPrint.Pos flattenComms = \case [] -> finalPure lcomm@(L (Anchor span _) _) : commRest -> \lastSpanEnd -> do case ExactPrint.ss2delta lastSpanEnd span of SameLine i -> do finalYield $ SamelineComment (i, lcomm) DifferentLine l c -> do finalYield $ EmptyLines (l - 1) finalYield $ NewlineComment (c - 1, lcomm) flattenComms commRest (ExactPrint.ss2posEnd span) flattenDecls :: [LImportDecl GhcPs] -> ExactPrint.Pos -> FinalList ImportLine ExactPrint.Pos flattenDecls = \case [] -> finalPure (L (SrcSpanAnn epAnn srcSpan@(RealSrcSpan declSpan _)) decl : declRest) -> \lastSpanEnd -> let (commsBefore, commsAfter, cleanEpAnn) = case epAnn of EpAnn anch s (EpaComments cs) -> ([], reverse cs, EpAnn anch s (EpaComments [])) EpAnn anch s (EpaCommentsBalanced cs1 cs2) -> (reverse cs1, reverse cs2, EpAnn anch s (EpaComments [])) EpAnnNotUsed -> ([], [], EpAnnNotUsed) in do span1 <- flattenComms commsBefore lastSpanEnd let newlines = case ExactPrint.ss2delta span1 declSpan of SameLine _ -> 0 DifferentLine i _ -> i - 1 finalYield $ EmptyLines newlines finalYield $ ImportStatement ImportStatementRecord { commentsBefore = [] , importStatement = L (SrcSpanAnn cleanEpAnn srcSpan) decl , commentsSameline = [] , commentsAfter = [] } span2 <- flattenComms commsAfter (ExactPrint.ss2posEnd declSpan) flattenDecls declRest span2 (L (SrcSpanAnn _epAnn UnhelpfulSpan{}) _decl : _declRest) -> error "UnhelpfulSpan" in flattenDecls is startPos data Partial = PartialCommsOnly [(Int, LEpaComment)] | PartialImport ImportStatementRecord groupifyImportLines :: [ImportLine] -> [ImportLine] groupifyImportLines ls = finalToList_ $ go (PartialCommsOnly []) ls where go acc [] = case acc of PartialCommsOnly comms -> reverse comms `forM_` \comm -> finalYield $ NewlineComment comm PartialImport partialRecord -> finalYield $ ImportStatement $ unpartial partialRecord go acc (EmptyLines 0 : lineR) = go acc lineR go acc (line1 : lineR) = do newAcc <- case acc of PartialCommsOnly comms -> case line1 of e@EmptyLines{} -> do reverse comms `forM_` \comm -> finalYield $ NewlineComment comm finalYield e pure $ PartialCommsOnly [] SamelineComment comm -> do pure $ PartialCommsOnly (comm : comms) NewlineComment comm -> pure $ PartialCommsOnly (comm : comms) ImportStatement record -> pure $ PartialImport $ record { commentsBefore = comms } PartialImport partialRecord -> case line1 of e@EmptyLines{} -> do finalYield $ ImportStatement $ unpartial partialRecord finalYield e pure $ PartialCommsOnly [] SamelineComment comm -> do if (null $ commentsAfter partialRecord) then pure $ PartialImport partialRecord { commentsSameline = tokenOnly comm : commentsSameline partialRecord } else pure $ PartialImport partialRecord { commentsAfter = comm : commentsAfter partialRecord } NewlineComment comm -> pure $ PartialImport $ partialRecord { commentsAfter = comm : commentsAfter partialRecord } ImportStatement record -> do let contestedComments = commentsAfter partialRecord finalYield $ ImportStatement $ unpartial $ partialRecord { commentsAfter = [] } pure $ PartialImport $ record { commentsBefore = contestedComments } -- comments in between will stay connected to the following decl go newAcc lineR tokenOnly :: (Int, LEpaComment) -> (Int, EpaCommentTok) tokenOnly (ind, L _ (EpaComment tok _)) = (ind, tok) unpartial :: ImportStatementRecord -> ImportStatementRecord unpartial partialRecord = ImportStatementRecord { commentsBefore = reverse (commentsBefore partialRecord) , importStatement = importStatement partialRecord , commentsSameline = reverse (commentsSameline partialRecord) , commentsAfter = reverse (commentsAfter partialRecord) } sortCommentedImports :: [ImportLine] -> [ImportLine] sortCommentedImports = mergeGroups . map (fmap (combineImports . sortGroups)) . groupify where mergeGroups :: [Either ImportLine [ImportStatementRecord]] -> [ImportLine] mergeGroups xs = xs >>= \case Left x -> [x] Right y -> ImportStatement <$> y sortGroups :: [ImportStatementRecord] -> [ImportStatementRecord] sortGroups = List.sortOn (moduleNameString . unLoc . ideclName . unLoc . importStatement) combineImports :: [ImportStatementRecord] -> [ImportStatementRecord] combineImports = go Nothing where go Nothing [] = [] go (Just r1) [] = [r1] go Nothing (r2 : rs) = go (Just r2) rs go (Just r1) (r2 : rs) = case (unpack r1, unpack r2) of (Nothing, _) -> r1 : go (Just r2) rs (_, Nothing) -> r1 : go (Just r2) rs (Just u1, Just u2) -> let (modName1, pkg1, src1, safe1, q1, alias1, mllies1) = u1 (modName2, pkg2, src2, safe2, q2, alias2, mllies2) = u2 inner1 = GHC.unLoc $ importStatement r1 mostThingsEqual = modName1 == modName2 && pkg1 == pkg2 && src1 == src2 && safe1 == safe2 && ((q1 == GHC.NotQualified) == (q2 == GHC.NotQualified)) && (unLoc <$> alias1) == (unLoc <$> alias2) merged explicits = go (Just ImportStatementRecord { commentsBefore = commentsBefore r1 ++ commentsBefore r2 , importStatement = L (getLoc $ importStatement r1) GHC.ImportDecl { GHC.ideclExt = GHC.ideclExt inner1 , GHC.ideclSourceSrc = NoSourceText , GHC.ideclName = GHC.ideclName inner1 , GHC.ideclPkgQual = pkg1 , GHC.ideclSource = src1 , GHC.ideclSafe = safe1 , GHC.ideclQualified = q1 , GHC.ideclImplicit = False , GHC.ideclAs = alias1 , GHC.ideclHiding = explicits } , commentsSameline = (commentsSameline r1 ++ commentsSameline r2) , commentsAfter = commentsAfter r1 ++ commentsAfter r2 } ) rs in case (mostThingsEqual, mllies1, mllies2) of (True, Nothing, _) -> merged Nothing (True, _, Nothing) -> merged Nothing (True, Just (False, l1), Just (False, l2)) -> merged (Just (False, L (getLoc l1) (unLoc l1 ++ unLoc l2))) _ -> r1 : go (Just r2) rs unpack x = case unLoc $ importStatement x of GHC.ImportDecl _ _ (L _ modName) pkg src safe q False alias mllies -> case mllies of Nothing -> Just (modName, pkg, src, safe, q, alias, Nothing) Just (_, (L ann _)) -> case GHC.comments $ GHC.ann ann of EpaComments [] -> Just (modName, pkg, src, safe, q, alias, mllies) EpaCommentsBalanced [] [] -> Just (modName, pkg, src, safe, q, alias, mllies) _ -> Nothing _ -> Nothing groupify :: [ImportLine] -> [Either ImportLine [ImportStatementRecord]] groupify cs = go [] cs where go [] = \case (l@EmptyLines{} : rest) -> Left l : go [] rest (l@NewlineComment{} : rest) -> Left l : go [] rest (l@SamelineComment{} : rest) -> Left l : go [] rest (ImportStatement r : rest) -> go [r] rest [] -> [] go acc = \case (l@EmptyLines{} : rest) -> Right (reverse acc) : Left l : go [] rest (l@NewlineComment{} : rest) -> Left l : Right (reverse acc) : go [] rest (l@SamelineComment{} : rest) -> Left l : Right (reverse acc) : go [] rest (ImportStatement r : rest) -> go (r : acc) rest [] -> [Right (reverse acc)] rdrNameToText :: RdrName -> Text -- rdrNameToText = Text.pack . show . flip runSDoc unsafeGlobalDynFlags . ppr rdrNameToText (Unqual occname) = Text.pack $ occNameString occname rdrNameToText (Qual mname occname) = Text.pack $ moduleNameString mname ++ "." ++ occNameString occname rdrNameToText (Orig modul occname) = Text.pack $ moduleNameString (moduleName modul) ++ occNameString occname rdrNameToText (Exact name) = Text.pack $ getOccString name getDeclBindingNames :: GHC.LHsDecl GhcPs -> [String] getDeclBindingNames (L _ decl) = case decl of GHC.SigD _ (GHC.TypeSig _ ns _) -> ns <&> \(L _ n) -> Text.unpack (rdrNameToText n) GHC.ValD _ (GHC.FunBind _ (L _ n) _ _) -> [Text.unpack $ rdrNameToText n] _ -> []