{-|
Module      : Z.Data.Vector.Sort
Description : Sorting vectors
Copyright   : (c) 2008-2011 Dan Doel, (c) Dong Han, 2017-2018
License     : BSD
Maintainer  : winterland1989@gmail.com
Stability   : experimental
Portability : non-portable

This module provide three stable sorting algorithms, which are:

  * 'mergeSort', a /O(log(n))/ general-purpose sorting algorithms for all different size vectors.

  * 'insertSort' a /O(n^2)/ sorting algorithms suitable for very small vectors.

  * 'radixSort' a /O(n)/ sorting algorithms based on 'Radix' instance, which is prefered on large vectors.

Sorting is always performed in ascending order. To reverse the order, either use @XXSortBy@ or use 'Down', 'RadixDown' newtypes. In general changing comparing functions can be done by creating auxiliary newtypes and 'Ord' instances (make sure you inline instance's method for performence!). Or 'Radix' instances in 'radixSort' case, for example:

@
data Foo = Foo { key :: Int16, ... }

instance Radix Foo where
    -- You should add INLINE pragmas to following methods
    bucketSize = bucketSize . key
    passes = passes . key
    radixLSB = radixLSB . key
    radix i = radix i . key
    radixMSB = radixMSB . key
@

-}

module Z.Data.Vector.Sort (
  -- * Sort
    mergeSort
  , mergeSortBy
  , mergeTileSize
  , insertSort
  , insertSortBy
  , Down(..)
  , radixSort
  , Radix(..)
  , RadixDown(..)
  -- * merge duplicated
  , mergeDupAdjacent
  , mergeDupAdjacentLeft
  , mergeDupAdjacentRight
  , mergeDupAdjacentBy
  ) where

import           Control.Monad.ST
import           Data.Bits
import           Data.Int
import           Data.Ord               (Down (..))
import           Data.Primitive         (sizeOf)
import           Data.Word
import           Prelude                hiding (splitAt)
import           Z.Data.Array
import           Z.Data.Array.Unaligned
import           Z.Data.Vector.Base
import           Z.Data.Vector.Extra

--------------------------------------------------------------------------------
-- Comparison Sort

-- | /O(n*log(n))/ Sort vector based on element's 'Ord' instance with classic
-- <https://en.wikipedia.org/wiki/Merge_sort mergesort> algorithm.
--
-- This is a stable sort, During sorting two O(n) worker arrays are needed, one of
-- them will be freezed into the result vector. The merge sort only begin at tile
-- size larger than 'mergeTileSize', each tile will be sorted with 'insertSort', then
-- iteratively merged into larger array, until all elements are sorted.
mergeSort :: forall v a. (Vec v a, Ord a) => v a -> v a
{-# INLINABLE mergeSort #-}
mergeSort :: v a -> v a
mergeSort = (a -> a -> Ordering) -> v a -> v a
forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Ordering) -> v a -> v a
mergeSortBy a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare

mergeSortBy :: forall v a. Vec v a => (a -> a -> Ordering) -> v a -> v a
{-# INLINE mergeSortBy #-}
mergeSortBy :: (a -> a -> Ordering) -> v a -> v a
mergeSortBy a -> a -> Ordering
cmp vec :: v a
vec@(Vec IArray v a
_ Int
_ Int
l)
    | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
mergeTileSize = (a -> a -> Ordering) -> v a -> v a
forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Ordering) -> v a -> v a
insertSortBy a -> a -> Ordering
cmp v a
vec
    | Bool
otherwise = (forall s. ST s (v a)) -> v a
forall a. (forall s. ST s a) -> a
runST (do
        -- create two worker array
        MArr (IArray v) s a
w1 <- Int -> ST s (MArr (IArray v) s a)
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
Int -> m (MArr arr s a)
newArr Int
l
        MArr (IArray v) s a
w2 <- Int -> ST s (MArr (IArray v) s a)
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
Int -> m (MArr arr s a)
newArr Int
l
        v a -> Int -> MArr (IArray v) s a -> ST s ()
forall s. v a -> Int -> MArr (IArray v) s a -> ST s ()
firstPass v a
vec Int
0 MArr (IArray v) s a
w1
        IArray v a
w <- MArr (IArray v) s a
-> MArr (IArray v) s a -> Int -> ST s (IArray v a)
forall s.
MArr (IArray v) s a
-> MArr (IArray v) s a -> Int -> ST s (IArray v a)
mergePass MArr (IArray v) s a
w1 MArr (IArray v) s a
w2 Int
mergeTileSize
        v a -> ST s (v a)
forall (m :: * -> *) a. Monad m => a -> m a
return (v a -> ST s (v a)) -> v a -> ST s (v a)
forall a b. (a -> b) -> a -> b
$! IArray v a -> Int -> Int -> v a
forall (v :: * -> *) a. Vec v a => IArray v a -> Int -> Int -> v a
fromArr IArray v a
w Int
0 Int
l)
  where
    firstPass :: forall s. v a -> Int -> MArr (IArray v) s a -> ST s ()
    {-# INLINABLE firstPass #-}
    firstPass :: v a -> Int -> MArr (IArray v) s a -> ST s ()
firstPass !v a
v !Int
i !MArr (IArray v) s a
marr
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
l     = () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise = do
            let (v a
v',v a
rest) = Int -> v a -> (v a, v a)
forall (v :: * -> *) a. Vec v a => Int -> v a -> (v a, v a)
splitAt Int
mergeTileSize v a
v
            (a -> a -> Ordering)
-> v a -> Int -> MArr (IArray v) s a -> ST s ()
forall (v :: * -> *) a s.
Vec v a =>
(a -> a -> Ordering)
-> v a -> Int -> MArr (IArray v) s a -> ST s ()
insertSortToMArr a -> a -> Ordering
cmp v a
v' Int
i MArr (IArray v) s a
marr
            v a -> Int -> MArr (IArray v) s a -> ST s ()
forall s. v a -> Int -> MArr (IArray v) s a -> ST s ()
firstPass v a
rest (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
mergeTileSize) MArr (IArray v) s a
marr

    mergePass :: forall s. MArr (IArray v) s a -> MArr (IArray v) s a -> Int -> ST s (IArray v a)
    {-# INLINABLE mergePass #-}
    mergePass :: MArr (IArray v) s a
-> MArr (IArray v) s a -> Int -> ST s (IArray v a)
mergePass !MArr (IArray v) s a
w1 !MArr (IArray v) s a
w2 !Int
blockSiz
        | Int
blockSiz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
l = MArr (IArray v) s a -> ST s (IArray v a)
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> m (arr a)
unsafeFreezeArr MArr (IArray v) s a
w1
        | Bool
otherwise     = do
            MArr (IArray v) s a -> MArr (IArray v) s a -> Int -> Int -> ST s ()
forall s.
MArr (IArray v) s a -> MArr (IArray v) s a -> Int -> Int -> ST s ()
mergeLoop MArr (IArray v) s a
w1 MArr (IArray v) s a
w2 Int
blockSiz Int
0
            MArr (IArray v) s a
-> MArr (IArray v) s a -> Int -> ST s (IArray v a)
forall s.
MArr (IArray v) s a
-> MArr (IArray v) s a -> Int -> ST s (IArray v a)
mergePass MArr (IArray v) s a
w2 MArr (IArray v) s a
w1 (Int
blockSizInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
2) -- swap worker array and continue merging

    mergeLoop :: forall s. MArr (IArray v) s a -> MArr (IArray v) s a -> Int -> Int -> ST s ()
    {-# INLINABLE mergeLoop #-}
    mergeLoop :: MArr (IArray v) s a -> MArr (IArray v) s a -> Int -> Int -> ST s ()
mergeLoop !MArr (IArray v) s a
src !MArr (IArray v) s a
target !Int
blockSiz !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
blockSiz =                 -- remaining elements less than a block
            if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
l
            then () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            else MArr (IArray v) s a
-> Int -> MArr (IArray v) s a -> Int -> Int -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> MArr arr s a -> Int -> Int -> m ()
copyMutableArr MArr (IArray v) s a
target Int
i MArr (IArray v) s a
src Int
i (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i)
        | Bool
otherwise = do
            let !mergeEnd :: Int
mergeEnd = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
blockSizInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
blockSiz) Int
l
            MArr (IArray v) s a
-> MArr (IArray v) s a
-> Int
-> Int
-> Int
-> Int
-> Int
-> ST s ()
forall s.
MArr (IArray v) s a
-> MArr (IArray v) s a
-> Int
-> Int
-> Int
-> Int
-> Int
-> ST s ()
mergeBlock MArr (IArray v) s a
src MArr (IArray v) s a
target (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
blockSiz) Int
mergeEnd Int
i (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
blockSiz) Int
i
            MArr (IArray v) s a -> MArr (IArray v) s a -> Int -> Int -> ST s ()
forall s.
MArr (IArray v) s a -> MArr (IArray v) s a -> Int -> Int -> ST s ()
mergeLoop MArr (IArray v) s a
src MArr (IArray v) s a
target Int
blockSiz Int
mergeEnd

    mergeBlock :: forall s. MArr (IArray v) s a -> MArr (IArray v) s a -> Int -> Int -> Int -> Int -> Int -> ST s ()
    {-# INLINABLE mergeBlock #-}
    mergeBlock :: MArr (IArray v) s a
-> MArr (IArray v) s a
-> Int
-> Int
-> Int
-> Int
-> Int
-> ST s ()
mergeBlock !MArr (IArray v) s a
src !MArr (IArray v) s a
target !Int
leftEnd !Int
rightEnd !Int
i !Int
j !Int
k = do
        a
lv <- MArr (IArray v) s a -> Int -> ST s a
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MArr (IArray v) s a
src Int
i
        a
rv <- MArr (IArray v) s a -> Int -> ST s a
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MArr (IArray v) s a
src Int
j
        case a
rv a -> a -> Ordering
`cmp` a
lv of
            Ordering
LT -> do
                MArr (IArray v) s a -> Int -> a -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MArr (IArray v) s a
target Int
k a
rv
                let !j' :: Int
j' = Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                    !k' :: Int
k' = Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                if Int
j' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
rightEnd
                then MArr (IArray v) s a
-> Int -> MArr (IArray v) s a -> Int -> Int -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> MArr arr s a -> Int -> Int -> m ()
copyMutableArr MArr (IArray v) s a
target Int
k' MArr (IArray v) s a
src Int
i (Int
leftEnd Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i)
                else MArr (IArray v) s a
-> MArr (IArray v) s a
-> Int
-> Int
-> Int
-> Int
-> Int
-> ST s ()
forall s.
MArr (IArray v) s a
-> MArr (IArray v) s a
-> Int
-> Int
-> Int
-> Int
-> Int
-> ST s ()
mergeBlock MArr (IArray v) s a
src MArr (IArray v) s a
target Int
leftEnd Int
rightEnd Int
i Int
j' Int
k'
            Ordering
_ -> do
                MArr (IArray v) s a -> Int -> a -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MArr (IArray v) s a
target Int
k a
lv
                let !i' :: Int
i' = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                    !k' :: Int
k' = Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                if Int
i' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
leftEnd
                then MArr (IArray v) s a
-> Int -> MArr (IArray v) s a -> Int -> Int -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> MArr arr s a -> Int -> Int -> m ()
copyMutableArr MArr (IArray v) s a
target Int
k' MArr (IArray v) s a
src Int
j (Int
rightEnd Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j)
                else MArr (IArray v) s a
-> MArr (IArray v) s a
-> Int
-> Int
-> Int
-> Int
-> Int
-> ST s ()
forall s.
MArr (IArray v) s a
-> MArr (IArray v) s a
-> Int
-> Int
-> Int
-> Int
-> Int
-> ST s ()
mergeBlock MArr (IArray v) s a
src MArr (IArray v) s a
target Int
leftEnd Int
rightEnd Int
i' Int
j Int
k'

-- | The mergesort tile size, @mergeTileSize = 8@.
mergeTileSize :: Int
{-# INLINE mergeTileSize #-}
mergeTileSize :: Int
mergeTileSize = Int
8

-- | /O(n^2)/ Sort vector based on element's 'Ord' instance with simple
-- <https://en.wikipedia.org/wiki/Insertion_sort insertion-sort> algorithm.
--
-- This is a stable sort. O(n) extra space are needed,
-- which will be freezed into result vector.
insertSort :: (Vec v a, Ord a) => v a -> v a
{-# INLINE insertSort #-}
insertSort :: v a -> v a
insertSort = (a -> a -> Ordering) -> v a -> v a
forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Ordering) -> v a -> v a
insertSortBy a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare

insertSortBy :: Vec v a => (a -> a -> Ordering) -> v a -> v a
{-# INLINE insertSortBy #-}
insertSortBy :: (a -> a -> Ordering) -> v a -> v a
insertSortBy a -> a -> Ordering
cmp v :: v a
v@(Vec IArray v a
_ Int
_ Int
l) | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = v a
v
                               | Bool
otherwise = Int -> (forall s. MArr (IArray v) s a -> ST s ()) -> v a
forall (v :: * -> *) a.
Vec v a =>
Int -> (forall s. MArr (IArray v) s a -> ST s ()) -> v a
create Int
l ((a -> a -> Ordering)
-> v a -> Int -> MArr (IArray v) s a -> ST s ()
forall (v :: * -> *) a s.
Vec v a =>
(a -> a -> Ordering)
-> v a -> Int -> MArr (IArray v) s a -> ST s ()
insertSortToMArr a -> a -> Ordering
cmp v a
v Int
0)

insertSortToMArr  :: Vec v a
                  => (a -> a -> Ordering)
                  -> v a            -- the original vector
                  -> Int            -- writing offset in the mutable array
                  -> MArr (IArray v) s a   -- writing mutable array, must have enough space!
                  -> ST s ()
{-# INLINE insertSortToMArr #-}
insertSortToMArr :: (a -> a -> Ordering)
-> v a -> Int -> MArr (IArray v) s a -> ST s ()
insertSortToMArr a -> a -> Ordering
cmp (Vec IArray v a
arr Int
s Int
l) Int
moff MArr (IArray v) s a
marr = Int -> ST s ()
go Int
s
  where
    !end :: Int
end = Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l
    !doff :: Int
doff = Int
moffInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
s
    go :: Int -> ST s ()
go !Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
end  = () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
          | Bool
otherwise = case IArray v a -> Int -> (# a #)
forall (arr :: * -> *) a. Arr arr a => arr a -> Int -> (# a #)
indexArr' IArray v a
arr Int
i of
               (# a
x #) -> do a -> Int -> ST s ()
insert a
x (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
doff)
                             Int -> ST s ()
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
    insert :: a -> Int -> ST s ()
insert !a
temp !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
moff = do
            MArr (IArray v) s a -> Int -> a -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MArr (IArray v) s a
marr Int
moff a
temp
        | Bool
otherwise = do
            a
x <- MArr (IArray v) s a -> Int -> ST s a
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MArr (IArray v) s a
marr (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
            case a
temp a -> a -> Ordering
`cmp` a
x of
                Ordering
LT -> do
                    MArr (IArray v) s a -> Int -> a -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MArr (IArray v) s a
marr Int
i a
x
                    a -> Int -> ST s ()
insert a
temp (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                Ordering
_ -> MArr (IArray v) s a -> Int -> a -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MArr (IArray v) s a
marr Int
i a
temp

--------------------------------------------------------------------------------
-- Radix Sort

-- | Types contain radixs, which can be inspected with 'radix' during different 'passes'.
--
-- The default instances share a same 'bucketSize' 256, which seems to be a good default.
class Radix a where
    -- | The size of an auxiliary array, i.e. the counting bucket
    bucketSize :: a -> Int
    -- | The number of passes necessary to sort an array of es,
    --   it equals to the key's byte number.
    passes :: a -> Int
    -- | The radix function used in the first pass, works on the least significant bit.
    radixLSB  :: a -> Int
    -- | The radix function parameterized by the current pass (0 < pass < passes e-1).
    radix  :: Int -> a -> Int
    -- | The radix function used in the last pass, works on the most significant bit.
    radixMSB  :: a -> Int

instance Radix Int8 where
    {-# INLINE bucketSize #-};
    bucketSize :: Int8 -> Int
bucketSize Int8
_ = Int
256
    {-# INLINE passes #-}
    passes :: Int8 -> Int
passes Int8
_ = Int
1
    {-# INLINE radixLSB #-}
    radixLSB :: Int8 -> Int
radixLSB Int8
a =  Int
255 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
a Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` Int
128
    {-# INLINE radix #-}
    radix :: Int -> Int8 -> Int
radix Int
_ Int8
a =  Int
255 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
a Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` Int
128
    {-# INLINE radixMSB #-}
    radixMSB :: Int8 -> Int
radixMSB Int8
a =  Int
255 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
a Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` Int
128

#define MULTI_BYTES_INT_RADIX(T) \
    {-# INLINE bucketSize #-}; \
    bucketSize _ = 256; \
    {-# INLINE passes #-}; \
    passes _ = sizeOf (undefined :: T); \
    {-# INLINE radixLSB #-}; \
    radixLSB a = fromIntegral (255 .&. a); \
    {-# INLINE radix #-}; \
    radix i a = fromIntegral (a `unsafeShiftR` (i `unsafeShiftL` 3)) .&. 255; \
    {-# INLINE radixMSB #-}; \
    radixMSB a = fromIntegral ((a `xor` minBound) `unsafeShiftR` ((passes a-1) `unsafeShiftL` 3)) .&. 255

instance Radix Int where MULTI_BYTES_INT_RADIX(Int)
instance Radix Int16 where MULTI_BYTES_INT_RADIX(Int16)
instance Radix Int32 where MULTI_BYTES_INT_RADIX(Int32)
instance Radix Int64 where MULTI_BYTES_INT_RADIX(Int64)

instance Radix Word8 where
    {-# INLINE bucketSize #-};
    bucketSize :: Word8 -> Int
bucketSize Word8
_ = Int
256
    {-# INLINE passes #-}
    passes :: Word8 -> Int
passes Word8
_ = Int
1
    {-# INLINE radixLSB #-}
    radixLSB :: Word8 -> Int
radixLSB = Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral
    {-# INLINE radix #-}
    radix :: Int -> Word8 -> Int
radix Int
_  = Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral
    {-# INLINE radixMSB #-}
    radixMSB :: Word8 -> Int
radixMSB = Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral

#define MULTI_BYTES_WORD_RADIX(T) \
    {-# INLINE bucketSize #-}; \
    bucketSize _ = 256; \
    {-# INLINE passes #-}; \
    passes _ = sizeOf (undefined :: T); \
    {-# INLINE radixLSB #-}; \
    radixLSB a = fromIntegral (255 .&. a); \
    {-# INLINE radix #-}; \
    radix i a = fromIntegral (a `unsafeShiftR` (i `unsafeShiftL` 3)) .&. 255; \
    {-# INLINE radixMSB #-}; \
    radixMSB a = fromIntegral (a `unsafeShiftR` ((passes a-1) `unsafeShiftL` 3)) .&. 255

instance Radix Word where MULTI_BYTES_INT_RADIX(Word)
instance Radix Word16 where MULTI_BYTES_INT_RADIX(Word16)
instance Radix Word32 where MULTI_BYTES_INT_RADIX(Word32)
instance Radix Word64 where MULTI_BYTES_INT_RADIX(Word64)

-- | Similar to 'Down' newtype for 'Ord', this newtype can inverse the order of a 'Radix'
-- instance when used in 'radixSort'.
newtype RadixDown a = RadixDown a deriving (Int -> RadixDown a -> ShowS
[RadixDown a] -> ShowS
RadixDown a -> String
(Int -> RadixDown a -> ShowS)
-> (RadixDown a -> String)
-> ([RadixDown a] -> ShowS)
-> Show (RadixDown a)
forall a. Show a => Int -> RadixDown a -> ShowS
forall a. Show a => [RadixDown a] -> ShowS
forall a. Show a => RadixDown a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [RadixDown a] -> ShowS
$cshowList :: forall a. Show a => [RadixDown a] -> ShowS
show :: RadixDown a -> String
$cshow :: forall a. Show a => RadixDown a -> String
showsPrec :: Int -> RadixDown a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> RadixDown a -> ShowS
Show, RadixDown a -> RadixDown a -> Bool
(RadixDown a -> RadixDown a -> Bool)
-> (RadixDown a -> RadixDown a -> Bool) -> Eq (RadixDown a)
forall a. Eq a => RadixDown a -> RadixDown a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: RadixDown a -> RadixDown a -> Bool
$c/= :: forall a. Eq a => RadixDown a -> RadixDown a -> Bool
== :: RadixDown a -> RadixDown a -> Bool
$c== :: forall a. Eq a => RadixDown a -> RadixDown a -> Bool
Eq)
                                    deriving newtype (Addr# -> Int# -> RadixDown a
Addr# -> Int# -> Int# -> RadixDown a -> State# s -> State# s
Addr# -> Int# -> State# s -> (# State# s, RadixDown a #)
Addr# -> Int# -> RadixDown a -> State# s -> State# s
ByteArray# -> Int# -> RadixDown a
MutableByteArray# s
-> Int# -> State# s -> (# State# s, RadixDown a #)
MutableByteArray# s -> Int# -> RadixDown a -> State# s -> State# s
MutableByteArray# s
-> Int# -> Int# -> RadixDown a -> State# s -> State# s
RadixDown a -> Int#
(RadixDown a -> Int#)
-> (RadixDown a -> Int#)
-> (ByteArray# -> Int# -> RadixDown a)
-> (forall s.
    MutableByteArray# s
    -> Int# -> State# s -> (# State# s, RadixDown a #))
-> (forall s.
    MutableByteArray# s -> Int# -> RadixDown a -> State# s -> State# s)
-> (forall s.
    MutableByteArray# s
    -> Int# -> Int# -> RadixDown a -> State# s -> State# s)
-> (Addr# -> Int# -> RadixDown a)
-> (forall s.
    Addr# -> Int# -> State# s -> (# State# s, RadixDown a #))
-> (forall s. Addr# -> Int# -> RadixDown a -> State# s -> State# s)
-> (forall s.
    Addr# -> Int# -> Int# -> RadixDown a -> State# s -> State# s)
-> Prim (RadixDown a)
forall s.
Addr# -> Int# -> Int# -> RadixDown a -> State# s -> State# s
forall s. Addr# -> Int# -> State# s -> (# State# s, RadixDown a #)
forall s. Addr# -> Int# -> RadixDown a -> State# s -> State# s
forall s.
MutableByteArray# s
-> Int# -> Int# -> RadixDown a -> State# s -> State# s
forall s.
MutableByteArray# s
-> Int# -> State# s -> (# State# s, RadixDown a #)
forall s.
MutableByteArray# s -> Int# -> RadixDown a -> State# s -> State# s
forall a. Prim a => Addr# -> Int# -> RadixDown a
forall a. Prim a => ByteArray# -> Int# -> RadixDown a
forall a. Prim a => RadixDown a -> Int#
forall a s.
Prim a =>
Addr# -> Int# -> Int# -> RadixDown a -> State# s -> State# s
forall a s.
Prim a =>
Addr# -> Int# -> State# s -> (# State# s, RadixDown a #)
forall a s.
Prim a =>
Addr# -> Int# -> RadixDown a -> State# s -> State# s
forall a s.
Prim a =>
MutableByteArray# s
-> Int# -> Int# -> RadixDown a -> State# s -> State# s
forall a s.
Prim a =>
MutableByteArray# s
-> Int# -> State# s -> (# State# s, RadixDown a #)
forall a s.
Prim a =>
MutableByteArray# s -> Int# -> RadixDown a -> State# s -> State# s
forall a.
(a -> Int#)
-> (a -> Int#)
-> (ByteArray# -> Int# -> a)
-> (forall s.
    MutableByteArray# s -> Int# -> State# s -> (# State# s, a #))
-> (forall s.
    MutableByteArray# s -> Int# -> a -> State# s -> State# s)
-> (forall s.
    MutableByteArray# s -> Int# -> Int# -> a -> State# s -> State# s)
-> (Addr# -> Int# -> a)
-> (forall s. Addr# -> Int# -> State# s -> (# State# s, a #))
-> (forall s. Addr# -> Int# -> a -> State# s -> State# s)
-> (forall s. Addr# -> Int# -> Int# -> a -> State# s -> State# s)
-> Prim a
setOffAddr# :: Addr# -> Int# -> Int# -> RadixDown a -> State# s -> State# s
$csetOffAddr# :: forall a s.
Prim a =>
Addr# -> Int# -> Int# -> RadixDown a -> State# s -> State# s
writeOffAddr# :: Addr# -> Int# -> RadixDown a -> State# s -> State# s
$cwriteOffAddr# :: forall a s.
Prim a =>
Addr# -> Int# -> RadixDown a -> State# s -> State# s
readOffAddr# :: Addr# -> Int# -> State# s -> (# State# s, RadixDown a #)
$creadOffAddr# :: forall a s.
Prim a =>
Addr# -> Int# -> State# s -> (# State# s, RadixDown a #)
indexOffAddr# :: Addr# -> Int# -> RadixDown a
$cindexOffAddr# :: forall a. Prim a => Addr# -> Int# -> RadixDown a
setByteArray# :: MutableByteArray# s
-> Int# -> Int# -> RadixDown a -> State# s -> State# s
$csetByteArray# :: forall a s.
Prim a =>
MutableByteArray# s
-> Int# -> Int# -> RadixDown a -> State# s -> State# s
writeByteArray# :: MutableByteArray# s -> Int# -> RadixDown a -> State# s -> State# s
$cwriteByteArray# :: forall a s.
Prim a =>
MutableByteArray# s -> Int# -> RadixDown a -> State# s -> State# s
readByteArray# :: MutableByteArray# s
-> Int# -> State# s -> (# State# s, RadixDown a #)
$creadByteArray# :: forall a s.
Prim a =>
MutableByteArray# s
-> Int# -> State# s -> (# State# s, RadixDown a #)
indexByteArray# :: ByteArray# -> Int# -> RadixDown a
$cindexByteArray# :: forall a. Prim a => ByteArray# -> Int# -> RadixDown a
alignment# :: RadixDown a -> Int#
$calignment# :: forall a. Prim a => RadixDown a -> Int#
sizeOf# :: RadixDown a -> Int#
$csizeOf# :: forall a. Prim a => RadixDown a -> Int#
Prim, ByteArray# -> Int# -> RadixDown a
ByteArray# -> Int -> RadixDown a
MutableByteArray# s
-> Int# -> State# s -> (# State# s, RadixDown a #)
MutableByteArray# s -> Int# -> RadixDown a -> State# s -> State# s
MutableByteArray# RealWorld -> Int -> IO (RadixDown a)
MutableByteArray# RealWorld -> Int -> RadixDown a -> IO ()
RadixDown a -> Int#
RadixDown a -> Int
(RadixDown a -> Int)
-> (RadixDown a -> Int#)
-> (ByteArray# -> Int# -> RadixDown a)
-> (forall s.
    MutableByteArray# s
    -> Int# -> State# s -> (# State# s, RadixDown a #))
-> (forall s.
    MutableByteArray# s -> Int# -> RadixDown a -> State# s -> State# s)
-> (MutableByteArray# RealWorld -> Int -> IO (RadixDown a))
-> (MutableByteArray# RealWorld -> Int -> RadixDown a -> IO ())
-> (ByteArray# -> Int -> RadixDown a)
-> Unaligned (RadixDown a)
forall s.
MutableByteArray# s
-> Int# -> State# s -> (# State# s, RadixDown a #)
forall s.
MutableByteArray# s -> Int# -> RadixDown a -> State# s -> State# s
forall a. Unaligned a => ByteArray# -> Int# -> RadixDown a
forall a. Unaligned a => ByteArray# -> Int -> RadixDown a
forall a.
Unaligned a =>
MutableByteArray# RealWorld -> Int -> IO (RadixDown a)
forall a.
Unaligned a =>
MutableByteArray# RealWorld -> Int -> RadixDown a -> IO ()
forall a. Unaligned a => RadixDown a -> Int#
forall a. Unaligned a => RadixDown a -> Int
forall a s.
Unaligned a =>
MutableByteArray# s
-> Int# -> State# s -> (# State# s, RadixDown a #)
forall a s.
Unaligned a =>
MutableByteArray# s -> Int# -> RadixDown a -> State# s -> State# s
forall a.
(a -> Int)
-> (a -> Int#)
-> (ByteArray# -> Int# -> a)
-> (forall s.
    MutableByteArray# s -> Int# -> State# s -> (# State# s, a #))
-> (forall s.
    MutableByteArray# s -> Int# -> a -> State# s -> State# s)
-> (MutableByteArray# RealWorld -> Int -> IO a)
-> (MutableByteArray# RealWorld -> Int -> a -> IO ())
-> (ByteArray# -> Int -> a)
-> Unaligned a
indexBA :: ByteArray# -> Int -> RadixDown a
$cindexBA :: forall a. Unaligned a => ByteArray# -> Int -> RadixDown a
pokeMBA :: MutableByteArray# RealWorld -> Int -> RadixDown a -> IO ()
$cpokeMBA :: forall a.
Unaligned a =>
MutableByteArray# RealWorld -> Int -> RadixDown a -> IO ()
peekMBA :: MutableByteArray# RealWorld -> Int -> IO (RadixDown a)
$cpeekMBA :: forall a.
Unaligned a =>
MutableByteArray# RealWorld -> Int -> IO (RadixDown a)
writeWord8ArrayAs# :: MutableByteArray# s -> Int# -> RadixDown a -> State# s -> State# s
$cwriteWord8ArrayAs# :: forall a s.
Unaligned a =>
MutableByteArray# s -> Int# -> RadixDown a -> State# s -> State# s
readWord8ArrayAs# :: MutableByteArray# s
-> Int# -> State# s -> (# State# s, RadixDown a #)
$creadWord8ArrayAs# :: forall a s.
Unaligned a =>
MutableByteArray# s
-> Int# -> State# s -> (# State# s, RadixDown a #)
indexWord8ArrayAs# :: ByteArray# -> Int# -> RadixDown a
$cindexWord8ArrayAs# :: forall a. Unaligned a => ByteArray# -> Int# -> RadixDown a
unalignedSize# :: RadixDown a -> Int#
$cunalignedSize# :: forall a. Unaligned a => RadixDown a -> Int#
unalignedSize :: RadixDown a -> Int
$cunalignedSize :: forall a. Unaligned a => RadixDown a -> Int
Unaligned)

instance Radix a => Radix (RadixDown a) where
    {-# INLINE bucketSize #-}
    bucketSize :: RadixDown a -> Int
bucketSize (RadixDown a
a) = a -> Int
forall a. Radix a => a -> Int
bucketSize a
a
    {-# INLINE passes #-}
    passes :: RadixDown a -> Int
passes (RadixDown a
a)  = a -> Int
forall a. Radix a => a -> Int
passes a
a
    {-# INLINE radixLSB #-}
    radixLSB :: RadixDown a -> Int
radixLSB (RadixDown a
a) = a -> Int
forall a. Radix a => a -> Int
bucketSize a
a Int -> Int -> Int
forall a. Num a => a -> a -> a
- a -> Int
forall a. Radix a => a -> Int
radixLSB a
a Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1
    {-# INLINE radix #-}
    radix :: Int -> RadixDown a -> Int
radix Int
i (RadixDown a
a) = a -> Int
forall a. Radix a => a -> Int
bucketSize a
a Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> a -> Int
forall a. Radix a => Int -> a -> Int
radix Int
i a
a Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1
    {-# INLINE radixMSB #-}
    radixMSB :: RadixDown a -> Int
radixMSB (RadixDown a
a) = a -> Int
forall a. Radix a => a -> Int
bucketSize a
a Int -> Int -> Int
forall a. Num a => a -> a -> a
- a -> Int
forall a. Radix a => a -> Int
radixMSB a
a Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1

-- | /O(n)/ Sort vector based on element's 'Radix' instance with
-- <https://en.wikipedia.org/wiki/Radix_sort radix-sort>,
-- (Least significant digit radix sorts variation).
--
-- This is a stable sort, one or two extra O(n) worker array are need
-- depend on how many 'passes' shall be performed, and a 'bucketSize'
-- counting bucket are also needed. This sort algorithms performed extremly
-- well on small byte size types such as 'Int8' or 'Word8', while on larger
-- type, constant passes may render this algorithm not suitable for small
-- vectors (turning point around 2^(2*passes)).
radixSort :: forall v a. (Vec v a, Radix a) => v a -> v a
{-# INLINABLE radixSort #-}
radixSort :: v a -> v a
radixSort v :: v a
v@(Vec IArray v a
arr Int
s Int
l)
    | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = v a
v
    | Bool
otherwise = (forall s. ST s (v a)) -> v a
forall a. (forall s. ST s a) -> a
runST (do
        MutablePrimArray s Int
bucket <- Int -> Int -> ST s (MArr PrimArray s Int)
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
Int -> a -> m (MArr arr s a)
newArrWith Int
buktSiz Int
0 :: ST s (MutablePrimArray s Int)
        MArr (IArray v) s a
w1 <- Int -> ST s (MArr (IArray v) s a)
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
Int -> m (MArr arr s a)
newArr Int
l :: ST s (MArr (IArray v) s a)
        IArray v a -> MutablePrimArray s Int -> Int -> ST s ()
forall s. IArray v a -> MutablePrimArray s Int -> Int -> ST s ()
firstCountPass IArray v a
arr MutablePrimArray s Int
bucket Int
s
        MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
forall s. MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
accumBucket MutablePrimArray s Int
bucket Int
buktSiz Int
0 Int
0
        IArray v a
-> Int -> MutablePrimArray s Int -> MArr (IArray v) s a -> ST s ()
forall s.
IArray v a
-> Int -> MutablePrimArray s Int -> MArr (IArray v) s a -> ST s ()
firstMovePass IArray v a
arr Int
s MutablePrimArray s Int
bucket MArr (IArray v) s a
w1
        IArray v a
w <- if Int
passSiz Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
            then MArr (IArray v) s a -> ST s (IArray v a)
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> m (arr a)
unsafeFreezeArr MArr (IArray v) s a
w1
            else do
                MArr (IArray v) s a
w2 <- Int -> ST s (MArr (IArray v) s a)
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
Int -> m (MArr arr s a)
newArr Int
l :: ST s (MArr (IArray v) s a)
                MArr (IArray v) s a
-> MArr (IArray v) s a
-> MutablePrimArray s Int
-> Int
-> Int
-> ST s (IArray v a)
forall s.
MArr (IArray v) s a
-> MArr (IArray v) s a
-> MutablePrimArray s Int
-> Int
-> Int
-> ST s (IArray v a)
radixLoop MArr (IArray v) s a
w1 MArr (IArray v) s a
w2 MutablePrimArray s Int
bucket Int
buktSiz Int
1
        v a -> ST s (v a)
forall (m :: * -> *) a. Monad m => a -> m a
return (v a -> ST s (v a)) -> v a -> ST s (v a)
forall a b. (a -> b) -> a -> b
$! IArray v a -> Int -> Int -> v a
forall (v :: * -> *) a. Vec v a => IArray v a -> Int -> Int -> v a
fromArr IArray v a
w Int
0 Int
l)
  where
    passSiz :: Int
passSiz = a -> Int
forall a. Radix a => a -> Int
passes (a
forall a. HasCallStack => a
undefined :: a)
    buktSiz :: Int
buktSiz = a -> Int
forall a. Radix a => a -> Int
bucketSize (a
forall a. HasCallStack => a
undefined :: a)
    !end :: Int
end = Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l

    {-# INLINABLE firstCountPass #-}
    firstCountPass :: forall s. IArray v a -> MutablePrimArray s Int -> Int -> ST s ()
    firstCountPass :: IArray v a -> MutablePrimArray s Int -> Int -> ST s ()
firstCountPass !IArray v a
arr' !MutablePrimArray s Int
bucket !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
end  = () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise = case IArray v a -> Int -> (# a #)
forall (arr :: * -> *) a. Arr arr a => arr a -> Int -> (# a #)
indexArr' IArray v a
arr' Int
i of
            (# a
x #) -> do
                let !r :: Int
r = a -> Int
forall a. Radix a => a -> Int
radixLSB a
x
                Int
c <- MArr PrimArray s Int -> Int -> ST s Int
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
r
                MArr PrimArray s Int -> Int -> Int -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
r (Int
cInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                IArray v a -> MutablePrimArray s Int -> Int -> ST s ()
forall s. IArray v a -> MutablePrimArray s Int -> Int -> ST s ()
firstCountPass IArray v a
arr' MutablePrimArray s Int
bucket (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

    {-# INLINABLE accumBucket #-}
    accumBucket :: forall s. MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
    accumBucket :: MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
accumBucket !MutablePrimArray s Int
bucket !Int
bsiz !Int
i !Int
acc
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
bsiz = () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise = do
            Int
c <- MArr PrimArray s Int -> Int -> ST s Int
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
i
            MArr PrimArray s Int -> Int -> Int -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
i Int
acc
            MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
forall s. MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
accumBucket MutablePrimArray s Int
bucket Int
bsiz (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
accInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
c)

    {-# INLINABLE firstMovePass #-}
    firstMovePass :: forall s. IArray v a -> Int -> MutablePrimArray s Int -> MArr (IArray v) s a -> ST s ()
    firstMovePass :: IArray v a
-> Int -> MutablePrimArray s Int -> MArr (IArray v) s a -> ST s ()
firstMovePass !IArray v a
arr' !Int
i !MutablePrimArray s Int
bucket !MArr (IArray v) s a
w
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
end  = () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise = case IArray v a -> Int -> (# a #)
forall (arr :: * -> *) a. Arr arr a => arr a -> Int -> (# a #)
indexArr' IArray v a
arr' Int
i of
            (# a
x #) -> do
                let !r :: Int
r = a -> Int
forall a. Radix a => a -> Int
radixLSB a
x
                Int
c <- MArr PrimArray s Int -> Int -> ST s Int
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
r
                MArr PrimArray s Int -> Int -> Int -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
r (Int
cInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                MArr (IArray v) s a -> Int -> a -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MArr (IArray v) s a
w Int
c a
x
                IArray v a
-> Int -> MutablePrimArray s Int -> MArr (IArray v) s a -> ST s ()
forall s.
IArray v a
-> Int -> MutablePrimArray s Int -> MArr (IArray v) s a -> ST s ()
firstMovePass IArray v a
arr' (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) MutablePrimArray s Int
bucket MArr (IArray v) s a
w

    {-# INLINABLE radixLoop #-}
    radixLoop :: forall s. MArr (IArray v) s a -> MArr (IArray v) s a -> MutablePrimArray s Int -> Int -> Int -> ST s ((IArray v) a)
    radixLoop :: MArr (IArray v) s a
-> MArr (IArray v) s a
-> MutablePrimArray s Int
-> Int
-> Int
-> ST s (IArray v a)
radixLoop !MArr (IArray v) s a
w1 !MArr (IArray v) s a
w2 !MutablePrimArray s Int
bucket !Int
bsiz !Int
pass
        | Int
pass Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
passSizInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1 = do
            MArr PrimArray s Int -> Int -> Int -> Int -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> Int -> a -> m ()
setArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
0 Int
bsiz Int
0   -- clear the counting bucket
            MArr (IArray v) s a -> MutablePrimArray s Int -> Int -> ST s ()
forall s.
MArr (IArray v) s a -> MutablePrimArray s Int -> Int -> ST s ()
lastCountPass MArr (IArray v) s a
w1 MutablePrimArray s Int
bucket Int
0
            MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
forall s. MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
accumBucket MutablePrimArray s Int
bucket Int
bsiz Int
0 Int
0
            MArr (IArray v) s a
-> MutablePrimArray s Int -> MArr (IArray v) s a -> Int -> ST s ()
forall s.
MArr (IArray v) s a
-> MutablePrimArray s Int -> MArr (IArray v) s a -> Int -> ST s ()
lastMovePass MArr (IArray v) s a
w1 MutablePrimArray s Int
bucket MArr (IArray v) s a
w2 Int
0
            MArr (IArray v) s a -> ST s (IArray v a)
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> m (arr a)
unsafeFreezeArr MArr (IArray v) s a
w2
        | Bool
otherwise = do
            MArr PrimArray s Int -> Int -> Int -> Int -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> Int -> a -> m ()
setArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
0 Int
bsiz Int
0   -- clear the counting bucket
            MArr (IArray v) s a
-> MutablePrimArray s Int -> Int -> Int -> ST s ()
forall s.
MArr (IArray v) s a
-> MutablePrimArray s Int -> Int -> Int -> ST s ()
countPass MArr (IArray v) s a
w1 MutablePrimArray s Int
bucket Int
pass Int
0
            MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
forall s. MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
accumBucket MutablePrimArray s Int
bucket Int
bsiz Int
0 Int
0
            MArr (IArray v) s a
-> MutablePrimArray s Int
-> Int
-> MArr (IArray v) s a
-> Int
-> ST s ()
forall s.
MArr (IArray v) s a
-> MutablePrimArray s Int
-> Int
-> MArr (IArray v) s a
-> Int
-> ST s ()
movePass MArr (IArray v) s a
w1 MutablePrimArray s Int
bucket Int
pass MArr (IArray v) s a
w2 Int
0
            MArr (IArray v) s a
-> MArr (IArray v) s a
-> MutablePrimArray s Int
-> Int
-> Int
-> ST s (IArray v a)
forall s.
MArr (IArray v) s a
-> MArr (IArray v) s a
-> MutablePrimArray s Int
-> Int
-> Int
-> ST s (IArray v a)
radixLoop MArr (IArray v) s a
w2 MArr (IArray v) s a
w1 MutablePrimArray s Int
bucket Int
bsiz (Int
passInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

    {-# INLINABLE countPass #-}
    countPass :: forall s. MArr (IArray v) s a -> MutablePrimArray s Int -> Int -> Int -> ST s ()
    countPass :: MArr (IArray v) s a
-> MutablePrimArray s Int -> Int -> Int -> ST s ()
countPass !MArr (IArray v) s a
marr !MutablePrimArray s Int
bucket !Int
pass !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
l  = () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise = do
                a
x <- MArr (IArray v) s a -> Int -> ST s a
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MArr (IArray v) s a
marr Int
i
                let !r :: Int
r = Int -> a -> Int
forall a. Radix a => Int -> a -> Int
radix Int
pass a
x
                Int
c <- MArr PrimArray s Int -> Int -> ST s Int
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
r
                MArr PrimArray s Int -> Int -> Int -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
r (Int
cInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                MArr (IArray v) s a
-> MutablePrimArray s Int -> Int -> Int -> ST s ()
forall s.
MArr (IArray v) s a
-> MutablePrimArray s Int -> Int -> Int -> ST s ()
countPass MArr (IArray v) s a
marr MutablePrimArray s Int
bucket Int
pass (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

    {-# INLINABLE movePass #-}
    movePass :: forall s. MArr (IArray v) s a -> MutablePrimArray s Int -> Int -> MArr (IArray v) s a -> Int -> ST s ()
    movePass :: MArr (IArray v) s a
-> MutablePrimArray s Int
-> Int
-> MArr (IArray v) s a
-> Int
-> ST s ()
movePass !MArr (IArray v) s a
src !MutablePrimArray s Int
bucket !Int
pass !MArr (IArray v) s a
target !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
l  = () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise = do
                a
x <- MArr (IArray v) s a -> Int -> ST s a
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MArr (IArray v) s a
src Int
i
                let !r :: Int
r = Int -> a -> Int
forall a. Radix a => Int -> a -> Int
radix Int
pass a
x
                Int
c <- MArr PrimArray s Int -> Int -> ST s Int
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
r
                MArr PrimArray s Int -> Int -> Int -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
r (Int
cInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                MArr (IArray v) s a -> Int -> a -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MArr (IArray v) s a
target Int
c a
x
                MArr (IArray v) s a
-> MutablePrimArray s Int
-> Int
-> MArr (IArray v) s a
-> Int
-> ST s ()
forall s.
MArr (IArray v) s a
-> MutablePrimArray s Int
-> Int
-> MArr (IArray v) s a
-> Int
-> ST s ()
movePass MArr (IArray v) s a
src MutablePrimArray s Int
bucket Int
pass MArr (IArray v) s a
target (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

    {-# INLINABLE lastCountPass #-}
    lastCountPass :: forall s. MArr (IArray v) s a -> MutablePrimArray s Int -> Int -> ST s ()
    lastCountPass :: MArr (IArray v) s a -> MutablePrimArray s Int -> Int -> ST s ()
lastCountPass !MArr (IArray v) s a
marr !MutablePrimArray s Int
bucket !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
l  = () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise = do
                a
x <- MArr (IArray v) s a -> Int -> ST s a
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MArr (IArray v) s a
marr Int
i
                let !r :: Int
r = a -> Int
forall a. Radix a => a -> Int
radixMSB a
x
                Int
c <- MArr PrimArray s Int -> Int -> ST s Int
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
r
                MArr PrimArray s Int -> Int -> Int -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
r (Int
cInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                MArr (IArray v) s a -> MutablePrimArray s Int -> Int -> ST s ()
forall s.
MArr (IArray v) s a -> MutablePrimArray s Int -> Int -> ST s ()
lastCountPass MArr (IArray v) s a
marr MutablePrimArray s Int
bucket (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

    {-# INLINABLE lastMovePass #-}
    lastMovePass :: forall s. MArr (IArray v) s a -> MutablePrimArray s Int -> MArr (IArray v) s a -> Int -> ST s ()
    lastMovePass :: MArr (IArray v) s a
-> MutablePrimArray s Int -> MArr (IArray v) s a -> Int -> ST s ()
lastMovePass !MArr (IArray v) s a
src !MutablePrimArray s Int
bucket !MArr (IArray v) s a
target !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
l  = () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise = do
                a
x <- MArr (IArray v) s a -> Int -> ST s a
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MArr (IArray v) s a
src Int
i
                let !r :: Int
r = a -> Int
forall a. Radix a => a -> Int
radixMSB a
x
                Int
c <- MArr PrimArray s Int -> Int -> ST s Int
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> m a
readArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
r
                MArr PrimArray s Int -> Int -> Int -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MutablePrimArray s Int
MArr PrimArray s Int
bucket Int
r (Int
cInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                MArr (IArray v) s a -> Int -> a -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MArr (IArray v) s a
target Int
c a
x
                MArr (IArray v) s a
-> MutablePrimArray s Int -> MArr (IArray v) s a -> Int -> ST s ()
forall s.
MArr (IArray v) s a
-> MutablePrimArray s Int -> MArr (IArray v) s a -> Int -> ST s ()
lastMovePass MArr (IArray v) s a
src MutablePrimArray s Int
bucket MArr (IArray v) s a
target (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

{- In fact IEEE float can be radix sorted like following:

newtype RadixDouble = RadixDouble Int64 deriving (Show, Eq, Prim)
instance Cast RadixDouble Double where cast (RadixDouble a) = cast a
instance Cast Double RadixDouble where cast a = RadixDouble (cast a)
instance Radix RadixDouble where
    {-# INLINE bucketSize #-}
    bucketSize (RadixDouble _) = 256
    {-# INLINE passes #-}
    passes (RadixDouble _)  = 8
    {-# INLINE radixLSB #-}
    radixLSB (RadixDouble a) | a > 0 = r
                             | otherwise = 255 - r
      where r = radixLSB a
    {-# INLINE radix #-}
    radix i (RadixDouble a) | a > 0 = r
                            | otherwise = 255 - r
      where r = radix i a
    {-# INLINE radixMSB #-}
    radixMSB (RadixDouble a) | r < 128  = r + 128
                             | otherwise = 255 - r
      where r = radixMSB (fromIntegral a :: Word64)

radixSortDouble :: PrimVector Double -> PrimVector Double
radixSortDouble v =  castVector (radixSort (castVector v :: PrimVector RadixDouble))

newtype RadixFloat = RadixFloat Int32 deriving (Show, Eq, Prim)
instance Cast RadixFloat Float where cast (RadixFloat a) = cast a
instance Cast Float RadixFloat where cast a = RadixFloat (cast a)
instance Radix RadixFloat where
    {-# INLINE bucketSize #-}
    bucketSize (RadixFloat _) = 256
    {-# INLINE passes #-}
    passes (RadixFloat _)  = 4
    {-# INLINE radixLSB #-}
    radixLSB (RadixFloat a) | a > 0 = r
                            | otherwise = 255 - r
      where r = radixLSB a
    {-# INLINE radix #-}
    radix i (RadixFloat a) | a > 0 = r
                           | otherwise = 255 - r
      where r = radix i a
    {-# INLINE radixMSB #-}
    radixMSB (RadixFloat a) | r < 128  = r + 128
                            | otherwise = 255 - r
      where r = radixMSB (fromIntegral a :: Word32)

radixSortFloat :: PrimVector Float -> PrimVector Float
radixSortFloat v =  castVector (radixSort (castVector v :: PrimVector RadixFloat))
-}

--------------------------------------------------------------------------------
-- | merge duplicated adjacent element, prefer left element.
--
-- Use this function on a sorted vector will have the same effects as 'nub'.
mergeDupAdjacent :: forall v a. (Vec v a, Eq a) => v a -> v a
{-# INLINE mergeDupAdjacent #-}
mergeDupAdjacent :: v a -> v a
mergeDupAdjacent = (a -> a -> Bool) -> (a -> a -> a) -> v a -> v a
forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Bool) -> (a -> a -> a) -> v a -> v a
mergeDupAdjacentBy a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==) a -> a -> a
forall a b. a -> b -> a
const

-- | Merge duplicated adjacent element, prefer left element.
mergeDupAdjacentLeft :: forall v a. Vec v a
                     => (a -> a -> Bool)   -- ^ equality tester, @\ left right -> eq left right@
                     -> v a
                     -> v a
mergeDupAdjacentLeft :: (a -> a -> Bool) -> v a -> v a
mergeDupAdjacentLeft a -> a -> Bool
eq = (a -> a -> Bool) -> (a -> a -> a) -> v a -> v a
forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Bool) -> (a -> a -> a) -> v a -> v a
mergeDupAdjacentBy a -> a -> Bool
eq a -> a -> a
forall a b. a -> b -> a
const
{-# INLINE mergeDupAdjacentLeft #-}

-- | Merge duplicated adjacent element, prefer right element.
mergeDupAdjacentRight :: forall v a. Vec v a
                      => (a -> a -> Bool)  -- ^ equality tester, @\ left right -> eq left right@
                      -> v a
                      -> v a
{-# INLINE mergeDupAdjacentRight #-}
mergeDupAdjacentRight :: (a -> a -> Bool) -> v a -> v a
mergeDupAdjacentRight a -> a -> Bool
eq = (a -> a -> Bool) -> (a -> a -> a) -> v a -> v a
forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Bool) -> (a -> a -> a) -> v a -> v a
mergeDupAdjacentBy a -> a -> Bool
eq (\ a
_ a
x -> a
x)

-- | Merge duplicated adjacent element, based on a equality tester and a merger function.
mergeDupAdjacentBy :: forall v a. Vec v a
                   => (a -> a -> Bool)  -- ^ equality tester, @\ left right -> eq left right@
                   -> (a -> a -> a)     -- ^ the merger, @\ left right -> merge left right@
                   -> v a -> v a
{-# INLINABLE mergeDupAdjacentBy #-}
mergeDupAdjacentBy :: (a -> a -> Bool) -> (a -> a -> a) -> v a -> v a
mergeDupAdjacentBy a -> a -> Bool
eq a -> a -> a
merger v :: v a
v@(Vec IArray v a
arr Int
s Int
l)
    | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = v a
forall (v :: * -> *) a. Vec v a => v a
empty
    | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = v a
v
    | Bool
otherwise = Int -> (forall s. MArr (IArray v) s a -> ST s Int) -> v a
forall (v :: * -> *) a.
(Vec v a, HasCallStack) =>
Int -> (forall s. MArr (IArray v) s a -> ST s Int) -> v a
createN Int
l ((forall s. MArr (IArray v) s a -> ST s Int) -> v a)
-> (forall s. MArr (IArray v) s a -> ST s Int) -> v a
forall a b. (a -> b) -> a -> b
$ \ MArr (IArray v) s a
marr -> do
        a
x0 <- IArray v a -> Int -> ST s a
forall (arr :: * -> *) a (m :: * -> *).
(Arr arr a, Monad m) =>
arr a -> Int -> m a
indexArrM IArray v a
arr Int
0
        MArr (IArray v) s a -> Int -> a -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MArr (IArray v) s a
marr Int
0 a
x0
        IArray v a -> MArr (IArray v) s a -> Int -> Int -> a -> ST s Int
forall s.
IArray v a -> MArr (IArray v) s a -> Int -> Int -> a -> ST s Int
go IArray v a
arr MArr (IArray v) s a
marr Int
s Int
1 a
x0
  where
    !end :: Int
end = Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l
    go :: forall s. IArray v a -> MArr (IArray v) s a -> Int -> Int -> a -> ST s Int
    go :: IArray v a -> MArr (IArray v) s a -> Int -> Int -> a -> ST s Int
go !IArray v a
arr' !MArr (IArray v) s a
marr !Int
i !Int
j !a
x
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
end  = Int -> ST s Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
j
        | Bool
otherwise = do
            a
x' <- IArray v a -> Int -> ST s a
forall (arr :: * -> *) a (m :: * -> *).
(Arr arr a, Monad m) =>
arr a -> Int -> m a
indexArrM IArray v a
arr' Int
i
            if a
x a -> a -> Bool
`eq` a
x'
            then do
                let !x'' :: a
x'' = a -> a -> a
merger a
x a
x'
                MArr (IArray v) s a -> Int -> a -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MArr (IArray v) s a
marr (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a
x''
                IArray v a -> MArr (IArray v) s a -> Int -> Int -> a -> ST s Int
forall s.
IArray v a -> MArr (IArray v) s a -> Int -> Int -> a -> ST s Int
go IArray v a
arr' MArr (IArray v) s a
marr (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
j a
x''
            else do
                MArr (IArray v) s a -> Int -> a -> ST s ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MArr (IArray v) s a
marr Int
j a
x'
                IArray v a -> MArr (IArray v) s a -> Int -> Int -> a -> ST s Int
forall s.
IArray v a -> MArr (IArray v) s a -> Int -> Int -> a -> ST s Int
go IArray v a
arr' MArr (IArray v) s a
marr (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) a
x'