module Data.Vector.Storable.Buffer (
Buffer,
newBuffer,
pushNextElement,
toVector,
mapBufferM, mapBufferM_,
) where
import Data.IORef
import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as M
import Foreign hiding(new)
data Buffer a = B { end :: !(IORef Int)
, dat :: !(M.IOVector a)
}
newBuffer :: Storable a
=> Int
-> IO (Buffer a)
newBuffer n = do
v <- M.new n
o <- newIORef 0
return $ B o v
pushNextElement :: Storable a => Buffer a -> a -> IO ()
pushNextElement b@(B o v) e = do
let n = M.length v
i <- readIORef o
M.unsafeWrite v i e
if i == (n1)
then writeIORef o 0
else writeIORef o (i+1)
toVector :: Storable a => Buffer a -> V.Vector a
toVector (B o v) = unsafePerformIO $ do
let n = M.length v
w <- M.new n
i <- readIORef o
M.unsafeWith v $ \p ->
M.unsafeWith w $ \q -> do
copyArray q (p `advancePtr` i) (ni)
if i /= 0
then copyArray (q `advancePtr` (ni)) p i
else return ()
V.unsafeFreeze w
mapBufferM :: (Storable a, Storable b) => (a -> IO b) -> Buffer a -> IO (V.Vector b)
mapBufferM f (B o v) = do
let n = M.length v
w <- M.new n
i <- readIORef o
go w 0 i n
V.unsafeFreeze w
where go w' i' o' n'
| i' + 1 == n' = do
x <- M.unsafeRead v (if i'+o' >= n' then i'+o'n' else i'+o')
y <- f x
M.unsafeWrite w' i' y
| otherwise = do
x <- M.unsafeRead v (if i'+o' >= n' then i'+o'n' else i'+o')
y <- f x
M.unsafeWrite w' i' y
let i'' = if i' + 1 == n' then 0 else i' + 1
go w' i'' o' n'
mapBufferM_ :: (Storable a) => (a -> IO b) -> Buffer a -> IO ()
mapBufferM_ f (B o v) = do
let n = M.length v
i <- readIORef o
go 0 i n
where go i' o' n'
| i' + 1 == n' = do
x <- M.unsafeRead v (if i'+o' >= n' then i'+o'n' else i'+o')
_ <- f x
return ()
| otherwise = do
x <- M.unsafeRead v (if i'+o' >= n' then i'+o'n' else i'+o')
_ <- f x
let i'' = if i' + 1 == n' then 0 else i' + 1
go i'' o' n'