From a59df1f3913344768e4d5d2b7835652e787eecef Mon Sep 17 00:00:00 2001
From: sniperrifle2004 <sniperrifle2004+github@gmail.com>
Date: Tue, 19 Dec 2017 14:28:22 +0100
Subject: [PATCH] Fix comments!! :hooray:

---
 src-literatetests/10-tests.blt                | 23 +++++++++++
 src-literatetests/tests-context-free.blt      | 23 +++++++++++
 .../Brittany/Internal/Layouters/Import.hs     | 41 +++++++++++++++----
 3 files changed, 79 insertions(+), 8 deletions(-)

diff --git a/src-literatetests/10-tests.blt b/src-literatetests/10-tests.blt
index 8c5b547..6b49e57 100644
--- a/src-literatetests/10-tests.blt
+++ b/src-literatetests/10-tests.blt
@@ -709,6 +709,29 @@ import qualified Data.List                               as L
 -- Test
 import           Test                                     ( test )
 
+#test import-with-comments-2
+
+import           Test                                     ( abc
+                                                          , def
+                                                          -- comment
+                                                          )
+
+#test import-with-comments-3
+
+import           Test                                     ( abc
+                                                          -- comment
+                                                          )
+
+#test import-with-comments-4
+import           Test                                     ( abc
+                                                          -- comment
+                                                          , def
+                                                          , ghi
+                                                          {- comment -}
+                                                          , jkl
+                                                          -- comment
+                                                          )
+
 #test preamble full-preamble
 {-# LANGUAGE BangPatterns #-}
 
diff --git a/src-literatetests/tests-context-free.blt b/src-literatetests/tests-context-free.blt
index 5f21502..1bc25ac 100644
--- a/src-literatetests/tests-context-free.blt
+++ b/src-literatetests/tests-context-free.blt
@@ -741,6 +741,29 @@ import           Data.List                                ( nub ) -- Test
 import qualified Data.List                               as L
                                                           ( foldl' ) {- Test -}
 
+#test import-with-comments-2
+
+import           Test                                     ( abc
+                                                          , def
+                                                          -- comment
+                                                          )
+
+#test import-with-comments-3
+
+import           Test                                     ( abc
+                                                          -- comment
+                                                          )
+
+#test import-with-comments-4
+import           Test                                     ( abc
+                                                          -- comment
+                                                          , def
+                                                          , ghi
+                                                          {- comment -}
+                                                          , jkl
+                                                          -- comment
+                                                          )
+
 -- Test
 import           Test                                     ( test )
 
diff --git a/src/Language/Haskell/Brittany/Internal/Layouters/Import.hs b/src/Language/Haskell/Brittany/Internal/Layouters/Import.hs
index 4f00f00..e2ba394 100644
--- a/src/Language/Haskell/Brittany/Internal/Layouters/Import.hs
+++ b/src/Language/Haskell/Brittany/Internal/Layouters/Import.hs
@@ -76,25 +76,50 @@ layoutImport limportD@(L _ importD) = docWrapNode limportD $ case importD of
         docEnsureIndent (BrIndentSpecial qLength) $ appSep $ docLit modNameT
       hidDoc =
         if hiding then appSep $ docLit $ Text.pack "hiding" else docEmpty
-      importHead = docSeq [importQualifiers, modNameD]
-      Just lies  = mlies
+      importHead      = docSeq [importQualifiers, modNameD]
+      Just lies       = mlies
       Just (_, llies) = mllies
-      (ieH:ieT)  = map layoutIE lies
+      (ieH:ieT)       = map layoutIE lies
       makeIENode ie = docSeq [docCommaSep, ie]
       bindings@(bindingsH:bindingsT) =
         docSeq [docParenLSep, ieH]
-          :  map makeIENode ieT
+          :  bindingsT'
           ++ [docSeq [docSeparator, docParenR]]
+        where
+          -- Handle the last element with docWrapNode llies
+          bindingsT' =
+            case ieT of
+              [] -> []
+              [ie] -> [makeIENode $ docWrapNode llies $ ie]
+              _ -> map makeIENode (List.init ieT) ++ [makeIENode $ docWrapNode llies $ List.last ieT]
       bindingsD = case mlies of
-        Nothing  -> docSeq [docEmpty]
+        Nothing -> docSeq [docEmpty]
         -- ..[hiding].( )
-        Just []  -> docSeq [hidDoc, docParenLSep, docParenR]
+        Just [] -> do
+          hasComments <- hasAnyCommentsBelow llies
+          if hasComments
+          then
+            docWrapNodeRest llies
+            $ docPar
+                (docSeq [hidDoc, docParenLSep, docWrapNode llies docEmpty])
+            $ docLines [docParenR]
+          else
+            docWrapNodeRest llies $ docSeq [hidDoc, docParenLSep, docParenR]
         -- ..[hiding].( b )
-        Just [_] -> docSeq $ hidDoc : bindings
+        Just [_] -> do
+          hasComments <- hasAnyCommentsBelow llies
+          if hasComments
+          then
+            docWrapNodeRest llies $ docPar (docSeq [hidDoc, docWrapNode llies $ bindingsH]) $ docLines
+              bindingsT
+          else
+            docWrapNodeRest llies $ docSeq $ hidDoc : bindings
         -- ..[hiding].( b
         --            , b'
         --            )
-        Just _ -> docWrapNode llies $ docPar (docSeq [hidDoc, docSetBaseY $ bindingsH]) $ docLines bindingsT
+        Just _ ->
+          docWrapNodeRest llies $ docPar (docSeq [hidDoc, docSetBaseY $ bindingsH])
+            $ docLines bindingsT
       bindingLine =
         docEnsureIndent (BrIndentSpecial (importCol - bindingCost)) bindingsD
     case asT of