Haskell in-place quicksort, Ord a => [a] -> IO [a]

Question

As noted below, this is more or less a direct translation from http://rosettacode.org/wiki/Sorting_algorithms/Quicksort#C. I'm fairly satisfied with its readability right now compared to the C code; but still wondering how to make it more succinct and clear. Any ideas/opinions?

{-
In-place quicksort algorithm.
Adapted from http://rosettacode.org/wiki/Sorting_algorithms/Quicksort#C.

Main considerations:

  - Type should be as close to [a] -> [a] as possible.

  - Comparatively most of the time should be spent in comparing and
    swapping members, and little in copying the list onto and out of the
    function's stack frame. So don't worry about forcing user to deal
    with mutable arrays/vectors; instead accept and return immutable
    lists, and mutate internally.

  - Since in-place, need to account for effect of changing members.
    Therefore return IO [a].

  - Work in IO monad rather than State/ST/others because want to be as
    close to the bottom of the monad stack as possible--we will end up
    at IO anyway ultimately.

  - Internally use as an Array of IORefs, but we can avoid exposing that
    to the user by sequencing all the effects into an IO [a].

  - Internal functions all 'close over' a single shared Array Int (IORef
    a). Is this a bad practice? My intuition is no, because we don't
    expose that to the user.

  - Haskell's ability to have let statements and monadic bindings spread
    anywhere inside a 'do' block is really awesome.
-}

import Control.Monad (forM_, mapM, when)
import Data.Array ((!), Array, bounds, elems, listArray)
import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef)

