module Data.Vector.Compress where

import Data.Bits
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

--
import Data.Vector.Sort.Radix

compress :: U.Vector Int -> U.Vector Int
compress :: Vector Int -> Vector Int
compress Vector Int
vec = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
U.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
  mvec <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
UM.unsafeNew (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector Int
vec)
  U.mapM_ (\(Int
i, Int
x) -> MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
mvec (Int
x Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
0xffffffff) Int
i)
    . U.postscanl'
      ( \(!Int
i, !Int
x) Int
y ->
          if Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR Int
x Int
32 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR Int
y Int
32
            then (Int
i, Int
y)
            else (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
y)
      )
      (-1, -1)
    . radixSortInt
    $ U.imap (\Int
i Int
x -> Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftL Int
x Int
32 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
i) vec
  return mvec
{-# INLINE compress #-}