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
  MVector s Int
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)
  ((Int, Int) -> ST s ()) -> Vector (Int, Int) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
(a -> m b) -> Vector a -> m ()
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)
    (Vector (Int, Int) -> ST s ())
-> (Vector Int -> Vector (Int, Int)) -> Vector Int -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Int) -> Int -> (Int, Int))
-> (Int, Int) -> Vector Int -> Vector (Int, Int)
forall a b.
(Unbox a, Unbox b) =>
(a -> b -> a) -> a -> Vector b -> Vector a
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)
      )
      (-Int
1, -Int
1)
    (Vector Int -> Vector (Int, Int))
-> (Vector Int -> Vector Int) -> Vector Int -> Vector (Int, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Vector Int
radixSortInt
    (Vector Int -> ST s ()) -> Vector Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> Vector Int -> Vector Int
forall a b.
(Unbox a, Unbox b) =>
(Int -> a -> b) -> Vector a -> Vector b
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) Vector Int
vec
  MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return MVector s Int
mvec
{-# INLINE compress #-}