Fix quadratic behaviour (fixes #34)

Split up annotations by top-level elements in one
go, instead of doing the filtering per top-level
element (which necessarily makes things quadratic,
or rather O(n*m) with n top-level elements and m
size of annotation map). The fixed version should
be O(log n * m), and log n is negligible.

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
+    layoutBriDoc briDoc
-  _         -> briDocMToPPM (briDocByExactNoComment d) >>= layoutBriDoc d
+  _         -> 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 :: BriDocNumbered -> PPMLocal ()
-layoutBriDoc ast briDoc = do
+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
         }

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 () ()))

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) =

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]