Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
{-# 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment