Support alignment of cases with a single guard

2 files changed, 48 insertions, 11 deletions

diff --git a/lib/Language/Haskell/Stylish/Step/SimpleAlign.hs b/lib/Language/Haskell/Stylish/Step/SimpleAlign.hs
index c83c482..924d6c5 100644
--- a/lib/Language/Haskell/Stylish/Step/SimpleAlign.hs
+++ b/lib/Language/Haskell/Stylish/Step/SimpleAlign.hs
@@ -8,6 +8,7 @@ module Language.Haskell.Stylish.Step.SimpleAlign
 
 --------------------------------------------------------------------------------
 import           Data.Data                       (Data)
+import           Data.List                       (foldl')
 import           Data.Maybe                      (maybeToList)
 import qualified Language.Haskell.Exts           as H
 
@@ -42,14 +43,30 @@ 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 "-> "
-    }
+-- | For this to work well, we require a way to merge annotations.  This merge
+-- operation should follow the semigroup laws.
+altToAlignable :: (l -> l -> l) -> H.Alt l -> Maybe (Alignable l)
+altToAlignable _ (H.Alt _   _   _   (Just _)) = Nothing
+altToAlignable _ (H.Alt ann pat rhs@(H.UnGuardedRhs _ _) Nothing) = Just $
+    Alignable
+        { aContainer = ann
+        , aLeft      = H.ann pat
+        , aRight     = H.ann rhs
+        , aRightLead = length "-> "
+        }
+altToAlignable
+        merge
+        (H.Alt ann pat (H.GuardedRhss _ [H.GuardedRhs _ guards rhs]) Nothing) =
+    -- We currently only support the case where an alternative has a single
+    -- guarded RHS.  If there are more, we would need to return multiple
+    -- `Alignable`s from this function, which would be a significant change.
+    Just $ Alignable
+        { aContainer = ann
+        , aLeft      = foldl' merge (H.ann pat) (map H.ann guards)
+        , aRight     = H.ann rhs
+        , aRightLead = length "-> "
+        }
+altToAlignable _ _ = Nothing
 
 
 --------------------------------------------------------------------------------
@@ -101,9 +118,11 @@ step maxColumns config = makeStep "Cases" $ \ls (module', _) ->
             , change_ <- align maxColumns aligns
             ]
 
-        configured             = concat $
-            [changes cases   altToAlignable       | cCases config] ++
-            [changes tlpats  matchToAlignable     | cTopLevelPatterns config] ++
+        configured = concat $
+            [ changes cases (altToAlignable H.mergeSrcSpan)
+            | cCases config
+            ] ++
+            [changes tlpats  matchToAlignable | cTopLevelPatterns config] ++
             [changes records fieldDeclToAlignable | cRecords 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 a57e6e9..b8afab4 100644
--- a/tests/Language/Haskell/Stylish/Step/SimpleAlign/Tests.hs
+++ b/tests/Language/Haskell/Stylish/Step/SimpleAlign/Tests.hs
@@ -25,6 +25,7 @@ tests = testGroup "Language.Haskell.Stylish.Step.SimpleAlign.Tests"
     , testCase "case 05" case05
     , testCase "case 06" case06
     , testCase "case 07" case07
+    , testCase "case 08" case08
     ]
 
 
@@ -148,3 +149,20 @@ case07 =
         , "    , barqux :: Int"
         , "    }"
         ]
+
+
+--------------------------------------------------------------------------------
+case08 :: Assertion
+case08 = expected @=? testStep (step 80 defaultConfig) input
+  where
+    input = unlines
+        [ "canDrink mbAge = case mbAge of"
+        , "    Just age | age > 18 -> True"
+        , "    _ -> False"
+        ]
+
+    expected = unlines
+        [ "canDrink mbAge = case mbAge of"
+        , "    Just age | age > 18 -> True"
+        , "    _                   -> False"
+        ]