Align options for imports

3 files changed, 153 insertions, 59 deletions

diff --git a/src/Language/Haskell/Stylish/Config.hs b/src/Language/Haskell/Stylish/Config.hs
index 0304ae5..bda5593 100644
--- a/src/Language/Haskell/Stylish/Config.hs
+++ b/src/Language/Haskell/Stylish/Config.hs
@@ -10,23 +10,25 @@ module Language.Haskell.Stylish.Config
 
 
 --------------------------------------------------------------------------------
-import           Control.Applicative                    (pure, (<$>), (<*>))
-import           Control.Monad                          (forM, mzero)
-import           Data.Aeson                             (FromJSON (..))
-import qualified Data.Aeson                             as A
-import qualified Data.Aeson.Types                       as A
-import qualified Data.ByteString                        as B
-import           Data.List                              (inits, intercalate)
-import           Data.Map                               (Map)
-import qualified Data.Map                               as M
-import           Data.Yaml                              (decodeEither)
+import           Control.Applicative                              (pure, (<$>),
+                                                                   (<*>))
+import           Control.Monad                                    (forM, mzero)
+import           Data.Aeson                                       (FromJSON (..))
+import qualified Data.Aeson                                       as A
+import qualified Data.Aeson.Types                                 as A
+import qualified Data.ByteString                                  as B
+import           Data.List                                        (inits,
+                                                                   intercalate)
+import           Data.Map                                         (Map)
+import qualified Data.Map                                         as M
+import           Data.Yaml                                        (decodeEither)
 import           System.Directory
-import           System.FilePath                        (joinPath, splitPath,
-                                                         (</>))
+import           System.FilePath                                  (joinPath,
+                                                                   splitPath,
+                                                                   (</>))
 
 
 --------------------------------------------------------------------------------
-import           Paths_stylish_haskell                  (getDataFileName)
 import           Language.Haskell.Stylish.Step
 import qualified Language.Haskell.Stylish.Step.Imports            as Imports
 import qualified Language.Haskell.Stylish.Step.LanguagePragmas    as LanguagePragmas
@@ -35,6 +37,7 @@ import qualified Language.Haskell.Stylish.Step.Tabs               as Tabs
 import qualified Language.Haskell.Stylish.Step.TrailingWhitespace as TrailingWhitespace
 import qualified Language.Haskell.Stylish.Step.UnicodeSyntax      as UnicodeSyntax
 import           Language.Haskell.Stylish.Verbose
+import           Paths_stylish_haskell                            (getDataFileName)
 
 
 --------------------------------------------------------------------------------
@@ -161,7 +164,11 @@ parseEnum strs _   (Just k) = case lookup k strs of
 parseImports :: Config -> A.Object -> A.Parser Step
 parseImports config o = Imports.step
     <$> pure (configColumns config)
-    <*> (o A..:? "align" >>= parseEnum aligns Imports.Global)
+    <*> (Imports.Align
+        <$> (o A..:? "align" >>= parseEnum aligns Imports.Global)
+        <*> (o A..:? "list_align" >>= parseEnum listAligns Imports.SameLine)
+        <*> (o A..:? "long_list_align"
+            >>= parseEnum longListAligns Imports.Inline))
   where
     aligns =
         [ ("global", Imports.Global)
@@ -170,6 +177,16 @@ parseImports config o = Imports.step
         , ("none",   Imports.None)
         ]
 
+    listAligns =
+        [ ("new line",  Imports.NewLine)
+        , ("same line", Imports.SameLine)
+        ]
+
+    longListAligns =
+        [ ("inline",    Imports.Inline)
+        , ("multiline", Imports.Multiline)
+        ]
+
 
 --------------------------------------------------------------------------------
 parseLanguagePragmas :: Config -> A.Object -> A.Parser Step
@@ -181,7 +198,8 @@ parseLanguagePragmas config o = LanguagePragmas.step
     styles =
         [ ("vertical",     LanguagePragmas.Vertical)
         , ("compact",      LanguagePragmas.Compact)
-        , ("compact_line", LanguagePragmas.CompactLine)]
+        , ("compact_line", LanguagePragmas.CompactLine)
+        ]
 
 
 --------------------------------------------------------------------------------
