1--------------------------------------------------------------------------------
2{-# LANGUAGE RecordWildCards #-}
3{-# LANGUAGE TypeFamilies    #-}
4module Language.Haskell.Stylish.Step.SimpleAlign
5    ( Config (..)
6    , Align (..)
7    , defaultConfig
8    , step
9    ) where
10
11
12--------------------------------------------------------------------------------
13import           Data.Either                     (partitionEithers)
14import           Data.Foldable                   (toList)
15import           Data.List                       (foldl', foldl1', sortOn)
16import           Data.Maybe                      (fromMaybe)
17import qualified GHC.Hs                          as Hs
18import qualified SrcLoc                          as S
19
20
21--------------------------------------------------------------------------------
22import           Language.Haskell.Stylish.Align
23import           Language.Haskell.Stylish.Editor
24import           Language.Haskell.Stylish.Module
25import           Language.Haskell.Stylish.Step
26import           Language.Haskell.Stylish.Util
27
28
29--------------------------------------------------------------------------------
30data Config = Config
31    { cCases            :: Align
32    , cTopLevelPatterns :: Align
33    , cRecords          :: Align
34    , cMultiWayIf       :: Align
35    } deriving (Show)
36
37data Align
38    = Always
39    | Adjacent
40    | Never
41    deriving (Eq, Show)
42
43defaultConfig :: Config
44defaultConfig = Config
45    { cCases            = Always
46    , cTopLevelPatterns = Always
47    , cRecords          = Always
48    , cMultiWayIf       = Always
49    }
50
51groupAlign :: Align -> [Alignable S.RealSrcSpan] -> [[Alignable S.RealSrcSpan]]
52groupAlign a xs = case a of
53    Never    -> []
54    Adjacent -> byLine . sortOn (S.srcSpanStartLine . aLeft) $ xs
55    Always   -> [xs]
56  where
57    byLine = map toList . groupByLine aLeft
58
59
60--------------------------------------------------------------------------------
61type Record = [S.Located (Hs.ConDeclField Hs.GhcPs)]
62
63
64--------------------------------------------------------------------------------
65records :: S.Located (Hs.HsModule Hs.GhcPs) -> [Record]
66records modu = do
67  let decls           = map S.unLoc (Hs.hsmodDecls (S.unLoc modu))
68      tyClDecls       = [ tyClDecl | Hs.TyClD _ tyClDecl <- decls ]
69      dataDecls       = [ d | d@(Hs.DataDecl _ _ _ _ _)  <- tyClDecls ]
70      dataDefns       = map Hs.tcdDataDefn dataDecls
71  d@Hs.ConDeclH98 {} <- concatMap getConDecls dataDefns
72  case Hs.con_args d of
73      Hs.RecCon rec -> [S.unLoc rec]
74      _             -> []
75 where
76  getConDecls :: Hs.HsDataDefn Hs.GhcPs -> [Hs.ConDecl Hs.GhcPs]
77  getConDecls d@Hs.HsDataDefn {} = map S.unLoc $ Hs.dd_cons d
78  getConDecls (Hs.XHsDataDefn x) = Hs.noExtCon x
79
80
81--------------------------------------------------------------------------------
82recordToAlignable :: Config -> Record -> [[Alignable S.RealSrcSpan]]
83recordToAlignable conf = groupAlign (cRecords conf) . fromMaybe [] . traverse fieldDeclToAlignable
84
85
86--------------------------------------------------------------------------------
87fieldDeclToAlignable
88    :: S.Located (Hs.ConDeclField Hs.GhcPs) -> Maybe (Alignable S.RealSrcSpan)
89fieldDeclToAlignable (S.L _ (Hs.XConDeclField x)) = Hs.noExtCon x
90fieldDeclToAlignable (S.L matchLoc (Hs.ConDeclField _ names ty _)) = do
91  matchPos <- toRealSrcSpan matchLoc
92  leftPos  <- toRealSrcSpan $ S.getLoc $ last names
93  tyPos    <- toRealSrcSpan $ S.getLoc ty
94  Just $ Alignable
95    { aContainer = matchPos
96    , aLeft      = leftPos
97    , aRight     = tyPos
98    , aRightLead = length ":: "
99    }
100
101
102--------------------------------------------------------------------------------
103matchGroupToAlignable
104    :: Config
105    -> Hs.MatchGroup Hs.GhcPs (Hs.LHsExpr Hs.GhcPs)
106    -> [[Alignable S.RealSrcSpan]]
107matchGroupToAlignable _conf (Hs.XMatchGroup x) = Hs.noExtCon x
108matchGroupToAlignable conf (Hs.MG _ alts _) = cases' ++ patterns'
109  where
110    (cases, patterns) = partitionEithers . fromMaybe [] $ traverse matchToAlignable (S.unLoc alts)
111    cases' = groupAlign (cCases conf) cases
112    patterns' = groupAlign (cTopLevelPatterns conf) patterns
113
114
115--------------------------------------------------------------------------------
116matchToAlignable
117    :: S.Located (Hs.Match Hs.GhcPs (Hs.LHsExpr Hs.GhcPs))
118    -> Maybe (Either (Alignable S.RealSrcSpan) (Alignable S.RealSrcSpan))
119matchToAlignable (S.L matchLoc m@(Hs.Match _ Hs.CaseAlt pats@(_ : _) grhss)) = do
120  let patsLocs   = map S.getLoc pats
121      pat        = last patsLocs
122      guards     = getGuards m
123      guardsLocs = map S.getLoc guards
124      left       = foldl' S.combineSrcSpans pat guardsLocs
125  body     <- rhsBody grhss
126  matchPos <- toRealSrcSpan matchLoc
127  leftPos  <- toRealSrcSpan left
128  rightPos <- toRealSrcSpan $ S.getLoc body
129  Just . Left $ Alignable
130    { aContainer = matchPos
131    , aLeft      = leftPos
132    , aRight     = rightPos
133    , aRightLead = length "-> "
134    }
135matchToAlignable (S.L matchLoc (Hs.Match _ (Hs.FunRhs name _ _) pats@(_ : _) grhss)) = do
136  body <- unguardedRhsBody grhss
137  let patsLocs = map S.getLoc pats
138      nameLoc  = S.getLoc name
139      left     = last (nameLoc : patsLocs)
140      bodyLoc  = S.getLoc body
141  matchPos <- toRealSrcSpan matchLoc
142  leftPos  <- toRealSrcSpan left
143  bodyPos  <- toRealSrcSpan bodyLoc
144  Just . Right $ Alignable
145    { aContainer = matchPos
146    , aLeft      = leftPos
147    , aRight     = bodyPos
148    , aRightLead = length "= "
149    }
150matchToAlignable (S.L _ (Hs.XMatch x))      = Hs.noExtCon x
151matchToAlignable (S.L _ (Hs.Match _ _ _ _)) = Nothing
152
153
154--------------------------------------------------------------------------------
155multiWayIfToAlignable
156    :: Config
157    -> Hs.LHsExpr Hs.GhcPs
158    -> [[Alignable S.RealSrcSpan]]
159multiWayIfToAlignable conf (S.L _ (Hs.HsMultiIf _ grhss)) =
160    groupAlign (cMultiWayIf conf) as
161  where
162    as = fromMaybe [] $ traverse grhsToAlignable grhss
163multiWayIfToAlignable _conf _ = []
164
165
166--------------------------------------------------------------------------------
167grhsToAlignable
168    :: S.Located (Hs.GRHS Hs.GhcPs (Hs.LHsExpr Hs.GhcPs))
169    -> Maybe (Alignable S.RealSrcSpan)
170grhsToAlignable (S.L grhsloc (Hs.GRHS _ guards@(_ : _) body)) = do
171    let guardsLocs = map S.getLoc guards
172        bodyLoc    = S.getLoc body
173        left       = foldl1' S.combineSrcSpans guardsLocs
174    matchPos <- toRealSrcSpan grhsloc
175    leftPos  <- toRealSrcSpan left
176    bodyPos  <- toRealSrcSpan bodyLoc
177    Just $ Alignable
178        { aContainer = matchPos
179        , aLeft      = leftPos
180        , aRight     = bodyPos
181        , aRightLead = length "-> "
182        }
183grhsToAlignable (S.L _ (Hs.XGRHS x)) = Hs.noExtCon x
184grhsToAlignable (S.L _ _)            = Nothing
185
186
187--------------------------------------------------------------------------------
188step :: Maybe Int -> Config -> Step
189step maxColumns config@(Config {..}) = makeStep "Cases" $ \ls module' ->
190    let changes
191            :: (S.Located (Hs.HsModule Hs.GhcPs) -> [a])
192            -> (a -> [[Alignable S.RealSrcSpan]])
193            -> [Change String]
194        changes search toAlign =
195            (concatMap . concatMap) (align maxColumns) . map toAlign $ search (parsedModule module')
196
197        configured :: [Change String]
198        configured = concat $
199            [changes records (recordToAlignable config)] ++
200            [changes everything (matchGroupToAlignable config)] ++
201            [changes everything (multiWayIfToAlignable config)] in
202    applyChanges configured ls
203