    {-# LANGUAGE FlexibleContexts #-}
    {-# LANGUAGE TypeFamilies #-}
    {-# LANGUAGE BangPatterns #-}
    {-# LANGUAGE ScopedTypeVariables #-}

    -- ---------------------------------------------------------------------------
    -- |
    -- Module      : Data.Vector.Algorithms.Common
    -- Copyright   : (c) 2008-2011 Dan Doel
    -- Maintainer  : Dan Doel
    -- Stability   : Experimental
    -- Portability : Portable
    --
    -- Common operations and utility functions for all sorts

    module Data.Vector.Algorithms.Common
      ( type Comparison
      , copyOffset
      , inc
      , countLoop
      , midPoint
      , uniqueMutableBy
      )
      where

    import Prelude hiding (read, length)

    import Control.Monad.Primitive

    import Data.Vector.Generic.Mutable
    import Data.Word (Word)

    import qualified Data.Vector.Primitive.Mutable as PV

    -- | A type of comparisons between two values of a given type.
    type Comparison e = e -> e -> Ordering

    copyOffset :: (PrimMonad m, MVector v e)
               => v (PrimState m) e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
    copyOffset :: forall (m :: * -> *) (v :: * -> * -> *) e.
    (PrimMonad m, MVector v e) =>
    v (PrimState m) e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
    copyOffset v (PrimState m) e
    from v (PrimState m) e
    to Int
    iFrom Int
    iTo Int
    len =
      v (PrimState m) e -> v (PrimState m) e -> m ()
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    v (PrimState m) a -> v (PrimState m) a -> m ()
    unsafeCopy (Int -> Int -> v (PrimState m) e -> v (PrimState m) e
    forall (v :: * -> * -> *) a s.
    MVector v a =>
    Int -> Int -> v s a -> v s a
    unsafeSlice Int
    iTo Int
    len v (PrimState m) e
    to) (Int -> Int -> v (PrimState m) e -> v (PrimState m) e
    forall (v :: * -> * -> *) a s.
    MVector v a =>
    Int -> Int -> v s a -> v s a
    unsafeSlice Int
    iFrom Int
    len v (PrimState m) e
    from)
    {-# INLINE copyOffset #-}

    inc :: (PrimMonad m, MVector v Int) => v (PrimState m) Int -> Int -> m Int
    inc :: forall (m :: * -> *) (v :: * -> * -> *).
    (PrimMonad m, MVector v Int) =>
    v (PrimState m) Int -> Int -> m Int
    inc v (PrimState m) Int
    arr Int
    i = v (PrimState m) Int -> Int -> m Int
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    v (PrimState m) a -> Int -> m a
    unsafeRead v (PrimState m) Int
    arr Int
    i m Int -> (Int -> m Int) -> m Int
    forall a b. m a -> (a -> m b) -> m b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= \Int
    e -> v (PrimState m) Int -> Int -> Int -> m ()
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    v (PrimState m) a -> Int -> a -> m ()
    unsafeWrite v (PrimState m) Int
    arr Int
    i (Int
    eInt -> Int -> Int
    forall a. Num a => a -> a -> a
    +Int
    1) m () -> m Int -> m Int
    forall a b. m a -> m b -> m b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Int -> m Int
    forall a. a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return Int
    e
    {-# INLINE inc #-}

    -- shared bucket sorting stuff
    countLoop :: (PrimMonad m, MVector v e)
              => (e -> Int)
              -> v (PrimState m) e -> PV.MVector (PrimState m) Int -> m ()
    countLoop :: forall (m :: * -> *) (v :: * -> * -> *) e.
    (PrimMonad m, MVector v e) =>
    (e -> Int)
    -> v (PrimState m) e -> MVector (PrimState m) Int -> m ()
    countLoop e -> Int
    rdx v (PrimState m) e
    src MVector (PrimState m) Int
    count = MVector (PrimState m) Int -> Int -> m ()
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    v (PrimState m) a -> a -> m ()
    set MVector (PrimState m) Int
    count Int
    0 m () -> m () -> m ()
    forall a b. m a -> m b -> m b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Int -> m ()
    go Int
    0
     where
     len :: Int
    len = v (PrimState m) e -> Int
    forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
    length v (PrimState m) e
    src
     go :: Int -> m ()
    go Int
    i
       | Int
    i Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    len    = v (PrimState m) e -> Int -> m e
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    v (PrimState m) a -> Int -> m a
    unsafeRead v (PrimState m) e
    src Int
    i m e -> (e -> m Int) -> m Int
    forall a b. m a -> (a -> m b) -> m b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= MVector (PrimState m) Int -> Int -> m Int
    forall (m :: * -> *) (v :: * -> * -> *).
    (PrimMonad m, MVector v Int) =>
    v (PrimState m) Int -> Int -> m Int
    inc MVector (PrimState m) Int
    count (Int -> m Int) -> (e -> Int) -> e -> m Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . e -> Int
    rdx m Int -> m () -> m ()
    forall a b. m a -> m b -> m b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Int -> m ()
    go (Int
    iInt -> Int -> Int
    forall a. Num a => a -> a -> a
    +Int
    1)
       | Bool
    otherwise  = () -> m ()
    forall a. a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ()
    {-# INLINE countLoop #-}

    midPoint :: Int -> Int -> Int
    midPoint :: Int -> Int -> Int
    midPoint Int
    a Int
    b =
      Word -> Int
    toInt (Word -> Int) -> Word -> Int
    forall a b. (a -> b) -> a -> b
    $ (Int -> Word
    toWord Int
    a Word -> Word -> Word
    forall a. Num a => a -> a -> a
    + Int -> Word
    toWord Int
    b) Word -> Word -> Word
    forall a. Integral a => a -> a -> a
    `div` Word
    2
      where
        toWord :: Int -> Word
        toWord :: Int -> Word
    toWord = Int -> Word
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral

        toInt :: Word -> Int
        toInt :: Word -> Int
    toInt = Word -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral
    {-# INLINE midPoint #-}

    -- Adapted from Andrew Martin's uniquqMutable in the primitive-sort package
    uniqueMutableBy :: forall m v a . (PrimMonad m, MVector v a)
      => Comparison a -> v (PrimState m) a -> m (v (PrimState m) a)
    uniqueMutableBy :: forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    Comparison a -> v (PrimState m) a -> m (v (PrimState m) a)
    uniqueMutableBy Comparison a
    cmp v (PrimState m) a
    mv = do
      let !len :: Int
    len = v (PrimState m) a -> Int
    forall s. v s a -> Int
    forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
    basicLength v (PrimState m) a
    mv
      if Int
    len Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    1
        then do
          !a0 <- v (PrimState m) a -> Int -> m a
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    v (PrimState m) a -> Int -> m a
    unsafeRead v (PrimState m) a
    mv Int
    0
          let findFirstDuplicate :: a -> Int -> m Int
              findFirstDuplicate !a
    prev !Int
    ix = if Int
    ix Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    len
                then do
                  a <- v (PrimState m) a -> Int -> m a
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    v (PrimState m) a -> Int -> m a
    unsafeRead v (PrimState m) a
    mv Int
    ix
                  if cmp a prev == EQ
                    then return ix
                    else findFirstDuplicate a (ix + 1)
                else Int -> m Int
    forall a. a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return Int
    ix
          dupIx <- findFirstDuplicate a0 1
          if dupIx == len
            then return mv
            else do
              let deduplicate :: a -> Int -> Int -> m Int
                  deduplicate !a
    prev !Int
    srcIx !Int
    dstIx = if Int
    srcIx Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    len
                    then do
                      a <- v (PrimState m) a -> Int -> m a
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    v (PrimState m) a -> Int -> m a
    unsafeRead v (PrimState m) a
    mv Int
    srcIx
                      if cmp a prev == EQ
                        then deduplicate a (srcIx + 1) dstIx
                        else do
                          unsafeWrite mv dstIx a
                          deduplicate a (srcIx + 1) (dstIx + 1)
                    else Int -> m Int
    forall a. a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return Int
    dstIx
              !a <- unsafeRead mv dupIx
              !reducedLen <- deduplicate a (dupIx + 1) dupIx
              resizeVector mv reducedLen
        else v (PrimState m) a -> m (v (PrimState m) a)
    forall a. a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return v (PrimState m) a
    mv
    {-# INLINABLE uniqueMutableBy #-}

    -- Used internally in uniqueMutableBy: copies the elements of a vector to one
    -- of a smaller size.
    resizeVector
      :: (MVector v a, PrimMonad m)
      =>  v (PrimState m) a -> Int -> m (v (PrimState m) a)
    resizeVector :: forall (v :: * -> * -> *) a (m :: * -> *).
    (MVector v a, PrimMonad m) =>
    v (PrimState m) a -> Int -> m (v (PrimState m) a)
    resizeVector !v (PrimState m) a
    src !Int
    sz = do
      dst <- Int -> m (v (PrimState m) a)
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    Int -> m (v (PrimState m) a)
    unsafeNew Int
    sz
      copyToSmaller dst src
      pure dst
    {-# inline resizeVector #-}

    -- Used internally in resizeVector: copy a vector from a larger to
    -- smaller vector. Should not be used if the source vector
    -- is smaller than the target vector.
    copyToSmaller
      :: (MVector v a, PrimMonad m)
      => v (PrimState m) a -> v (PrimState m) a -> m ()
    copyToSmaller :: forall (v :: * -> * -> *) a (m :: * -> *).
    (MVector v a, PrimMonad m) =>
    v (PrimState m) a -> v (PrimState m) a -> m ()
    copyToSmaller !v (PrimState m) a
    dst !v (PrimState m) a
    src = ST (PrimState m) () -> m ()
    forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
    stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
    forall a b. (a -> b) -> a -> b
    $ Int -> ST (PrimState m) ()
    do_copy Int
    0
        where
          !n :: Int
    n = v (PrimState m) a -> Int
    forall s. v s a -> Int
    forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
    basicLength v (PrimState m) a
    dst

          do_copy :: Int -> ST (PrimState m) ()
    do_copy Int
    i | Int
    i Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    n = do
                                x <- v (PrimState m) a -> Int -> ST (PrimState m) a
    forall s. v s a -> Int -> ST s a
    forall (v :: * -> * -> *) a s.
    MVector v a =>
    v s a -> Int -> ST s a
    basicUnsafeRead v (PrimState m) a
    src Int
    i
                                basicUnsafeWrite dst i x
                                do_copy (i+1)
                    | Bool
    otherwise = () -> ST (PrimState m) ()
    forall a. a -> ST (PrimState m) a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ()