{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}

-- | Vector conversions and utilities.

module System.IO.Streams.Vector
 ( -- * Vector conversions
   fromVector
 , toVector
 , toVectorSized
 , outputToVector
 , outputToVectorSized
 , toMutableVector
 , toMutableVectorSized
 , outputToMutableVector
 , outputToMutableVectorSized
 , writeVector

   -- * Utility
 , chunkVector
 , vectorOutputStream
 , vectorOutputStreamSized
 , mutableVectorOutputStream
 , mutableVectorOutputStreamSized
 ) where

------------------------------------------------------------------------------
import           Control.Concurrent.MVar     (modifyMVar, modifyMVar_, newMVar)
import           Control.Monad               (liftM, (>=>))
import           Control.Monad.IO.Class      (MonadIO (..))
import           Control.Monad.Primitive     (PrimState (..))
import           Data.IORef                  (IORef, newIORef, readIORef, writeIORef)
import           Data.Vector.Generic         (Vector (..))
import qualified Data.Vector.Generic         as V
import           Data.Vector.Generic.Mutable (MVector)
import qualified Data.Vector.Generic.Mutable as VM
import           System.IO.Streams.Internal  (InputStream, OutputStream, fromGenerator, yield)
import qualified System.IO.Streams.Internal  as S


------------------------------------------------------------------------------
-- | Transforms a vector into an 'InputStream' that yields each of the values
-- in the vector in turn.
--
-- @
-- ghci> import "Control.Monad"
-- ghci> import qualified "System.IO.Streams" as Streams
-- ghci> import qualified "Data.Vector" as V
-- ghci> let v = V.'Data.Vector.fromList' [1, 2]
-- ghci> is <- Streams.'fromVector' v
-- ghci> 'Control.Monad.replicateM' 3 (Streams.'read' is)
-- ['Just' 1,'Just' 2,'Nothing']
-- @
fromVector :: Vector v a => v a -> IO (InputStream a)
fromVector :: v a -> IO (InputStream a)
fromVector = Generator a () -> IO (InputStream a)
forall r a. Generator r a -> IO (InputStream r)
fromGenerator (Generator a () -> IO (InputStream a))
-> (v a -> Generator a ()) -> v a -> IO (InputStream a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Generator a ()) -> v a -> Generator a ()
forall (m :: * -> *) (v :: * -> *) a b.
(Monad m, Vector v a) =>
(a -> m b) -> v a -> m ()
V.mapM_ a -> Generator a ()
forall r. r -> Generator r ()
yield
{-# INLINE fromVector #-}


------------------------------------------------------------------------------
-- | Drains an 'InputStream', converting it to a vector. Note that this
-- function reads the entire 'InputStream' strictly into memory and as such is
-- not recommended for streaming applications or where the size of the input is
-- not bounded or known.
--
-- @
-- ghci> is <- Streams.'System.IO.Streams.fromList' [(1::Int)..4]
-- ghci> Streams.'toVector' is :: 'IO' (V.'Vector' Int)
-- fromList [1,2,3,4]
-- @
toVector :: Vector v a => InputStream a -> IO (v a)
toVector :: InputStream a -> IO (v a)
toVector = Int -> InputStream a -> IO (v a)
forall (v :: * -> *) a.
Vector v a =>
Int -> InputStream a -> IO (v a)
toVectorSized Int
dEFAULT_BUFSIZ
{-# INLINE toVector #-}


------------------------------------------------------------------------------
-- | Like 'toVector', but allows control over how large the vector buffer is to
-- start with.
toVectorSized :: Vector v a => Int -> InputStream a -> IO (v a)
toVectorSized :: Int -> InputStream a -> IO (v a)
toVectorSized Int
n = Int -> InputStream a -> IO (Mutable v (PrimState IO) a)
forall (v :: * -> * -> *) a.
MVector v a =>
Int -> InputStream a -> IO (v (PrimState IO) a)
toMutableVectorSized Int
n (InputStream a -> IO (Mutable v RealWorld a))
-> (Mutable v RealWorld a -> IO (v a)) -> InputStream a -> IO (v a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Mutable v RealWorld a -> IO (v a)
forall (v :: * -> *) a (m :: * -> *).
(Vector v a, PrimMonad m) =>
Mutable v (PrimState m) a -> m (v a)
V.basicUnsafeFreeze
{-# INLINE toVectorSized #-}


------------------------------------------------------------------------------
-- | Drains an 'InputStream', converting it to a mutable vector. Note that this
-- function reads the entire 'InputStream' strictly into memory and as such is
-- not recommended for streaming applications or where the size of the input is
-- not bounded or known.
toMutableVector :: VM.MVector v a => InputStream a -> IO (v (PrimState IO) a)
toMutableVector :: InputStream a -> IO (v (PrimState IO) a)
toMutableVector = Int -> InputStream a -> IO (v (PrimState IO) a)
forall (v :: * -> * -> *) a.
MVector v a =>
Int -> InputStream a -> IO (v (PrimState IO) a)
toMutableVectorSized Int
dEFAULT_BUFSIZ


------------------------------------------------------------------------------
-- | Like 'toMutableVector', but allows control over how large the vector
-- buffer is to start with.
toMutableVectorSized :: VM.MVector v a =>
                        Int            -- ^ initial size of the vector buffer
                     -> InputStream a
                     -> IO (v (PrimState IO) a)
toMutableVectorSized :: Int -> InputStream a -> IO (v (PrimState IO) a)
toMutableVectorSized Int
initialSize InputStream a
input = Int -> IO (VectorFillInfo v a)
forall (v :: * -> * -> *) a.
MVector v a =>
Int -> IO (VectorFillInfo v a)
vfNew Int
initialSize IO (VectorFillInfo v a)
-> (VectorFillInfo v a -> IO (v RealWorld a)) -> IO (v RealWorld a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= VectorFillInfo v a -> IO (v RealWorld a)
forall (v :: * -> * -> *).
MVector v a =>
VectorFillInfo v a -> IO (v RealWorld a)
go
  where
    go :: VectorFillInfo v a -> IO (v RealWorld a)
go VectorFillInfo v a
vfi = InputStream a -> IO (Maybe a)
forall a. InputStream a -> IO (Maybe a)
S.read InputStream a
input IO (Maybe a)
-> (Maybe a -> IO (v RealWorld a)) -> IO (v RealWorld a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO (v RealWorld a)
-> (a -> IO (v RealWorld a)) -> Maybe a -> IO (v RealWorld a)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (VectorFillInfo v a -> IO (v (PrimState IO) a)
forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> IO (v (PrimState IO) a)
vfFinish VectorFillInfo v a
vfi) (VectorFillInfo v a -> a -> IO (VectorFillInfo v a)
forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> a -> IO (VectorFillInfo v a)
vfAppend VectorFillInfo v a
vfi (a -> IO (VectorFillInfo v a))
-> (VectorFillInfo v a -> IO (v RealWorld a))
-> a
-> IO (v RealWorld a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> VectorFillInfo v a -> IO (v RealWorld a)
go)
{-# INLINE toMutableVectorSized #-}


------------------------------------------------------------------------------
-- | 'vectorOutputStream' returns an 'OutputStream' which stores values fed
-- into it and an action which flushes all stored values to a vector.
--
-- The flush action resets the store.
--
-- Note that this function /will/ buffer any input sent to it on the heap.
-- Please don't use this unless you're sure that the amount of input provided
-- is bounded and will fit in memory without issues.
--
-- @
-- ghci> (os, flush) <- Streams.'vectorOutputStream' :: IO ('OutputStream' Int, IO (V.'Vector' Int))
-- ghci> Streams.'System.IO.Streams.write' (Just 1) os
-- ghci> Streams.'System.IO.Streams.write' (Just 2) os
-- ghci> flush
-- fromList [1,2]
-- ghci> Streams.'System.IO.Streams.write' (Just 3) os
-- ghci> Streams.'System.IO.Streams.write' Nothing  os
-- ghci> Streams.'System.IO.Streams.write' (Just 4) os
-- ghci> flush
-- fromList [3]
-- @
vectorOutputStream :: Vector v c => IO (OutputStream c, IO (v c))
vectorOutputStream :: IO (OutputStream c, IO (v c))
vectorOutputStream = Int -> IO (OutputStream c, IO (v c))
forall (v :: * -> *) c.
Vector v c =>
Int -> IO (OutputStream c, IO (v c))
vectorOutputStreamSized Int
dEFAULT_BUFSIZ
{-# INLINE vectorOutputStream #-}


------------------------------------------------------------------------------
-- | Like 'vectorOutputStream', but allows control over how large the vector
-- buffer is to start with.
vectorOutputStreamSized :: Vector v c => Int -> IO (OutputStream c, IO (v c))
vectorOutputStreamSized :: Int -> IO (OutputStream c, IO (v c))
vectorOutputStreamSized Int
n = do
    (OutputStream c
os, IO (Mutable v RealWorld c)
flush) <- Int -> IO (OutputStream c, IO (Mutable v (PrimState IO) c))
forall (v :: * -> * -> *) c.
MVector v c =>
Int -> IO (OutputStream c, IO (v (PrimState IO) c))
mutableVectorOutputStreamSized Int
n
    (OutputStream c, IO (v c)) -> IO (OutputStream c, IO (v c))
forall (m :: * -> *) a. Monad m => a -> m a
return ((OutputStream c, IO (v c)) -> IO (OutputStream c, IO (v c)))
-> (OutputStream c, IO (v c)) -> IO (OutputStream c, IO (v c))
forall a b. (a -> b) -> a -> b
$! (OutputStream c
os, IO (Mutable v RealWorld c)
flush IO (Mutable v RealWorld c)
-> (Mutable v RealWorld c -> IO (v c)) -> IO (v c)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Mutable v RealWorld c -> IO (v c)
forall (v :: * -> *) a (m :: * -> *).
(Vector v a, PrimMonad m) =>
Mutable v (PrimState m) a -> m (v a)
V.basicUnsafeFreeze)


------------------------------------------------------------------------------
data VectorFillInfo v c = VectorFillInfo {
      VectorFillInfo v c -> v (PrimState IO) c
_vec :: !(v (PrimState IO) c)
    , VectorFillInfo v c -> IORef Int
_idx :: {-# UNPACK #-} !(IORef Int)

    -- TODO: vector contains its own size
    , VectorFillInfo v c -> IORef Int
_sz  :: {-# UNPACK #-} !(IORef Int)
    }


------------------------------------------------------------------------------
vfNew :: MVector v a => Int -> IO (VectorFillInfo v a)
vfNew :: Int -> IO (VectorFillInfo v a)
vfNew Int
initialSize = do
    v RealWorld a
v  <- Int -> IO (v (PrimState IO) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
VM.unsafeNew Int
initialSize
    IORef Int
i  <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0
    IORef Int
sz <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
initialSize
    VectorFillInfo v a -> IO (VectorFillInfo v a)
forall (m :: * -> *) a. Monad m => a -> m a
return (VectorFillInfo v a -> IO (VectorFillInfo v a))
-> VectorFillInfo v a -> IO (VectorFillInfo v a)
forall a b. (a -> b) -> a -> b
$! v (PrimState IO) a -> IORef Int -> IORef Int -> VectorFillInfo v a
forall (v :: * -> * -> *) c.
v (PrimState IO) c -> IORef Int -> IORef Int -> VectorFillInfo v c
VectorFillInfo v RealWorld a
v (PrimState IO) a
v IORef Int
i IORef Int
sz


------------------------------------------------------------------------------
vfFinish :: MVector v a =>
            VectorFillInfo v a
         -> IO (v (PrimState IO) a)
vfFinish :: VectorFillInfo v a -> IO (v (PrimState IO) a)
vfFinish VectorFillInfo v a
vfi = (Int -> v RealWorld a) -> IO Int -> IO (v RealWorld a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM ((Int -> v RealWorld a -> v RealWorld a)
-> v RealWorld a -> Int -> v RealWorld a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> v RealWorld a -> v RealWorld a
forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
VM.unsafeTake v RealWorld a
v) (IO Int -> IO (v RealWorld a)) -> IO Int -> IO (v RealWorld a)
forall a b. (a -> b) -> a -> b
$ IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
i
  where
    v :: v (PrimState IO) a
v = VectorFillInfo v a -> v (PrimState IO) a
forall (v :: * -> * -> *) c.
VectorFillInfo v c -> v (PrimState IO) c
_vec VectorFillInfo v a
vfi
    i :: IORef Int
i = VectorFillInfo v a -> IORef Int
forall (v :: * -> * -> *) c. VectorFillInfo v c -> IORef Int
_idx VectorFillInfo v a
vfi


------------------------------------------------------------------------------
vfAppend :: MVector v a =>
         VectorFillInfo v a
      -> a
      -> IO (VectorFillInfo v a)
vfAppend :: VectorFillInfo v a -> a -> IO (VectorFillInfo v a)
vfAppend VectorFillInfo v a
vfi !a
x = do
    Int
i  <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
iRef
    Int
sz <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
szRef
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz then Int -> IO (VectorFillInfo v a)
add Int
i else Int -> IO (VectorFillInfo v a)
grow Int
sz
  where
    v :: v (PrimState IO) a
v     = VectorFillInfo v a -> v (PrimState IO) a
forall (v :: * -> * -> *) c.
VectorFillInfo v c -> v (PrimState IO) c
_vec VectorFillInfo v a
vfi
    iRef :: IORef Int
iRef  = VectorFillInfo v a -> IORef Int
forall (v :: * -> * -> *) c. VectorFillInfo v c -> IORef Int
_idx VectorFillInfo v a
vfi
    szRef :: IORef Int
szRef = VectorFillInfo v a -> IORef Int
forall (v :: * -> * -> *) c. VectorFillInfo v c -> IORef Int
_sz VectorFillInfo v a
vfi

    add :: Int -> IO (VectorFillInfo v a)
add Int
i = do
        v (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VM.unsafeWrite v RealWorld a
v (PrimState IO) a
v Int
i a
x
        IORef Int -> Int -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Int
iRef (Int -> IO ()) -> Int -> IO ()
forall a b. (a -> b) -> a -> b
$! Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        VectorFillInfo v a -> IO (VectorFillInfo v a)
forall (m :: * -> *) a. Monad m => a -> m a
return VectorFillInfo v a
vfi

    grow :: Int -> IO (VectorFillInfo v a)
grow Int
sz = do
        let !sz' :: Int
sz' = Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2
        v RealWorld a
v' <- v (PrimState IO) a -> Int -> IO (v (PrimState IO) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (v (PrimState m) a)
VM.unsafeGrow v RealWorld a
v (PrimState IO) a
v Int
sz
        IORef Int -> Int -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Int
szRef Int
sz'
        VectorFillInfo v a -> a -> IO (VectorFillInfo v a)
forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> a -> IO (VectorFillInfo v a)
vfAppend (VectorFillInfo v a
vfi { _vec :: v (PrimState IO) a
_vec = v RealWorld a
v (PrimState IO) a
v' }) a
x


------------------------------------------------------------------------------
-- | 'mutableVectorOutputStream' returns an 'OutputStream' which stores values
-- fed into it and an action which flushes all stored values to a vector.
--
-- The flush action resets the store.
--
-- Note that this function /will/ buffer any input sent to it on the heap.
-- Please don't use this unless you're sure that the amount of input provided
-- is bounded and will fit in memory without issues.
mutableVectorOutputStream :: VM.MVector v c =>
                             IO (OutputStream c, IO (v (PrimState IO) c))
mutableVectorOutputStream :: IO (OutputStream c, IO (v (PrimState IO) c))
mutableVectorOutputStream = Int -> IO (OutputStream c, IO (v (PrimState IO) c))
forall (v :: * -> * -> *) c.
MVector v c =>
Int -> IO (OutputStream c, IO (v (PrimState IO) c))
mutableVectorOutputStreamSized Int
dEFAULT_BUFSIZ


------------------------------------------------------------------------------
-- | Like 'mutableVectorOutputStream', but allows control over how large the
-- vector buffer is to start with.
mutableVectorOutputStreamSized :: VM.MVector v c =>
                                  Int
                               -> IO (OutputStream c, IO (v (PrimState IO) c))
mutableVectorOutputStreamSized :: Int -> IO (OutputStream c, IO (v (PrimState IO) c))
mutableVectorOutputStreamSized Int
initialSize = do
    MVar (VectorFillInfo v c)
r <- Int -> IO (VectorFillInfo v c)
forall (v :: * -> * -> *) a.
MVector v a =>
Int -> IO (VectorFillInfo v a)
vfNew Int
initialSize IO (VectorFillInfo v c)
-> (VectorFillInfo v c -> IO (MVar (VectorFillInfo v c)))
-> IO (MVar (VectorFillInfo v c))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= VectorFillInfo v c -> IO (MVar (VectorFillInfo v c))
forall a. a -> IO (MVar a)
newMVar
    OutputStream c
c <- Consumer c () -> IO (OutputStream c)
forall r a. Consumer r a -> IO (OutputStream r)
S.fromConsumer (Consumer c () -> IO (OutputStream c))
-> Consumer c () -> IO (OutputStream c)
forall a b. (a -> b) -> a -> b
$ MVar (VectorFillInfo v c) -> Consumer c ()
forall (v :: * -> * -> *) a.
MVector v a =>
MVar (VectorFillInfo v a) -> Consumer a ()
consumer MVar (VectorFillInfo v c)
r
    (OutputStream c, IO (v RealWorld c))
-> IO (OutputStream c, IO (v RealWorld c))
forall (m :: * -> *) a. Monad m => a -> m a
return (OutputStream c
c, MVar (VectorFillInfo v c) -> IO (v RealWorld c)
forall (v :: * -> * -> *) a.
MVector v a =>
MVar (VectorFillInfo v a) -> IO (v RealWorld a)
flush MVar (VectorFillInfo v c)
r)

  where
    consumer :: MVar (VectorFillInfo v a) -> Consumer a ()
consumer MVar (VectorFillInfo v a)
r = Consumer a ()
go
      where
        go :: Consumer a ()
go = Consumer a (Maybe a)
forall r. Consumer r (Maybe r)
S.await Consumer a (Maybe a) -> (Maybe a -> Consumer a ()) -> Consumer a ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
             (Consumer a () -> (a -> Consumer a ()) -> Maybe a -> Consumer a ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> Consumer a ()
forall (m :: * -> *) a. Monad m => a -> m a
return (() -> Consumer a ()) -> () -> Consumer a ()
forall a b. (a -> b) -> a -> b
$! ()) ((a -> Consumer a ()) -> Maybe a -> Consumer a ())
-> (a -> Consumer a ()) -> Maybe a -> Consumer a ()
forall a b. (a -> b) -> a -> b
$ \a
c -> do
                 IO () -> Consumer a ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> Consumer a ()) -> IO () -> Consumer a ()
forall a b. (a -> b) -> a -> b
$ MVar (VectorFillInfo v a)
-> (VectorFillInfo v a -> IO (VectorFillInfo v a)) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar (VectorFillInfo v a)
r ((VectorFillInfo v a -> IO (VectorFillInfo v a)) -> IO ())
-> (VectorFillInfo v a -> IO (VectorFillInfo v a)) -> IO ()
forall a b. (a -> b) -> a -> b
$ (VectorFillInfo v a -> a -> IO (VectorFillInfo v a))
-> a -> VectorFillInfo v a -> IO (VectorFillInfo v a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip VectorFillInfo v a -> a -> IO (VectorFillInfo v a)
forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> a -> IO (VectorFillInfo v a)
vfAppend a
c
                 Consumer a ()
go)

    flush :: MVar (VectorFillInfo v a) -> IO (v RealWorld a)
flush MVar (VectorFillInfo v a)
r = MVar (VectorFillInfo v a)
-> (VectorFillInfo v a -> IO (VectorFillInfo v a, v RealWorld a))
-> IO (v RealWorld a)
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar (VectorFillInfo v a)
r ((VectorFillInfo v a -> IO (VectorFillInfo v a, v RealWorld a))
 -> IO (v RealWorld a))
-> (VectorFillInfo v a -> IO (VectorFillInfo v a, v RealWorld a))
-> IO (v RealWorld a)
forall a b. (a -> b) -> a -> b
$ \VectorFillInfo v a
vfi -> do
                                !v RealWorld a
v   <- VectorFillInfo v a -> IO (v (PrimState IO) a)
forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> IO (v (PrimState IO) a)
vfFinish VectorFillInfo v a
vfi
                                VectorFillInfo v a
vfi' <- Int -> IO (VectorFillInfo v a)
forall (v :: * -> * -> *) a.
MVector v a =>
Int -> IO (VectorFillInfo v a)
vfNew Int
initialSize
                                (VectorFillInfo v a, v RealWorld a)
-> IO (VectorFillInfo v a, v RealWorld a)
forall (m :: * -> *) a. Monad m => a -> m a
return ((VectorFillInfo v a, v RealWorld a)
 -> IO (VectorFillInfo v a, v RealWorld a))
-> (VectorFillInfo v a, v RealWorld a)
-> IO (VectorFillInfo v a, v RealWorld a)
forall a b. (a -> b) -> a -> b
$! (VectorFillInfo v a
vfi', v RealWorld a
v)
{-# INLINE mutableVectorOutputStreamSized #-}


------------------------------------------------------------------------------
-- | Given an IO action that requires an 'OutputStream', creates one and
-- captures all the output the action sends to it as a mutable vector.
--
-- Example:
--
-- @
-- ghci> import "Control.Applicative"
-- ghci> ('connect' \<\$\> 'System.IO.Streams.fromList' [1, 2, 3::'Int'])
--        \>\>= 'outputToMutableVector'
--        \>\>= V.'Data.Vector.freeze'
-- fromList [1,2,3]
-- @
outputToMutableVector :: MVector v a =>
                         (OutputStream a -> IO b)
                      -> IO (v (PrimState IO) a)
outputToMutableVector :: (OutputStream a -> IO b) -> IO (v (PrimState IO) a)
outputToMutableVector = Int -> (OutputStream a -> IO b) -> IO (v (PrimState IO) a)
forall (v :: * -> * -> *) a b.
MVector v a =>
Int -> (OutputStream a -> IO b) -> IO (v (PrimState IO) a)
outputToMutableVectorSized Int
dEFAULT_BUFSIZ
{-# INLINE outputToMutableVector #-}


------------------------------------------------------------------------------
-- | Like 'outputToMutableVector', but allows control over how large the vector
-- buffer is to start with.
outputToMutableVectorSized :: MVector v a =>
                              Int
                           -> (OutputStream a -> IO b)
                           -> IO (v (PrimState IO) a)
outputToMutableVectorSized :: Int -> (OutputStream a -> IO b) -> IO (v (PrimState IO) a)
outputToMutableVectorSized Int
n OutputStream a -> IO b
f = do
    (OutputStream a
os, IO (v RealWorld a)
getVec) <- Int -> IO (OutputStream a, IO (v (PrimState IO) a))
forall (v :: * -> * -> *) c.
MVector v c =>
Int -> IO (OutputStream c, IO (v (PrimState IO) c))
mutableVectorOutputStreamSized Int
n
    b
_ <- OutputStream a -> IO b
f OutputStream a
os
    IO (v RealWorld a)
getVec
{-# INLINE outputToMutableVectorSized #-}


------------------------------------------------------------------------------
-- | Given an IO action that requires an 'OutputStream', creates one and
-- captures all the output the action sends to it as a vector.
--
-- Example:
--
-- @
-- ghci> (('connect' <$> 'System.IO.Streams.fromList' [1, 2, 3]) >>= 'outputToVector')
--           :: IO ('Data.Vector.Vector' Int)
-- fromList [1,2,3]
-- @
outputToVector :: Vector v a => (OutputStream a -> IO b) -> IO (v a)
outputToVector :: (OutputStream a -> IO b) -> IO (v a)
outputToVector = Int -> (OutputStream a -> IO b) -> IO (v a)
forall (v :: * -> *) a b.
Vector v a =>
Int -> (OutputStream a -> IO b) -> IO (v a)
outputToVectorSized Int
dEFAULT_BUFSIZ
{-# INLINE outputToVector #-}


------------------------------------------------------------------------------
-- | Like 'outputToVector', but allows control over how large the vector buffer
-- is to start with.
outputToVectorSized :: Vector v a =>
                       Int
                    -> (OutputStream a -> IO b)
                    -> IO (v a)
outputToVectorSized :: Int -> (OutputStream a -> IO b) -> IO (v a)
outputToVectorSized Int
n = Int -> (OutputStream a -> IO b) -> IO (Mutable v (PrimState IO) a)
forall (v :: * -> * -> *) a b.
MVector v a =>
Int -> (OutputStream a -> IO b) -> IO (v (PrimState IO) a)
outputToMutableVectorSized Int
n ((OutputStream a -> IO b) -> IO (Mutable v RealWorld a))
-> (Mutable v RealWorld a -> IO (v a))
-> (OutputStream a -> IO b)
-> IO (v a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Mutable v RealWorld a -> IO (v a)
forall (v :: * -> *) a (m :: * -> *).
(Vector v a, PrimMonad m) =>
Mutable v (PrimState m) a -> m (v a)
V.basicUnsafeFreeze
{-# INLINE outputToVectorSized #-}


------------------------------------------------------------------------------
-- | Splits an input stream into chunks of at most size @n@.
--
-- Example:
--
-- @
-- ghci> ('System.IO.Streams.fromList' [1..14::Int] >>= 'chunkVector' 4 >>= 'System.IO.Streams.toList')
--          :: IO ['Data.Vector.Vector' Int]
-- [fromList [1,2,3,4],fromList [5,6,7,8],fromList [9,10,11,12],fromList [13,14]]
-- @
chunkVector :: Vector v a => Int -> InputStream a -> IO (InputStream (v a))
chunkVector :: Int -> InputStream a -> IO (InputStream (v a))
chunkVector Int
n InputStream a
input = if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
                        then [Char] -> IO (InputStream (v a))
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO (InputStream (v a)))
-> [Char] -> IO (InputStream (v a))
forall a b. (a -> b) -> a -> b
$ [Char]
"chunkVector: bad size: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n
                        else Int -> IO (VectorFillInfo (Mutable v) a)
forall (v :: * -> * -> *) a.
MVector v a =>
Int -> IO (VectorFillInfo v a)
vfNew Int
n IO (VectorFillInfo (Mutable v) a)
-> (VectorFillInfo (Mutable v) a -> IO (InputStream (v a)))
-> IO (InputStream (v a))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Generator (v a) () -> IO (InputStream (v a))
forall r a. Generator r a -> IO (InputStream r)
fromGenerator (Generator (v a) () -> IO (InputStream (v a)))
-> (VectorFillInfo (Mutable v) a -> Generator (v a) ())
-> VectorFillInfo (Mutable v) a
-> IO (InputStream (v a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> VectorFillInfo (Mutable v) a -> Generator (v a) ()
forall (v :: * -> *).
Vector v a =>
Int -> VectorFillInfo (Mutable v) a -> Generator (v a) ()
go Int
n
  where
    doneChunk :: VectorFillInfo (Mutable v) a -> Generator (v a) ()
doneChunk !VectorFillInfo (Mutable v) a
vfi = do
        IO (v a) -> Generator (v a) (v a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (VectorFillInfo (Mutable v) a -> IO (Mutable v (PrimState IO) a)
forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> IO (v (PrimState IO) a)
vfFinish VectorFillInfo (Mutable v) a
vfi IO (Mutable v RealWorld a)
-> (Mutable v RealWorld a -> IO (v a)) -> IO (v a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Mutable v RealWorld a -> IO (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
V.unsafeFreeze) Generator (v a) (v a)
-> (v a -> Generator (v a) ()) -> Generator (v a) ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= v a -> Generator (v a) ()
forall r. r -> Generator r ()
yield
        !VectorFillInfo (Mutable v) a
vfi' <- IO (VectorFillInfo (Mutable v) a)
-> Generator (v a) (VectorFillInfo (Mutable v) a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (VectorFillInfo (Mutable v) a)
 -> Generator (v a) (VectorFillInfo (Mutable v) a))
-> IO (VectorFillInfo (Mutable v) a)
-> Generator (v a) (VectorFillInfo (Mutable v) a)
forall a b. (a -> b) -> a -> b
$ Int -> IO (VectorFillInfo (Mutable v) a)
forall (v :: * -> * -> *) a.
MVector v a =>
Int -> IO (VectorFillInfo v a)
vfNew Int
n
        Int -> VectorFillInfo (Mutable v) a -> Generator (v a) ()
go Int
n VectorFillInfo (Mutable v) a
vfi'

    go :: Int -> VectorFillInfo (Mutable v) a -> Generator (v a) ()
go !Int
k !VectorFillInfo (Mutable v) a
vfi | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0    = VectorFillInfo (Mutable v) a -> Generator (v a) ()
doneChunk VectorFillInfo (Mutable v) a
vfi
               | Bool
otherwise = IO (Maybe a) -> Generator (v a) (Maybe a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (InputStream a -> IO (Maybe a)
forall a. InputStream a -> IO (Maybe a)
S.read InputStream a
input) Generator (v a) (Maybe a)
-> (Maybe a -> Generator (v a) ()) -> Generator (v a) ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Generator (v a) ()
-> (a -> Generator (v a) ()) -> Maybe a -> Generator (v a) ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Generator (v a) ()
finish a -> Generator (v a) ()
chunk
      where
        finish :: Generator (v a) ()
finish = do
            v a
v <- IO (v a) -> Generator (v a) (v a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (VectorFillInfo (Mutable v) a -> IO (Mutable v (PrimState IO) a)
forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> IO (v (PrimState IO) a)
vfFinish VectorFillInfo (Mutable v) a
vfi IO (Mutable v RealWorld a)
-> (Mutable v RealWorld a -> IO (v a)) -> IO (v a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Mutable v RealWorld a -> IO (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
V.unsafeFreeze)
            if v a -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
V.null v a
v then () -> Generator (v a) ()
forall (m :: * -> *) a. Monad m => a -> m a
return (() -> Generator (v a) ()) -> () -> Generator (v a) ()
forall a b. (a -> b) -> a -> b
$! () else v a -> Generator (v a) ()
forall r. r -> Generator r ()
yield v a
v

        chunk :: a -> Generator (v a) ()
chunk a
x = do
            !VectorFillInfo (Mutable v) a
vfi' <- IO (VectorFillInfo (Mutable v) a)
-> Generator (v a) (VectorFillInfo (Mutable v) a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (VectorFillInfo (Mutable v) a)
 -> Generator (v a) (VectorFillInfo (Mutable v) a))
-> IO (VectorFillInfo (Mutable v) a)
-> Generator (v a) (VectorFillInfo (Mutable v) a)
forall a b. (a -> b) -> a -> b
$ VectorFillInfo (Mutable v) a
-> a -> IO (VectorFillInfo (Mutable v) a)
forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> a -> IO (VectorFillInfo v a)
vfAppend VectorFillInfo (Mutable v) a
vfi a
x
            Int -> VectorFillInfo (Mutable v) a -> Generator (v a) ()
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) VectorFillInfo (Mutable v) a
vfi'
{-# INLINE chunkVector #-}


------------------------------------------------------------------------------
-- | Feeds a vector to an 'OutputStream'. Does /not/ write an end-of-stream to
-- the stream.
--
-- @
-- ghci> let v = V.'fromList' [1..4] :: V.'Vector' Int
-- ghci> os \<- Streams.'unlines' Streams.'stdout' >>= Streams.'System.IO.Streams.contramap' (S.pack . show) :: IO ('OutputStream' Int)
-- ghci> Streams.'writeVector' v os
-- 1
-- 2
-- 3
-- 4
-- @
writeVector :: Vector v a => v a -> OutputStream a -> IO ()
writeVector :: v a -> OutputStream a -> IO ()
writeVector v a
v OutputStream a
out = (a -> IO ()) -> v a -> IO ()
forall (m :: * -> *) (v :: * -> *) a b.
(Monad m, Vector v a) =>
(a -> m b) -> v a -> m ()
V.mapM_ ((Maybe a -> OutputStream a -> IO ())
-> OutputStream a -> Maybe a -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip Maybe a -> OutputStream a -> IO ()
forall a. Maybe a -> OutputStream a -> IO ()
S.write OutputStream a
out (Maybe a -> IO ()) -> (a -> Maybe a) -> a -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Maybe a
forall a. a -> Maybe a
Just) v a
v
{-# INLINE writeVector #-}


------------------------------------------------------------------------------
dEFAULT_BUFSIZ :: Int
dEFAULT_BUFSIZ :: Int
dEFAULT_BUFSIZ = Int
64