From 39c48b33f1a45fe2265719ff735ae5d22d3b63be Mon Sep 17 00:00:00 2001
From: Lennart Spitzner <hexagoxel@hexagoxel.de>
Date: Fri, 8 Dec 2017 19:57:16 +0100
Subject: [PATCH] Fix error in Annotation filtering (fixes #70)

---
 src-literatetests/15-regressions.blt              |  4 ++++
 src/Language/Haskell/Brittany/Internal.hs         |  8 ++++----
 .../Haskell/Brittany/Internal/ExactPrintUtils.hs  | 15 +++++++++++----
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/src-literatetests/15-regressions.blt b/src-literatetests/15-regressions.blt
index 029caa1..c2290ba 100644
--- a/src-literatetests/15-regressions.blt
+++ b/src-literatetests/15-regressions.blt
@@ -509,3 +509,7 @@ fmapuv f xs = G.generate (G.length xs) (f . (xs G.!))
 
 #test parallellistcomp-workaround
 cs0 = 0 : [ c / Interval n | c <- cs | n <- [1..] ]
+
+#test issue 70
+{-# LANGUAGE TemplateHaskell #-}
+deriveFromJSON (unPrefix "assignPost") ''AssignmentPost
diff --git a/src/Language/Haskell/Brittany/Internal.hs b/src/Language/Haskell/Brittany/Internal.hs
index 64c139a..e6256ec 100644
--- a/src/Language/Haskell/Brittany/Internal.hs
+++ b/src/Language/Haskell/Brittany/Internal.hs
@@ -320,8 +320,8 @@ withTransformedAnns ast m = do
 
 ppDecl :: LHsDecl RdrName -> PPMLocal ()
 ppDecl d@(L loc decl) = case decl of
-  SigD sig  -> -- trace (_sigHead sig) $
-               withTransformedAnns d $ do
+  SigD sig -> -- trace (_sigHead sig) $
+              withTransformedAnns d $ do
     -- runLayouter $ Old.layoutSig (L loc sig)
     briDoc <- briDocMToPPM $ layoutSig (L loc sig)
     layoutBriDoc briDoc
@@ -332,9 +332,9 @@ ppDecl d@(L loc decl) = case decl of
       eitherNode <- layoutBind (L loc bind)
       case eitherNode of
         Left  ns -> docLines $ return <$> ns
-        Right n -> return n
+        Right n  -> return n
     layoutBriDoc briDoc
-  _         -> briDocMToPPM (briDocByExactNoComment d) >>= layoutBriDoc
+  _ -> briDocMToPPM (briDocByExactNoComment d) >>= layoutBriDoc
 
 _sigHead :: Sig RdrName -> String
 _sigHead = \case
diff --git a/src/Language/Haskell/Brittany/Internal/ExactPrintUtils.hs b/src/Language/Haskell/Brittany/Internal/ExactPrintUtils.hs
index 74ed50d..081032d 100644
--- a/src/Language/Haskell/Brittany/Internal/ExactPrintUtils.hs
+++ b/src/Language/Haskell/Brittany/Internal/ExactPrintUtils.hs
@@ -217,11 +217,18 @@ extractToplevelAnns
 extractToplevelAnns lmod anns = output
  where
   (L _ (HsModule _ _ _ ldecls _ _)) = lmod
-  declMap :: Map ExactPrint.AnnKey ExactPrint.AnnKey
-  declMap = Map.unions $ ldecls <&> \ldecl ->
+  declMap1 :: Map ExactPrint.AnnKey ExactPrint.AnnKey
+  declMap1 = 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
+  declMap2 :: Map ExactPrint.AnnKey ExactPrint.AnnKey
+  declMap2 =
+    Map.fromList
+      $ [ (captured, declMap1 Map.! k)
+        | (k, ExactPrint.Ann _ _ _ _ _ (Just captured)) <- Map.toList anns
+        ]
+  declMap = declMap1 `Map.union` declMap2
+  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)