QuickSort with STUArray and Vector(s)

gistfile1.hs

{-# LANGUAGE BangPatterns #-}

module Main where

import Control.Applicative
import Control.Monad
import Control.Monad.ST
import Data.Array.ST

import Data.List (sort)

import qualified Data.Vector.Algorithms.Intro as Intro
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Storable as VS

import System.Random
import Criterion.Main

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

main :: IO ()
main = do
    !r1 <- randomList 10000 10000
    !r2 <- randomList 100000 100000
    defaultMain $ [
        bgroup "" [
             bench "rnd 10^4 list" $ nf qs r1
           , bench "rnd 10^5 list" $ nf qs r2
           , bench "rnd 10^4 list (Data.List.sort)" $ nf sort r1
           , bench "rnd 10^5 list (Data.List.sort)" $ nf sort r2
           , bench "rnd 10^4 STUArray" $ nf quickSort r1
           , bench "rnd 10^5 STUArray" $ nf quickSort r2
           , bench "rnd 10^4 Vector.Unboxed Intro" $ nf introSortVU r1
           , bench "rnd 10^5 Vector.Unboxed Intro" $ nf introSortVU r2
           , bench "rnd 10^4 Vector.Storable Intro" $ nf introSortVS r1
           , bench "rnd 10^5 Vector.Storable Intro" $ nf introSortVS r2
           ]]

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

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

qs :: Ord a => [a] -> [a]
qs [] = []
qs (x:xs) = qs lt ++ [x] ++ qs gt
  where
    lt = filter (<x) xs
    gt = filter (>=x) xs

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

type Index = Int
type Value = Int
type SUA s = STUArray s Index Value
type PRED = Value -> Value -> Bool

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

quickSort :: [Value] -> [Value]
quickSort xs = runST $ do
    arr <- newListArray (beg,end) xs
    quicksort beg end arr
    toList beg end arr
  where
    beg = 0
    end = length xs - 1

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

loop :: Value -> Index -> Index -> Index -> SUA s -> ST s Index
loop 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
        swap ni nj arr
        loop t u ni nj arr
  where
    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

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

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

-----

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