{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
module Data.ArrayList.Generic
(
ArrayList
, new, fromVector
, vector, size
, unsafeVector
, push, grow
, freeze
) where
import Control.Monad (when)
import Control.Monad.Primitive
import Control.Monad.ST (ST)
import Data.Primitive.MutVar
import Data.Vector.Generic (Mutable, Vector)
import qualified Data.Vector.Generic as GV
import Data.Vector.Generic.Mutable (MVector)
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed.Mutable as VUM
import GHC.Prim (RealWorld)
data ArrayList v s a = ArrayList
{ arrayListSize :: !(MutVar s Int)
, arrayListVector :: !(MutVar s (v s a))
}
new :: (PrimMonad m, MVector v a) => Int -> m (ArrayList v (PrimState m) a)
new len = ArrayList <$> newMutVar 0 <*> (newMutVar =<< GM.new len)
{-# INLINEABLE new #-}
fromVector :: (PrimMonad m, Vector v a) => v a -> m (ArrayList (Mutable v) (PrimState m) a)
fromVector vec = ArrayList <$> newMutVar (GV.length vec) <*> (newMutVar =<< GV.thaw vec)
{-# INLINEABLE fromVector #-}
vector :: (PrimMonad m, MVector v a) => ArrayList v (PrimState m) a -> m (v (PrimState m) a)
vector (ArrayList sizeRef mvecRef) = do
!size_ <- readMutVar sizeRef
GM.unsafeTake size_ <$> readMutVar mvecRef
{-# INLINEABLE vector #-}
unsafeVector :: (PrimMonad m, MVector v a) => ArrayList v (PrimState m) a -> m (v (PrimState m) a)
unsafeVector (ArrayList _ mvecRef) = readMutVar mvecRef
{-# INLINEABLE unsafeVector #-}
size :: (PrimMonad m, MVector v a) => ArrayList v (PrimState m) a -> m Int
size (ArrayList sizeRef _) = readMutVar sizeRef
{-# INLINEABLE size #-}
push :: (PrimMonad m, MVector v a) => ArrayList v (PrimState m) a -> a -> m ()
push al@(ArrayList sizeRef mvecRef) a = do
grow al 1
!size_ <- readMutVar sizeRef
!mvec <- readMutVar mvecRef
GM.unsafeWrite mvec (size_ - 1) a
{-# INLINEABLE push #-}
{-# SPECIALIZE push :: ArrayList VM.MVector RealWorld a -> a -> IO () #-}
{-# SPECIALIZE push :: ArrayList VM.MVector s a -> a -> ST s () #-}
{-# SPECIALIZE push :: (GM.MVector VUM.MVector a) => ArrayList VUM.MVector RealWorld a -> a -> IO () #-}
{-# SPECIALIZE push :: (GM.MVector VUM.MVector a) => ArrayList VUM.MVector s a -> a -> ST s () #-}
{-# SPECIALIZE push :: (MVector v a) => ArrayList v RealWorld a -> a -> IO () #-}
{-# SPECIALIZE push :: (MVector v a) => ArrayList v s a -> a -> ST s () #-}
grow :: (PrimMonad m, MVector v a) => ArrayList v (PrimState m) a -> Int -> m ()
grow (ArrayList sizeRef mvecRef) extraSize = do
!size_ <- readMutVar sizeRef
!mvec <- readMutVar mvecRef
let
capacity = GM.length mvec
newSize = size_ + extraSize
newCapacity = floor $ fromIntegral capacity * factor
when (capacity < newSize) $ do
!mvec' <- if newSize > newCapacity
then
{-# SCC "ArrayList.grow.toNewSize" #-} GM.unsafeGrow mvec (newSize - capacity)
else
{-# SCC "ArrayList.grow.toNewCapacity" #-} GM.unsafeGrow mvec (newCapacity - capacity)
writeMutVar mvecRef mvec'
writeMutVar sizeRef newSize
{-# INLINEABLE grow #-}
{-# SPECIALIZE grow :: ArrayList VM.MVector RealWorld a -> Int -> IO () #-}
{-# SPECIALIZE grow :: ArrayList VM.MVector s a -> Int -> ST s () #-}
{-# SPECIALIZE grow :: (GM.MVector VUM.MVector a) => ArrayList VUM.MVector RealWorld a -> Int -> IO () #-}
{-# SPECIALIZE grow :: (GM.MVector VUM.MVector a) => ArrayList VUM.MVector s a -> Int -> ST s () #-}
{-# SPECIALIZE grow :: (MVector v a) => ArrayList v RealWorld a -> Int -> IO () #-}
{-# SPECIALIZE grow :: (MVector v a) => ArrayList v s a -> Int -> ST s () #-}
factor :: Double
factor = 1.5
freeze :: (PrimMonad m, Vector v a) => ArrayList (Mutable v) (PrimState m) a -> m (v a)
freeze (ArrayList sizeRef mvecRef) = do
!size_ <- readMutVar sizeRef
!mvec <- readMutVar mvecRef
let sizedMVec = GM.unsafeTake size_ mvec
GV.freeze sizedMVec
{-# INLINEABLE freeze #-}