module Data.Vector.Buffer (
Buffer,
newBuffer,
pushNextElement,
toVector,
mapBufferM, mapBufferM_,
) where
import Data.IORef
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as M
import Foreign hiding(new)
data Buffer a = B { nxt :: !(IORef Int)
, siz :: !Int
, dat :: !(M.IOVector a)
}
newBuffer :: Int
-> IO (Buffer a)
newBuffer n = do
v <- M.new n
o <- newIORef 0
return $ B o n v
pushNextElement :: Buffer a -> a -> IO ()
pushNextElement (B o n v) e = do
i <- readIORef o
M.unsafeWrite v i e
if i == (n1)
then writeIORef o 0
else writeIORef o (i+1)
toVector :: Buffer a -> IO (V.Vector a)
toVector (B o n v) = do
w <- M.new n
i <- readIORef o
let n' = ni
let vs = M.slice i n' v
let ws = M.slice 0 n' w
M.unsafeCopy ws vs
if i /= 0
then do
let vf = M.slice 0 i v
let wf = M.slice n' i w
M.unsafeCopy wf vf
else return ()
V.unsafeFreeze w
mapBufferM :: (a -> IO b) -> Buffer a -> IO (V.Vector b)
mapBufferM f (B o n v) = do
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
let j = i'+o'
x <- M.unsafeRead v (j `mod` n')
y <- f x
M.unsafeWrite w' i' y
| otherwise = do
let j = i'+o'
x <- M.unsafeRead v (j `mod` n')
y <- f x
M.unsafeWrite w' i' y
go w' ((i'+1) `mod` n') o' n'
mapBufferM_ :: (a -> IO b) -> Buffer a -> IO ()
mapBufferM_ f (B o n v) = do
i <- readIORef o
go 0 i n
where go !i' !o' !n'
| i' + 1 == n' = do
let j = i'+o'
x <- M.unsafeRead v (j `mod` n')
_ <- f x
return ()
| otherwise = do
let j = i'+o'
x <- M.unsafeRead v (j `mod` n')
_ <- f x
go ((i'+1) `mod` n') o' n'