Simple alignment for top-level cases

2 files changed, 59 insertions, 6 deletions

diff --git a/lib/Language/Haskell/Stylish/Step/Cases.hs b/lib/Language/Haskell/Stylish/Step/Cases.hs
index e71c6c9..5ea30f4 100644
--- a/lib/Language/Haskell/Stylish/Step/Cases.hs
+++ b/lib/Language/Haskell/Stylish/Step/Cases.hs
@@ -34,13 +34,34 @@ altToAlignable (H.Alt ann pat rhs Nothing)  = Just $ Alignable
 
 
 --------------------------------------------------------------------------------
+tlpats :: Data l => H.Module l -> [[H.Match l]]
+tlpats modu = [matches | H.FunBind _ matches <- everything modu]
+
+
+--------------------------------------------------------------------------------
+matchToAlignable :: H.Match l -> Maybe (Alignable l)
+matchToAlignable (H.InfixMatch _ _ _ _ _ _)           = Nothing
+matchToAlignable (H.Match _   _    []   _   _)        = Nothing
+matchToAlignable (H.Match _   _    _    _   (Just _)) = Nothing
+matchToAlignable (H.Match ann name pats rhs Nothing)  = Just $ Alignable
+    { aContainer = ann
+    , aLeft      = last (H.ann name : map H.ann pats)
+    , aRight     = H.ann rhs
+    , aRightLead = length "= "
+    }
+
+
+--------------------------------------------------------------------------------
 step :: Int -> Step
 step maxColumns = makeStep "Cases" $ \ls (module', _) ->
-    let module'' = fmap H.srcInfoSpan module' in
+    let module''               = fmap H.srcInfoSpan module'
+        changes search toAlign =
+            [ change_
+            | case_   <- search module''
+            , aligns  <- maybeToList (mapM toAlign case_)
+            , change_ <- align maxColumns aligns
+            ] in
+
     applyChanges
-        [ change_
-        | case_   <- cases module''
-        , aligns  <- maybeToList (mapM altToAlignable case_)
-        , change_ <- align maxColumns aligns
-        ]
+        (changes cases altToAlignable ++ changes tlpats matchToAlignable)
         ls
diff --git a/tests/Language/Haskell/Stylish/Step/Cases/Tests.hs b/tests/Language/Haskell/Stylish/Step/Cases/Tests.hs
index cab671e..6a104f8 100644
--- a/tests/Language/Haskell/Stylish/Step/Cases/Tests.hs
+++ b/tests/Language/Haskell/Stylish/Step/Cases/Tests.hs
@@ -19,6 +19,8 @@ import           Language.Haskell.Stylish.Tests.Util
 tests :: Test
 tests = testGroup "Language.Haskell.Stylish.Step.Records.Tests"
     [ testCase "case 01" case01
+    , testCase "case 02" case02
+    , testCase "case 03" case03
     ]
 
 
@@ -37,3 +39,33 @@ case01 = expected @=? testStep (step 80) input
         , "    Left _  -> Nothing"
         , "    Right x -> Just x"
         ]
+
+
+--------------------------------------------------------------------------------
+case02 :: Assertion
+case02 = expected @=? testStep (step 80) input
+  where
+    input = unlines
+        [ "eitherToMaybe (Left _) = Nothing"
+        , "eitherToMaybe (Right x) = Just x"
+        ]
+
+    expected = unlines
+        [ "eitherToMaybe (Left _)  = Nothing"
+        , "eitherToMaybe (Right x) = Just x"
+        ]
+
+
+--------------------------------------------------------------------------------
+case03 :: Assertion
+case03 = expected @=? testStep (step 80) input
+  where
+    input = unlines
+        [ "heady def [] = def"
+        , "heady _ (x : _) = x"
+        ]
+
+    expected = unlines
+        [ "heady def []    = def"
+        , "heady _ (x : _) = x"
+        ]