Add support for aligning multi way ifs

2 files changed, 50 insertions, 2 deletions

diff --git a/lib/Language/Haskell/Stylish/Step/SimpleAlign.hs b/lib/Language/Haskell/Stylish/Step/SimpleAlign.hs
index e02c270..e03f665 100644
--- a/lib/Language/Haskell/Stylish/Step/SimpleAlign.hs
+++ b/lib/Language/Haskell/Stylish/Step/SimpleAlign.hs
@@ -9,7 +9,7 @@ module Language.Haskell.Stylish.Step.SimpleAlign
 
 --------------------------------------------------------------------------------
 import           Control.Monad                   (guard)
-import           Data.List                       (foldl')
+import           Data.List                       (foldl', foldl1')
 import           Data.Maybe                      (fromMaybe)
 import qualified GHC.Hs                          as Hs
 import qualified SrcLoc                          as S
@@ -135,6 +135,39 @@ matchToAlignable _conf (S.L _ (Hs.Match _ _ _ _))  = Nothing
 
 
 --------------------------------------------------------------------------------
+multiWayIfToAlignable
+    :: Config
+    -> Hs.LHsExpr Hs.GhcPs
+    -> [Alignable S.RealSrcSpan]
+multiWayIfToAlignable conf (S.L _ (Hs.HsMultiIf _ grhss)) =
+  fromMaybe [] $ traverse (grhsToAlignable conf) grhss
+multiWayIfToAlignable _conf _ = []
+
+
+--------------------------------------------------------------------------------
+grhsToAlignable
+    :: Config
+    -> S.Located (Hs.GRHS Hs.GhcPs (Hs.LHsExpr Hs.GhcPs))
+    -> Maybe (Alignable S.RealSrcSpan)
+grhsToAlignable conf (S.L grhsloc (Hs.GRHS _ guards@(_ : _) body)) = do
+    guard $ cCases conf
+    let guardsLocs = map S.getLoc guards
+        bodyLoc    = S.getLoc body
+        left       = foldl1' S.combineSrcSpans guardsLocs
+    matchPos <- toRealSrcSpan grhsloc
+    leftPos  <- toRealSrcSpan left
+    bodyPos  <- toRealSrcSpan bodyLoc
+    Just $ Alignable
+        { aContainer = matchPos
+        , aLeft      = leftPos
+        , aRight     = bodyPos
+        , aRightLead = length "-> "
+        }
+grhsToAlignable _conf (S.L _ (Hs.XGRHS x))  = Hs.noExtCon x
+grhsToAlignable _conf (S.L _ _)             = Nothing
+
+
+--------------------------------------------------------------------------------
 step :: Maybe Int -> Config -> Step
 step maxColumns config = makeStep "Cases" $ \ls module' ->
     let changes
@@ -147,5 +180,6 @@ step maxColumns config = makeStep "Cases" $ \ls module' ->
         configured :: [Change String]
         configured = concat $
             [changes records recordToAlignable | cRecords config] ++
-            [changes everything (matchGroupToAlignable config)] in
+            [changes everything (matchGroupToAlignable config)] ++
+            [changes everything (multiWayIfToAlignable config)] in
     applyChanges configured ls
diff --git a/tests/Language/Haskell/Stylish/Step/SimpleAlign/Tests.hs b/tests/Language/Haskell/Stylish/Step/SimpleAlign/Tests.hs
index fa17784..5b502d1 100644
--- a/tests/Language/Haskell/Stylish/Step/SimpleAlign/Tests.hs
+++ b/tests/Language/Haskell/Stylish/Step/SimpleAlign/Tests.hs
@@ -31,6 +31,7 @@ tests = testGroup "Language.Haskell.Stylish.Step.SimpleAlign.Tests"
     , testCase "case 10" case10
     , testCase "case 11" case11
     , testCase "case 12" case12
+    , testCase "case 13" case13
     ]
 
 
@@ -199,3 +200,16 @@ case12 = assertSnippet (step Nothing defaultConfig {cCases = False}) input input
         , "  Just y -> 1"
         , "  Nothing -> 2"
         ]
+
+
+--------------------------------------------------------------------------------
+case13 :: Assertion
+case13 = assertSnippet (step Nothing defaultConfig)
+    [ "cond n = if"
+    , "    | n < 10, x <- 1 -> x"
+    , "    | otherwise -> 2"
+    ]
+    [ "cond n = if"
+    , "    | n < 10, x <- 1 -> x"
+    , "    | otherwise      -> 2"
+    ]