{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}

module Data.RRBVector.Internal.Sorting
    ( sort
    , sortBy
    , sortOn
    ) where

import Data.Foldable (toList)
import Data.Foldable.WithIndex (ifor_)
import Data.Primitive.Array
import Data.SamSort (sortArrayBy)
import Data.Semigroup (Arg(..))

import Data.RRBVector.Internal

uninitialized :: a
uninitialized :: forall a. a
uninitialized = [Char] -> a
forall a. [Char] -> a
errorWithoutStackTrace [Char]
"uninitialized"

-- | \(O(n \log n)\). Sort the vector in ascending order.
-- The sort is stable, meaning the order of equal elements is preserved.
--
-- @since 0.2.2.0
sort :: (Ord a) => Vector a -> Vector a
sort :: forall a. Ord a => Vector a -> Vector a
sort = (a -> a -> Ordering) -> Vector a -> Vector a
forall a. (a -> a -> Ordering) -> Vector a -> Vector a
sortBy a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare

-- | \(O(n \log n)\). Sort the vector in ascending order according to the specified comparison function.
-- The sort is stable, meaning the order of equal elements is preserved.
--
-- @since 0.2.2.0
sortBy :: (a -> a -> Ordering) -> Vector a -> Vector a
sortBy :: forall a. (a -> a -> Ordering) -> Vector a -> Vector a
sortBy a -> a -> Ordering
cmp Vector a
v =
    let sortedArr :: Array a
sortedArr = Int -> a -> (forall s. MutableArray s a -> ST s ()) -> Array a
forall a.
Int -> a -> (forall s. MutableArray s a -> ST s ()) -> Array a
createArray (Vector a -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
v) a
forall a. a
uninitialized ((forall s. MutableArray s a -> ST s ()) -> Array a)
-> (forall s. MutableArray s a -> ST s ()) -> Array a
forall a b. (a -> b) -> a -> b
$ \arr :: MutableArray s a
arr@(MutableArray MutableArray# s a
arr#) -> do
            Vector a -> (Int -> a -> ST s ()) -> ST s ()
forall i (t :: * -> *) (f :: * -> *) a b.
(FoldableWithIndex i t, Applicative f) =>
t a -> (i -> a -> f b) -> f ()
ifor_ Vector a
v (MutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s a
MutableArray (PrimState (ST s)) a
arr)
            (a -> a -> Ordering) -> MutableArray# s a -> Int -> Int -> ST s ()
forall a s.
(a -> a -> Ordering) -> MutableArray# s a -> Int -> Int -> ST s ()
sortArrayBy a -> a -> Ordering
cmp MutableArray# s a
arr# Int
0 (Vector a -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
v)
    in [a] -> Vector a
forall a. [a] -> Vector a
fromList ([a] -> Vector a) -> (Array a -> [a]) -> Array a -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> [a]
forall a. Array a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Array a -> Vector a) -> Array a -> Vector a
forall a b. (a -> b) -> a -> b
$ Array a
sortedArr

-- | \(O(n \log n)\). Sort the vector in ascending order by comparing the results of applying the key function to each element.
-- The sort is stable, meaning the order of equal elements is preserved.
-- @`sortOn` f@ is equivalent to @`sortBy` (`Data.Ord.comparing` f)@, but only evaluates @f@ once for each element.
--
-- @since 0.2.2.0
sortOn :: (Ord b) => (a -> b) -> Vector a -> Vector a
sortOn :: forall b a. Ord b => (a -> b) -> Vector a -> Vector a
sortOn a -> b
f Vector a
v =
    let sortedArr :: Array (Arg b a)
sortedArr = Int
-> Arg b a
-> (forall s. MutableArray s (Arg b a) -> ST s ())
-> Array (Arg b a)
forall a.
Int -> a -> (forall s. MutableArray s a -> ST s ()) -> Array a
createArray (Vector a -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
v) Arg b a
forall a. a
uninitialized ((forall s. MutableArray s (Arg b a) -> ST s ())
 -> Array (Arg b a))
-> (forall s. MutableArray s (Arg b a) -> ST s ())
-> Array (Arg b a)
forall a b. (a -> b) -> a -> b
$ \arr :: MutableArray s (Arg b a)
arr@(MutableArray MutableArray# s (Arg b a)
arr#) -> do
            Vector a -> (Int -> a -> ST s ()) -> ST s ()
forall i (t :: * -> *) (f :: * -> *) a b.
(FoldableWithIndex i t, Applicative f) =>
t a -> (i -> a -> f b) -> f ()
ifor_ Vector a
v ((Int -> a -> ST s ()) -> ST s ())
-> (Int -> a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i a
x -> let !y :: b
y = a -> b
f a
x in MutableArray (PrimState (ST s)) (Arg b a)
-> Int -> Arg b a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s (Arg b a)
MutableArray (PrimState (ST s)) (Arg b a)
arr Int
i (b -> a -> Arg b a
forall a b. a -> b -> Arg a b
Arg b
y a
x)
            (Arg b a -> Arg b a -> Ordering)
-> MutableArray# s (Arg b a) -> Int -> Int -> ST s ()
forall a s.
(a -> a -> Ordering) -> MutableArray# s a -> Int -> Int -> ST s ()
sortArrayBy Arg b a -> Arg b a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare MutableArray# s (Arg b a)
arr# Int
0 (Vector a -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
v)
    in [a] -> Vector a
forall a. [a] -> Vector a
fromList ([a] -> Vector a)
-> (Array (Arg b a) -> [a]) -> Array (Arg b a) -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Arg b a -> a) -> [Arg b a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Arg b
_ a
x) -> a
x) ([Arg b a] -> [a])
-> (Array (Arg b a) -> [Arg b a]) -> Array (Arg b a) -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array (Arg b a) -> [Arg b a]
forall a. Array a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Array (Arg b a) -> Vector a) -> Array (Arg b a) -> Vector a
forall a b. (a -> b) -> a -> b
$ Array (Arg b a)
sortedArr