Add new option for aligning groups of adjacent items

Co-authored-by: 1computer1 <onecomputer00@gmail.com>

4 files changed, 114 insertions, 83 deletions

diff --git a/lib/Language/Haskell/Stylish/Config.hs b/lib/Language/Haskell/Stylish/Config.hs
index 682d7d7..dde9d0d 100644
--- a/lib/Language/Haskell/Stylish/Config.hs
+++ b/lib/Language/Haskell/Stylish/Config.hs
@@ -1,5 +1,5 @@
 --------------------------------------------------------------------------------
-{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE BlockArguments    #-}
 {-# LANGUAGE LambdaCase        #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE TemplateHaskell   #-}
@@ -10,10 +10,12 @@ module Language.Haskell.Stylish.Config
     , defaultConfigBytes
     , configFilePath
     , loadConfig
+    , parseConfig
     ) where
 
 
 --------------------------------------------------------------------------------
+import           Control.Applicative                              ((<|>))
 import           Control.Monad                                    (forM, mzero)
 import           Data.Aeson                                       (FromJSON (..))
 import qualified Data.Aeson                                       as A
@@ -43,8 +45,8 @@ import           Language.Haskell.Stylish.Config.Internal
 import           Language.Haskell.Stylish.Step
 import qualified Language.Haskell.Stylish.Step.Data               as Data
 import qualified Language.Haskell.Stylish.Step.Imports            as Imports
-import qualified Language.Haskell.Stylish.Step.ModuleHeader       as ModuleHeader
 import qualified Language.Haskell.Stylish.Step.LanguagePragmas    as LanguagePragmas
+import qualified Language.Haskell.Stylish.Step.ModuleHeader       as ModuleHeader
 import qualified Language.Haskell.Stylish.Step.SimpleAlign        as SimpleAlign
 import qualified Language.Haskell.Stylish.Step.Squash             as Squash
 import qualified Language.Haskell.Stylish.Step.Tabs               as Tabs
@@ -74,7 +76,7 @@ data ExitCodeBehavior
   deriving (Eq)
 
 instance Show ExitCodeBehavior where
-  show NormalExitBehavior = "normal"
+  show NormalExitBehavior        = "normal"
   show ErrorOnFormatExitBehavior = "error_on_format"
 
 --------------------------------------------------------------------------------
@@ -206,12 +208,22 @@ parseSimpleAlign :: Config -> A.Object -> A.Parser Step
 parseSimpleAlign c o = SimpleAlign.step
     <$> pure (configColumns c)
     <*> (SimpleAlign.Config
-        <$> withDef SimpleAlign.cCases            "cases"
-        <*> withDef SimpleAlign.cTopLevelPatterns "top_level_patterns"
-        <*> withDef SimpleAlign.cRecords          "records"
-        <*> withDef SimpleAlign.cMultiWayIf       "multi_way_if")
+        <$> parseAlign "cases"              SimpleAlign.cCases
+        <*> parseAlign "top_level_patterns" SimpleAlign.cTopLevelPatterns
+        <*> parseAlign "records"            SimpleAlign.cRecords
+        <*> parseAlign "multi_way_if"       SimpleAlign.cMultiWayIf)
   where
-    withDef f k = fromMaybe (f SimpleAlign.defaultConfig) <$> (o A..:? k)
+    parseAlign key f =
+        (o A..:? key >>= parseEnum aligns (f SimpleAlign.defaultConfig)) <|>
+        (boolToAlign <$> o A..: key)
+    aligns =
+        [ ("always",   SimpleAlign.Always)
+        , ("adjacent", SimpleAlign.Adjacent)
+        , ("never",    SimpleAlign.Never)
+        ]
+    boolToAlign True  = SimpleAlign.Always
+    boolToAlign False = SimpleAlign.Never
+
 
 --------------------------------------------------------------------------------
 parseRecords :: Config -> A.Object -> A.Parser Step
