From 8d7b46b9e916de843aed5aeda47ecb6e6c3658de Mon Sep 17 00:00:00 2001
From: Taylor Fausak <taylor@fausak.me>
Date: Sun, 7 Nov 2021 13:01:54 +0000
Subject: [PATCH] Fix handling of comments

---
 data/10-tests.blt                             | 48 +++++++++++++++++++
 .../Language/Haskell/Brittany/Internal.hs     | 18 +++----
 2 files changed, 58 insertions(+), 8 deletions(-)

diff --git a/data/10-tests.blt b/data/10-tests.blt
index 75babb0..79d9a0a 100644
--- a/data/10-tests.blt
+++ b/data/10-tests.blt
@@ -1568,6 +1568,13 @@ type instance F Int = IO Int
 type family F a
 type instance F Int = IO Int -- x
 
+#test type-instance-with-module-header
+
+{-# language TypeFamilies #-}
+module M where
+type family F a
+type instance F Int = IO Int
+
 #test newtype-instance-without-comment
 
 {-# language TypeFamilies #-}
@@ -1580,6 +1587,13 @@ newtype instance F Int = N Int
 data family F a
 newtype instance F Int = N Int -- x
 
+#test newtype-instance-with-module-header
+
+{-# language TypeFamilies #-}
+module M where
+data family F a
+newtype instance F Int = N Int
+
 #test data-instance-without-comment
 
 {-# language TypeFamilies #-}
@@ -1592,6 +1606,13 @@ data instance F Int = D Int
 data family F a
 data instance F Int = D Int -- x
 
+#test data-instance-with-module-header
+
+{-# language TypeFamilies #-}
+module M where
+data family F a
+data instance F Int = D Int
+
 #test instance-type-without-comment
 
 {-# language TypeFamilies #-}
@@ -1608,6 +1629,15 @@ class C a where
 instance C Int where
   type F Int = IO Int -- x
 
+#test instance-type-with-module-header
+
+{-# language TypeFamilies #-}
+module M where
+class C a where
+  type family F a
+instance C Int where
+  type F Int = IO Int
+
 #test instance-newtype-without-comment
 
 {-# language TypeFamilies #-}
@@ -1624,6 +1654,15 @@ class C a where
 instance C Int where
   newtype F Int = N Int -- x
 
+#test instance-newtype-with-module-header
+
+{-# language TypeFamilies #-}
+module M where
+class C a where
+  data family F a
+instance C Int where
+  newtype F Int = N Int
+
 #test instance-data-without-comment
 
 {-# language TypeFamilies #-}
@@ -1640,6 +1679,15 @@ class C a where
 instance C Int where
   data F Int = D Int -- x
 
+#test instance-data-with-module-header
+
+{-# language TypeFamilies #-}
+module M where
+class C a where
+  data family F a
+instance C Int where
+  data F Int = D Int
+
 ###############################################################################
 ###############################################################################
 ###############################################################################
diff --git a/source/library/Language/Haskell/Brittany/Internal.hs b/source/library/Language/Haskell/Brittany/Internal.hs
index 71e885b..8d3e72e 100644
--- a/source/library/Language/Haskell/Brittany/Internal.hs
+++ b/source/library/Language/Haskell/Brittany/Internal.hs
@@ -400,7 +400,7 @@ parsePrintModuleTests conf filename input = do
         then return $ pPrintModule moduleConf perItemConf anns parsedModule
         else lift
           $ pPrintModuleAndCheck moduleConf perItemConf anns parsedModule
-      if all isErrorUnusedComment errs
+      if null errs
         then pure $ TextL.toStrict $ ltext
         else
           let
@@ -413,11 +413,6 @@ parsePrintModuleTests conf filename input = do
               ErrorOutputCheck       -> "Output is not syntactically valid."
           in  throwE $ "pretty printing error(s):\n" ++ List.unlines errStrs
 
-isErrorUnusedComment :: BrittanyError -> Bool
-isErrorUnusedComment x = case x of
-  ErrorUnusedComment _ -> True
-  _ -> False
-
 -- this approach would for if there was a pure GHC.parseDynamicFilePragma.
 -- Unfortunately that does not exist yet, so we cannot provide a nominally
 -- pure interface.
@@ -461,7 +456,14 @@ toLocal conf anns m = do
 
 ppModule :: GenLocated SrcSpan HsModule -> PPM ()
 ppModule lmod@(L _loc _m@(HsModule _ _name _exports _ decls _ _)) = do
-  let annKey = ExactPrint.mkAnnKey lmod
+  defaultAnns <- do
+    anns <- mAsk
+    let annKey = ExactPrint.mkAnnKey lmod
+    let annMap = Map.findWithDefault Map.empty annKey anns
+    let isEof = (== ExactPrint.AnnEofPos)
+    let overAnnsDP f a = a { ExactPrint.annsDP = f $ ExactPrint.annsDP a }
+    pure $ fmap (overAnnsDP . filter $ isEof . fst) annMap
+
   post <- ppPreamble lmod
   decls `forM_` \decl -> do
     let declAnnKey       = ExactPrint.mkAnnKey decl
@@ -472,7 +474,7 @@ ppModule lmod@(L _loc _m@(HsModule _ _name _exports _ decls _ _)) = do
           declBindingNames <&> \n -> Map.lookup n $ _icd_perBinding inlineConf
     filteredAnns <- mAsk
       <&> \annMap ->
-        Map.union (Map.findWithDefault Map.empty annKey annMap) $
+        Map.union defaultAnns $
         Map.findWithDefault Map.empty declAnnKey annMap
 
     traceIfDumpConf "bridoc annotations filtered/transformed"