{-# 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]
_ -> []