module Control.Quiver.Sort (
spsort
, spsortBy
, spsortOn
, spfilesort
, spfilesortBy
) where
import Control.Quiver.Binary
import Control.Quiver.ByteString
import Control.Quiver.Group
import Control.Quiver.Instances ()
import Control.Quiver.Interleave
import Control.Quiver.SP
import Control.Applicative (liftA2)
import Control.Exception (IOException, finally)
import Control.Monad (join)
import Control.Monad.Catch (MonadCatch (..), MonadMask)
import Control.Monad.IO.Class (MonadIO (..))
import Control.Monad.Trans.Resource (MonadResource, allocate)
import Data.Bool (bool)
import Data.Function (on)
import Data.List (sortBy)
import Data.Maybe (fromMaybe)
import System.Directory (doesDirectoryExist, getPermissions,
getTemporaryDirectory,
removeDirectoryRecursive, writable)
import System.IO (hClose, openTempFile)
import System.IO.Temp (createTempDirectory)
spsort :: (Ord a, Monad m) => SP a a m ()
spsort = spsortBy compare
spsortBy :: (Monad m) => (a -> a -> Ordering) -> SP a a m ()
spsortBy f = (sortBy f <$> spfoldr (:) []) >>= spevery
spsortOn :: (Ord b, Monad m) => (a -> b) -> SP a a m ()
spsortOn f = sppure ((,) <*> f)
>->> spsortBy (compare `on` snd)
>->> sppure fst >&> snd
spfilesort :: (Binary a, Ord a, MonadResource m, MonadMask m) => Maybe Int -> Maybe FilePath
-> P () a a () m (SPResult IOException)
spfilesort = spfilesortBy compare
spfilesortBy :: (Binary a, MonadResource m, MonadMask m) => (a -> a -> Ordering) -> Maybe Int -> Maybe FilePath
-> P () a a () m (SPResult IOException)
spfilesortBy cmp mchunks mdir = do mdir' <- join <$> liftIO (traverse checkDir mdir)
getTmpDir mdir' "quiver-sort" pipeline
where
checkDir dir = do ex <- liftA2 (liftA2 (&&)) doesDirectoryExist (fmap writable . getPermissions) dir
return (bool Nothing (Just dir) ex)
getTmpDir = maybe withSystemTempDirectory withTempDirectory
pipeline tmpDir = toFiles tmpDir >>= either spfailed (sortFromFiles cmp)
toFiles tmpDir = sortToFiles chunkSize cmp tmpDir >->> spToList >&> uncurry (flip checkFailed)
chunkSize = fromMaybe 10000 mchunks
sortToFiles :: (Binary a, MonadIO m) => Int -> (a -> a -> Ordering) -> FilePath
-> SP a FilePath m IOException
sortToFiles chunkSize cmp tmpDir = spchunks chunkSize
>->> spTraverseUntil sortChunk
>&> snd
where
sortChunk as = liftIO $ do (fl,h) <- openTempFile tmpDir "quiver-sort-chunk"
finally (checkFailed fl <$> sprun (pipeline h)) (hClose h)
where
pipeline h = spevery (sortBy cmp as) >->> spencode >->> qPut h >&> snd
sortFromFiles :: (Binary a, MonadIO m) => (a -> a -> Ordering) -> [FilePath]
-> Producer a () m (SPResult IOException)
sortFromFiles cmp fls = spinterleave cmp (map readFl fls)
where
readFl fl = qhoist liftIO (qReadFile fl readSize) >->> spdecode >&> fst
readSize = 4096
spTraverseUntil :: (Monad m) => (a -> m (Either e b)) -> SP a b m e
spTraverseUntil k = loop
where
loop = spconsume loop' spcomplete
loop' a = qlift (k a) >>= either spfailed (>:> loop)
checkFailed :: r -> SPResult e -> Either e r
checkFailed _ (Just (Just e)) = Left e
checkFailed r _ = Right r
spToList :: SQ a x f [a]
spToList = spfoldr (:) []
withSystemTempDirectory :: (MonadResource m) =>
String
-> (FilePath -> m a)
-> m a
withSystemTempDirectory template action = liftIO getTemporaryDirectory >>= \tmpDir -> withTempDirectory tmpDir template action
withTempDirectory :: (MonadResource m) =>
FilePath
-> String
-> (FilePath -> m a)
-> m a
withTempDirectory targetDir template withTmp = do
(_release, tmpDir) <- allocate (createTempDirectory targetDir template)
(ignoringIOErrors . removeDirectoryRecursive)
withTmp tmpDir
ignoringIOErrors :: (MonadCatch m) => m () -> m ()
ignoringIOErrors ioe = ioe `catch` (\(_ :: IOError) -> return ())