{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Data.Vector.Growable
  ( Growable
  -- * Type synonyms
  , GrowableVector
  , GrowableUnboxedVector
  , GrowableStorableVector
  , GrowablePrimitiveVector
  , GrowableIOVector
  , GrowableUnboxedIOVector
  , GrowableStorableIOVector
  , GrowablePrimitiveIOVector
  -- * Operations
  , new
  , withCapacity
  , replicate
  , replicateM
  , push
  , pop
  , length
  , null
  , read
  , write
  , modify
  , thaw
  , freeze
  , unsafeFreeze
  , fromGrowable
  , toGrowable
  -- * Atomic operation
  , CASVector(..)
  , atomicPush
  , atomicPop
  ) where

import Prelude hiding (read, length, replicate, null)
import Control.Exception
import Control.Monad.Catch
import Control.Monad.Primitive
import Data.Atomics
import Data.Bits
import Data.Primitive.MutVar
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as MV
import Data.Vector.Generic.Mutable (MVector)
import qualified Data.Vector.Mutable as V
import qualified Data.Vector.Unboxed.Mutable as U
import qualified Data.Vector.Primitive.Mutable as P
import qualified Data.Vector.Storable.Mutable as S
import Unsafe.Coerce

data Pending a = Pending !Int (Ticket a) !a | Complete

data GVState v s a = GVState !Int !(v s a) (MutVar s (Pending a))

-- | 'Growable' is a dynamic vector based on mutable vector @v@.
newtype Growable v s a = Growable (MutVar s (GVState v s a))

type GrowableVector = Growable V.MVector
type GrowableUnboxedVector = Growable U.MVector
type GrowableStorableVector = Growable S.MVector
type GrowablePrimitiveVector = Growable S.MVector
type GrowableIOVector = Growable V.MVector RealWorld
type GrowableUnboxedIOVector = Growable U.MVector RealWorld
type GrowableStorableIOVector = Growable S.MVector RealWorld
type GrowablePrimitiveIOVector = Growable P.MVector RealWorld

-- | Create an empty vector with the given number of pre-allocated elements.
withCapacity :: (PrimMonad m, MVector v a) => Int -> m (Growable v (PrimState m) a) 
withCapacity :: Int -> m (Growable v (PrimState m) a)
withCapacity Int
cap = do
  v (PrimState m) a
vec <- Int -> m (v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MV.new Int
cap
  MutVar (PrimState m) (Pending a)
ref <- Pending a -> m (MutVar (PrimState m) (Pending a))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar Pending a
forall a. Pending a
Complete
  MutVar (PrimState m) (GVState v (PrimState m) a)
-> Growable v (PrimState m) a
forall (v :: * -> * -> *) s a.
MutVar s (GVState v s a) -> Growable v s a
Growable (MutVar (PrimState m) (GVState v (PrimState m) a)
 -> Growable v (PrimState m) a)
-> m (MutVar (PrimState m) (GVState v (PrimState m) a))
-> m (Growable v (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GVState v (PrimState m) a
-> m (MutVar (PrimState m) (GVState v (PrimState m) a))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar (Int
-> v (PrimState m) a
-> MutVar (PrimState m) (Pending a)
-> GVState v (PrimState m) a
forall (v :: * -> * -> *) s a.
Int -> v s a -> MutVar s (Pending a) -> GVState v s a
GVState Int
0 v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
ref)
{-# INlINE withCapacity #-}

-- | Create an empty vector
new :: (PrimMonad m, MVector v a) => m (Growable v (PrimState m) a) 
new :: m (Growable v (PrimState m) a)
new = Int -> m (Growable v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (Growable v (PrimState m) a)
withCapacity Int
0
{-# INLINE new #-}

-- | Create a vector and fill with the initial value.
replicate :: (PrimMonad m, MVector v a) => Int -> a -> m (Growable v (PrimState m) a)
replicate :: Int -> a -> m (Growable v (PrimState m) a)
replicate Int
len a
a = do
  v (PrimState m) a
vec <- Int -> a -> m (v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> a -> m (v (PrimState m) a)
MV.replicate Int
len a
a
  MutVar (PrimState m) (Pending a)
ref <- Pending a -> m (MutVar (PrimState m) (Pending a))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar Pending a
forall a. Pending a
Complete
  MutVar (PrimState m) (GVState v (PrimState m) a)
-> Growable v (PrimState m) a
forall (v :: * -> * -> *) s a.
MutVar s (GVState v s a) -> Growable v s a
Growable (MutVar (PrimState m) (GVState v (PrimState m) a)
 -> Growable v (PrimState m) a)
-> m (MutVar (PrimState m) (GVState v (PrimState m) a))
-> m (Growable v (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GVState v (PrimState m) a
-> m (MutVar (PrimState m) (GVState v (PrimState m) a))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar (Int
-> v (PrimState m) a
-> MutVar (PrimState m) (Pending a)
-> GVState v (PrimState m) a
forall (v :: * -> * -> *) s a.
Int -> v s a -> MutVar s (Pending a) -> GVState v s a
GVState Int
len v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
ref)
{-# INLINE replicate #-}

-- | Like 'replicate', but initialises the elements by running the action repeatedly
replicateM :: (PrimMonad m, MVector v a) => Int -> m a -> m (Growable v (PrimState m) a)
replicateM :: Int -> m a -> m (Growable v (PrimState m) a)
replicateM Int
len m a
a = do
  v (PrimState m) a
vec <- Int -> m a -> m (v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m a -> m (v (PrimState m) a)
MV.replicateM Int
len m a
a
  MutVar (PrimState m) (Pending a)
ref <- Pending a -> m (MutVar (PrimState m) (Pending a))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar Pending a
forall a. Pending a
Complete
  MutVar (PrimState m) (GVState v (PrimState m) a)
-> Growable v (PrimState m) a
forall (v :: * -> * -> *) s a.
MutVar s (GVState v s a) -> Growable v s a
Growable (MutVar (PrimState m) (GVState v (PrimState m) a)
 -> Growable v (PrimState m) a)
-> m (MutVar (PrimState m) (GVState v (PrimState m) a))
-> m (Growable v (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GVState v (PrimState m) a
-> m (MutVar (PrimState m) (GVState v (PrimState m) a))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar (Int
-> v (PrimState m) a
-> MutVar (PrimState m) (Pending a)
-> GVState v (PrimState m) a
forall (v :: * -> * -> *) s a.
Int -> v s a -> MutVar s (Pending a) -> GVState v s a
GVState Int
len v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
ref)
{-# INLINE replicateM #-}

-- | Append an element to the vector (not atomic).
push :: (PrimMonad m, MVector v a) => Growable v (PrimState m) a -> a -> m ()
push :: Growable v (PrimState m) a -> a -> m ()
push (Growable MutVar (PrimState m) (GVState v (PrimState m) a)
ref) a
val = do
  GVState Int
len v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
pending <- MutVar (PrimState m) (GVState v (PrimState m) a)
-> m (GVState v (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (GVState v (PrimState m) a)
ref
  v (PrimState m) a
vec' <- if v (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
MV.length v (PrimState m) a
vec Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len
    then v (PrimState m) a -> Int -> m (v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (v (PrimState m) a)
MV.unsafeGrow v (PrimState m) a
vec (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
len)
    else v (PrimState m) a -> m (v (PrimState m) a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure v (PrimState m) a
vec
  MutVar (PrimState m) (GVState v (PrimState m) a)
-> GVState v (PrimState m) a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> a -> m ()
writeMutVar MutVar (PrimState m) (GVState v (PrimState m) a)
ref (GVState v (PrimState m) a -> m ())
-> GVState v (PrimState m) a -> m ()
forall a b. (a -> b) -> a -> b
$ Int
-> v (PrimState m) a
-> MutVar (PrimState m) (Pending a)
-> GVState v (PrimState m) a
forall (v :: * -> * -> *) s a.
Int -> v s a -> MutVar s (Pending a) -> GVState v s a
GVState (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) v (PrimState m) a
vec' MutVar (PrimState m) (Pending a)
pending
  v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MV.write v (PrimState m) a
vec' Int
len a
val
{-# INLINE push #-}

-- | Pop the last element. Returns 'Nothing' if the vector is empty.
pop :: (PrimMonad m, MVector v a) => Growable v (PrimState m) a -> m (Maybe a)
pop :: Growable v (PrimState m) a -> m (Maybe a)
pop (Growable MutVar (PrimState m) (GVState v (PrimState m) a)
ref) = do
  GVState Int
len v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
pending <- MutVar (PrimState m) (GVState v (PrimState m) a)
-> m (GVState v (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (GVState v (PrimState m) a)
ref
  if Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then Maybe a -> m (Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
    else do
      MutVar (PrimState m) (GVState v (PrimState m) a)
-> GVState v (PrimState m) a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> a -> m ()
writeMutVar MutVar (PrimState m) (GVState v (PrimState m) a)
ref (GVState v (PrimState m) a -> m ())
-> GVState v (PrimState m) a -> m ()
forall a b. (a -> b) -> a -> b
$ Int
-> v (PrimState m) a
-> MutVar (PrimState m) (Pending a)
-> GVState v (PrimState m) a
forall (v :: * -> * -> *) s a.
Int -> v s a -> MutVar s (Pending a) -> GVState v s a
GVState (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
pending
      a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> m a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MV.unsafeRead v (PrimState m) a
vec (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
{-# INLINE pop #-}

-- | Get the length of the vector.
length :: (PrimMonad m) => Growable v (PrimState m) a -> m Int
length :: Growable v (PrimState m) a -> m Int
length (Growable MutVar (PrimState m) (GVState v (PrimState m) a)
ref) = do
  GVState Int
len v (PrimState m) a
_ MutVar (PrimState m) (Pending a)
_ <- MutVar (PrimState m) (GVState v (PrimState m) a)
-> m (GVState v (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (GVState v (PrimState m) a)
ref
  Int -> m Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
len
{-# INLINE length #-}

-- | Returns 'True' if the vector is empty
null :: (PrimMonad m) => Growable v (PrimState m) a -> m Bool
null :: Growable v (PrimState m) a -> m Bool
null (Growable MutVar (PrimState m) (GVState v (PrimState m) a)
ref) = do
  GVState Int
len v (PrimState m) a
_ MutVar (PrimState m) (Pending a)
_ <- MutVar (PrimState m) (GVState v (PrimState m) a)
-> m (GVState v (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (GVState v (PrimState m) a)
ref
  Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
{-# INLINE null #-}

-- | May throw 'IndexOutOfBounds'
read :: (PrimMonad m, MVector v a, MonadThrow m) => Growable v (PrimState m) a -> Int -> m a
read :: Growable v (PrimState m) a -> Int -> m a
read (Growable MutVar (PrimState m) (GVState v (PrimState m) a)
ref) Int
i = do
  GVState Int
len v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
_ <- MutVar (PrimState m) (GVState v (PrimState m) a)
-> m (GVState v (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (GVState v (PrimState m) a)
ref
  if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len
    then v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MV.unsafeRead v (PrimState m) a
vec Int
i
    else ArrayException -> m a
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (ArrayException -> m a) -> ArrayException -> m a
forall a b. (a -> b) -> a -> b
$ String -> ArrayException
IndexOutOfBounds (String -> ArrayException) -> String -> ArrayException
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
len String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" <= " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
i
{-# INLINE read #-}

-- | Throws 'IndexOutOfBounds' if the index is larger than the size.
write :: (PrimMonad m, MVector v a, MonadThrow m) => Growable v (PrimState m) a -> Int -> a -> m ()
write :: Growable v (PrimState m) a -> Int -> a -> m ()
write (Growable MutVar (PrimState m) (GVState v (PrimState m) a)
ref) Int
i a
val = do
  GVState Int
len v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
pending <- MutVar (PrimState m) (GVState v (PrimState m) a)
-> m (GVState v (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (GVState v (PrimState m) a)
ref
  case Int
len Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Int
i of
    Ordering
LT -> ArrayException -> m ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (ArrayException -> m ()) -> ArrayException -> m ()
forall a b. (a -> b) -> a -> b
$ String -> ArrayException
IndexOutOfBounds (String -> ArrayException) -> String -> ArrayException
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
len String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" < " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
i
    Ordering
EQ -> do
      let amount :: Int
amount = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` (Int -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros Int
i) Int -> Int -> Int
forall a. Num a => a -> a -> a
- v (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
MV.length v (PrimState m) a
vec
      v (PrimState m) a
vec' <- v (PrimState m) a -> Int -> m (v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (v (PrimState m) a)
MV.unsafeGrow v (PrimState m) a
vec Int
amount
      v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite v (PrimState m) a
vec' Int
i a
val
      MutVar (PrimState m) (GVState v (PrimState m) a)
-> GVState v (PrimState m) a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> a -> m ()
writeMutVar MutVar (PrimState m) (GVState v (PrimState m) a)
ref (GVState v (PrimState m) a -> m ())
-> GVState v (PrimState m) a -> m ()
forall a b. (a -> b) -> a -> b
$ Int
-> v (PrimState m) a
-> MutVar (PrimState m) (Pending a)
-> GVState v (PrimState m) a
forall (v :: * -> * -> *) s a.
Int -> v s a -> MutVar s (Pending a) -> GVState v s a
GVState (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) v (PrimState m) a
vec' MutVar (PrimState m) (Pending a)
pending
    Ordering
GT -> v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite v (PrimState m) a
vec Int
i a
val
{-# INLINE write #-}

modify :: (PrimMonad m, MVector v a, MonadThrow m) => Growable v (PrimState m) a -> (a -> a) -> Int -> m ()
modify :: Growable v (PrimState m) a -> (a -> a) -> Int -> m ()
modify (Growable MutVar (PrimState m) (GVState v (PrimState m) a)
ref) a -> a
f Int
i = do
  GVState Int
len v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
_ <- MutVar (PrimState m) (GVState v (PrimState m) a)
-> m (GVState v (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (GVState v (PrimState m) a)
ref
  if Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i
    then ArrayException -> m ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (ArrayException -> m ()) -> ArrayException -> m ()
forall a b. (a -> b) -> a -> b
$ String -> ArrayException
IndexOutOfBounds (String -> ArrayException) -> String -> ArrayException
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
len String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" <= " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
i
    else v (PrimState m) a -> (a -> a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MV.unsafeModify v (PrimState m) a
vec a -> a
f Int
i
{-# INLINE modify #-}

-- | Thaw an immutable vector and create a 'Growable' one.
thaw :: (G.Vector v a, PrimMonad m) => v a -> m (Growable (G.Mutable v) (PrimState m) a)
thaw :: v a -> m (Growable (Mutable v) (PrimState m) a)
thaw v a
v = do
  Mutable v (PrimState m) a
vec <- v a -> m (Mutable v (PrimState m) a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
G.thaw v a
v
  MutVar (PrimState m) (Pending a)
pending <- Pending a -> m (MutVar (PrimState m) (Pending a))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar Pending a
forall a. Pending a
Complete
  MutVar (PrimState m) (GVState (Mutable v) (PrimState m) a)
-> Growable (Mutable v) (PrimState m) a
forall (v :: * -> * -> *) s a.
MutVar s (GVState v s a) -> Growable v s a
Growable (MutVar (PrimState m) (GVState (Mutable v) (PrimState m) a)
 -> Growable (Mutable v) (PrimState m) a)
-> m (MutVar (PrimState m) (GVState (Mutable v) (PrimState m) a))
-> m (Growable (Mutable v) (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GVState (Mutable v) (PrimState m) a
-> m (MutVar (PrimState m) (GVState (Mutable v) (PrimState m) a))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar (Int
-> Mutable v (PrimState m) a
-> MutVar (PrimState m) (Pending a)
-> GVState (Mutable v) (PrimState m) a
forall (v :: * -> * -> *) s a.
Int -> v s a -> MutVar s (Pending a) -> GVState v s a
GVState (v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
v) Mutable v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
pending)
{-# INLINE thaw #-}

-- | Take a snapshot of a 'Growable' vector.
freeze :: (G.Vector v a, PrimMonad m) => Growable (G.Mutable v) (PrimState m) a -> m (v a)
freeze :: Growable (Mutable v) (PrimState m) a -> m (v a)
freeze (Growable MutVar (PrimState m) (GVState (Mutable v) (PrimState m) a)
ref) = do
  GVState Int
len Mutable v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
_ <- MutVar (PrimState m) (GVState (Mutable v) (PrimState m) a)
-> m (GVState (Mutable v) (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (GVState (Mutable v) (PrimState m) a)
ref
  v a
v <- Mutable v (PrimState m) a -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.freeze Mutable v (PrimState m) a
vec
  v a -> m (v a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (v a -> m (v a)) -> v a -> m (v a)
forall a b. (a -> b) -> a -> b
$! Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.unsafeTake Int
len v a
v
{-# INLINE freeze #-}

-- | Take a snapshot of a 'Growable' vector. The original vector may not be used.
unsafeFreeze :: (G.Vector v a, PrimMonad m) => Growable (G.Mutable v) (PrimState m) a -> m (v a)
unsafeFreeze :: Growable (Mutable v) (PrimState m) a -> m (v a)
unsafeFreeze (Growable MutVar (PrimState m) (GVState (Mutable v) (PrimState m) a)
ref) = do
  GVState Int
len Mutable v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
_ <- MutVar (PrimState m) (GVState (Mutable v) (PrimState m) a)
-> m (GVState (Mutable v) (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (GVState (Mutable v) (PrimState m) a)
ref
  v a
v <- Mutable v (PrimState m) a -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze Mutable v (PrimState m) a
vec
  v a -> m (v a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (v a -> m (v a)) -> v a -> m (v a)
forall a b. (a -> b) -> a -> b
$! Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.unsafeTake Int
len v a
v
{-# INLINE unsafeFreeze #-}

-- | Turn 'Growable' vector into a regular mutable vector.
fromGrowable :: (PrimMonad m, MVector v a) => Growable v (PrimState m) a -> m (v (PrimState m) a)
fromGrowable :: Growable v (PrimState m) a -> m (v (PrimState m) a)
fromGrowable (Growable MutVar (PrimState m) (GVState v (PrimState m) a)
ref) = do
  GVState Int
len v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
_ <- MutVar (PrimState m) (GVState v (PrimState m) a)
-> m (GVState v (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (GVState v (PrimState m) a)
ref
  v (PrimState m) a -> m (v (PrimState m) a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (v (PrimState m) a -> m (v (PrimState m) a))
-> v (PrimState m) a -> m (v (PrimState m) a)
forall a b. (a -> b) -> a -> b
$! Int -> v (PrimState m) a -> v (PrimState m) a
forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
MV.unsafeTake Int
len v (PrimState m) a
vec
{-# INLINE fromGrowable #-}

-- | Create a 'Growable' vector from a mutable vector.
toGrowable ::  (PrimMonad m, MVector v a) => v (PrimState m) a -> m (Growable v (PrimState m) a)
toGrowable :: v (PrimState m) a -> m (Growable v (PrimState m) a)
toGrowable v (PrimState m) a
vec = do
    MutVar (PrimState m) (Pending a)
pending <- Pending a -> m (MutVar (PrimState m) (Pending a))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar Pending a
forall a. Pending a
Complete
    MutVar (PrimState m) (GVState v (PrimState m) a)
-> Growable v (PrimState m) a
forall (v :: * -> * -> *) s a.
MutVar s (GVState v s a) -> Growable v s a
Growable (MutVar (PrimState m) (GVState v (PrimState m) a)
 -> Growable v (PrimState m) a)
-> m (MutVar (PrimState m) (GVState v (PrimState m) a))
-> m (Growable v (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GVState v (PrimState m) a
-> m (MutVar (PrimState m) (GVState v (PrimState m) a))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar (Int
-> v (PrimState m) a
-> MutVar (PrimState m) (Pending a)
-> GVState v (PrimState m) a
forall (v :: * -> * -> *) s a.
Int -> v s a -> MutVar s (Pending a) -> GVState v s a
GVState (v (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
MV.length v (PrimState m) a
vec) v (PrimState m) a
vec MutVar (PrimState m) (Pending a)
pending)
{-# INLINE toGrowable #-}

class MVector v a => CASVector v a where
  readVectorElem :: v RealWorld a -> Int -> IO (Ticket a)
  casVectorElem :: v RealWorld a -> Int -> Ticket a -> a -> IO (Bool, Ticket a)

instance CASVector V.MVector a where
  readVectorElem :: MVector RealWorld a -> Int -> IO (Ticket a)
readVectorElem (V.MVector Int
_ Int
_ MutableArray RealWorld a
arr) Int
i = MutableArray RealWorld a -> Int -> IO (Ticket a)
forall a. MutableArray RealWorld a -> Int -> IO (Ticket a)
readArrayElem MutableArray RealWorld a
arr Int
i
  casVectorElem :: MVector RealWorld a -> Int -> Ticket a -> a -> IO (Bool, Ticket a)
casVectorElem (V.MVector Int
_ Int
_ MutableArray RealWorld a
arr) Int
i Ticket a
x a
y = MutableArray RealWorld a
-> Int -> Ticket a -> a -> IO (Bool, Ticket a)
forall a.
MutableArray RealWorld a
-> Int -> Ticket a -> a -> IO (Bool, Ticket a)
casArrayElem MutableArray RealWorld a
arr Int
i Ticket a
x a
y

instance CASVector P.MVector Int where
  readVectorElem :: MVector RealWorld Int -> Int -> IO (Ticket Int)
readVectorElem MVector RealWorld Int
vec Int
i = Int -> Ticket Int
forgeIntTicket (Int -> Ticket Int) -> IO Int -> IO (Ticket Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MVector (PrimState m) a -> Int -> m a
P.unsafeRead MVector RealWorld Int
MVector (PrimState IO) Int
vec Int
i
  casVectorElem :: MVector RealWorld Int
-> Int -> Ticket Int -> Int -> IO (Bool, Ticket Int)
casVectorElem (P.MVector Int
_ Int
_ MutableByteArray RealWorld
arr) Int
i Ticket Int
x Int
y = do
    let old :: Int
old = Ticket Int -> Int
forall a. Ticket a -> a
peekTicket Ticket Int
x
    Int
old' <- MutableByteArray RealWorld -> Int -> Int -> Int -> IO Int
casByteArrayInt MutableByteArray RealWorld
arr Int
i Int
old Int
y
    (Bool, Ticket Int) -> IO (Bool, Ticket Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
old Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
old', Int -> Ticket Int
forgeIntTicket Int
y)

forgeIntTicket :: Int -> Ticket Int
forgeIntTicket :: Int -> Ticket Int
forgeIntTicket = Int -> Ticket Int
forall a b. a -> b
unsafeCoerce

complete :: CASVector v a => v RealWorld a -> MutVar RealWorld (Pending a) -> IO ()
complete :: v RealWorld a -> MutVar RealWorld (Pending a) -> IO ()
complete v RealWorld a
vec MutVar RealWorld (Pending a)
v = MutVar (PrimState IO) (Pending a) -> IO (Pending a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar RealWorld (Pending a)
MutVar (PrimState IO) (Pending a)
v IO (Pending a) -> (Pending a -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  Pending a
Complete -> () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  Pending Int
i Ticket a
oldVal a
newVal -> do
    (Bool
_done, Ticket a
_) <- v RealWorld a -> Int -> Ticket a -> a -> IO (Bool, Ticket a)
forall (v :: * -> * -> *) a.
CASVector v a =>
v RealWorld a -> Int -> Ticket a -> a -> IO (Bool, Ticket a)
casVectorElem v RealWorld a
vec Int
i Ticket a
oldVal a
newVal
    MutVar (PrimState IO) (Pending a) -> Pending a -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> a -> m ()
writeMutVar MutVar RealWorld (Pending a)
MutVar (PrimState IO) (Pending a)
v Pending a
forall a. Pending a
Complete

-- | Atomically push a value to the end of the vector.
-- | Based on <https://www.stroustrup.com/lock-free-vector.pdf Damian Dechev, Peter Pirkelbauer, and Bjarne Stroustrup - Lock-free Dynamically Resizable Arrays>
atomicPush :: CASVector v a => Growable v RealWorld a -> a -> IO ()
atomicPush :: Growable v RealWorld a -> a -> IO ()
atomicPush (Growable (MutVar MutVar# RealWorld (GVState v RealWorld a)
mut)) a
val = IO ()
go
  where
    go :: IO ()
go = do
      Ticket (GVState v RealWorld a)
old <- MutVar# RealWorld (GVState v RealWorld a)
-> IO (Ticket (GVState v RealWorld a))
forall a. MutVar# RealWorld a -> IO (Ticket a)
readMutVarForCAS MutVar# RealWorld (GVState v RealWorld a)
mut
      let GVState Int
len v RealWorld a
vec MutVar RealWorld (Pending a)
pending = Ticket (GVState v RealWorld a) -> GVState v RealWorld a
forall a. Ticket a -> a
peekTicket Ticket (GVState v RealWorld a)
old
      v RealWorld a -> MutVar RealWorld (Pending a) -> IO ()
forall (v :: * -> * -> *) a.
CASVector v a =>
v RealWorld a -> MutVar RealWorld (Pending a) -> IO ()
complete v RealWorld a
vec MutVar RealWorld (Pending a)
pending
      v RealWorld a
vec' <- if v RealWorld a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
MV.length v RealWorld a
vec Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len
        then v (PrimState IO) a -> Int -> IO (v (PrimState IO) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (v (PrimState m) a)
MV.unsafeGrow v RealWorld a
v (PrimState IO) a
vec (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
len)
        else v RealWorld a -> IO (v RealWorld a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure v RealWorld a
vec
      Ticket a
oldVal <- v RealWorld a -> Int -> IO (Ticket a)
forall (v :: * -> * -> *) a.
CASVector v a =>
v RealWorld a -> Int -> IO (Ticket a)
readVectorElem v RealWorld a
vec' Int
len
      MutVar RealWorld (Pending a)
pending' <- Pending a -> IO (MutVar (PrimState IO) (Pending a))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar (Pending a -> IO (MutVar (PrimState IO) (Pending a)))
-> Pending a -> IO (MutVar (PrimState IO) (Pending a))
forall a b. (a -> b) -> a -> b
$ Int -> Ticket a -> a -> Pending a
forall a. Int -> Ticket a -> a -> Pending a
Pending Int
len Ticket a
oldVal a
val
      (Bool
success, Ticket (GVState v RealWorld a)
_) <- MutVar# RealWorld (GVState v RealWorld a)
-> Ticket (GVState v RealWorld a)
-> GVState v RealWorld a
-> IO (Bool, Ticket (GVState v RealWorld a))
forall a.
MutVar# RealWorld a -> Ticket a -> a -> IO (Bool, Ticket a)
casMutVar MutVar# RealWorld (GVState v RealWorld a)
mut Ticket (GVState v RealWorld a)
old (GVState v RealWorld a
 -> IO (Bool, Ticket (GVState v RealWorld a)))
-> GVState v RealWorld a
-> IO (Bool, Ticket (GVState v RealWorld a))
forall a b. (a -> b) -> a -> b
$ Int
-> v RealWorld a
-> MutVar RealWorld (Pending a)
-> GVState v RealWorld a
forall (v :: * -> * -> *) s a.
Int -> v s a -> MutVar s (Pending a) -> GVState v s a
GVState (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) v RealWorld a
vec' MutVar RealWorld (Pending a)
pending'
      if Bool
success
        then v RealWorld a -> MutVar RealWorld (Pending a) -> IO ()
forall (v :: * -> * -> *) a.
CASVector v a =>
v RealWorld a -> MutVar RealWorld (Pending a) -> IO ()
complete v RealWorld a
vec' MutVar RealWorld (Pending a)
pending'
        else IO ()
go

-- | Pop the last element. Returns 'Nothing' if the vector is empty.
atomicPop :: CASVector v a => Growable v RealWorld a -> IO (Maybe a)
atomicPop :: Growable v RealWorld a -> IO (Maybe a)
atomicPop (Growable (MutVar MutVar# RealWorld (GVState v RealWorld a)
mut)) = IO (Maybe a)
go
  where
    go :: IO (Maybe a)
go = do
      Ticket (GVState v RealWorld a)
old <- MutVar# RealWorld (GVState v RealWorld a)
-> IO (Ticket (GVState v RealWorld a))
forall a. MutVar# RealWorld a -> IO (Ticket a)
readMutVarForCAS MutVar# RealWorld (GVState v RealWorld a)
mut
      let GVState Int
len v RealWorld a
vec MutVar RealWorld (Pending a)
pending = Ticket (GVState v RealWorld a) -> GVState v RealWorld a
forall a. Ticket a -> a
peekTicket Ticket (GVState v RealWorld a)
old
      v RealWorld a -> MutVar RealWorld (Pending a) -> IO ()
forall (v :: * -> * -> *) a.
CASVector v a =>
v RealWorld a -> MutVar RealWorld (Pending a) -> IO ()
complete v RealWorld a
vec MutVar RealWorld (Pending a)
pending
      if Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then Maybe a -> IO (Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
        else do
          a
result <- v (PrimState IO) a -> Int -> IO a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MV.unsafeRead v RealWorld a
v (PrimState IO) a
vec (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
          MutVar RealWorld (Pending a)
pending' <- Pending a -> IO (MutVar (PrimState IO) (Pending a))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar Pending a
forall a. Pending a
Complete
          (Bool
success, Ticket (GVState v RealWorld a)
_) <- MutVar# RealWorld (GVState v RealWorld a)
-> Ticket (GVState v RealWorld a)
-> GVState v RealWorld a
-> IO (Bool, Ticket (GVState v RealWorld a))
forall a.
MutVar# RealWorld a -> Ticket a -> a -> IO (Bool, Ticket a)
casMutVar MutVar# RealWorld (GVState v RealWorld a)
mut Ticket (GVState v RealWorld a)
old (GVState v RealWorld a
 -> IO (Bool, Ticket (GVState v RealWorld a)))
-> GVState v RealWorld a
-> IO (Bool, Ticket (GVState v RealWorld a))
forall a b. (a -> b) -> a -> b
$ Int
-> v RealWorld a
-> MutVar RealWorld (Pending a)
-> GVState v RealWorld a
forall (v :: * -> * -> *) s a.
Int -> v s a -> MutVar s (Pending a) -> GVState v s a
GVState (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) v RealWorld a
vec MutVar RealWorld (Pending a)
pending'
          if Bool
success
            then Maybe a -> IO (Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> IO (Maybe a)) -> Maybe a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
result
            else IO (Maybe a)
go
{-# INLINE atomicPop #-}