{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
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
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 #-}
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 #-}
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
!a
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 -> 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
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 Comparison a
cmp a
a a
prev Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ
then Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
ix
else a -> Int -> m Int
findFirstDuplicate a
a (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
else Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
ix
Int
dupIx <- a -> Int -> m Int
findFirstDuplicate a
a0 Int
1
if Int
dupIx Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len
then 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
else do
let deduplicate :: a -> Int -> Int -> m Int
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
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 Comparison a
cmp a
a a
prev Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ
then a -> Int -> Int -> m Int
deduplicate a
a (Int
srcIx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
dstIx
else do
v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite v (PrimState m) a
mv Int
dstIx a
a
a -> Int -> Int -> m Int
deduplicate a
a (Int
srcIx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
dstIx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
else Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
dstIx
!a
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
dupIx
!Int
reducedLen <- a -> Int -> Int -> m Int
deduplicate a
a (Int
dupIx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
dupIx
v (PrimState m) a -> Int -> m (v (PrimState m) a)
forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
v (PrimState m) a -> Int -> m (v (PrimState m) a)
resizeVector v (PrimState m) a
mv Int
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 #-}
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
v (PrimState m) a
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
v (PrimState m) a -> v (PrimState m) a -> m ()
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
v (PrimState m) a -> m (v (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure v (PrimState m) a
dst
{-# inline resizeVector #-}
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
a
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
v (PrimState m) a -> Int -> a -> ST (PrimState m) ()
forall s. v s a -> Int -> a -> ST s ()
forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> Int -> a -> ST s ()
basicUnsafeWrite v (PrimState m) a
dst Int
i a
x
Int -> ST (PrimState m) ()
do_copy (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
| Bool
otherwise = () -> ST (PrimState m) ()
forall a. a -> ST (PrimState m) a
forall (m :: * -> *) a. Monad m => a -> m a
return ()