QuickSort with Array

gistfile1.hs

{-# LANGUAGE BangPatterns #-}

module Main where

import Control.Applicative
import Control.Monad
import Control.Monad.ST
import Criterion.Main
import Data.Array.ST
import Data.List (sort)
import qualified Data.Vector.Algorithms.Intro as Intro
import qualified Data.Vector.Storable as VS
import Data.Vector.Unboxed.Mutable (STVector)
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import System.Random

----------------------------------------------------------------

type Index = Int
type Value = Int
type SUA s = STUArray s Index Value
type VUA s = STVector s Value

----------------------------------------------------------------

main :: IO ()
main = do
    !r1 <- randomList 10000 10000
    !r2 <- randomList 100000 100000
    defaultMain $ [
        bgroup "" [
             bench "sort          rnd 10^4" $ nf sort r1
           , bench "sort          rnd 10^5" $ nf sort r2             
           , bench "quickSortList rnd 10^4" $ nf quickSortList r1
           , bench "quickSortList rnd 10^5" $ nf quickSortList r2
           , bench "quickSortST   rnd 10^4" $ nf quickSortST   r1
           , bench "quickSortST   rnd 10^5" $ nf quickSortST   r2
           , bench "quickSortVec  rnd 10^4" $ nf quickSortVec  r1
           , bench "quickSortVec  rnd 10^5" $ nf quickSortVec  r2
           , bench "introSort VU  rnd 10^4" $ nf introSortVU   r1
           , bench "introSort VU  rnd 10^5" $ nf introSortVU   r2
           , bench "introSort VS  rnd 10^4" $ nf introSortVS   r1
           , bench "introSort VS  rnd 10^5" $ nf introSortVS   r2
           ]
        ]

randomList :: Index -> Value -> IO [Value]
randomList n boundary = replicateM n randomInt
  where
      randomInt :: IO Value
      randomInt = getStdRandom (randomR (0,boundary))

----------------------------------------------------------------
    
quickSortList :: [Value] -> [Value]
quickSortList [] = []
quickSortList (x:xs) = quickSortList lt ++ [x] ++ quickSortList gt
  where
    lt = filter (<x)  xs
    gt = filter (>=x) xs

----------------------------------------------------------------

quickSortST :: [Value] -> [Value]
quickSortST xs = runST $ do
    arr <- newListArray (beg,end) xs
    quicksortST beg end arr
    toListST beg end arr
  where
    beg = 0
    end = length xs - 1

quicksortST :: Index -> Index -> SUA s -> ST s ()
quicksortST l u arr
  | l >= u    = return ()
  | otherwise = do
    -- swapST l random arr -- good pivot
    t <- readArray arr l
    let i = l
        j = u + 1
    nj <- loopST t u i j arr
    swapST arr l nj
    quicksortST l (nj-1) arr
    quicksortST (nj+1) u arr

loopST :: Value -> Index -> Index -> Index -> SUA s -> ST s Index
loopST !t !u !i !j arr = do
    ni <- doWhile i (+1)         (< t)
    nj <- doWhile j (subtract 1) (> t)
    if ni > nj then
        return nj
      else do
        swapST arr ni nj
        loopST t u ni nj arr
  where
    {-# INLINE doWhile #-}
    doWhile k op p
      | nk > u    = return nk
      | otherwise = do
          x <- readArray arr nk
          if p x then
              doWhile nk op p
            else
              return nk
      where
        nk = op k

swapST :: SUA s -> Index ->Index -> ST s ()
swapST sua i j = do
    xi <- readArray sua i
    xj <- readArray sua j
    writeArray sua i xj
    writeArray sua j xi

toListST :: Index -> Index -> SUA s -> ST s [Value]
toListST i lim arr
  | i > lim   = return []
  | otherwise = (:) <$> readArray arr i <*> toListST (i+1) lim arr

----------------------------------------------------------------

quickSortVec :: [Value] -> [Value]
quickSortVec xs = runST $ do
    arr <- VU.unsafeThaw $ VU.fromList xs
    let beg = 0
        end = VUM.length arr - 1
    quicksortVec beg end arr
    VU.toList <$> VU.unsafeFreeze arr

quicksortVec :: Index -> Index -> VUA s -> ST s ()
quicksortVec l u arr
  | l >= u    = return ()
  | otherwise = do
    -- VUM.unsafeSwap arr l random -- good pivot
    t <- VUM.unsafeRead arr l
    let i = l
        j = u + 1
    nj <- loopVec t u i j arr
    VUM.unsafeSwap arr l nj
    quicksortVec l (nj-1) arr
    quicksortVec (nj+1) u arr

loopVec :: Value -> Index -> Index -> Index -> VUA s -> ST s Index
loopVec !t !u !i !j arr = do
    ni <- doWhile i (+1)         (< t)
    nj <- doWhile j (subtract 1) (> t)
    if ni > nj then
        return nj
      else do
        VUM.unsafeSwap arr ni nj
        loopVec t u ni nj arr
  where
    {-# INLINE doWhile #-}
    doWhile k op p
      | nk > u    = return nk
      | otherwise = do
          x <- VUM.unsafeRead arr nk
          if p x then
              doWhile nk op p
            else
              return nk
      where
        nk = op k

----------------------------------------------------------------

introSortVU :: [Value] -> [Value]
introSortVU xs = runST $ do
  vum <- VU.unsafeThaw $ VU.fromList xs
  Intro.sort vum
  VU.toList <$> VU.unsafeFreeze vum

introSortVS :: [Value] -> [Value]
introSortVS xs = runST $ do
  vum <- VS.unsafeThaw $ VS.fromList xs
  Intro.sort vum
  VS.toList <$> VS.unsafeFreeze vum

                 10^4       10^5
sort            15.09825   256.7437
quickSortList   12.97226   193.1537
quickSortST      7.502490   92.16319
quickSortVec     6.115590   71.58241
introSort VU     3.233967   35.04761
introSort VS     3.267389   35.36919