@@ -295,8 +307,8 @@ parseImports config o = fmap (Imports.step columns) $ Imports.Options
 
     parseListPadding = \case
         A.String "module_name" -> pure Imports.LPModuleName
-        A.Number n | n >= 1 -> pure $ Imports.LPConstant (truncate n)
-        v -> A.typeMismatch "'module_name' or >=1 number" v
+        A.Number n | n >= 1    -> pure $ Imports.LPConstant (truncate n)
+        v                      -> A.typeMismatch "'module_name' or >=1 number" v
 
 --------------------------------------------------------------------------------
 parseLanguagePragmas :: Config -> A.Object -> A.Parser Step
diff --git a/lib/Language/Haskell/Stylish/GHC.hs b/lib/Language/Haskell/Stylish/GHC.hs
index ee2d59f..c99d4bf 100644
--- a/lib/Language/Haskell/Stylish/GHC.hs
+++ b/lib/Language/Haskell/Stylish/GHC.hs
@@ -6,6 +6,7 @@ module Language.Haskell.Stylish.GHC
   , dropBeforeLocated
   , dropBeforeAndAfter
     -- * Unsafe getters
+  , unsafeGetRealSrcSpan
   , getEndLineUnsafe
   , getStartLineUnsafe
     -- * Standard settings
@@ -18,32 +19,33 @@ module Language.Haskell.Stylish.GHC
   ) where
 
 --------------------------------------------------------------------------------
-import           Data.Function (on)
+import           Data.Function     (on)
 
 --------------------------------------------------------------------------------
-import           DynFlags                        (Settings(..), defaultDynFlags)
-import qualified DynFlags                        as GHC
-import           FileSettings                    (FileSettings(..))
-import           GHC.Fingerprint                 (fingerprint0)
+import           DynFlags          (Settings (..), defaultDynFlags)
+import qualified DynFlags          as GHC
+import           FileSettings      (FileSettings (..))
+import           GHC.Fingerprint   (fingerprint0)
 import           GHC.Platform
-import           GHC.Version                     (cProjectVersion)
-import           GhcNameVersion                  (GhcNameVersion(..))
-import           PlatformConstants               (PlatformConstants(..))
-import           SrcLoc                          (GenLocated(..), SrcSpan(..))
-import           SrcLoc                          (Located, RealLocated)
-import           SrcLoc                          (srcSpanStartLine, srcSpanEndLine)
-import           ToolSettings                    (ToolSettings(..))
-import qualified Outputable                      as GHC
+import           GHC.Version       (cProjectVersion)
+import           GhcNameVersion    (GhcNameVersion (..))
+import qualified Outputable        as GHC
+import           PlatformConstants (PlatformConstants (..))
+import           SrcLoc            (GenLocated (..), Located, RealLocated,
+                                    RealSrcSpan, SrcSpan (..), srcSpanEndLine,
+                                    srcSpanStartLine)
+import           ToolSettings      (ToolSettings (..))
+
+unsafeGetRealSrcSpan :: Located a -> RealSrcSpan
+unsafeGetRealSrcSpan = \case
+  (L (RealSrcSpan s) _) -> s
+  _                     -> error "could not get source code location"
 
 getStartLineUnsafe :: Located a -> Int
-getStartLineUnsafe = \case
-  (L (RealSrcSpan s) _) -> srcSpanStartLine s
-  _ -> error "could not get start line of block"
+getStartLineUnsafe = srcSpanStartLine . unsafeGetRealSrcSpan
 
 getEndLineUnsafe :: Located a -> Int
-getEndLineUnsafe = \case
-  (L (RealSrcSpan s) _) -> srcSpanEndLine s
-  _ -> error "could not get end line of block"
+getEndLineUnsafe = srcSpanEndLine . unsafeGetRealSrcSpan
 
 dropAfterLocated :: Maybe (Located a) -> [RealLocated b] -> [RealLocated b]
 dropAfterLocated loc xs = case loc of
