From e67855fd9885c7fb1fefbb0dbf37708387e7e6e4 Mon Sep 17 00:00:00 2001
From: Lennart Spitzner <lsp@informatik.uni-kiel.de>
Date: Thu, 18 Aug 2016 00:59:12 +0200
Subject: [PATCH] Fix list comprehension comment placement issue

---
 src-unittests/IdentityTests.hs                |  9 +++++++
 src/Language/Haskell/Brittany/LayoutBasics.hs |  4 +++
 .../Haskell/Brittany/Layouters/Expr.hs        | 25 ++++++++++++++-----
 3 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/src-unittests/IdentityTests.hs b/src-unittests/IdentityTests.hs
index 8e3f818..3a4164e 100644
--- a/src-unittests/IdentityTests.hs
+++ b/src-unittests/IdentityTests.hs
@@ -639,3 +639,12 @@ regressionTests = do
             utargets''              = "foo"
         return ()
       |]
+  it "list comprehension comment placement" $ do
+    roundTripEqual $
+      [text|
+      func =
+        [ (thing, take 10 alts) --TODO: select best ones
+        | (thing, _got, alts@(_:_)) <- nosuchFooThing
+        , gast                      <- award
+        ]
+      |]
diff --git a/src/Language/Haskell/Brittany/LayoutBasics.hs b/src/Language/Haskell/Brittany/LayoutBasics.hs
index f142dd7..c7e8936 100644
--- a/src/Language/Haskell/Brittany/LayoutBasics.hs
+++ b/src/Language/Haskell/Brittany/LayoutBasics.hs
@@ -43,6 +43,7 @@ module Language.Haskell.Brittany.LayoutBasics
   , docEmpty
   , docLit
   , docAlt
+  , docAltFilter
   , docLines
   , docCols
   , docSeq
@@ -910,6 +911,9 @@ docExt x anns shouldAddComment = allocateNode $ BDFExternal
 docAlt :: [ToBriDocM BriDocNumbered] -> ToBriDocM BriDocNumbered
 docAlt l = allocateNode . BDFAlt =<< sequence l
 
+docAltFilter :: [(Bool, ToBriDocM BriDocNumbered)] -> ToBriDocM BriDocNumbered
+docAltFilter = docAlt . map snd . filter fst
+
 
 docSeq :: [ToBriDocM BriDocNumbered] -> ToBriDocM BriDocNumbered
 docSeq l = allocateNode . BDFSeq =<< sequence l
diff --git a/src/Language/Haskell/Brittany/Layouters/Expr.hs b/src/Language/Haskell/Brittany/Layouters/Expr.hs
index bfc337c..f219725 100644
--- a/src/Language/Haskell/Brittany/Layouters/Expr.hs
+++ b/src/Language/Haskell/Brittany/Layouters/Expr.hs
@@ -424,18 +424,31 @@ layoutExpr lexpr@(L _ expr) = docWrapNode lexpr $ case expr of
                                     ; MonadComp -> True
                                     ; _ -> False } -> do
     stmtDocs <- docSharedWrapper layoutStmt `mapM` stmts
-    docAlt
-      [ docSeq
-        [ appSep $ docLit $ Text.pack "["
-        , appSep $ docForceSingleline $ List.last stmtDocs
+    hasComments <- hasAnyCommentsBelow lexpr
+    docAltFilter
+      [ (,) (not hasComments)
+      $ docSeq
+        [ docNodeAnnKW lexpr Nothing
+        $ appSep
+        $ docLit
+        $ Text.pack "["
+        , docNodeAnnKW lexpr (Just AnnOpenS)
+        $ appSep
+        $ docForceSingleline
+        $ List.last stmtDocs
         , appSep $ docLit $ Text.pack "|"
         , docSeq $ List.intersperse docCommaSep
                 $ fmap docForceSingleline $ List.init stmtDocs
         , docLit $ Text.pack " ]"
         ]
-      , let
+      , (,) True
+      $ let
           start = docCols ColListComp
-                    [appSep $ docLit $ Text.pack "[", List.last stmtDocs]
+                    [ docNodeAnnKW lexpr Nothing
+                    $ appSep $ docLit $ Text.pack "["
+                    , docNodeAnnKW lexpr (Just AnnOpenS)
+                    $ List.last stmtDocs
+                    ]
           (s1:sM) = List.init stmtDocs
           line1 = docCols ColListComp
                     [appSep $ docLit $ Text.pack "|", s1]