Sort.hs

{-# 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 arr
  | isEmpty arr = setComp (getComp arr) empty
  | otherwise = scatMaybes $ sunfoldrN (liftSz2 (+) sz oneSz) count (0, 0, sorted ! 0)
  where
    sz@(Sz k) = size sorted
    count (!i, !n, !prev)
      | i < k =
          let !e' = unsafeLinearIndex sorted i
           in if prev == e'
                then Just (Nothing, (i + 1, n + 1, prev))
                else Just (Just (prev, n), (i + 1, 1, e'))
      | otherwise = Just (Just (prev, n), (i + 1, n, prev))
    {-# INLINE count #-}
    sorted = quicksort $ flatten 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 marr f start end = fromLeft start (end + 1)
  where
    fromLeft i j
      | i == j = pure i
      | otherwise = do
          e <- f =<< unsafeLinearRead marr i
          if e
            then fromLeft (i + 1) j
            else fromRight i (j - 1)
    fromRight i j
      | i == j = pure i
      | otherwise = do
          x <- unsafeLinearRead marr j
          e <- f x
          if e
            then do
              unsafeLinearWrite marr j =<< unsafeLinearRead marr i
              unsafeLinearWrite marr i x
              fromLeft (i + 1) j
            else fromRight i (j - 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 _ arr = unsafePerformIO $ withLoadMArray_ arr 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 _ f arr =
  unsafePerformIO $ withLoadMArray_ arr (quicksortByM_ (\x y -> pure $ f x 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 _ f arr =
  withRunInIO $ \run -> withLoadMArray_ arr (quicksortByM_ (\x y -> run (f x 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 arr = unsafePerformIO $ withMArray_ arr 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 f arr = withRunInIO $ \run -> withMArray_ arr (quicksortByM_ (\x y -> run (f x 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 f arr =
  unsafePerformIO $ withMArray_ arr (quicksortByM_ (\x y -> pure $ f x 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_ = quicksortInternalM_ (\e1 e2 -> pure $ e1 < e2) (\e1 e2 -> pure $ e1 == 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_ compareM =
  quicksortInternalM_ (\x y -> (LT ==) <$> compareM x y) (\x y -> (EQ ==) <$> compareM x 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_ fLT fEQ scheduler marr
  | numWorkers scheduler < 2 || depthPar <= 0 = qsortSeq 0 (k - 1)
  | otherwise = qsortPar depthPar 0 (k - 1)
  where
    -- How deep into the search tree should we continue scheduling jobs. Constants below
    -- were discovered empirically:
    depthPar = min (logNumWorkers + 4) (logSize - 10)
    k = unSz (sizeOfMArray marr)
    -- We must use log becuase decinding into a tree creates an exponential number of jobs
    logNumWorkers = 63 - countLeadingZeros (fromIntegral (numWorkers scheduler) :: Word64)
    -- Using many cores on small vectors only makes things slower
    logSize = 63 - countLeadingZeros (fromIntegral k :: Word64)
    ltSwap i j = do
      ei <- unsafeLinearRead marr i
      ej <- unsafeLinearRead marr j
      lt <- fLT ei ej
      if lt
        then do
          unsafeLinearWrite marr i ej
          unsafeLinearWrite marr j ei
          pure ei
        else pure ej
    {-# INLINE ltSwap #-}
    getPivot lo hi = do
      let !mid = (hi + lo) `div` 2
      _ <- ltSwap mid lo
      _ <- ltSwap hi lo
      ltSwap mid hi
    {-# INLINE getPivot #-}
    qsortPar !n !lo !hi =
      when (lo < hi) $ do
        p <- getPivot lo hi
        l <- unsafeUnstablePartitionRegionM marr (`fLT` p) lo (hi - 1)
        h <- unsafeUnstablePartitionRegionM marr (`fEQ` p) l hi
        if n > 0
          then do
            let !n' = n - 1
            scheduleWork scheduler $ qsortPar n' lo (l - 1)
            scheduleWork scheduler $ qsortPar n' h hi
          else do
            qsortSeq lo (l - 1)
            qsortSeq h hi
    qsortSeq !lo !hi =
      when (lo < hi) $ do
        p <- getPivot lo hi
        l <- unsafeUnstablePartitionRegionM marr (`fLT` p) lo (hi - 1)
        h <- unsafeUnstablePartitionRegionM marr (`fEQ` p) l hi
        qsortSeq lo (l - 1)
        qsortSeq h hi
{-# INLINE quicksortInternalM_ #-}