Add alignment for some case expressions

3 files changed, 62 insertions, 1 deletions

diff --git a/lib/Language/Haskell/Stylish.hs b/lib/Language/Haskell/Stylish.hs
index cfc7807..6e54d65 100644
--- a/lib/Language/Haskell/Stylish.hs
+++ b/lib/Language/Haskell/Stylish.hs
@@ -3,6 +3,7 @@ module Language.Haskell.Stylish
     ( -- * Run
       runSteps
       -- * Steps
+    , cases
     , imports
     , languagePragmas
     , records
@@ -32,6 +33,7 @@ import           Control.Monad                                    (foldM)
 import           Language.Haskell.Stylish.Config
 import           Language.Haskell.Stylish.Parse
 import           Language.Haskell.Stylish.Step
+import qualified Language.Haskell.Stylish.Step.Cases              as Cases
 import qualified Language.Haskell.Stylish.Step.Imports            as Imports
 import qualified Language.Haskell.Stylish.Step.LanguagePragmas    as LanguagePragmas
 import qualified Language.Haskell.Stylish.Step.Records            as Records
@@ -43,6 +45,12 @@ import           Paths_stylish_haskell                            (version)
 
 
 --------------------------------------------------------------------------------
+cases :: Int  -- ^ Columns
+      -> Step
+cases = Cases.step
+
+
+--------------------------------------------------------------------------------
 imports :: Int -- ^ columns
         -> Imports.Align
         -> Step
diff --git a/lib/Language/Haskell/Stylish/Config.hs b/lib/Language/Haskell/Stylish/Config.hs
index 2f72958..7994216 100644
--- a/lib/Language/Haskell/Stylish/Config.hs
+++ b/lib/Language/Haskell/Stylish/Config.hs
@@ -28,6 +28,7 @@ import           System.FilePath                                  (joinPath,
 
 --------------------------------------------------------------------------------
 import           Language.Haskell.Stylish.Step
+import qualified Language.Haskell.Stylish.Step.Cases              as Cases
 import qualified Language.Haskell.Stylish.Step.Imports            as Imports
 import qualified Language.Haskell.Stylish.Step.LanguagePragmas    as LanguagePragmas
 import qualified Language.Haskell.Stylish.Step.Records            as Records
@@ -126,7 +127,8 @@ parseConfig _            = mzero
 --------------------------------------------------------------------------------
 catalog :: Map String (Config -> A.Object -> A.Parser Step)
 catalog = M.fromList
-    [ ("imports",             parseImports)
+    [ ("cases",               parseCases)
+    , ("imports",             parseImports)
     , ("language_pragmas",    parseLanguagePragmas)
     , ("records",             parseRecords)
     , ("tabs",                parseTabs)
@@ -155,6 +157,11 @@ parseEnum strs _   (Just k) = case lookup k strs of
 
 
 --------------------------------------------------------------------------------
+parseCases :: Config -> A.Object -> A.Parser Step
+parseCases c _ = return (Cases.step $ configColumns c)
+
+
+--------------------------------------------------------------------------------
 parseImports :: Config -> A.Object -> A.Parser Step
 parseImports config o = Imports.step
     <$> pure (configColumns config)
diff --git a/lib/Language/Haskell/Stylish/Step/Cases.hs b/lib/Language/Haskell/Stylish/Step/Cases.hs
new file mode 100644
index 0000000..e71c6c9
--- /dev/null
+++ b/lib/Language/Haskell/Stylish/Step/Cases.hs
@@ -0,0 +1,46 @@
+--------------------------------------------------------------------------------
+module Language.Haskell.Stylish.Step.Cases
+    ( step
+    ) where
+
+
+--------------------------------------------------------------------------------
+import           Data.Data                       (Data)
+import           Data.Maybe                       (maybeToList)
+import qualified Language.Haskell.Exts.Annotated as H
+
+
+--------------------------------------------------------------------------------
+import           Language.Haskell.Stylish.Align
+import           Language.Haskell.Stylish.Editor
+import           Language.Haskell.Stylish.Step
+import           Language.Haskell.Stylish.Util
+
+
+--------------------------------------------------------------------------------
+cases :: Data l => H.Module l -> [[H.Alt l]]
+cases modu = [alts | H.Case _ _ alts <- everything modu]
+
+
+--------------------------------------------------------------------------------
+altToAlignable :: H.Alt l -> Maybe (Alignable l)
+altToAlignable (H.Alt _   _   _   (Just _)) = Nothing
+altToAlignable (H.Alt ann pat rhs Nothing)  = Just $ Alignable
+    { aContainer = ann
+    , aLeft      = H.ann pat
+    , aRight     = H.ann rhs
+    , aRightLead = length "-> "
+    }
+
+
+--------------------------------------------------------------------------------
+step :: Int -> Step
+step maxColumns = makeStep "Cases" $ \ls (module', _) ->
+    let module'' = fmap H.srcInfoSpan module' in
+    applyChanges
+        [ change_
+        | case_   <- cases module''
+        , aligns  <- maybeToList (mapM altToAlignable case_)
+        , change_ <- align maxColumns aligns
+        ]
+        ls