{-# LANGUAGE TupleSections #-}

module Floskell.Imports ( sortImports, groupImports, splitImports ) where

import           Control.Monad.Trans.State    ( State, execState, gets, modify )

import           Data.Function                ( on )
import           Data.List
                 ( groupBy, inits, intercalate, sortOn, sortOn, unfoldr )
import qualified Data.Map                     as M
import           Data.Monoid                  ( First(..) )

import           Floskell.Config
                 ( ImportsGroup(..), ImportsGroupOrder(..) )

import           Language.Haskell.Exts.Syntax ( ImportDecl(..), ModuleName(..) )

moduleName :: ImportDecl a -> String
moduleName :: ImportDecl a -> String
moduleName ImportDecl a
i = case ImportDecl a -> ModuleName a
forall l. ImportDecl l -> ModuleName l
importModule ImportDecl a
i of
    (ModuleName a
_ String
s) -> String
s

splitOn :: Char -> String -> [String]
splitOn :: Char -> String -> [String]
splitOn Char
c = (String -> Maybe (String, String)) -> String -> [String]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr String -> Maybe (String, String)
go
  where
    go :: String -> Maybe (String, String)
go [] = Maybe (String, String)
forall a. Maybe a
Nothing
    go String
x = (String, String) -> Maybe (String, String)
forall a. a -> Maybe a
Just ((String, String) -> Maybe (String, String))
-> (String, String) -> Maybe (String, String)
forall a b. (a -> b) -> a -> b
$ Int -> String -> String
forall a. Int -> [a] -> [a]
drop Int
1 (String -> String) -> (String, String) -> (String, String)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
break (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
c) String
x

modulePrefixes :: String -> [String]
modulePrefixes :: String -> [String]
modulePrefixes = ([String] -> String) -> [[String]] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
".") ([[String]] -> [String])
-> (String -> [[String]]) -> String -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[String]] -> [[String]]
forall a. [a] -> [a]
reverse ([[String]] -> [[String]])
-> (String -> [[String]]) -> String -> [[String]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> [[String]]
forall a. [a] -> [[a]]
inits ([String] -> [[String]])
-> (String -> [String]) -> String -> [[String]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> String -> [String]
splitOn Char
'.'

data St a = St { St a -> Map String Int
stIndex  :: M.Map String Int
               , St a -> Map Int (ImportsGroup, [ImportDecl a])
stGroups :: M.Map Int (ImportsGroup, [ImportDecl a])
               , St a -> [ImportDecl a]
stRest   :: [ImportDecl a]
               }

commonPrefixLength :: Eq a => [[a]] -> Int
commonPrefixLength :: [[a]] -> Int
commonPrefixLength = Int -> [[a]] -> Int
forall a t. (Eq a, Num t) => t -> [[a]] -> t
go Int
0
  where
    go :: t -> [[a]] -> t
go t
l [] = t
l
    go t
l ([] : [[a]]
_) = t
l
    go t
l ((a
x : [a]
xs) : [[a]]
ys) =
        if ([a] -> Bool) -> [[a]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([a] -> [a] -> Bool
forall a. Eq a => a -> a -> Bool
== [ a
x ]) ([a] -> Bool) -> ([a] -> [a]) -> [a] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
1) [[a]]
ys then t -> [[a]] -> t
go (t
l t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) ([a]
xs [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [[a]]
ys) else t
l

sortImports :: [ImportDecl a] -> [ImportDecl a]
sortImports :: [ImportDecl a] -> [ImportDecl a]
sortImports = (ImportDecl a -> String) -> [ImportDecl a] -> [ImportDecl a]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn ImportDecl a -> String
forall a. ImportDecl a -> String
moduleName

groupImports :: Int -> [ImportDecl a] -> [[ImportDecl a]]
groupImports :: Int -> [ImportDecl a] -> [[ImportDecl a]]
groupImports Int
n = (ImportDecl a -> ImportDecl a -> Bool)
-> [ImportDecl a] -> [[ImportDecl a]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy ([String] -> [String] -> Bool
forall a. Eq a => a -> a -> Bool
(==) ([String] -> [String] -> Bool)
-> (ImportDecl a -> [String])
-> ImportDecl a
-> ImportDecl a
-> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` Int -> ImportDecl a -> [String]
forall a. Int -> ImportDecl a -> [String]
prefix Int
n)
  where
    prefix :: Int -> ImportDecl a -> [String]
prefix Int
l = Int -> [String] -> [String]
forall a. Int -> [a] -> [a]
take Int
1 ([String] -> [String])
-> (ImportDecl a -> [String]) -> ImportDecl a -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [String] -> [String]
forall a. Int -> [a] -> [a]
drop Int
l ([String] -> [String])
-> (ImportDecl a -> [String]) -> ImportDecl a -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> String -> [String]
splitOn Char
'.' (String -> [String])
-> (ImportDecl a -> String) -> ImportDecl a -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ImportDecl a -> String
forall a. ImportDecl a -> String
moduleName

lookupFirst :: Ord a => [a] -> M.Map a b -> Maybe b
lookupFirst :: [a] -> Map a b -> Maybe b
lookupFirst [a]
ks Map a b
m = First b -> Maybe b
forall a. First a -> Maybe a
getFirst (First b -> Maybe b)
-> ([First b] -> First b) -> [First b] -> Maybe b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [First b] -> First b
forall a. Monoid a => [a] -> a
mconcat ([First b] -> Maybe b) -> [First b] -> Maybe b
forall a b. (a -> b) -> a -> b
$ (a -> First b) -> [a] -> [First b]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe b -> First b
forall a. Maybe a -> First a
First (Maybe b -> First b) -> (a -> Maybe b) -> a -> First b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Map a b -> Maybe b
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map a b
m)) [a]
ks

placeImport :: ImportDecl a -> State (St a) ()
placeImport :: ImportDecl a -> State (St a) ()
placeImport ImportDecl a
i = do
    Maybe Int
idx <- (St a -> Maybe Int) -> StateT (St a) Identity (Maybe Int)
forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
gets ([String] -> Map String Int -> Maybe Int
forall a b. Ord a => [a] -> Map a b -> Maybe b
lookupFirst (String -> [String]
modulePrefixes (String -> [String]) -> String -> [String]
forall a b. (a -> b) -> a -> b
$ ImportDecl a -> String
forall a. ImportDecl a -> String
moduleName ImportDecl a
i) (Map String Int -> Maybe Int)
-> (St a -> Map String Int) -> St a -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. St a -> Map String Int
forall a. St a -> Map String Int
stIndex)
    case Maybe Int
idx of
        Just Int
idx' -> (St a -> St a) -> State (St a) ()
forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify ((St a -> St a) -> State (St a) ())
-> (St a -> St a) -> State (St a) ()
forall a b. (a -> b) -> a -> b
$ \St a
s -> St a
s { stGroups :: Map Int (ImportsGroup, [ImportDecl a])
stGroups = Int
-> Map Int (ImportsGroup, [ImportDecl a])
-> Map Int (ImportsGroup, [ImportDecl a])
placeAt Int
idx' (St a -> Map Int (ImportsGroup, [ImportDecl a])
forall a. St a -> Map Int (ImportsGroup, [ImportDecl a])
stGroups St a
s) }
        Maybe Int
Nothing -> (St a -> St a) -> State (St a) ()
forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify ((St a -> St a) -> State (St a) ())
-> (St a -> St a) -> State (St a) ()
forall a b. (a -> b) -> a -> b
$ \St a
s -> St a
s { stRest :: [ImportDecl a]
stRest = St a -> [ImportDecl a]
forall a. St a -> [ImportDecl a]
stRest St a
s [ImportDecl a] -> [ImportDecl a] -> [ImportDecl a]
forall a. [a] -> [a] -> [a]
++ [ ImportDecl a
i ] }
  where
    placeAt :: Int
-> Map Int (ImportsGroup, [ImportDecl a])
-> Map Int (ImportsGroup, [ImportDecl a])
placeAt = ((ImportsGroup, [ImportDecl a]) -> (ImportsGroup, [ImportDecl a]))
-> Int
-> Map Int (ImportsGroup, [ImportDecl a])
-> Map Int (ImportsGroup, [ImportDecl a])
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (([ImportDecl a] -> [ImportDecl a])
-> (ImportsGroup, [ImportDecl a]) -> (ImportsGroup, [ImportDecl a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([ImportDecl a] -> [ImportDecl a] -> [ImportDecl a]
forall a. [a] -> [a] -> [a]
++ [ ImportDecl a
i ]))

splitImports :: [ImportsGroup] -> [ImportDecl a] -> [[ImportDecl a]]
splitImports :: [ImportsGroup] -> [ImportDecl a] -> [[ImportDecl a]]
splitImports [ImportsGroup]
groups [ImportDecl a]
imports = St a -> [[ImportDecl a]]
forall a. St a -> [[ImportDecl a]]
extract (St a -> [[ImportDecl a]]) -> St a -> [[ImportDecl a]]
forall a b. (a -> b) -> a -> b
$
    State (St a) () -> St a -> St a
forall s a. State s a -> s -> s
execState ((ImportDecl a -> State (St a) ())
-> [ImportDecl a] -> State (St a) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ImportDecl a -> State (St a) ()
forall a. ImportDecl a -> State (St a) ()
placeImport [ImportDecl a]
imports) St a
forall a. St a
initial
  where
    initial :: St a
initial = St :: forall a.
Map String Int
-> Map Int (ImportsGroup, [ImportDecl a]) -> [ImportDecl a] -> St a
St { stIndex :: Map String Int
stIndex  = [(String, Int)] -> Map String Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(String, Int)] -> Map String Int)
-> ([[(String, Int)]] -> [(String, Int)])
-> [[(String, Int)]]
-> Map String Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[(String, Int)]] -> [(String, Int)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(String, Int)]] -> Map String Int)
-> [[(String, Int)]] -> Map String Int
forall a b. (a -> b) -> a -> b
$
                       (Int -> ImportsGroup -> [(String, Int)])
-> [Int] -> [ImportsGroup] -> [[(String, Int)]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
n ImportsGroup
g -> (String -> (String, Int)) -> [String] -> [(String, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (, Int
n) (ImportsGroup -> [String]
importsPrefixes ImportsGroup
g))
                               [ Int
0 .. ]
                               [ImportsGroup]
groups
                 , stGroups :: Map Int (ImportsGroup, [ImportDecl a])
stGroups = [(Int, (ImportsGroup, [ImportDecl a]))]
-> Map Int (ImportsGroup, [ImportDecl a])
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Int, (ImportsGroup, [ImportDecl a]))]
 -> Map Int (ImportsGroup, [ImportDecl a]))
-> [(Int, (ImportsGroup, [ImportDecl a]))]
-> Map Int (ImportsGroup, [ImportDecl a])
forall a b. (a -> b) -> a -> b
$
                       (Int -> ImportsGroup -> (Int, (ImportsGroup, [ImportDecl a])))
-> [Int]
-> [ImportsGroup]
-> [(Int, (ImportsGroup, [ImportDecl a]))]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
n ImportsGroup
g -> (Int
n, (ImportsGroup
g, []))) [ Int
0 .. ] [ImportsGroup]
groups
                 , stRest :: [ImportDecl a]
stRest   = []
                 }

    extract :: St a -> [[ImportDecl a]]
extract St a
s = ([ImportDecl a] -> Bool) -> [[ImportDecl a]] -> [[ImportDecl a]]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool)
-> ([ImportDecl a] -> Bool) -> [ImportDecl a] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ImportDecl a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null) ([[ImportDecl a]] -> [[ImportDecl a]])
-> [[ImportDecl a]] -> [[ImportDecl a]]
forall a b. (a -> b) -> a -> b
$
        ((ImportsGroup, [ImportDecl a]) -> [[ImportDecl a]])
-> [(ImportsGroup, [ImportDecl a])] -> [[ImportDecl a]]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (ImportsGroup, [ImportDecl a]) -> [[ImportDecl a]]
forall a. (ImportsGroup, [ImportDecl a]) -> [[ImportDecl a]]
maybeSortAndGroup (Map Int (ImportsGroup, [ImportDecl a])
-> [(ImportsGroup, [ImportDecl a])]
forall k a. Map k a -> [a]
M.elems (Map Int (ImportsGroup, [ImportDecl a])
 -> [(ImportsGroup, [ImportDecl a])])
-> Map Int (ImportsGroup, [ImportDecl a])
-> [(ImportsGroup, [ImportDecl a])]
forall a b. (a -> b) -> a -> b
$ St a -> Map Int (ImportsGroup, [ImportDecl a])
forall a. St a -> Map Int (ImportsGroup, [ImportDecl a])
stGroups St a
s) [[ImportDecl a]] -> [[ImportDecl a]] -> [[ImportDecl a]]
forall a. [a] -> [a] -> [a]
++ [ St a -> [ImportDecl a]
forall a. St a -> [ImportDecl a]
stRest St a
s ]

    maybeSortAndGroup :: (ImportsGroup, [ImportDecl a]) -> [[ImportDecl a]]
maybeSortAndGroup (ImportsGroup
g, [ImportDecl a]
is) = case ImportsGroup -> ImportsGroupOrder
importsOrder ImportsGroup
g of
        ImportsGroupOrder
ImportsGroupKeep -> [ [ImportDecl a]
is ]
        ImportsGroupOrder
ImportsGroupSorted -> [ [ImportDecl a] -> [ImportDecl a]
forall a. [ImportDecl a] -> [ImportDecl a]
sortImports [ImportDecl a]
is ]
        ImportsGroupOrder
ImportsGroupGrouped -> Int -> [ImportDecl a] -> [[ImportDecl a]]
forall a. Int -> [ImportDecl a] -> [[ImportDecl a]]
groupImports ([String] -> Int
forall a. Eq a => [[a]] -> Int
commonPrefixLength ([String] -> Int) -> [String] -> Int
forall a b. (a -> b) -> a -> b
$
                                             ImportsGroup -> [String]
importsPrefixes ImportsGroup
g) ([ImportDecl a] -> [[ImportDecl a]])
-> [ImportDecl a] -> [[ImportDecl a]]
forall a b. (a -> b) -> a -> b
$ [ImportDecl a] -> [ImportDecl a]
forall a. [ImportDecl a] -> [ImportDecl a]
sortImports [ImportDecl a]
is