diff --git a/src/Language/Haskell/Stylish/Step/Imports.hs b/src/Language/Haskell/Stylish/Step/Imports.hs
index 486fba7..b1e8843 100644
--- a/src/Language/Haskell/Stylish/Step/Imports.hs
+++ b/src/Language/Haskell/Stylish/Step/Imports.hs
@@ -1,6 +1,10 @@
+{-# LANGUAGE RecordWildCards #-}
 --------------------------------------------------------------------------------
 module Language.Haskell.Stylish.Step.Imports
     ( Align (..)
+    , ImportAlign (..)
+    , ListAlign (..)
+    , LongListAlign (..)
     , step
     ) where
 
@@ -22,13 +26,29 @@ import           Language.Haskell.Stylish.Util
 
 
 --------------------------------------------------------------------------------
-data Align
+data Align = Align
+    { importAlign   :: ImportAlign
+    , listAlign     :: ListAlign
+    , longListAlign :: LongListAlign
+    }
+    deriving (Eq, Show)
+
+data ImportAlign
     = Global
     | File
     | Group
     | None
     deriving (Eq, Show)
 
+data ListAlign
+    = SameLine
+    | NewLine
+    deriving (Eq, Show)
+
+data LongListAlign
+    = Inline
+    | Multiline
+    deriving (Eq, Show)
 
 --------------------------------------------------------------------------------
 imports :: H.Module l -> [H.ImportDecl l]
@@ -95,28 +115,52 @@ prettyImportSpec x                      = H.prettyPrint x
 
 
 --------------------------------------------------------------------------------
-prettyImport :: Int -> Bool -> Bool -> Int -> H.ImportDecl l -> [String]
-prettyImport columns padQualified padName longest imp =
-    wrap columns base (length base + 2) $
-    (if hiding then ("hiding" :) else id) $
-    case importSpecs of
-        Nothing -> []     -- Import everything
-        Just [] -> ["()"] -- Instance only imports
-        Just is ->
-            withInit (++ ",") $
-            withHead ("(" ++) $
-            withLast (++ ")") $
-            map prettyImportSpec $
-            is
+prettyImport :: Int -> Align -> Bool -> Bool -> Int -> H.ImportDecl l
+             -> [String]
+prettyImport columns Align{..} padQualified padName longest imp =
+    case longListAlign of
+        Inline -> inlineWrap
+        Multiline -> if listAlign == NewLine || length inlineWrap > 1
+            then multilineWrap
+            else inlineWrap
   where
-    base = unwords $ concat
-         [ ["import"]
-         , qualified
-         , (fmap show $ maybeToList $ H.importPkg imp)
-         , [(if hasExtras && padName then padRight longest else id)
-            (importName imp)]
-         , ["as " ++ as | H.ModuleName _ as <- maybeToList $ H.importAs imp]
-         ]
+    inlineWrap = inlineWrapper
+        $ withInit (++ ",")
+        $ withHead ("(" ++)
+        $ withLast (++ ")")
+        $ specs
+
+    inlineWrapper = case listAlign of
+        SameLine -> wrap columns inlineBase (inlineBaseLength + 1)
+        NewLine  -> (inlineBase :) . wrapRest columns 4
+
+    multilineWrap = multilineBase : (wrapRest 0 4
+        $ (withHead ("( " ++)
+        $ withTail (", " ++)
+        $ specs) ++ [")"])
+
+    inlineBase = base $ padImport $ importName imp
+
+    multilineBase = base $ importName imp
+
+    padImport = if hasExtras && padName
+        then padRight longest
+        else id
+
+    base' baseName importAs hiding' = unwords $ concat $ filter (not . null)
+        [ ["import"]
+        , qualified
+        , (fmap show $ maybeToList $ H.importPkg imp)
+        , [baseName]
+        , importAs
+        , hiding'
+        ]
+
+    base baseName = base' baseName
+        ["as " ++ as | H.ModuleName _ as <- maybeToList $ H.importAs imp]
+        (if hiding then (["hiding"]) else [])
+
+    inlineBaseLength = length $ base' (padImport $ importName imp) [] []
 
     (hiding, importSpecs) = case H.importSpecs imp of
         Just (H.ImportSpecList _ h l) -> (h, Just l)
@@ -129,21 +173,29 @@ prettyImport columns padQualified padName longest imp =
         | padQualified          = ["         "]
         | otherwise             = []
 
+    specs = case importSpecs of
+        Nothing -> []     -- Import everything
+        Just [] -> ["()"] -- Instance only imports
+        Just is -> map prettyImportSpec is
+
 
 --------------------------------------------------------------------------------
-prettyImportGroup :: Int -> Align -> Bool -> Int -> [H.ImportDecl LineBlock]
+prettyImportGroup :: Int -> Align -> Bool -> Int
+                  -> [H.ImportDecl LineBlock]
                   -> Lines
 prettyImportGroup columns align fileAlign longest imps =
-    concatMap (prettyImport columns padQual padName longest') $
+    concatMap (prettyImport columns align padQual padName longest') $
     sortBy compareImports imps
   where
-    longest' = case align of
+    align' = importAlign align
+
+    longest' = case align' of
         Group -> longestImport imps
         _     -> longest
 
-    padName = align /= None
+    padName = align' /= None
 
-    padQual = case align of
+    padQual = case align' of
         Global -> True
         File   -> fileAlign
         Group  -> any H.importQualified imps
@@ -167,6 +219,6 @@ step' columns align ls (module', _) = flip applyChanges ls
     longest = longestImport imps
     groups  = groupAdjacent [(H.ann i, i) | i <- imps]
 
-    fileAlign = case align of
+    fileAlign = case importAlign align of
         File -> any H.importQualified imps
         _    -> False
diff --git a/src/Language/Haskell/Stylish/Util.hs b/src/Language/Haskell/Stylish/Util.hs
index 004c3f1..f94c356 100644
--- a/src/Language/Haskell/Stylish/Util.hs
+++ b/src/Language/Haskell/Stylish/Util.hs
@@ -6,10 +6,12 @@ module Language.Haskell.Stylish.Util
     , everything
     , infoPoints
     , wrap
+    , wrapRest
 
     , withHead
-    , withLast
     , withInit
+    , withTail
+    , withLast
     ) where
 
 
@@ -34,7 +36,12 @@ nameToString (H.Symbol _ str) = str
 
 --------------------------------------------------------------------------------
 indent :: Int -> String -> String
-indent len str = replicate len ' ' ++ str
+indent len = (indentPrefix len ++)
+
+
+--------------------------------------------------------------------------------
+indentPrefix :: Int -> String
+indentPrefix = (`replicate` ' ')
 
 
 --------------------------------------------------------------------------------
@@ -58,21 +65,33 @@ wrap :: Int       -- ^ Maximum line width
      -> Int       -- ^ Indentation
      -> [String]  -- ^ Strings to add/wrap
      -> Lines     -- ^ Resulting lines
-wrap maxWidth leading ind strs =
-    let (ls, curr, _) = foldl step ([], leading, length leading) strs
-    in ls ++ [curr]
+wrap maxWidth leading ind strs' = wrap' leading strs'
   where
-    -- TODO: In order to optimize this, use a difference list instead of a
-    -- regular list for 'ls'.
-    step (ls, curr, width) str
-        | nextLine  = (ls ++ [curr], indent ind str, ind + len)
-        | otherwise = (ls, curr ++ " " ++ str, width')
-      where
-        -- Put it on the next line if it would make the current line too long,
-        -- AND if it doesn't make the next line too long.
-        nextLine = width' > maxWidth && ind + len <= maxWidth
-        len      = length str
-        width'   = width + 1 + len
+    wrap' ss [] = [ss]
+    wrap' ss (str:strs)
+        | overflows ss str =
+            ss : wrapRest maxWidth ind (str:strs)
+        | otherwise = wrap' (ss ++ " " ++ str) strs
+
+    overflows ss str = (length ss + length str) >= maxWidth
+
+
+--------------------------------------------------------------------------------
+wrapRest :: Int
+         -> Int
+         -> [String]
+         -> Lines
+wrapRest maxWidth ind = reverse . wrapRest' [] ""
+  where
+    wrapRest' ls ss []
+        | null ss = ls
+        | otherwise = ss:ls
+    wrapRest' ls ss (str:strs)
+        | overflows ss str = wrapRest' (ss:ls) "" (str:strs)
+        | null ss = wrapRest' ls (indent ind str) strs
+        | otherwise = wrapRest' ls (ss ++ " " ++ str) strs
+
+    overflows ss str = (length ss + length str) >= maxWidth && length ss > ind
 
 
 --------------------------------------------------------------------------------
@@ -93,3 +112,8 @@ withInit :: (a -> a) -> [a] -> [a]
 withInit _ []       = []
 withInit _ (x : []) = [x]
 withInit f (x : xs) = f x : withInit f xs
+
+--------------------------------------------------------------------------------
+withTail :: (a -> a) -> [a] -> [a]
+withTail _ [] = []
+withTail f (x : xs) = x : map f xs