diff --git a/lib/Language/Haskell/Stylish/Module.hs b/lib/Language/Haskell/Stylish/Module.hs
index 2cc8f47..3dbebe0 100644
--- a/lib/Language/Haskell/Stylish/Module.hs
+++ b/lib/Language/Haskell/Stylish/Module.hs
@@ -24,6 +24,7 @@ module Language.Haskell.Stylish.Module
   , moduleComments
   , moduleLanguagePragmas
   , queryModule
+  , groupByLine
 
     -- * Imports
   , canMergeImport
@@ -192,22 +193,21 @@ moduleImports m
 
 -- | Get groups of imports from module
 moduleImportGroups :: Module -> [NonEmpty (Located Import)]
-moduleImportGroups = go [] Nothing . moduleImports
+moduleImportGroups = groupByLine unsafeGetRealSrcSpan . moduleImports
+
+-- The same logic as 'Language.Haskell.Stylish.Module.moduleImportGroups'.
+groupByLine :: (a -> RealSrcSpan) -> [a] -> [NonEmpty a]
+groupByLine f = go [] Nothing
   where
-    -- Run through all imports (assume they are sorted already in order of
-    -- appearance in the file) and group the ones that are on consecutive
-    -- lines.
-    go  :: [Located Import] -> Maybe Int -> [Located Import]
-        -> [NonEmpty (Located Import)]
-    go acc _        []                   = ne acc
-    go acc mbCurrentLine (imp : impRest) =
-        let
-          lStart = getStartLineUnsafe imp
-          lEnd = getEndLineUnsafe imp in
-        case mbCurrentLine of
-            Just lPrevEnd | lPrevEnd + 1 < lStart
-              -> ne acc ++ go [imp] (Just lEnd) impRest
-            _ -> go (acc ++ [imp]) (Just lEnd) impRest
+    go acc _ [] = ne acc
+    go acc mbCurrentLine (x:xs) =
+      let
+        lStart = GHC.srcSpanStartLine (f x)
+        lEnd = GHC.srcSpanEndLine (f x) in
+      case mbCurrentLine of
+        Just lPrevEnd | lPrevEnd + 1 < lStart
+          -> ne acc ++ go [x] (Just lEnd) xs
+        _ -> go (acc ++ [x]) (Just lEnd) xs
 
     ne []       = []
     ne (x : xs) = [x :| xs]
