From d0fcfc563e49a7d4f955922fdc27217df7c4d929 Mon Sep 17 00:00:00 2001
From: Lennart Spitzner <lsp@informatik.uni-kiel.de>
Date: Thu, 18 Aug 2016 01:14:11 +0200
Subject: [PATCH] Fix if-then-else comment placement issue

---
 src-unittests/IdentityTests.hs                |  9 ++++++
 .../Haskell/Brittany/Layouters/Expr.hs        | 30 ++++++++++++++-----
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/src-unittests/IdentityTests.hs b/src-unittests/IdentityTests.hs
index 3a4164e..fa35c15 100644
--- a/src-unittests/IdentityTests.hs
+++ b/src-unittests/IdentityTests.hs
@@ -648,3 +648,12 @@ regressionTests = do
         , gast                      <- award
         ]
       |]
+  it "if-then-else comment placement" $ do
+    roundTripEqual $
+      [text|
+      func = if x
+        then if y -- y is important
+          then foo
+          else bar
+        else Nothing
+      |]
diff --git a/src/Language/Haskell/Brittany/Layouters/Expr.hs b/src/Language/Haskell/Brittany/Layouters/Expr.hs
index f219725..ae1efde 100644
--- a/src/Language/Haskell/Brittany/Layouters/Expr.hs
+++ b/src/Language/Haskell/Brittany/Layouters/Expr.hs
@@ -296,8 +296,10 @@ layoutExpr lexpr@(L _ expr) = docWrapNode lexpr $ case expr of
     ifExprDoc   <- docSharedWrapper layoutExpr ifExpr
     thenExprDoc <- docSharedWrapper layoutExpr thenExpr
     elseExprDoc <- docSharedWrapper layoutExpr elseExpr
-    docAlt
-      [ docSeq
+    hasComments <- hasAnyCommentsBelow lexpr
+    docAltFilter
+      [ (,) (not hasComments)
+      $ docSeq
         [ appSep $ docLit $ Text.pack "if"
         , appSep $ docForceSingleline ifExprDoc
         , appSep $ docLit $ Text.pack "then"
@@ -305,11 +307,15 @@ layoutExpr lexpr@(L _ expr) = docWrapNode lexpr $ case expr of
         , appSep $ docLit $ Text.pack "else"
         , docForceSingleline elseExprDoc
         ]
-      , docSetParSpacing
+      , (,) True
+      $ docSetParSpacing
       $ docAddBaseY BrIndentRegular
       $ docPar
           ( docAddBaseY (BrIndentSpecial 3)
-          $ docSeq [appSep $ docLit $ Text.pack "if", docForceSingleline ifExprDoc])
+          $ docSeq
+            [ docNodeAnnKW lexpr Nothing $ appSep $ docLit $ Text.pack "if"
+            , docNodeAnnKW lexpr (Just AnnIf) $ docForceSingleline ifExprDoc
+            ])
           (docLines
             [ docAddBaseY BrIndentRegular
             $ docAlt
@@ -324,10 +330,14 @@ layoutExpr lexpr@(L _ expr) = docWrapNode lexpr $ case expr of
               $ docPar (docLit $ Text.pack "else") elseExprDoc
               ]
             ])
-      , docAddBaseY BrIndentRegular
+      , (,) True
+      $ docAddBaseY BrIndentRegular
       $ docPar
           ( docAddBaseY (BrIndentSpecial 3)
-          $ docSeq [appSep $ docLit $ Text.pack "if", ifExprDoc])
+          $ docSeq
+            [ docNodeAnnKW lexpr Nothing $ appSep $ docLit $ Text.pack "if"
+            , docNodeAnnKW lexpr (Just AnnIf) $ ifExprDoc
+            ])
           (docLines
             [ docAddBaseY BrIndentRegular
             $ docAlt
@@ -342,9 +352,13 @@ layoutExpr lexpr@(L _ expr) = docWrapNode lexpr $ case expr of
               $ docPar (docLit $ Text.pack "else") elseExprDoc
               ]
             ])
-      , docLines
+      , (,) True
+      $ docLines
         [ docAddBaseY (BrIndentSpecial 3)
-        $ docSeq [appSep $ docLit $ Text.pack "if", ifExprDoc]
+        $ docSeq
+          [ docNodeAnnKW lexpr Nothing $ appSep $ docLit $ Text.pack "if"
+          , docNodeAnnKW lexpr (Just AnnIf) $ ifExprDoc
+          ]
         , docAddBaseY BrIndentRegular
         $ docPar (docLit $ Text.pack "then") thenExprDoc
         , docAddBaseY BrIndentRegular