module Data.Vector.Mutable.PushBack where
import Prelude hiding (length, read)
import Control.Monad
import Data.IORef
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import System.IO.Unsafe
data IOVector a = IOVector !(IORef (VM.IOVector a)) !(VUM.IOVector Int)
new :: Int -> IO (IOVector a)
new p = new' (p + 10)
where new' p = IOVector <$> (newIORef =<< VM.new p) <*> (VUM.replicate 1 0)
read :: IOVector a -> Int -> IO a
read (IOVector vref _) k = readIORef vref >>= \vec -> VM.read vec k
safeLength :: IOVector a -> IO Int
safeLength (IOVector _ uvec) = VUM.read uvec 0
length :: IOVector a -> Int
length pvec = unsafePerformIO $ safeLength pvec
safeCapacity :: IOVector a -> IO Int
safeCapacity (IOVector vref _) = fmap VM.length $ readIORef vref
capacity :: IOVector a -> Int
capacity pvec = unsafePerformIO $ safeCapacity pvec
write :: IOVector a -> Int -> a -> IO ()
write (IOVector vref _) i v = do
vec <- readIORef vref
VM.write vec i v
insert
:: IOVector a
-> Int
-> a
-> IO ()
insert pvec i v = do
len <- safeLength pvec
read pvec (len - 1) >>= push pvec
forM_ (reverse [i .. len - 2]) $ \j -> read pvec j >>= write pvec (j + 1)
write pvec i v
delete :: IOVector a -> Int -> IO ()
delete pvec@(IOVector _ uvec) i = do
len <- safeLength pvec
forM_ [i + 1 .. len - 1] $ \j -> read pvec j >>= write pvec (j - 1)
VUM.modify uvec (\x -> x - 1) 0
push :: IOVector a -> a -> IO ()
push pvec@(IOVector vref uvec) v = do
vec <- readIORef vref
len <- safeLength pvec
cap <- safeCapacity pvec
when (len == cap) $ do
vec' <- VM.grow vec cap
writeIORef vref vec'
write pvec len v
VUM.modify uvec (+ 1) 0
fromList :: [a] -> IO (IOVector a)
fromList xs = do
vec <- V.thaw $ V.fromList xs
vec' <- VM.grow vec ((VM.length vec + 5) * 2)
vref <- newIORef vec'
uvec <- VU.thaw $ VU.fromList [VM.length vec]
return $ IOVector vref uvec
asIOVector :: IOVector a -> IO (VM.IOVector a)
asIOVector pvec@(IOVector vref _) = do
len <- safeLength pvec
readIORef vref >>= \vec -> return (VM.slice 0 len vec)
asUnsafeIOVector :: IOVector a -> VM.IOVector a
asUnsafeIOVector pvec = unsafePerformIO $ asIOVector pvec