diff --git a/src/Language/Haskell/Brittany/Internal.hs b/src/Language/Haskell/Brittany/Internal.hs index 980992f..4c4bbf0 100644 --- a/src/Language/Haskell/Brittany/Internal.hs +++ b/src/Language/Haskell/Brittany/Internal.hs @@ -140,6 +140,7 @@ pPrintModule conf anns parsedModule = $ MultiRWSS.withMultiWriterW $ MultiRWSS.withMultiReader anns $ MultiRWSS.withMultiReader conf + $ MultiRWSS.withMultiReader (extractToplevelAnns parsedModule anns) $ do traceIfDumpConf "bridoc annotations raw" _dconf_dump_annotations $ annsDoc anns @@ -266,7 +267,19 @@ ppModule lmod@(L loc m@(HsModule _name _exports _imports decls _ _)) = do anns' = Map.insert (ExactPrint.Types.mkAnnKey lmod) mAnn' anns in (anns', post) MultiRWSS.withMultiReader anns' $ processDefault emptyModule - decls `forM_` ppDecl + decls `forM_` \decl -> do + filteredAnns <- mAsk <&> \annMap -> + Map.findWithDefault Map.empty (ExactPrint.Types.mkAnnKey decl) annMap + + traceIfDumpConf "bridoc annotations filtered/transformed" + _dconf_dump_annotations + $ annsDoc filteredAnns + + config <- mAsk + + MultiRWSS.withoutMultiReader $ do + MultiRWSS.mPutRawR $ config :+: filteredAnns :+: HNil + ppDecl decl let finalComments = filter ( fst .> \case ExactPrint.Types.AnnComment{} -> True @@ -291,7 +304,7 @@ ppModule lmod@(L loc m@(HsModule _name _exports _imports decls _ _)) = do ppmMoveToExactLoc $ ExactPrint.Types.DP (eofX - cmX, eofY - cmY) _ -> return () -withTransformedAnns :: Data ast => ast -> PPM () -> PPM () +withTransformedAnns :: Data ast => ast -> PPMLocal () -> PPMLocal () withTransformedAnns ast m = do -- TODO: implement `local` for MultiReader/MultiRWS readers@(conf :+: anns :+: HNil) <- MultiRWSS.mGetRawR @@ -305,13 +318,13 @@ withTransformedAnns ast m = do in annsBalanced -ppDecl :: LHsDecl RdrName -> PPM () +ppDecl :: LHsDecl RdrName -> PPMLocal () ppDecl d@(L loc decl) = case decl of SigD sig -> -- trace (_sigHead sig) $ withTransformedAnns d $ do -- runLayouter $ Old.layoutSig (L loc sig) briDoc <- briDocMToPPM $ layoutSig (L loc sig) - layoutBriDoc d briDoc + layoutBriDoc briDoc ValD bind -> -- trace (_bindHead bind) $ withTransformedAnns d $ do -- Old.layoutBind (L loc bind) @@ -320,8 +333,8 @@ ppDecl d@(L loc decl) = case decl of case eitherNode of Left ns -> docLines $ return <$> ns Right n -> return n - layoutBriDoc d briDoc - _ -> briDocMToPPM (briDocByExactNoComment d) >>= layoutBriDoc d + layoutBriDoc briDoc + _ -> briDocMToPPM (briDocByExactNoComment d) >>= layoutBriDoc _sigHead :: Sig RdrName -> String _sigHead = \case @@ -337,8 +350,8 @@ _bindHead = \case -layoutBriDoc :: Data.Data.Data ast => ast -> BriDocNumbered -> PPM () -layoutBriDoc ast briDoc = do +layoutBriDoc :: BriDocNumbered -> PPMLocal () +layoutBriDoc briDoc = do -- first step: transform the briDoc. briDoc' <- MultiRWSS.withMultiStateS BDEmpty $ do -- Note that briDoc is BriDocNumbered, but state type is BriDoc. @@ -374,11 +387,6 @@ layoutBriDoc ast briDoc = do -- return simpl anns :: ExactPrint.Types.Anns <- mAsk - let filteredAnns = filterAnns ast anns - - traceIfDumpConf "bridoc annotations filtered/transformed" - _dconf_dump_annotations - $ annsDoc filteredAnns let state = LayoutState { _lstate_baseYs = [0] @@ -388,7 +396,7 @@ layoutBriDoc ast briDoc = do -- thing properly. , _lstate_indLevels = [0] , _lstate_indLevelLinger = 0 - , _lstate_comments = filteredAnns + , _lstate_comments = anns , _lstate_commentCol = Nothing , _lstate_addSepSpace = Nothing } diff --git a/src/Language/Haskell/Brittany/Internal/ExactPrintUtils.hs b/src/Language/Haskell/Brittany/Internal/ExactPrintUtils.hs index c3a7d72..420ac24 100644 --- a/src/Language/Haskell/Brittany/Internal/ExactPrintUtils.hs +++ b/src/Language/Haskell/Brittany/Internal/ExactPrintUtils.hs @@ -5,6 +5,8 @@ module Language.Haskell.Brittany.Internal.ExactPrintUtils , parseModuleFromString , commentAnnFixTransform , commentAnnFixTransformGlob + , extractToplevelAnns + , foldedAnnKeys ) where @@ -14,7 +16,6 @@ where import Language.Haskell.Brittany.Internal.Types import Language.Haskell.Brittany.Internal.Config.Types -import Language.Haskell.Brittany.Internal.LayouterBasics import Language.Haskell.Brittany.Internal.Utils import DynFlags ( getDynFlags ) @@ -36,6 +37,7 @@ import qualified Language.Haskell.GHC.ExactPrint.Preprocess as ExactPrint import qualified Language.Haskell.GHC.ExactPrint.Delta as ExactPrint import qualified Data.Generics as SYB +-- import Data.Generics.Schemes @@ -209,3 +211,40 @@ moveTrailingComments astFrom astTo = do ans' = Map.insert k1 an1' $ Map.insert k2 an2' ans ExactPrint.modifyAnnsT moveComments + +-- | split a set of annotations in a module into a map from top-level module +-- elements to the relevant annotations. Avoids quadratic behaviour a trivial +-- implementation would have. +extractToplevelAnns + :: Located (HsModule RdrName) + -> ExactPrint.Anns + -> Map ExactPrint.AnnKey ExactPrint.Anns +extractToplevelAnns lmod anns = output + where + (L _ (HsModule _ _ _ ldecls _ _)) = lmod + declMap :: Map ExactPrint.AnnKey ExactPrint.AnnKey + declMap = Map.unions $ ldecls <&> \ldecl -> + Map.fromSet (const (ExactPrint.mkAnnKey ldecl)) (foldedAnnKeys ldecl) + modKey = ExactPrint.mkAnnKey lmod + output = groupMap (\k _ -> Map.findWithDefault modKey k declMap) anns + +groupMap :: (Ord k, Ord l) => (k -> a -> l) -> Map k a -> Map l (Map k a) +groupMap f = Map.foldlWithKey' (\m k a -> Map.alter (insert k a) (f k a) m) + Map.empty + where + insert k a Nothing = Just (Map.singleton k a) + insert k a (Just m) = Just (Map.insert k a m) + +foldedAnnKeys :: Data.Data.Data ast => ast -> Set ExactPrint.AnnKey +foldedAnnKeys ast = SYB.everything + Set.union + ( \x -> maybe + Set.empty + Set.singleton + [ SYB.gmapQi 1 (\t -> ExactPrint.mkAnnKey $ L l t) x + | locTyCon == SYB.typeRepTyCon (SYB.typeOf x) + , l <- SYB.gmapQi 0 SYB.cast x + ] + ) + ast + where locTyCon = SYB.typeRepTyCon (SYB.typeOf (L () ())) diff --git a/src/Language/Haskell/Brittany/Internal/LayouterBasics.hs b/src/Language/Haskell/Brittany/Internal/LayouterBasics.hs index 5314fcf..cffcad7 100644 --- a/src/Language/Haskell/Brittany/Internal/LayouterBasics.hs +++ b/src/Language/Haskell/Brittany/Internal/LayouterBasics.hs @@ -67,6 +67,7 @@ import qualified Data.Text.Lazy.Builder as Text.Builder import Language.Haskell.Brittany.Internal.Config.Types import Language.Haskell.Brittany.Internal.Types import Language.Haskell.Brittany.Internal.Utils +import Language.Haskell.Brittany.Internal.ExactPrintUtils import RdrName ( RdrName(..) ) import GHC ( Located, runGhc, GenLocated(L), moduleNameString ) @@ -81,6 +82,8 @@ import Data.Generics.Schemes import DataTreePrint +import Data.HList.HList + processDefault @@ -222,20 +225,6 @@ extractAllComments ann = _ -> [] ) -foldedAnnKeys :: Data.Data.Data ast => ast -> Set ExactPrint.AnnKey -foldedAnnKeys ast = everything - Set.union - ( \x -> maybe Set.empty - Set.singleton - [ gmapQi 1 (\t -> ExactPrint.mkAnnKey $ L l t) x - | locTyCon == typeRepTyCon (typeOf x) - , l <- gmapQi 0 cast x - ] - ) - ast - where - locTyCon = typeRepTyCon (typeOf (L () ())) - filterAnns :: Data.Data.Data ast => ast -> ExactPrint.Anns -> ExactPrint.Anns filterAnns ast anns = Map.filterWithKey (\k _ -> k `Set.member` foldedAnnKeys ast) anns @@ -614,7 +603,7 @@ spacifyDocs :: [ToBriDocM BriDocNumbered] -> [ToBriDocM BriDocNumbered] spacifyDocs [] = [] spacifyDocs ds = fmap appSep (List.init ds) ++ [List.last ds] -briDocMToPPM :: ToBriDocM a -> PPM a +briDocMToPPM :: ToBriDocM a -> PPMLocal a briDocMToPPM m = do readers <- MultiRWSS.mGetRawR let ((x, errs), debugs) = diff --git a/src/Language/Haskell/Brittany/Internal/Types.hs b/src/Language/Haskell/Brittany/Internal/Types.hs index e5c5e2c..557f9b3 100644 --- a/src/Language/Haskell/Brittany/Internal/Types.hs +++ b/src/Language/Haskell/Brittany/Internal/Types.hs @@ -28,7 +28,15 @@ import Data.Generics.Uniplate.Direct as Uniplate -type PPM a = MultiRWSS.MultiRWS '[Config, ExactPrint.Anns] '[Text.Builder.Builder, [BrittanyError], Seq String] '[] a +type PPM = MultiRWSS.MultiRWS + '[Map ExactPrint.AnnKey ExactPrint.Anns, Config, ExactPrint.Anns] + '[Text.Builder.Builder, [BrittanyError], Seq String] + '[] + +type PPMLocal = MultiRWSS.MultiRWS + '[Config, ExactPrint.Anns] + '[Text.Builder.Builder, [BrittanyError], Seq String] + '[] data LayoutState = LayoutState { _lstate_baseYs :: [Int]