From 8861f16624947b6c75a2c1b1c85e8f26b746c1ec Mon Sep 17 00:00:00 2001
From: Lennart Spitzner <hexagoxel@hexagoxel.de>
Date: Mon, 2 Sep 2019 17:02:20 +0200
Subject: [PATCH] Fix comment handling with let-in

---
 src-literatetests/10-tests.blt                | 33 ++++++++
 .../Brittany/Internal/Layouters/Expr.hs       | 75 ++++++++++---------
 2 files changed, 72 insertions(+), 36 deletions(-)

diff --git a/src-literatetests/10-tests.blt b/src-literatetests/10-tests.blt
index f833847..7516fd4 100644
--- a/src-literatetests/10-tests.blt
+++ b/src-literatetests/10-tests.blt
@@ -612,6 +612,39 @@ func =
           _ -> True
       ]
 
+###############################################################################
+###############################################################################
+###############################################################################
+#group expression.let
+###############################################################################
+###############################################################################
+###############################################################################
+
+#test single-bind-comment-long
+testMethod foo bar baz qux =
+  let x = undefined foo bar baz qux qux baz bar :: String
+  -- some comment explaining the in expression
+  in  undefined foo x :: String
+
+#test single-bind-comment-short
+testMethod foo bar baz qux =
+  let x = undefined :: String
+  -- some comment explaining the in expression
+  in  undefined :: String
+
+#test single-bind-comment-before
+testMethod foo bar baz qux =
+  -- some comment explaining the in expression
+  let x = undefined :: String in undefined :: String
+
+#test multiple-binds-comment
+foo foo bar baz qux =
+  let a = 1
+      b = 2
+      c = 3
+  -- some comment explaining the in expression
+  in  undefined :: String
+
 
 ###############################################################################
 ###############################################################################
diff --git a/src/Language/Haskell/Brittany/Internal/Layouters/Expr.hs b/src/Language/Haskell/Brittany/Internal/Layouters/Expr.hs
index 6fad40b..3ecf133 100644
--- a/src/Language/Haskell/Brittany/Internal/Layouters/Expr.hs
+++ b/src/Language/Haskell/Brittany/Internal/Layouters/Expr.hs
@@ -725,10 +725,10 @@ layoutExpr lexpr@(L _ expr) = do
 #else
     HsLet binds exp1 -> do
 #endif
-      expDoc1   <- docSharedWrapper layoutExpr exp1
+      expDoc1     <- docSharedWrapper layoutExpr exp1
       -- We jump through some ugly hoops here to ensure proper sharing.
-      mBindDocs <- mapM (fmap (fmap return) . docWrapNodeRest lexpr . return)
-               =<< layoutLocalBinds binds
+      hasComments <- hasAnyCommentsBelow lexpr
+      mBindDocs   <- fmap (fmap (fmap pure)) $ layoutLocalBinds binds
       let
         ifIndentFreeElse :: a -> a -> a
         ifIndentFreeElse x y =
@@ -745,37 +745,38 @@ layoutExpr lexpr@(L _ expr) = do
       -- if "let" is moved horizontally as part of the transformation, as the
       -- comments before the first let item are moved horizontally with it.
       docSetBaseAndIndent $ case mBindDocs of
-        Just [bindDoc] -> docAlt
-          [ docSeq
-              [ appSep $ docLit $ Text.pack "let"
-              , appSep $ docForceSingleline bindDoc
-              , appSep $ docLit $ Text.pack "in"
-              , docForceSingleline expDoc1
-              ]
-          , docLines
-              [ docAlt
-                  [ docSeq
-                      [ appSep $ docLit $ Text.pack "let"
-                      , ifIndentFreeElse docSetBaseAndIndent docForceSingleline
-                      $ bindDoc
-                      ]
-                  , docAddBaseY BrIndentRegular
-                  $ docPar
-                    (docLit $ Text.pack "let")
-                    (docSetBaseAndIndent bindDoc)
-                  ]
-              , docAlt
-                  [ docSeq
-                      [ appSep $ docLit $ Text.pack $ ifIndentFreeElse "in " "in"
-                      , ifIndentFreeElse docSetBaseAndIndent docForceSingleline expDoc1
-                      ]
-                  , docAddBaseY BrIndentRegular
-                  $ docPar
-                    (docLit $ Text.pack "in")
-                    (docSetBaseY expDoc1)
-                  ]
-              ]
-          ]
+        Just [bindDoc] -> runFilteredAlternative $ do
+          addAlternativeCond (not hasComments) $ docSeq
+            [ appSep $ docLit $ Text.pack "let"
+            , docNodeAnnKW lexpr (Just AnnLet)
+            $ appSep $ docForceSingleline bindDoc
+            , appSep $ docLit $ Text.pack "in"
+            , docForceSingleline expDoc1
+            ]
+          addAlternative $ docLines
+            [ docNodeAnnKW lexpr (Just AnnLet)
+            $ docAlt
+                [ docSeq
+                    [ appSep $ docLit $ Text.pack "let"
+                    , ifIndentFreeElse docSetBaseAndIndent docForceSingleline
+                    $ bindDoc
+                    ]
+                , docAddBaseY BrIndentRegular
+                $ docPar
+                  (docLit $ Text.pack "let")
+                  (docSetBaseAndIndent bindDoc)
+                ]
+            , docAlt
+                [ docSeq
+                    [ appSep $ docLit $ Text.pack $ ifIndentFreeElse "in " "in"
+                    , ifIndentFreeElse docSetBaseAndIndent docForceSingleline expDoc1
+                    ]
+                , docAddBaseY BrIndentRegular
+                $ docPar
+                  (docLit $ Text.pack "in")
+                  (docSetBaseY expDoc1)
+                ]
+            ]
         Just bindDocs@(_:_) -> runFilteredAlternative $ do
           --either
           --  let
@@ -805,7 +806,8 @@ layoutExpr lexpr@(L _ expr) = do
             IndentPolicyLeft     -> docLines noHangingBinds
             IndentPolicyMultiple -> docLines noHangingBinds
             IndentPolicyFree     -> docLines
-              [ docSeq
+              [ docNodeAnnKW lexpr (Just AnnLet)
+              $ docSeq
                 [ appSep $ docLit $ Text.pack "let"
                 , docSetBaseAndIndent $ docLines bindDocs
                 ]
@@ -816,7 +818,8 @@ layoutExpr lexpr@(L _ expr) = do
               ]
           addAlternative
             $ docLines
-            [ docAddBaseY BrIndentRegular
+            [ docNodeAnnKW lexpr (Just AnnLet)
+            $ docAddBaseY BrIndentRegular
             $ docPar
               (docLit $ Text.pack "let")
               (docSetBaseAndIndent $ docLines $ bindDocs)