    {-# language BangPatterns, RankNTypes, ScopedTypeVariables #-}
    module Data.Vector.Algorithms where

    import Prelude hiding (length)
    import Control.Monad
    import Control.Monad.Primitive
    import Control.Monad.ST (runST)

    import Data.Vector.Generic.Mutable
    import qualified Data.Vector.Generic as V
    import qualified Data.Vector.Unboxed.Mutable as UMV
    import qualified Data.Bit as Bit

    import Data.Vector.Algorithms.Common (Comparison)
    import Data.Vector.Algorithms.Intro (sortUniqBy)
    import qualified Data.Vector.Algorithms.Search  as S

    -- | The `nub` function which removes duplicate elements from a vector.
    nub :: forall v e . (V.Vector v e, Ord e) => v e -> v e
    nub :: forall (v :: * -> *) e. (Vector v e, Ord e) => v e -> v e
    nub = Comparison e -> v e -> v e
    forall (v :: * -> *) e. Vector v e => Comparison e -> v e -> v e
    nubBy Comparison e
    forall a. Ord a => a -> a -> Ordering
    compare
    {-# INLINE nub #-}

    -- | A version of `nub` with a custom comparison predicate.
    --
    -- /Note:/ This function makes use of `sortByUniq` using the intro
    -- sort algorithm.
    nubBy ::
      forall v e . (V.Vector v e) =>
      Comparison e -> v e -> v e
    nubBy :: forall (v :: * -> *) e. Vector v e => Comparison e -> v e -> v e
    nubBy Comparison e
    cmp v e
    vec = (forall s. ST s (v e)) -> v e
    forall a. (forall s. ST s a) -> a
    runST ((forall s. ST s (v e)) -> v e) -> (forall s. ST s (v e)) -> v e
    forall a b. (a -> b) -> a -> b
    $ do
      mv <- v e -> ST s (Mutable v (PrimState (ST s)) e)
    forall (m :: * -> *) (v :: * -> *) a.
    (PrimMonad m, Vector v a) =>
    v a -> m (Mutable v (PrimState m) a)
    V.unsafeThaw v e
    vec -- safe as the nubByMut algorithm copies the input
      destMV <- nubByMut sortUniqBy cmp mv
      v <- V.unsafeFreeze destMV
      pure (V.force v)
    {-# INLINE nubBy #-}

    -- | The `nubByMut` function takes in an in-place sort algorithm
    -- and uses it to do a de-deduplicated sort. It then uses this to
    -- remove duplicate elements from the input.
    --
    -- /Note:/ Since this algorithm needs the original input and so
    -- copies before sorting in-place. As such, it is safe to use on
    -- immutable inputs.
    nubByMut ::
      forall m v e . (PrimMonad m, MVector v e) =>
      (Comparison e -> v (PrimState m) e -> m (v (PrimState m) e))
      -> Comparison e -> v (PrimState m) e -> m (v (PrimState m) e)
    nubByMut :: forall (m :: * -> *) (v :: * -> * -> *) e.
    (PrimMonad m, MVector v e) =>
    (Comparison e -> v (PrimState m) e -> m (v (PrimState m) e))
    -> Comparison e -> v (PrimState m) e -> m (v (PrimState m) e)
    nubByMut Comparison e -> v (PrimState m) e -> m (v (PrimState m) e)
    alg Comparison e
    cmp v (PrimState m) e
    inp = do
      let len :: Int
    len = v (PrimState m) e -> Int
    forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
    length v (PrimState m) e
    inp
      inp' <- v (PrimState m) e -> m (v (PrimState m) e)
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    v (PrimState m) a -> m (v (PrimState m) a)
    clone v (PrimState m) e
    inp
      sortUniqs <- alg cmp inp'
      let uniqLen = v (PrimState m) e -> Int
    forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
    length v (PrimState m) e
    sortUniqs
      bitmask <- UMV.replicate uniqLen (Bit.Bit False) -- bitmask to track which elements have
                                                       -- already been seen.
      dest ::  v (PrimState m) e <- unsafeNew uniqLen  -- return vector
      let
        go :: Int -> Int -> m ()
        go !Int
    srcInd !Int
    destInd
          | Int
    srcInd Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    len = () -> m ()
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
          | Int
    destInd Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    uniqLen = () -> m ()
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
          | Bool
    otherwise = do
              curr    <- 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
    inp Int
    srcInd                -- read current element
              sortInd <- S.binarySearchBy cmp sortUniqs curr  -- find sorted index
              bit <- UMV.unsafeRead bitmask sortInd           -- check if we have already seen
                                                              -- this element in bitvector
              case bit of
                -- if we have seen it then iterate
                Bit.Bit Bool
    True -> Int -> Int -> m ()
    go (Int
    srcInd Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1) Int
    destInd
                -- if we haven't then write it into output
                -- and mark that it has been seen
                Bit.Bit Bool
    False -> do
                  MVector (PrimState m) Bit -> Int -> Bit -> m ()
    forall (m :: * -> *) a.
    (PrimMonad m, Unbox a) =>
    MVector (PrimState m) a -> Int -> a -> m ()
    UMV.unsafeWrite MVector (PrimState m) Bit
    bitmask Int
    sortInd (Bool -> Bit
    Bit.Bit Bool
    True)
                  v (PrimState m) e -> Int -> e -> m ()
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    v (PrimState m) a -> Int -> a -> m ()
    unsafeWrite v (PrimState m) e
    dest Int
    destInd e
    curr
                  Int -> Int -> m ()
    go (Int
    srcInd Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1) (Int
    destInd Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1)
      go 0 0
      pure dest
    {-# INLINABLE nubByMut #-}