Fix records with comments (#257)

* Format records where comments are in the same line as the field name

* Fix records format, records with comments will now be formatted

* Fix formatting of comments below

Co-authored-by: Łukasz Gołębiewski <lukasz.golebiewski@gmail.com>

2 files changed, 64 insertions, 33 deletions

diff --git a/lib/Language/Haskell/Stylish/Step/Data.hs b/lib/Language/Haskell/Stylish/Step/Data.hs
index 9acd22b..94aaf22 100644
--- a/lib/Language/Haskell/Stylish/Step/Data.hs
+++ b/lib/Language/Haskell/Stylish/Step/Data.hs
@@ -25,12 +25,18 @@ step' indentSize ls (module', allComments) = applyChanges changes ls
     datas' = datas $ fmap linesFromSrcSpan module'
     changes = datas' >>= maybeToList . changeDecl allComments indentSize
 
-findComment :: LineBlock -> [Comment] -> Maybe Comment
-findComment lb = find commentOnLine
+findCommentOnLine :: LineBlock -> [Comment] -> Maybe Comment
+findCommentOnLine lb = find commentOnLine
   where
     commentOnLine (Comment _ (H.SrcSpan _ start _ end _) _) =
       blockStart lb == start && blockEnd lb == end
 
+findCommentBelowLine :: LineBlock -> [Comment] -> Maybe Comment
+findCommentBelowLine lb = find commentOnLine
+  where
+    commentOnLine (Comment _ (H.SrcSpan _ start _ end _) _) =
+      blockStart lb == start - 1 && blockEnd lb == end - 1
+
 commentsWithin :: LineBlock -> [Comment] -> [Comment]
 commentsWithin lb = filter within
   where
@@ -39,9 +45,8 @@ commentsWithin lb = filter within
 
 changeDecl :: [Comment] -> Int -> H.Decl LineBlock -> Maybe ChangeLine
 changeDecl _ _ (H.DataDecl _ (H.DataType _) Nothing _ [] _) = Nothing
-changeDecl allComments indentSize (H.DataDecl block (H.DataType _) Nothing dhead decls derivings)
-  | null $ commentsWithin block allComments = Just $ change block (const $ concat newLines)
-  | otherwise = Nothing
+changeDecl allComments indentSize (H.DataDecl block (H.DataType _) Nothing dhead decls derivings) =
+  Just $ change block (const $ concat newLines)
   where
     newLines = fmap constructors zipped ++ [fmap (indented . H.prettyPrint) derivings]
     zipped = zip decls ([1..] ::[Int])
@@ -53,14 +58,17 @@ changeDecl _ _ _ = Nothing
 
 processConstructor :: [Comment] -> String -> Int -> H.QualConDecl LineBlock -> [String]
 processConstructor allComments init indentSize (H.QualConDecl _ _ _ (H.RecDecl _ dname fields)) = do
-  init <> H.prettyPrint dname : n1 : ns ++ [indented "}"]
+  init <> H.prettyPrint dname : n1 ++ ns ++ [indented "}"]
   where
     n1 = processName "{ " ( extractField $ head fields)
-    ns = fmap (processName ", " . extractField) (tail fields)
-    processName prefix (fnames, _type, Nothing) =
-      indented prefix <> intercalate ", " (fmap H.prettyPrint fnames) <> " :: " <> H.prettyPrint _type
-    processName prefix (fnames, _type, (Just (Comment _ _ c))) =
-      indented prefix <> intercalate ", " (fmap H.prettyPrint fnames) <> " :: " <> H.prettyPrint _type <> " --" <> c
-    extractField (H.FieldDecl lb names _type) = (names, _type, findComment lb allComments)
+    ns = tail fields >>= (processName ", " . extractField)
+    processName prefix (fnames, _type, lineComment, commentBelowLine) =
+      [indented prefix <> intercalate ", " (fmap H.prettyPrint fnames) <> " :: " <> H.prettyPrint _type <> addLineComment lineComment] ++ addCommentBelow commentBelowLine
+    addLineComment (Just (Comment _ _ c)) = " --" <> c
+    addLineComment Nothing                = ""
+    addCommentBelow Nothing                = []
+    addCommentBelow (Just (Comment _ _ c)) = [indented "--" <> c]
+    extractField (H.FieldDecl lb names _type) =
+      (names, _type, findCommentOnLine lb allComments, findCommentBelowLine lb allComments)
     indented = indent indentSize
 processConstructor _ init _ decl = [init <> trimLeft (H.prettyPrint decl)]
diff --git a/tests/Language/Haskell/Stylish/Step/Data/Tests.hs b/tests/Language/Haskell/Stylish/Step/Data/Tests.hs
index b152819..712ffae 100644
--- a/tests/Language/Haskell/Stylish/Step/Data/Tests.hs
+++ b/tests/Language/Haskell/Stylish/Step/Data/Tests.hs
@@ -29,6 +29,7 @@ tests = testGroup "Language.Haskell.Stylish.Step.Data.Tests"
     , testCase "case 16" case16
     , testCase "case 17" case17
     , testCase "case 18" case18
+    , testCase "case 19" case19
     ]
 
 case00 :: Assertion
@@ -289,6 +290,26 @@ case15 :: Assertion
 case15 = expected @=? testStep (step 2) input
   where
     input = unlines
+       [ "module Herp where"
+       , ""
+       , "data Foo a = Foo"
+       , "  { a :: a, -- comment"
+       , "   a2 :: String"
+       , "  }"
+       ]
+    expected = unlines
+       [ "module Herp where"
+       , ""
+       , "data Foo a = Foo"
+       , "  { a :: a -- comment"
+       , "  , a2 :: String"
+       , "  }"
+       ]
+
+case16 :: Assertion
+case16 = expected @=? testStep (step 2) input
+  where
+    input = unlines
       [ "module Herp where"
       , ""
       , "data Foo = Foo {"
@@ -298,20 +319,20 @@ case15 = expected @=? testStep (step 2) input
     expected = unlines
       [ "module Herp where"
       , ""
-      , "data Foo = Foo {"
-      , "   a :: Int -- ^ comment"
+      , "data Foo = Foo"
+      , "  { a :: Int -- ^ comment"
       , "  }"
       ]
 
-case16 :: Assertion
-case16 = expected @=? testStep (step 2) input
+case17 :: Assertion
+case17 = expected @=? testStep (step 2) input
   where
     input = unlines
        [ "module Herp where"
        , ""
        , "data Foo a = Foo"
        , "  { a :: a,"
-       , "-- ^ comment"
+       , "-- comment"
        , "   a2 :: String"
        , "  }"
        ]
@@ -319,20 +340,21 @@ case16 = expected @=? testStep (step 2) input
        [ "module Herp where"
        , ""
        , "data Foo a = Foo"
-       , "  { a :: a,"
-       , "-- ^ comment"
-       , "   a2 :: String"
+       , "  { a :: a"
+       , "  -- comment"
+       , "  , a2 :: String"
        , "  }"
        ]
 
-case17 :: Assertion
-case17 = expected @=? testStep (step 2) input
+case18 :: Assertion
+case18 = expected @=? testStep (step 2) input
   where
     input = unlines
        [ "module Herp where"
        , ""
        , "data Foo a = Foo"
-       , "  { a :: a, -- comment"
+       , "  { a :: a,"
+       , "-- ^ comment"
        , "   a2 :: String"
        , "  }"
        ]
@@ -340,29 +362,30 @@ case17 = expected @=? testStep (step 2) input
        [ "module Herp where"
        , ""
        , "data Foo a = Foo"
-       , "  { a :: a, -- comment"
-       , "   a2 :: String"
+       , "  { a :: a"
+       , "  -- ^ comment"
+       , "  , a2 :: String"
        , "  }"
        ]
 
-case18 :: Assertion
-case18 = expected @=? testStep (step 2) input
+case19 :: Assertion
+case19 = expected @=? testStep (step 2) input
   where
     input = unlines
        [ "module Herp where"
        , ""
        , "data Foo a = Foo"
-       , "  { a :: a,"
-       , "-- comment "
-       , "   a2 :: String"
+       , "  { firstName, lastName :: String,"
+       , "-- ^ names"
+       , "   age :: Int"
        , "  }"
        ]
     expected = unlines
        [ "module Herp where"
        , ""
        , "data Foo a = Foo"
-       , "  { a :: a,"
-       , "-- comment "
-       , "   a2 :: String"
+       , "  { firstName, lastName :: String"
+       , "  -- ^ names"
+       , "  , age :: Int"
        , "  }"
        ]