module Data.Vector.Conduit
(
sourceVector,
sourceMVector,
consumeVector,
consumeMVector
)
where
import Control.Monad.Primitive
import Control.Monad.ST.Safe
import Control.Monad.Trans.Class
import Data.Conduit
import qualified Data.DList as D
import qualified Data.Vector.Generic as V
import qualified Data.Vector.Generic.Mutable as M
sourceVector :: (Resource m, V.Vector v a) => v a -> Source m a
sourceVector v = sourceState 0 f
where f index | index == V.length v = return StateClosed
| otherwise = return $ StateOpen (index + 1) (v V.! index)
sourceMVector :: (PrimMonad m, Resource m, M.MVector v a)
=> v (PrimState m) a
-> Source m a
sourceMVector v = sourceState 0 f
where f index | index == M.length v = return StateClosed
| otherwise = do x <- lift $ M.read v index
return $ StateOpen (index + 1) x
consumeVector :: (Resource m, V.Vector v a) => Sink a m (v a)
consumeVector = sinkState D.empty push close
where push xs x = return . StateProcessing $ D.snoc xs x
close xs = return . V.fromList . D.toList $ xs
consumeMVector :: (PrimMonad m, Resource m, M.MVector v a)
=> Sink a m (v (PrimState m) a)
consumeMVector = sinkState (Nothing, 0) push close
where push (v, index) x = do v' <- case v of
Nothing -> lift $ M.new 10
Just vec -> return vec
let len = M.length v'
v'' <- if index >= len
then lift $ M.grow v' len
else return v'
lift $ M.write v'' index x
return $ StateProcessing (Just v'', index + 1)
close (Nothing, index) = lift $ M.new 0
close (Just v, index) = return $ M.take index v