module Pipes.Vector (
toVector,
runToVectorP,
runToVector,
fromProducer,
ToVector
) where
import Control.Applicative
import Control.Monad
import Control.Monad.Trans.State.Strict as S
import Control.Monad.Primitive
import Control.Monad.Primitive.Class
import Pipes
import Pipes.Internal (unsafeHoist)
import Pipes.Lift
import qualified Data.Vector.Generic as V
import qualified Data.Vector.Generic.Mutable as M
data ToVectorState v e m = ToVecS { result :: V.Mutable v (PrimState (BasePrimMonad m)) e
, idx :: Int
}
newtype ToVector v e m r = TV {unTV :: S.StateT (ToVectorState v e m) m r}
deriving (Functor, Applicative, Monad, MonadIO)
instance MonadTrans (ToVector v e) where
lift = TV . lift
instance MonadPrim m => MonadPrim (Proxy a' a b' b m) where
type BasePrimMonad (Proxy a' a b' b m) = BasePrimMonad m
liftPrim = lift . liftPrim
instance MonadPrim m => MonadPrim (ToVector v e m) where
type BasePrimMonad (ToVector v e m) = BasePrimMonad m
liftPrim = TV . liftPrim
maxChunkSize :: Int
maxChunkSize = 8*1024*1024
toVector
:: (MonadPrim m, M.MVector (V.Mutable v) e)
=> Consumer e (ToVector v e m) r
toVector = forever $ do
length <- M.length . result <$> lift (TV get)
pos <- idx `liftM` lift (TV get)
lift $ TV $ when (pos >= length) $ do
v <- result `liftM` get
v' <- liftPrim $ M.unsafeGrow v (min length maxChunkSize)
modify $ \(ToVecS r i) -> ToVecS v' i
r <- await
lift $ TV $ do
v <- result `liftM` get
liftPrim $ M.unsafeWrite v pos r
modify $ \(ToVecS r i) -> ToVecS r (pos+1)
runToVectorP
:: (MonadPrim m, V.Vector v e)
=> Proxy a' a b' b (ToVector v e m) r
-> Proxy a' a b' b m (v e)
runToVectorP x = do
v <- liftPrim $ M.new 10
s <- execStateP (ToVecS v 0) (hoist unTV x)
frozen <- liftPrim $ V.freeze (result s)
return $ V.take (idx s) frozen
runToVector :: (MonadPrim m, V.Vector v e)
=> ToVector v e m r -> m (v e)
runToVector (TV a) = do
v <- liftPrim $ M.new 10
s <- execStateT a (ToVecS v 0)
frozen <- liftPrim $ V.freeze (result s)
return $ V.take (idx s) frozen
fromProducer :: (V.Vector v e, MonadPrim m) => Producer e (ToVector v e m) r -> m (v e)
fromProducer p = runEffect $ runToVectorP (p >-> toVector)