content: {-# LANGUAGE FlexibleContexts    #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  --
  -- Radix sort for a subclass of element types
  --
  
  module Radix where
  
  import Random
  
  import Prelude                  as P
  import Data.Array.Accelerate    as A
  
  import Data.Bits
  import Data.List                ( sort )
  import Data.Array.Unboxed       ( IArray, UArray, listArray, bounds, elems )
  import System.Random.MWC
  
  
  -- Radix sort
  -- ----------
  
  class Elt e => Radix e where
    passes :: e {- dummy -} -> Int
    radix  :: Exp Int -> Exp e -> Exp Int
  
  instance Radix Int32 where
    passes    = bitSize
    radix i e = i ==* (passes' - 1) ? (radix' (e `xor` minBound), radix' e)
      where
        radix' x = A.fromIntegral $ (x `A.shiftR` i) .&. 1
        passes'  = constant (passes (undefined :: Int32))
  
  -- For IEEE-754 floating-point representation. Unsafe, but widely supported.
  -- TLM: unsafeCoerce does not work in the CUDA backend.
  --
  -- 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 => Acc (Vector a) -> Acc (Vector a)
  sortAcc = sortAccBy id
  
  sortAccBy
      :: forall a r. (Elt a, Radix r)
      => (Exp a -> Exp r)
      -> Acc (Vector a)
      -> Acc (Vector a)
  sortAccBy rdx arr = foldr1 (>->) (P.map radixPass [0..p-1]) arr
    where
      p = passes (undefined :: r)
      --
      deal f x      = let (a,b) = unlift x in (f ==* 0) ? (a,b)
      radixPass k v = let k'    = unit (constant k)
                          flags = A.map (radix (the k') . rdx) v
                          idown = prescanl (+) 0 . A.map (xor 1)        $ flags
                          iup   = A.map (size v - 1 -) . prescanr (+) 0 $ flags
                          index = A.zipWith deal flags (A.zip idown iup)
                      in
                      permute const v (\ix -> index1 (index!ix)) v
  
  
  sortRef :: (Ord a, IArray UArray a) => UArray Int a -> UArray Int a
  sortRef xs = listArray (bounds xs) $ sort (elems xs)
  
  
  -- Main
  -- ----
  
  run :: Int -> IO (() -> UArray Int Int32, () -> Acc (Vector Int32))
  run n = withSystemRandom $ \gen -> do
    vec  <- randomUArrayR (minBound,maxBound) gen n
    vec' <- use `fmap` convertUArray vec
    --
    return (run_ref vec, run_acc vec')
    where
      {-# NOINLINE run_ref #-}
      run_ref xs () = sortRef xs
      run_acc xs () = sortAcc xs
  
  