real-quick-sort.hs

real-quick-sort.hs

#!/usr/bin/env stack
{-stack script
--resolver=lts-13.18
--package=base-unicode-symbols
--package=array
--package=MonadRandom
--package=hspec
-}

-- Author:
--   Viacheslav Lotsmanov <lotsmanov89@gmail.com>
--   https://github.com/unclechu
--   April 2019

{-# LANGUAGE UnicodeSyntax, BangPatterns, MultiWayIf #-}
{-# LANGUAGE ScopedTypeVariables, FlexibleContexts #-}

import Prelude.Unicode
import Data.Function (fix)
import Data.Array.IO
import Control.Monad (when)

-- just for tests
import Data.List (sort)
import Control.Monad (forM_)
import Control.Monad.Random.Class (getRandoms)
import Test.Hspec


quickSort ∷ (MArray IOUArray α IO, Ord α) ⇒ [α] → IO [α]
quickSort [] = pure []
quickSort xs = go where
  go = do
    let !maxIx = pred $ length xs
    !(arr ∷ IOUArray Int a) ← newListArray (0, maxIx) xs
    qsort arr 0 maxIx >> getElems arr

  qsort !arr = fix $ \again !leftIx !rightIx →
    let !diff = rightIx - leftIx in if

    | diff ≤ 1 →
        when (diff ≡ 1) $ do
          !a ← readArray arr leftIx
          !b ← readArray arr rightIx
          when (a > b) $ do
            writeArray arr leftIx  b
            writeArray arr rightIx a

    | otherwise → do
        let !mid = (rightIx + leftIx) `div` 2
        !pivot ← swap arr mid rightIx
        !newPivotIx ← f arr pivot rightIx leftIx $ pred rightIx
        again leftIx newPivotIx
        again (succ newPivotIx) rightIx

  f !arr !pivot !pivotIx = fix $ \again !leftIx !rightIx → do
    !lGt ← firstGreater arr pivot rightIx leftIx  -- left  (greater than)
    !rLt ← firstLess    arr pivot leftIx  rightIx -- right (less than)

    if | lGt > rLt  → lGt <$ swap arr lGt pivotIx
       | lGt ≡ (-1) → pure rightIx
       | rLt ≡ (-1) → leftIx <$ swap arr leftIx pivotIx
       | otherwise  → swap arr lGt rLt >> again (succ lGt) (pred rLt)

  firstGreater !arr !pivot !rightIx = fix $ \next !ix → if
    | ix > rightIx → pure (-1)
    | otherwise → readArray arr ix >>= \ !x →
        if x > pivot then pure ix else next $ succ ix

  firstLess !arr !pivot !leftIx = fix $ \next !ix → if
    | ix < leftIx → pure (-1)
    | otherwise → readArray arr ix >>= \ !x →
        if x ≤ pivot then pure ix else next $ pred ix

  swap !arr !aIx !bIx = do
    !aEl ← readArray arr aIx
    !bEl ← readArray arr bIx
    writeArray arr aIx bEl
    writeArray arr bIx aEl
    pure aEl


main ∷ IO ()
main = hspec $ parallel $ do
  let amount = [1..100∷Word]

  do
    let ref = [] ∷ [Int]

    it "empty list" $
      forM_ amount (const $ quickSort ref >>= (`shouldBe` ref))

  it "some example" $ do
    let src = [ 9817899481354916781,  14888124000582553526
              , 12438762095863496482, 15327772538303920662
              , 13875875630416247750 ∷ Word ]

    quickSort src >>= (`shouldBe` sort src)

  forM_ [1..100∷Int] $ \n →
    it (show n ⧺ " element(s)") $
      forM_ amount $ const $ do
        !(ref1 ∷ [Int])  ← take n <$> getRandoms
        !(ref2 ∷ [Word]) ← take n <$> getRandoms
        quickSort ref1 >>= (`shouldBe` sort ref1)
        quickSort ref2 >>= (`shouldBe` sort ref2)