{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MonoLocalBinds #-}

-- |
-- Module      : Data.Massiv.Array.Ops.Sort
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
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

-- | Count number of occurrences of each element in the array. Results will be
-- sorted in ascending order of the element.
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array as A
-- >>> xs = fromList Seq [2, 4, 3, 2, 4, 5, 2, 1] :: Array P Ix1 Int
-- >>> xs
-- Array P Seq (Sz1 8)
--   [ 2, 4, 3, 2, 4, 5, 2, 1 ]
-- >>> tally xs
-- Array DS Seq (Sz1 5)
--   [ (1,1), (2,3), (3,1), (4,2), (5,1) ]
--
-- @since 0.4.4
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
  | Array r ix e -> Bool
forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
isEmpty Array r ix e
arr = Comp -> Array DS Int (e, Int) -> Array DS Int (e, Int)
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
forall ix e. Comp -> Array DS ix e -> Array DS ix e
setComp (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix e
arr) Array DS Int (e, Int)
forall r ix e. Load r ix e => Array r ix e
empty
  | Bool
otherwise = Array DS Int (Maybe (e, Int)) -> Array DS Int (e, Int)
forall r ix a.
Stream r ix (Maybe a) =>
Array r ix (Maybe a) -> Vector DS a
scatMaybes (Array DS Int (Maybe (e, Int)) -> Array DS Int (e, Int))
-> Array DS Int (Maybe (e, Int)) -> Array DS Int (e, Int)
forall a b. (a -> b) -> a -> b
$ Sz1
-> ((Int, Int, e) -> Maybe (Maybe (e, Int), (Int, Int, e)))
-> (Int, Int, e)
-> Array DS Int (Maybe (e, Int))
forall e s. Sz1 -> (s -> Maybe (e, s)) -> s -> Vector DS e
sunfoldrN ((Int -> Int -> Int) -> Sz1 -> Sz1 -> Sz1
forall ix.
Index ix =>
(Int -> Int -> Int) -> Sz ix -> Sz ix -> Sz ix
liftSz2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Sz1
sz Sz1
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 Vector r e -> Int -> e
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) = Vector r e -> Sz1
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. 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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k =
          let !e' :: e
e' = Vector r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> 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 e -> e -> Bool
forall a. Eq a => a -> a -> Bool
== e
e'
                then (Maybe (e, Int), (Int, Int, e))
-> Maybe (Maybe (e, Int), (Int, Int, e))
forall a. a -> Maybe a
Just (Maybe (e, Int)
forall a. Maybe a
Nothing, (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, e
prev))
                else (Maybe (e, Int), (Int, Int, e))
-> Maybe (Maybe (e, Int), (Int, Int, e))
forall a. a -> Maybe a
Just ((e, Int) -> Maybe (e, Int)
forall a. a -> Maybe a
Just (e
prev, Int
n), (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
1, e
e'))
      | Bool
otherwise = (Maybe (e, Int), (Int, Int, e))
-> Maybe (Maybe (e, Int), (Int, Int, e))
forall a. a -> Maybe a
Just ((e, Int) -> Maybe (e, Int)
forall a. a -> Maybe a
Just (e
prev, Int
n), (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
n, e
prev))
    {-# INLINE count #-}
    sorted :: Vector r e
sorted = Vector r e -> Vector r e
forall r e. (Manifest r e, Ord e) => Vector r e -> Vector r e
quicksort (Vector r e -> Vector r e) -> Vector r e -> Vector r e
forall a b. (a -> b) -> a -> b
$ Array r ix e -> Vector r e
forall r ix e. (Index ix, Size r) => Array r ix e -> Vector r e
flatten Array r ix e
arr
{-# INLINE tally #-}

-- | Partition a segment of a vector. Starting and ending indices are unchecked.
--
-- @since 1.0.0
unsafeUnstablePartitionRegionM
  :: forall r e m
   . (Manifest r e, PrimMonad m)
  => MVector (PrimState m) r e
  -> (e -> m Bool)
  -> Ix1
  -- ^ Start index of the region
  -> Ix1
  -- ^ End index of the region
  -> 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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  where
    fromLeft :: Int -> Int -> m Int
fromLeft Int
i Int
j
      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j = Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i
      | Bool
otherwise = do
          Bool
e <- e -> m Bool
f (e -> m Bool) -> m e -> m Bool
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState m) r e -> Int -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
j
            else Int -> Int -> m Int
fromRight Int
i (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    fromRight :: Int -> Int -> m Int
fromRight Int
i Int
j
      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j = Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i
      | Bool
otherwise = do
          e
x <- MVector (PrimState m) r e -> Int -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(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
              MVector (PrimState m) r e -> Int -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MVector (PrimState m) r e
marr Int
j (e -> m ()) -> m e -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState m) r e -> Int -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MVector (PrimState m) r e
marr Int
i
              MVector (PrimState m) r e -> Int -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
j
            else Int -> Int -> m Int
fromRight Int
i (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
{-# INLINE unsafeUnstablePartitionRegionM #-}

-- | Same as `quicksort` except it accepts any array that is computable.
--
-- @since 1.0.2
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 = IO (Vector r' e) -> Vector r' e
forall a. IO a -> a
unsafePerformIO (IO (Vector r' e) -> Vector r' e)
-> IO (Vector r' e) -> Vector r' e
forall a b. (a -> b) -> a -> b
$ Vector r e
-> (Scheduler RealWorld () -> MArray RealWorld r' Int e -> IO ())
-> IO (Vector r' e)
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 Scheduler RealWorld () -> MArray RealWorld r' Int e -> IO ()
forall e r s (m :: * -> *).
(Ord e, Manifest r e, MonadPrimBase s m) =>
Scheduler s () -> MVector s r e -> m ()
quicksortM_
{-# INLINE quicksortAs #-}

-- | Same as `quicksortBy` except it accepts any array that is computable.
--
-- @since 1.0.2
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 =
  IO (Vector r' e) -> Vector r' e
forall a. IO a -> a
unsafePerformIO (IO (Vector r' e) -> Vector r' e)
-> IO (Vector r' e) -> Vector r' e
forall a b. (a -> b) -> a -> b
$ Vector r e
-> (Scheduler RealWorld () -> MArray RealWorld r' Int e -> IO ())
-> IO (Vector r' e)
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 ((e -> e -> IO Ordering)
-> Scheduler RealWorld () -> MArray RealWorld r' Int e -> IO ()
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 -> Ordering -> IO Ordering
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ordering -> IO Ordering) -> Ordering -> IO Ordering
forall a b. (a -> b) -> a -> b
$ e -> e -> Ordering
f e
x e
y))
{-# INLINE quicksortAsBy #-}

-- | Same as `quicksortByM` except it accepts any array that is computable.
--
-- @since 1.0.2
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 a. m a -> IO a) -> IO (Vector r' e)) -> m (Vector r' e)
forall b. ((forall a. m a -> IO a) -> IO b) -> m b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO (Vector r' e)) -> m (Vector r' e))
-> ((forall a. m a -> IO a) -> IO (Vector r' e)) -> m (Vector r' e)
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> Vector r e
-> (Scheduler RealWorld () -> MArray RealWorld r' Int e -> IO ())
-> IO (Vector r' e)
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 ((e -> e -> IO Ordering)
-> Scheduler RealWorld () -> MArray RealWorld r' Int e -> IO ()
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 -> m Ordering -> IO Ordering
forall a. m a -> IO a
run (e -> e -> m Ordering
f e
x e
y)))
{-# INLINE quicksortAsByM #-}

-- | This is an implementation of
-- [Quicksort](https://en.wikipedia.org/wiki/Quicksort), which is an efficient,
-- but unstable sort. This implementation uses Median-of-three for pivot
-- choosing, as such it performs very well not only for random values, but also
-- for common edge cases like already sorted, reversed sorted and arrays with
-- many duplicate elements. It will also respect the computation strategy and
-- will result in a nice speed up for systems with multiple CPUs.
--
-- @since 0.3.2
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 = IO (Vector r e) -> Vector r e
forall a. IO a -> a
unsafePerformIO (IO (Vector r e) -> Vector r e) -> IO (Vector r e) -> Vector r e
forall a b. (a -> b) -> a -> b
$ Vector r e
-> (Scheduler RealWorld () -> MArray RealWorld r Int e -> IO ())
-> IO (Vector r e)
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 Scheduler RealWorld () -> MArray RealWorld r Int e -> IO ()
forall e r s (m :: * -> *).
(Ord e, Manifest r e, MonadPrimBase s m) =>
Scheduler s () -> MVector s r e -> m ()
quicksortM_
{-# INLINE quicksort #-}

-- | Same as `quicksortBy`, but instead of `Ord` constraint expects a custom `Ordering`.
--
-- @since 0.6.1
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 a. m a -> IO a) -> IO (Vector r e)) -> m (Vector r e)
forall b. ((forall a. m a -> IO a) -> IO b) -> m b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO (Vector r e)) -> m (Vector r e))
-> ((forall a. m a -> IO a) -> IO (Vector r e)) -> m (Vector r e)
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> Vector r e
-> (Scheduler RealWorld () -> MArray RealWorld r Int e -> IO ())
-> IO (Vector r e)
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 ((e -> e -> IO Ordering)
-> Scheduler RealWorld () -> MArray RealWorld r Int e -> IO ()
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 -> m Ordering -> IO Ordering
forall a. m a -> IO a
run (e -> e -> m Ordering
f e
x e
y)))
{-# INLINE quicksortByM #-}

-- | Same as `quicksortBy`, but instead of `Ord` constraint expects a custom `Ordering`.
--
-- @since 0.6.1
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 =
  IO (Vector r e) -> Vector r e
forall a. IO a -> a
unsafePerformIO (IO (Vector r e) -> Vector r e) -> IO (Vector r e) -> Vector r e
forall a b. (a -> b) -> a -> b
$ Vector r e
-> (Scheduler RealWorld () -> MArray RealWorld r Int e -> IO ())
-> IO (Vector r e)
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 ((e -> e -> IO Ordering)
-> Scheduler RealWorld () -> MArray RealWorld r Int e -> IO ()
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 -> Ordering -> IO Ordering
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ordering -> IO Ordering) -> Ordering -> IO Ordering
forall a b. (a -> b) -> a -> b
$ e -> e -> Ordering
f e
x e
y))
{-# INLINE quicksortBy #-}

-- | Manifest version of `quicksort`
--
-- @since 0.3.2
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_ = (e -> e -> m Bool)
-> (e -> e -> m Bool) -> Scheduler s () -> MVector s r e -> m ()
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 -> Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ e
e1 e -> e -> Bool
forall a. Ord a => a -> a -> Bool
< e
e2) (\e
e1 e
e2 -> Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ e
e1 e -> e -> Bool
forall a. Eq a => a -> a -> Bool
== e
e2)
{-# INLINE quicksortM_ #-}

-- | Same as `quicksortM_`, but instead of `Ord` constraint expects a custom `Ordering`.
--
-- @since 0.6.1
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 =
  (e -> e -> m Bool)
-> (e -> e -> m Bool) -> Scheduler s () -> MVector s r e -> m ()
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 Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
==) (Ordering -> Bool) -> m Ordering -> m 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 Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
==) (Ordering -> Bool) -> m Ordering -> m 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
  | Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 Bool -> Bool -> Bool
|| Int
depthPar Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Int -> Int -> m ()
qsortSeq Int
0 (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  | Bool
otherwise = Int -> Int -> Int -> m ()
qsortPar Int
depthPar Int
0 (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  where
    -- How deep into the search tree should we continue scheduling jobs. Constants below
    -- were discovered imperically:
    depthPar :: Int
depthPar = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
logNumWorkers Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4) (Int
logSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
10)
    k :: Int
k = Sz1 -> Int
forall ix. Sz ix -> ix
unSz (MVector s r e -> Sz1
forall ix s. Index ix => MArray s r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MVector s r e
marr)
    -- We must use log becuase decinding into a tree creates an exponential number of jobs
    logNumWorkers :: Int
logNumWorkers = Int
63 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Word64 -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) :: Word64)
    -- Using many cores on small vectors only makes things slower
    logSize :: Int
logSize = Int
63 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Word64 -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros (Int -> Word64
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 <- MArray (PrimState m) r Int e -> Int -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MVector s r e
MArray (PrimState m) r Int e
marr Int
i
      e
ej <- MArray (PrimState m) r Int e -> Int -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MVector s r e
MArray (PrimState m) r Int e
marr Int
j
      Bool
lt <- e -> e -> m Bool
fLT e
ei e
ej
      if Bool
lt
        then do
          MArray (PrimState m) r Int e -> Int -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MVector s r e
MArray (PrimState m) r Int e
marr Int
i e
ej
          MArray (PrimState m) r Int e -> Int -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MVector s r e
MArray (PrimState m) r Int e
marr Int
j e
ei
          e -> m e
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure e
ei
        else e -> m e
forall a. a -> m a
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lo) Int -> Int -> Int
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 =
      Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
lo Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
hi) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        e
p <- Int -> Int -> m e
getPivot Int
lo Int
hi
        Int
l <- MArray (PrimState m) r Int e
-> (e -> m Bool) -> Int -> Int -> m Int
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
MArray (PrimState m) r Int e
marr (e -> e -> m Bool
`fLT` e
p) Int
lo (Int
hi Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        Int
h <- MArray (PrimState m) r Int e
-> (e -> m Bool) -> Int -> Int -> m Int
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
MArray (PrimState m) r Int e
marr (e -> e -> m Bool
`fEQ` e
p) Int
l Int
hi
        if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
          then do
            let !n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
            Scheduler s () -> m () -> m ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> m ()
qsortPar Int
n' Int
lo (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            Scheduler s () -> m () -> m ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler (m () -> m ()) -> m () -> m ()
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 Int -> Int -> Int
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 =
      Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
lo Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
hi) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        e
p <- Int -> Int -> m e
getPivot Int
lo Int
hi
        Int
l <- MArray (PrimState m) r Int e
-> (e -> m Bool) -> Int -> Int -> m Int
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
MArray (PrimState m) r Int e
marr (e -> e -> m Bool
`fLT` e
p) Int
lo (Int
hi Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        Int
h <- MArray (PrimState m) r Int e
-> (e -> m Bool) -> Int -> Int -> m Int
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
MArray (PrimState m) r Int e
marr (e -> e -> m Bool
`fEQ` e
p) Int
l Int
hi
        Int -> Int -> m ()
qsortSeq Int
lo (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        Int -> Int -> m ()
qsortSeq Int
h Int
hi
{-# INLINE quicksortInternalM_ #-}