{-# 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 :: forall a. ImportDecl a -> String
moduleName ImportDecl a
i = case 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 = forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr String -> Maybe (String, String)
go
  where
    go :: String -> Maybe (String, String)
go [] = forall a. Maybe a
Nothing
    go String
x = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
drop Int
1 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. (a -> Bool) -> [a] -> ([a], [a])
break (forall a. Eq a => a -> a -> Bool
== Char
c) String
x

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

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

commonPrefixLength :: Eq a => [[a]] -> Int
commonPrefixLength :: forall a. Eq a => [[a]] -> Int
commonPrefixLength = 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 forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((forall a. Eq a => a -> a -> Bool
== [ a
x ]) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> [a] -> [a]
take Int
1) [[a]]
ys then t -> [[a]] -> t
go (t
l forall a. Num a => a -> a -> a
+ t
1) ([a]
xs forall a. a -> [a] -> [a]
: [[a]]
ys) else t
l

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

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

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

splitImports :: [ImportsGroup] -> [ImportDecl a] -> [[ImportDecl a]]
splitImports :: forall a. [ImportsGroup] -> [ImportDecl a] -> [[ImportDecl a]]
splitImports [ImportsGroup]
groups [ImportDecl a]
imports = forall {a}. St a -> [[ImportDecl a]]
extract forall a b. (a -> b) -> a -> b
$
    forall s a. State s a -> s -> s
execState (forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall a. ImportDecl a -> State (St a) ()
placeImport [ImportDecl a]
imports) forall {a}. St a
initial
  where
    initial :: St a
initial = St { stIndex :: Map String Int
stIndex  = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$
                       forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
n ImportsGroup
g -> 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 = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$
                       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 = forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Bool
null) forall a b. (a -> b) -> a -> b
$
        forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall {a}. (ImportsGroup, [ImportDecl a]) -> [[ImportDecl a]]
maybeSortAndGroup (forall k a. Map k a -> [a]
M.elems forall a b. (a -> b) -> a -> b
$ forall a. St a -> Map Int (ImportsGroup, [ImportDecl a])
stGroups St a
s) forall a. [a] -> [a] -> [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 -> [ forall a. [ImportDecl a] -> [ImportDecl a]
sortImports [ImportDecl a]
is ]
        ImportsGroupOrder
ImportsGroupGrouped -> forall a. Int -> [ImportDecl a] -> [[ImportDecl a]]
groupImports (forall a. Eq a => [[a]] -> Int
commonPrefixLength forall a b. (a -> b) -> a -> b
$
                                             ImportsGroup -> [String]
importsPrefixes ImportsGroup
g) forall a b. (a -> b) -> a -> b
$ forall a. [ImportDecl a] -> [ImportDecl a]
sortImports [ImportDecl a]
is