--
-- Radix sort for a subclass of element types
--

module Radix where

import Random

import qualified Prelude
import Prelude               hiding (zip, map, scanl, scanr, zipWith, fst)
import Data.Bits             hiding (shiftL, shiftR, bit, testBit)
import Data.Array.Accelerate as Acc

import Data.List             (sort)
import Data.Array.Unboxed    (IArray, UArray, listArray, bounds, elems)
import System.Random.MWC
import Unsafe.Coerce


-- Radix sort
-- ----------

class Elt e => Radix e where
    passes :: Exp e   -> Int                -- Haskell-side control needs to know this
    radix  :: Exp Int -> Exp e -> Exp Int

instance Radix Int where                    -- may be 32- or 64-bit
    passes    = bitSize . (undefined :: Exp t -> t)
    radix i e = i ==* (passes' e - 1) ? (radix' (e `xor` minBound), radix' e)
      where
        radix' x = (x `shiftR` i) .&. 1
        passes'  = constant . passes

-- For IEEE-754 floating-point representation. Unsafe, but widely supported.
--
instance Radix Float where
    passes _   = 32
    radix i e  = let x = (unsafeCoerce e :: Exp Int32)
                 in  i ==* 31 ? (radix' (x `xor` minBound), radix' (floatFlip x))
      where
        floatFlip x = x `testBit` 31 ? (complement x, x)  -- twos-complement negative numbers
        radix'    x = x `testBit` i  ? (1,0)


--
-- A simple (parallel) radix sort implementation [1].
--
-- [1] G. E. Blelloch. "Prefix sums and their applications." Technical Report
--     CMU-CS-90-190. Carnegie Mellon University. 1990.
--
sortAcc :: Radix a => Vector a -> Acc (Vector a)
sortAcc = sortAccBy id

sortAccBy :: (Elt a, Radix r) => (Exp a -> Exp r) -> Vector a -> Acc (Vector a)
sortAccBy rdx arr = foldr1 (>->) (Prelude.map radixPass [0..p-1]) (use arr)
  where
    n = constant $ (arraySize $ arrayShape arr) - 1
    p = passes . rdx . (undefined :: Vector e -> Exp e) $ arr

    deal f x      = let (a,b) = unlift x in (f ==* 0) ? (a,b)
    radixPass k v = let flags = map (radix (constant k) . rdx) v
                        idown = prescanl (+) 0 . map (xor 1) $ flags
                        iup   = map (n-) . prescanr (+) 0    $ flags
                        index = zipWith deal flags (zip idown iup)
                    in
                    permute const v (\ix -> index1 (index!ix)) v


sortRef :: UArray Int Int -> UArray Int Int
sortRef xs = listArray (bounds xs) $ sort (elems xs)


-- Main
-- ----

run :: Int -> IO (() -> UArray Int Int, () -> Acc (Vector Int))
run n = withSystemRandom $ \gen -> do
  vec  <- randomUArrayR (minBound,maxBound) gen n
  vec' <- convertUArray vec
  --
  return (run_ref vec, run_acc vec')
  where
    {-# NOINLINE run_ref #-}
    run_ref xs () = sortRef xs
    run_acc xs () = sortAcc xs