{-# 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.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 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 (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 ()
qsort (forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) Int
0 (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) forall a. Num a => a -> a -> a
- Int
1)
where
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 #-}
qsort :: Int -> Int -> Int -> m ()
qsort !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 ()
qsort 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 ()
qsort Int
n' Int
h Int
hi
else do
Int -> Int -> Int -> m ()
qsort Int
n Int
lo (Int
l forall a. Num a => a -> a -> a
- Int
1)
Int -> Int -> Int -> m ()
qsort Int
n Int
h Int
hi
{-# INLINE quicksortInternalM_ #-}