qsortInPlace :: Ord a => [a] -> IO [a]
qsortInPlace xs = do
  xs' <- mapM newIORef xs

  let
    upperBound = length xs - 1
    ys = listArray (0, upperBound) xs'

    swap i1 i2 = do
      let
        r1 = ys ! i1
        r2 = ys ! i2
      t <- readIORef r1
      readIORef r2 >>= writeIORef r1
      writeIORef r2 t

    go lo hi
      | hi <= lo = return ()
      | otherwise = do
        let pivotIndex = (lo + hi) `div` 2

        pivot <- readIORef $ ys ! pivotIndex
        chg <- newIORef lo
        swap pivotIndex hi

        forM_ [lo..(hi - 1)] $ \i -> do
          ys_i <- readIORef $ ys ! i
          when (ys_i < pivot) $ do
            chg' <- readIORef chg
            swap i chg'
            modifyIORef' chg (+1)

        chg' <- readIORef chg
        swap chg' hi

        go lo $ chg' - 1
        go (chg' + 1) hi

  go 0 upperBound
  mapM readIORef . elems $ ys

main :: IO ()
main = do
  xs <- qsortInPlace [4, 5, 2, 3, 1, 2, 7, 1, 10]
  print xs

Attempt 2

Taking the advice of Feuerbach and others on the Haskell IRC channel (http://tunes.org/~nef/logs/haskell/15.01.02, 16:03:52 onwards), I've reimplemented using mutable vectors and the ST monad. The code is somewhat simplified and I'm also struck by how similarly the IORef and STRef types work.

{-
In-place quicksort algorithm.
Adapted from http://rosettacode.org/wiki/Sorting_algorithms/Quicksort#C.

Main considerations:

  - Type should be as close to [a] -> [a] as possible.

  - Comparatively most of the time should be spent in comparing and
    swapping members, and little in copying the list onto and out of the
    function's stack frame. So don't worry about forcing user to deal
    with mutable arrays/vectors; instead accept and return immutable
    lists, and mutate internally.

  - Since in-place, need to account for effect of changing members.
    Therefore internally use ST monad. The nice thing about ST is that
    we can run it and return [a]. So we can actually have type [a] ->
    [a]. The user never needs to know we did stateful computations.

  - Internally use a mutable vector, which out of the box provides the
    ability to swap its elements.

  - Internal 'go' function 'closes over' a single shared mutable vector
    (MVector) so we don't have to keep passing it back and forth between
    function stack frames.

  - Haskell's ability to have let statements and monadic bindings spread
    anywhere inside a 'do' block is really awesome. Not to mention its
    ability to encapsulate a complex series of monadic actions inside a
    single 'variable'.
-}

import Control.Monad (forM_, mapM, when)
import Control.Monad.ST (runST)
import Data.STRef (modifySTRef', newSTRef, readSTRef )
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM

qsortInPlace :: Ord a => [a] -> [a]
qsortInPlace xs =
  let
    vAction = do
      v <- V.thaw . V.fromList $ xs

      let
        go lo hi
          | hi <= lo = return ()
          | otherwise = do
            let pivotIndex = (lo + hi) `div` 2

            pivot <- VM.read v pivotIndex
            chg <- newSTRef lo
            VM.swap v pivotIndex hi

            forM_ [lo..(hi - 1)] $ \i -> do
              v_i <- VM.read v i
              when (v_i < pivot) $ do
                chg' <- readSTRef chg
                VM.swap v i chg'
                modifySTRef' chg (+1)

            chg' <- readSTRef chg
            VM.swap v chg' hi

            go lo $ chg' - 1
            go (chg' + 1) hi

      go 0 $ length xs - 1
      V.freeze v >>= return . V.toList

  in runST vAction

main :: IO ()
main = print $ qsortInPlace [4, 5, 2, 3, 1, 2, 7, 1, 10]

Answer 1

Nice and well documented code. Some ideas:

In this case you could use unsafeThaw, as you create a vector only to thaw it. This will save you copying it during thaw, but of course you must be careful and aware of the consequences. If using arrays, another option is newListArray.

When converting back to a list, for arrays you could either use runSTArray which again saves copying the entire array (as opposed to using freeze). Or instead of freeze followed by toList, use getElems. For vectors, there is unsafeFreeze too.

While using a STRef counter for the middle part certainly works, it's a bit imperative style for Haskell. The more common approach would be to have a recursive function where you pass it as an argument and return it back.

Somewhat related is looping using forM_ over a list. Probably passing it as an argument would be faster, as producing/consuming the list is likely to cause memory (de)allocation inside the loop, although less idiomatic. So it depends on what are your goals.

Finally, declaring functions using let inside do is indeed often convenient, but can impact readability, if they're long. I'd prefer to declare them separately so that the main function is just "thaw - go - freeze" and then the definition of go follows (or precedes). If the helper functions aren't recursive (that is, the recursion is hidden inside), they'll be properly inlined as needed by the compiler.

Also x >>= return . f is equivalent to liftM f x or just f <$> x.

Here is a possible variant with the ideas above, and a few minor more:

{-# LANGUAGE BangPatterns #-}
import Control.Monad (forM_, mapM, when)
import Control.Monad.ST (runST)
import Data.Functor ((<$>))
import Data.STRef (modifySTRef', newSTRef, readSTRef )
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM

qsortInPlace :: Ord a => [a] -> [a]
qsortInPlace xs = runST $ do
      v <- V.unsafeThaw . V.fromList $ xs
      pass v 0 (length xs - 1)
      V.toList <$> V.unsafeFreeze v
  where
    -- Hiding the recursion into the inner 'go' function is not just
    -- convenient, it allows 'split' to be non-recursive, which allows
    -- its inlining.
    split v pivot lo hi = go lo lo
      where
        -- Bang patterns should help the GHC optimizer here.
        go !chg !i | i >= hi = return chg
                   | otherwise = do
          v_i <- VM.read v i
          if (v_i < pivot) then do
              VM.swap v i chg
              go (chg + 1) (i + 1)
            else
              go chg (i + 1)

    pass v lo hi
      | hi <= lo = return ()
      | otherwise = do
        let pivotIndex = (lo + hi) `div` 2

        pivot <- VM.read v pivotIndex
        VM.swap v pivotIndex hi

        chg <- split v pivot lo hi
        VM.swap v chg hi

        pass v lo (chg - 1)
        pass v (chg + 1) hi

main :: IO ()
main = print $ qsortInPlace [4, 5, 2, 3, 1, 2, 7, 1, 10]

Also a good thing would be to add a QuickCheck test.