{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MonoLocalBinds #-}
module Data.Massiv.Array.Ops.Sort (
tally,
quicksort,
quicksortBy,
quicksortByM,
quicksortAs,
quicksortAsBy,
quicksortAsByM,
quicksortM_,
quicksortByM_,
unsafeUnstablePartitionRegionM,
) where
import Control.Monad (when)
import Control.Monad.IO.Unlift
import Control.Monad.Primitive
import Control.Scheduler
import Data.Bits (countLeadingZeros)
import Data.Massiv.Array.Delayed.Stream
import Data.Massiv.Array.Mutable
import Data.Massiv.Array.Ops.Transform
import Data.Massiv.Core.Common
import Data.Massiv.Vector (scatMaybes, sunfoldrN)
import Data.Word (Word64)
import System.IO.Unsafe
tally :: (Manifest r e, Load r ix e, Ord e) => Array r ix e -> Vector DS (e, Int)
tally :: forall r e ix.
(Manifest r e, Load r ix e, Ord e) =>
Array r ix e -> Vector DS (e, Int)
tally Array r ix e
arr
| forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
isEmpty Array r ix e
arr = forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) forall r ix e. Load r ix e => Array r ix e
empty
| Bool
otherwise = forall r ix a.
Stream r ix (Maybe a) =>
Array r ix (Maybe a) -> Vector DS a
scatMaybes forall a b. (a -> b) -> a -> b
$ forall e s. Sz1 -> (s -> Maybe (e, s)) -> s -> Vector DS e
sunfoldrN (forall ix.
Index ix =>
(Int -> Int -> Int) -> Sz ix -> Sz ix -> Sz ix
liftSz2 forall a. Num a => a -> a -> a
(+) Sz1
sz forall ix. Index ix => Sz ix
oneSz) (Int, Int, e) -> Maybe (Maybe (e, Int), (Int, Int, e))
count (Int
0, Int
0, Vector r e
sorted forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
! Int
0)
where
sz :: Sz1
sz@(Sz Int
k) = forall r ix e. Size r => Array r ix e -> Sz ix
size Vector r e
sorted
count :: (Int, Int, e) -> Maybe (Maybe (e, Int), (Int, Int, e))
count (!Int
i, !Int
n, !e
prev)
| Int
i forall a. Ord a => a -> a -> Bool
< Int
k =
let !e' :: e
e' = forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Vector r e
sorted Int
i
in if e
prev forall a. Eq a => a -> a -> Bool
== e
e'
then forall a. a -> Maybe a
Just (forall a. Maybe a
Nothing, (Int
i forall a. Num a => a -> a -> a
+ Int
1, Int
n forall a. Num a => a -> a -> a
+ Int
1, e
prev))
else forall a. a -> Maybe a
Just (forall a. a -> Maybe a
Just (e
prev, Int
n), (Int
i forall a. Num a => a -> a -> a
+ Int
1, Int
1, e
e'))
| Bool
otherwise = forall a. a -> Maybe a
Just (forall a. a -> Maybe a
Just (e
prev, Int
n), (Int
i forall a. Num a => a -> a -> a
+ Int
1, Int
n, e
prev))
{-# INLINE count #-}
sorted :: Vector r e
sorted = forall r e. (Manifest r e, Ord e) => Vector r e -> Vector r e
quicksort forall a b. (a -> b) -> a -> b
$ forall r ix e. (Index ix, Size r) => Array r ix e -> Vector r e
flatten Array r ix e
arr
{-# INLINE tally #-}
unsafeUnstablePartitionRegionM
:: forall r e m
. (Manifest r e, PrimMonad m)
=> MVector (PrimState m) r e
-> (e -> m Bool)
-> Ix1
-> Ix1
-> m Ix1
unsafeUnstablePartitionRegionM :: forall r e (m :: * -> *).
(Manifest r e, PrimMonad m) =>
MVector (PrimState m) r e -> (e -> m Bool) -> Int -> Int -> m Int
unsafeUnstablePartitionRegionM MVector (PrimState m) r e
marr e -> m Bool
f Int
start Int
end = Int -> Int -> m Int
fromLeft Int
start (Int
end forall a. Num a => a -> a -> a
+ Int
1)
where
fromLeft :: Int -> Int -> m Int
fromLeft Int
i Int
j
| Int
i forall a. Eq a => a -> a -> Bool
== Int
j = forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i
| Bool
otherwise = do
Bool
e <- e -> m Bool
f forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MVector (PrimState m) r e
marr Int
i
if Bool
e
then Int -> Int -> m Int
fromLeft (Int
i forall a. Num a => a -> a -> a
+ Int
1) Int
j
else Int -> Int -> m Int
fromRight Int
i (Int
j forall a. Num a => a -> a -> a
- Int
1)
fromRight :: Int -> Int -> m Int
fromRight Int
i Int
j
| Int
i forall a. Eq a => a -> a -> Bool
== Int
j = forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i
| Bool
otherwise = do
e
x <- forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MVector (PrimState m) r e
marr Int
j
Bool
e <- e -> m Bool
f e
x
if Bool
e
then do
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MVector (PrimState m) r e
marr Int
j forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MVector (PrimState m) r e
marr Int
i
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MVector (PrimState m) r e
marr Int
i e
x
Int -> Int -> m Int
fromLeft (Int
i forall a. Num a => a -> a -> a
+ Int
1) Int
j
else Int -> Int -> m Int
fromRight Int
i (Int
j forall a. Num a => a -> a -> a
- Int
1)
{-# INLINE unsafeUnstablePartitionRegionM #-}
quicksortAs
:: (Load r Ix1 e, Manifest r' e, Ord e) => r' -> Vector r e -> Vector r' e
quicksortAs :: forall r e r'.
(Load r Int e, Manifest r' e, Ord e) =>
r' -> Vector r e -> Vector r' e
quicksortAs r'
_ Vector r e
arr = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall r ix e r' (m :: * -> *) b.
(Load r' ix e, Manifest r e, MonadUnliftIO m) =>
Array r' ix e
-> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m b)
-> m (Array r ix e)
withLoadMArray_ Vector r e
arr forall e r s (m :: * -> *).
(Ord e, Manifest r e, MonadPrimBase s m) =>
Scheduler s () -> MVector s r e -> m ()
quicksortM_
{-# INLINE quicksortAs #-}
quicksortAsBy
:: (Load r Ix1 e, Manifest r' e) => r' -> (e -> e -> Ordering) -> Vector r e -> Vector r' e
quicksortAsBy :: forall r e r'.
(Load r Int e, Manifest r' e) =>
r' -> (e -> e -> Ordering) -> Vector r e -> Vector r' e
quicksortAsBy r'
_ e -> e -> Ordering
f Vector r e
arr =
forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall r ix e r' (m :: * -> *) b.
(Load r' ix e, Manifest r e, MonadUnliftIO m) =>
Array r' ix e
-> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m b)
-> m (Array r ix e)
withLoadMArray_ Vector r e
arr (forall r e s (m :: * -> *).
(Manifest r e, MonadPrimBase s m) =>
(e -> e -> m Ordering) -> Scheduler s () -> MVector s r e -> m ()
quicksortByM_ (\e
x e
y -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ e -> e -> Ordering
f e
x e
y))
{-# INLINE quicksortAsBy #-}
quicksortAsByM
:: (Load r Ix1 e, Manifest r' e, MonadUnliftIO m)
=> r'
-> (e -> e -> m Ordering)
-> Vector r e
-> m (Vector r' e)
quicksortAsByM :: forall r e r' (m :: * -> *).
(Load r Int e, Manifest r' e, MonadUnliftIO m) =>
r' -> (e -> e -> m Ordering) -> Vector r e -> m (Vector r' e)
quicksortAsByM r'
_ e -> e -> m Ordering
f Vector r e
arr =
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> forall r ix e r' (m :: * -> *) b.
(Load r' ix e, Manifest r e, MonadUnliftIO m) =>
Array r' ix e
-> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m b)
-> m (Array r ix e)
withLoadMArray_ Vector r e
arr (forall r e s (m :: * -> *).
(Manifest r e, MonadPrimBase s m) =>
(e -> e -> m Ordering) -> Scheduler s () -> MVector s r e -> m ()
quicksortByM_ (\e
x e
y -> forall a. m a -> IO a
run (e -> e -> m Ordering
f e
x e
y)))
{-# INLINE quicksortAsByM #-}
quicksort
:: (Manifest r e, Ord e) => Vector r e -> Vector r e
quicksort :: forall r e. (Manifest r e, Ord e) => Vector r e -> Vector r e
quicksort Vector r e
arr = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall r e ix (m :: * -> *) a.
(Manifest r e, Index ix, MonadUnliftIO m) =>
Array r ix e
-> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a)
-> m (Array r ix e)
withMArray_ Vector r e
arr forall e r s (m :: * -> *).
(Ord e, Manifest r e, MonadPrimBase s m) =>
Scheduler s () -> MVector s r e -> m ()
quicksortM_
{-# INLINE quicksort #-}
quicksortByM
:: (Manifest r e, MonadUnliftIO m) => (e -> e -> m Ordering) -> Vector r e -> m (Vector r e)
quicksortByM :: forall r e (m :: * -> *).
(Manifest r e, MonadUnliftIO m) =>
(e -> e -> m Ordering) -> Vector r e -> m (Vector r e)
quicksortByM e -> e -> m Ordering
f Vector r e
arr = forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> forall r e ix (m :: * -> *) a.
(Manifest r e, Index ix, MonadUnliftIO m) =>
Array r ix e
-> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a)
-> m (Array r ix e)
withMArray_ Vector r e
arr (forall r e s (m :: * -> *).
(Manifest r e, MonadPrimBase s m) =>
(e -> e -> m Ordering) -> Scheduler s () -> MVector s r e -> m ()
quicksortByM_ (\e
x e
y -> forall a. m a -> IO a
run (e -> e -> m Ordering
f e
x e
y)))
{-# INLINE quicksortByM #-}
quicksortBy :: Manifest r e => (e -> e -> Ordering) -> Vector r e -> Vector r e
quicksortBy :: forall r e.
Manifest r e =>
(e -> e -> Ordering) -> Vector r e -> Vector r e
quicksortBy e -> e -> Ordering
f Vector r e
arr =
forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall r e ix (m :: * -> *) a.
(Manifest r e, Index ix, MonadUnliftIO m) =>
Array r ix e
-> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a)
-> m (Array r ix e)
withMArray_ Vector r e
arr (forall r e s (m :: * -> *).
(Manifest r e, MonadPrimBase s m) =>
(e -> e -> m Ordering) -> Scheduler s () -> MVector s r e -> m ()
quicksortByM_ (\e
x e
y -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ e -> e -> Ordering
f e
x e
y))
{-# INLINE quicksortBy #-}
quicksortM_
:: (Ord e, Manifest r e, MonadPrimBase s m)
=> Scheduler s ()
-> MVector s r e
-> m ()
quicksortM_ :: forall e r s (m :: * -> *).
(Ord e, Manifest r e, MonadPrimBase s m) =>
Scheduler s () -> MVector s r e -> m ()
quicksortM_ = forall r e s (m :: * -> *).
(Manifest r e, MonadPrimBase s m) =>
(e -> e -> m Bool)
-> (e -> e -> m Bool) -> Scheduler s () -> MVector s r e -> m ()
quicksortInternalM_ (\e
e1 e
e2 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ e
e1 forall a. Ord a => a -> a -> Bool
< e
e2) (\e
e1 e
e2 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ e
e1 forall a. Eq a => a -> a -> Bool
== e
e2)
{-# INLINE quicksortM_ #-}
quicksortByM_
:: (Manifest r e, MonadPrimBase s m)
=> (e -> e -> m Ordering)
-> Scheduler s ()
-> MVector s r e
-> m ()
quicksortByM_ :: forall r e s (m :: * -> *).
(Manifest r e, MonadPrimBase s m) =>
(e -> e -> m Ordering) -> Scheduler s () -> MVector s r e -> m ()
quicksortByM_ e -> e -> m Ordering
compareM =
forall r e s (m :: * -> *).
(Manifest r e, MonadPrimBase s m) =>
(e -> e -> m Bool)
-> (e -> e -> m Bool) -> Scheduler s () -> MVector s r e -> m ()
quicksortInternalM_ (\e
x e
y -> (Ordering
LT forall a. Eq a => a -> a -> Bool
==) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> e -> e -> m Ordering
compareM e
x e
y) (\e
x e
y -> (Ordering
EQ forall a. Eq a => a -> a -> Bool
==) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> e -> e -> m Ordering
compareM e
x e
y)
{-# INLINE quicksortByM_ #-}
quicksortInternalM_
:: (Manifest r e, MonadPrimBase s m)
=> (e -> e -> m Bool)
-> (e -> e -> m Bool)
-> Scheduler s ()
-> MVector s r e
-> m ()
quicksortInternalM_ :: forall r e s (m :: * -> *).
(Manifest r e, MonadPrimBase s m) =>
(e -> e -> m Bool)
-> (e -> e -> m Bool) -> Scheduler s () -> MVector s r e -> m ()
quicksortInternalM_ e -> e -> m Bool
fLT e -> e -> m Bool
fEQ Scheduler s ()
scheduler MVector s r e
marr
| forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler forall a. Ord a => a -> a -> Bool
< Int
2 Bool -> Bool -> Bool
|| Int
depthPar forall a. Ord a => a -> a -> Bool
<= Int
0 = Int -> Int -> m ()
qsortSeq Int
0 (Int
k forall a. Num a => a -> a -> a
- Int
1)
| Bool
otherwise = Int -> Int -> Int -> m ()
qsortPar Int
depthPar Int
0 (Int
k forall a. Num a => a -> a -> a
- Int
1)
where
depthPar :: Int
depthPar = forall a. Ord a => a -> a -> a
min (Int
logNumWorkers forall a. Num a => a -> a -> a
+ Int
4) (Int
logSize forall a. Num a => a -> a -> a
- Int
10)
k :: Int
k = forall ix. Sz ix -> ix
unSz (forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MVector s r e
marr)
logNumWorkers :: Int
logNumWorkers = Int
63 forall a. Num a => a -> a -> a
- forall b. FiniteBits b => b -> Int
countLeadingZeros (forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) :: Word64)
logSize :: Int
logSize = Int
63 forall a. Num a => a -> a -> a
- forall b. FiniteBits b => b -> Int
countLeadingZeros (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k :: Word64)
ltSwap :: Int -> Int -> m e
ltSwap Int
i Int
j = do
e
ei <- forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MVector s r e
marr Int
i
e
ej <- forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MVector s r e
marr Int
j
Bool
lt <- e -> e -> m Bool
fLT e
ei e
ej
if Bool
lt
then do
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MVector s r e
marr Int
i e
ej
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MVector s r e
marr Int
j e
ei
forall (f :: * -> *) a. Applicative f => a -> f a
pure e
ei
else forall (f :: * -> *) a. Applicative f => a -> f a
pure e
ej
{-# INLINE ltSwap #-}
getPivot :: Int -> Int -> m e
getPivot Int
lo Int
hi = do
let !mid :: Int
mid = (Int
hi forall a. Num a => a -> a -> a
+ Int
lo) forall a. Integral a => a -> a -> a
`div` Int
2
e
_ <- Int -> Int -> m e
ltSwap Int
mid Int
lo
e
_ <- Int -> Int -> m e
ltSwap Int
hi Int
lo
Int -> Int -> m e
ltSwap Int
mid Int
hi
{-# INLINE getPivot #-}
qsortPar :: Int -> Int -> Int -> m ()
qsortPar !Int
n !Int
lo !Int
hi =
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
lo forall a. Ord a => a -> a -> Bool
< Int
hi) forall a b. (a -> b) -> a -> b
$ do
e
p <- Int -> Int -> m e
getPivot Int
lo Int
hi
Int
l <- forall r e (m :: * -> *).
(Manifest r e, PrimMonad m) =>
MVector (PrimState m) r e -> (e -> m Bool) -> Int -> Int -> m Int
unsafeUnstablePartitionRegionM MVector s r e
marr (e -> e -> m Bool
`fLT` e
p) Int
lo (Int
hi forall a. Num a => a -> a -> a
- Int
1)
Int
h <- forall r e (m :: * -> *).
(Manifest r e, PrimMonad m) =>
MVector (PrimState m) r e -> (e -> m Bool) -> Int -> Int -> m Int
unsafeUnstablePartitionRegionM MVector s r e
marr (e -> e -> m Bool
`fEQ` e
p) Int
l Int
hi
if Int
n forall a. Ord a => a -> a -> Bool
> Int
0
then do
let !n' :: Int
n' = Int
n forall a. Num a => a -> a -> a
- Int
1
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> m ()
qsortPar Int
n' Int
lo (Int
l forall a. Num a => a -> a -> a
- Int
1)
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> m ()
qsortPar Int
n' Int
h Int
hi
else do
Int -> Int -> m ()
qsortSeq Int
lo (Int
l forall a. Num a => a -> a -> a
- Int
1)
Int -> Int -> m ()
qsortSeq Int
h Int
hi
qsortSeq :: Int -> Int -> m ()
qsortSeq !Int
lo !Int
hi =
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
lo forall a. Ord a => a -> a -> Bool
< Int
hi) forall a b. (a -> b) -> a -> b
$ do
e
p <- Int -> Int -> m e
getPivot Int
lo Int
hi
Int
l <- forall r e (m :: * -> *).
(Manifest r e, PrimMonad m) =>
MVector (PrimState m) r e -> (e -> m Bool) -> Int -> Int -> m Int
unsafeUnstablePartitionRegionM MVector s r e
marr (e -> e -> m Bool
`fLT` e
p) Int
lo (Int
hi forall a. Num a => a -> a -> a
- Int
1)
Int
h <- forall r e (m :: * -> *).
(Manifest r e, PrimMonad m) =>
MVector (PrimState m) r e -> (e -> m Bool) -> Int -> Int -> m Int
unsafeUnstablePartitionRegionM MVector s r e
marr (e -> e -> m Bool
`fEQ` e
p) Int
l Int
hi
Int -> Int -> m ()
qsortSeq Int
lo (Int
l forall a. Num a => a -> a -> a
- Int
1)
Int -> Int -> m ()
qsortSeq Int
h Int
hi
{-# INLINE quicksortInternalM_ #-}