diff --git a/lib/Language/Haskell/Stylish/Step/SimpleAlign.hs b/lib/Language/Haskell/Stylish/Step/SimpleAlign.hs
index 523a6fb..f8aea50 100644
--- a/lib/Language/Haskell/Stylish/Step/SimpleAlign.hs
+++ b/lib/Language/Haskell/Stylish/Step/SimpleAlign.hs
@@ -3,14 +3,16 @@
 {-# LANGUAGE TypeFamilies    #-}
 module Language.Haskell.Stylish.Step.SimpleAlign
     ( Config (..)
+    , Align (..)
     , defaultConfig
     , step
     ) where
 
 
 --------------------------------------------------------------------------------
-import           Control.Monad                   (guard)
-import           Data.List                       (foldl', foldl1')
+import           Data.Either                     (partitionEithers)
+import           Data.Foldable                   (toList)
+import           Data.List                       (foldl', foldl1', sortOn)
 import           Data.Maybe                      (fromMaybe)
 import qualified GHC.Hs                          as Hs
 import qualified SrcLoc                          as S
@@ -26,22 +28,34 @@ import           Language.Haskell.Stylish.Util
 
 --------------------------------------------------------------------------------
 data Config = Config
-    { cCases            :: !Bool
-    , cTopLevelPatterns :: !Bool
-    , cRecords          :: !Bool
-    , cMultiWayIf       :: !Bool
+    { cCases            :: Align
+    , cTopLevelPatterns :: Align
+    , cRecords          :: Align
+    , cMultiWayIf       :: Align
     } deriving (Show)
 
+data Align
+    = Always
+    | Adjacent
+    | Never
+    deriving (Eq, Show)
 
---------------------------------------------------------------------------------
 defaultConfig :: Config
 defaultConfig = Config
-    { cCases            = True
-    , cTopLevelPatterns = True
-    , cRecords          = True
-    , cMultiWayIf       = True
+    { cCases            = Always
+    , cTopLevelPatterns = Always
+    , cRecords          = Always
+    , cMultiWayIf       = Always
     }
 
+groupAlign :: Align -> [Alignable S.RealSrcSpan] -> [[Alignable S.RealSrcSpan]]
+groupAlign a xs = case a of
+    Never    -> []
+    Adjacent -> byLine . sortOn (S.srcSpanStartLine . aLeft) $ xs
+    Always   -> [xs]
+  where
+    byLine = map toList . groupByLine aLeft
+
 
 --------------------------------------------------------------------------------
 type Record = [S.Located (Hs.ConDeclField Hs.GhcPs)]
@@ -65,8 +79,8 @@ records modu = do
 
 
 --------------------------------------------------------------------------------
-recordToAlignable :: Record -> [Alignable S.RealSrcSpan]
-recordToAlignable = fromMaybe [] . traverse fieldDeclToAlignable
+recordToAlignable :: Config -> Record -> [[Alignable S.RealSrcSpan]]
+recordToAlignable conf = groupAlign (cRecords conf) . fromMaybe [] . traverse fieldDeclToAlignable
 
 
 --------------------------------------------------------------------------------
@@ -89,36 +103,36 @@ fieldDeclToAlignable (S.L matchLoc (Hs.ConDeclField _ names ty _)) = do
 matchGroupToAlignable
     :: Config
     -> Hs.MatchGroup Hs.GhcPs (Hs.LHsExpr Hs.GhcPs)
-    -> [Alignable S.RealSrcSpan]
+    -> [[Alignable S.RealSrcSpan]]
 matchGroupToAlignable _conf (Hs.XMatchGroup x) = Hs.noExtCon x
-matchGroupToAlignable conf (Hs.MG _ alts _) =
-  fromMaybe [] $ traverse (matchToAlignable conf) (S.unLoc alts)
+matchGroupToAlignable conf (Hs.MG _ alts _) = cases' ++ patterns'
+  where
+    (cases, patterns) = partitionEithers . fromMaybe [] $ traverse matchToAlignable (S.unLoc alts)
+    cases' = groupAlign (cCases conf) cases
+    patterns' = groupAlign (cTopLevelPatterns conf) patterns
 
 
 --------------------------------------------------------------------------------
 matchToAlignable
-    :: Config
-    -> S.Located (Hs.Match Hs.GhcPs (Hs.LHsExpr Hs.GhcPs))
-    -> Maybe (Alignable S.RealSrcSpan)
-matchToAlignable conf (S.L matchLoc m@(Hs.Match _ Hs.CaseAlt pats@(_ : _) grhss)) = do
+    :: S.Located (Hs.Match Hs.GhcPs (Hs.LHsExpr Hs.GhcPs))
+    -> Maybe (Either (Alignable S.RealSrcSpan) (Alignable S.RealSrcSpan))
+matchToAlignable (S.L matchLoc m@(Hs.Match _ Hs.CaseAlt pats@(_ : _) grhss)) = do
   let patsLocs   = map S.getLoc pats
       pat        = last patsLocs
       guards     = getGuards m
       guardsLocs = map S.getLoc guards
       left       = foldl' S.combineSrcSpans pat guardsLocs
-  guard $ cCases conf
   body     <- rhsBody grhss
   matchPos <- toRealSrcSpan matchLoc
   leftPos  <- toRealSrcSpan left
   rightPos <- toRealSrcSpan $ S.getLoc body
-  Just $ Alignable
+  Just . Left $ Alignable
     { aContainer = matchPos
     , aLeft      = leftPos
     , aRight     = rightPos
     , aRightLead = length "-> "
     }
-matchToAlignable conf (S.L matchLoc (Hs.Match _ (Hs.FunRhs name _ _) pats@(_ : _) grhss)) = do
-  guard $ cTopLevelPatterns conf
+matchToAlignable (S.L matchLoc (Hs.Match _ (Hs.FunRhs name _ _) pats@(_ : _) grhss)) = do
   body <- unguardedRhsBody grhss
   let patsLocs = map S.getLoc pats
       nameLoc  = S.getLoc name
@@ -127,23 +141,26 @@ matchToAlignable conf (S.L matchLoc (Hs.Match _ (Hs.FunRhs name _ _) pats@(_ : _
   matchPos <- toRealSrcSpan matchLoc
   leftPos  <- toRealSrcSpan left
   bodyPos  <- toRealSrcSpan bodyLoc
-  Just $ Alignable
+  Just . Right $ Alignable
     { aContainer = matchPos
     , aLeft      = leftPos
     , aRight     = bodyPos
     , aRightLead = length "= "
     }
-matchToAlignable _conf (S.L _ (Hs.XMatch x))       = Hs.noExtCon x
-matchToAlignable _conf (S.L _ (Hs.Match _ _ _ _))  = Nothing
+matchToAlignable (S.L _ (Hs.XMatch x))      = Hs.noExtCon x
+matchToAlignable (S.L _ (Hs.Match _ _ _ _)) = Nothing
 
 
 --------------------------------------------------------------------------------
 multiWayIfToAlignable
-    :: Hs.LHsExpr Hs.GhcPs
-    -> [Alignable S.RealSrcSpan]
-multiWayIfToAlignable  (S.L _ (Hs.HsMultiIf _ grhss)) =
-  fromMaybe [] $ traverse grhsToAlignable grhss
-multiWayIfToAlignable _ = []
+    :: Config
+    -> Hs.LHsExpr Hs.GhcPs
+    -> [[Alignable S.RealSrcSpan]]
+multiWayIfToAlignable conf (S.L _ (Hs.HsMultiIf _ grhss)) =
+    groupAlign (cMultiWayIf conf) as
+  where
+    as = fromMaybe [] $ traverse grhsToAlignable grhss
+multiWayIfToAlignable _conf _ = []
 
 
 --------------------------------------------------------------------------------
@@ -163,8 +180,8 @@ grhsToAlignable (S.L grhsloc (Hs.GRHS _ guards@(_ : _) body)) = do
         , aRight     = bodyPos
         , aRightLead = length "-> "
         }
-grhsToAlignable (S.L _ (Hs.XGRHS x))  = Hs.noExtCon x
-grhsToAlignable (S.L _ _)             = Nothing
+grhsToAlignable (S.L _ (Hs.XGRHS x)) = Hs.noExtCon x
+grhsToAlignable (S.L _ _)            = Nothing
 
 
 --------------------------------------------------------------------------------
@@ -172,14 +189,14 @@ step :: Maybe Int -> Config -> Step
 step maxColumns config@(Config {..}) = makeStep "Cases" $ \ls module' ->
     let changes
             :: (S.Located (Hs.HsModule Hs.GhcPs) -> [a])
-            -> (a -> [Alignable S.RealSrcSpan])
+            -> (a -> [[Alignable S.RealSrcSpan]])
             -> [Change String]
-        changes search toAlign = concat $
-            map (align maxColumns) . map toAlign $ search (parsedModule module')
+        changes search toAlign =
+            (concatMap . concatMap) (align maxColumns) . map toAlign $ search (parsedModule module')
 
         configured :: [Change String]
         configured = concat $
-            [changes records recordToAlignable | cRecords ] ++
+            [changes records (recordToAlignable config)] ++
             [changes everything (matchGroupToAlignable config)] ++
-            [changes everything multiWayIfToAlignable | cMultiWayIf] in
+            [changes everything (multiWayIfToAlignable config)] in
     applyChanges configured ls