VecSort.hs

{-# LANGUAGE BangPatterns #-}
module VecSort(sort, sortBy) where
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV

{-# INLINE sort #-}
sort :: Ord a => [a] -> [a]
sort = sortBy compare

sortBy :: (a -> a -> Ordering) -> [a] -> [a]
sortBy cmp = V.toList . mergeSortAux . V.fromList
    where
    mergeSortAux l
        | n <= 1 = l
        | otherwise = merge (n1, l1') (n2, l2')
             where
                n  = V.length l
                n1 = div n 2
                n2 = n - n1
                l1 = V.unsafeSlice 0  n1 l
                l2 = V.unsafeSlice n1 n2 l
                l1' = mergeSortAux l1
                l2' = mergeSortAux l2
    merge (!n1,!as) (!n2,!bs) = V.create $ do
        res <- MV.unsafeNew (n1+n2)
        let go i1 i2 | i1 >= n1 && i2 >= n2 = return res
                     | i1 >= n1 = tick2 i1 i2
                     | i2 >= n2 = tick1 i1 i2
                     | otherwise = do
                         v1 <- V.unsafeIndexM as i1
                         v2 <- V.unsafeIndexM bs i2
                         case cmp v1 v2 of
                            GT  -> do
                                 MV.unsafeWrite res (i1 + i2) v2
                                 go i1 (i2 + 1)
                            _   -> do
                                 MV.unsafeWrite res (i1 + i2) v1
                                 go (i1 + 1) i2
            tick1 i1 i2 = do
                V.unsafeIndexM as i1 >>= MV.unsafeWrite res (i1 + i2)
                go (i1 + 1) i2
            tick2 i1 i2 = do
                V.unsafeIndexM bs i2 >>= MV.unsafeWrite res (i1 + i2)
                go i1 (i2 + 1)
        go 0 0