{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE CPP #-}
module System.IO.Streams.Vector
(
fromVector
, toVector
, toVectorSized
, outputToVector
, outputToVectorSized
, toMutableVector
, toMutableVectorSized
, outputToMutableVector
, outputToMutableVectorSized
, writeVector
, 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 (..), RealWorld)
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
#if MIN_VERSION_vector(0,13,0)
import Control.Monad.ST (stToIO)
#endif
basicUnsafeFreezeCompat :: Vector v a => V.Mutable v RealWorld a -> IO (v a)
#if MIN_VERSION_vector(0,13,0)
basicUnsafeFreezeCompat = stToIO . V.basicUnsafeFreeze
#else
basicUnsafeFreezeCompat :: forall (v :: * -> *) a.
Vector v a =>
Mutable v RealWorld a -> IO (v a)
basicUnsafeFreezeCompat = forall (v :: * -> *) a (m :: * -> *).
(Vector v a, PrimMonad m) =>
Mutable v (PrimState m) a -> m (v a)
V.basicUnsafeFreeze
#endif
fromVector :: Vector v a => v a -> IO (InputStream a)
fromVector :: forall (v :: * -> *) a. Vector v a => v a -> IO (InputStream a)
fromVector = forall r a. Generator r a -> IO (InputStream r)
fromGenerator forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) (v :: * -> *) a b.
(Monad m, Vector v a) =>
(a -> m b) -> v a -> m ()
V.mapM_ forall r. r -> Generator r ()
yield
{-# INLINE fromVector #-}
toVector :: Vector v a => InputStream a -> IO (v a)
toVector :: forall (v :: * -> *) a. Vector v a => InputStream a -> IO (v a)
toVector = forall (v :: * -> *) a.
Vector v a =>
Int -> InputStream a -> IO (v a)
toVectorSized Int
dEFAULT_BUFSIZ
{-# INLINE toVector #-}
toVectorSized :: Vector v a => Int -> InputStream a -> IO (v a)
toVectorSized :: forall (v :: * -> *) a.
Vector v a =>
Int -> InputStream a -> IO (v a)
toVectorSized Int
n = forall (v :: * -> * -> *) a.
MVector v a =>
Int -> InputStream a -> IO (v (PrimState IO) a)
toMutableVectorSized Int
n forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> forall (v :: * -> *) a.
Vector v a =>
Mutable v RealWorld a -> IO (v a)
basicUnsafeFreezeCompat
{-# INLINE toVectorSized #-}
toMutableVector :: VM.MVector v a => InputStream a -> IO (v (PrimState IO) a)
toMutableVector :: forall (v :: * -> * -> *) a.
MVector v a =>
InputStream a -> IO (v (PrimState IO) a)
toMutableVector = forall (v :: * -> * -> *) a.
MVector v a =>
Int -> InputStream a -> IO (v (PrimState IO) a)
toMutableVectorSized Int
dEFAULT_BUFSIZ
toMutableVectorSized :: VM.MVector v a =>
Int
-> InputStream a
-> IO (v (PrimState IO) a)
toMutableVectorSized :: forall (v :: * -> * -> *) a.
MVector v a =>
Int -> InputStream a -> IO (v (PrimState IO) a)
toMutableVectorSized Int
initialSize InputStream a
input = forall (v :: * -> * -> *) a.
MVector v a =>
Int -> IO (VectorFillInfo v a)
vfNew Int
initialSize forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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 = forall a. InputStream a -> IO (Maybe a)
S.read InputStream a
input forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> IO (v (PrimState IO) a)
vfFinish VectorFillInfo v a
vfi) (forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> a -> IO (VectorFillInfo v a)
vfAppend VectorFillInfo v a
vfi 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 :: Vector v c => IO (OutputStream c, IO (v c))
vectorOutputStream :: forall (v :: * -> *) c. Vector v c => IO (OutputStream c, IO (v c))
vectorOutputStream = forall (v :: * -> *) c.
Vector v c =>
Int -> IO (OutputStream c, IO (v c))
vectorOutputStreamSized Int
dEFAULT_BUFSIZ
{-# INLINE vectorOutputStream #-}
vectorOutputStreamSized :: Vector v c => Int -> IO (OutputStream c, IO (v c))
vectorOutputStreamSized :: forall (v :: * -> *) c.
Vector v c =>
Int -> IO (OutputStream c, IO (v c))
vectorOutputStreamSized Int
n = do
(OutputStream c
os, IO (Mutable v RealWorld c)
flush) <- forall (v :: * -> * -> *) c.
MVector v c =>
Int -> IO (OutputStream c, IO (v (PrimState IO) c))
mutableVectorOutputStreamSized Int
n
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! (OutputStream c
os, IO (Mutable v RealWorld c)
flush forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (v :: * -> *) a.
Vector v a =>
Mutable v RealWorld a -> IO (v a)
basicUnsafeFreezeCompat)
data VectorFillInfo v c = VectorFillInfo {
forall (v :: * -> * -> *) c.
VectorFillInfo v c -> v (PrimState IO) c
_vec :: !(v (PrimState IO) c)
, forall (v :: * -> * -> *) c. VectorFillInfo v c -> IORef Int
_idx :: {-# UNPACK #-} !(IORef Int)
, forall (v :: * -> * -> *) c. VectorFillInfo v c -> IORef Int
_sz :: {-# UNPACK #-} !(IORef Int)
}
vfNew :: MVector v a => Int -> IO (VectorFillInfo v a)
vfNew :: forall (v :: * -> * -> *) a.
MVector v a =>
Int -> IO (VectorFillInfo v a)
vfNew Int
initialSize = do
v RealWorld a
v <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
VM.unsafeNew Int
initialSize
IORef Int
i <- forall a. a -> IO (IORef a)
newIORef Int
0
IORef Int
sz <- forall a. a -> IO (IORef a)
newIORef Int
initialSize
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall (v :: * -> * -> *) c.
v (PrimState IO) c -> IORef Int -> IORef Int -> VectorFillInfo v c
VectorFillInfo v RealWorld a
v IORef Int
i IORef Int
sz
vfFinish :: MVector v a =>
VectorFillInfo v a
-> IO (v (PrimState IO) a)
vfFinish :: forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> IO (v (PrimState IO) a)
vfFinish VectorFillInfo v a
vfi = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
VM.unsafeTake v RealWorld a
v) forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef Int
i
where
v :: v (PrimState IO) a
v = forall (v :: * -> * -> *) c.
VectorFillInfo v c -> v (PrimState IO) c
_vec VectorFillInfo v a
vfi
i :: IORef Int
i = 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 :: forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> a -> IO (VectorFillInfo v a)
vfAppend VectorFillInfo v a
vfi !a
x = do
Int
i <- forall a. IORef a -> IO a
readIORef IORef Int
iRef
Int
sz <- forall a. IORef a -> IO a
readIORef IORef Int
szRef
if Int
i 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 = forall (v :: * -> * -> *) c.
VectorFillInfo v c -> v (PrimState IO) c
_vec VectorFillInfo v a
vfi
iRef :: IORef Int
iRef = forall (v :: * -> * -> *) c. VectorFillInfo v c -> IORef Int
_idx VectorFillInfo v a
vfi
szRef :: IORef Int
szRef = forall (v :: * -> * -> *) c. VectorFillInfo v c -> IORef Int
_sz VectorFillInfo v a
vfi
add :: Int -> IO (VectorFillInfo v a)
add Int
i = do
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VM.unsafeWrite v RealWorld a
v Int
i a
x
forall a. IORef a -> a -> IO ()
writeIORef IORef Int
iRef forall a b. (a -> b) -> a -> b
$! Int
i forall a. Num a => a -> a -> a
+ Int
1
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 forall a. Num a => a -> a -> a
* Int
2
v RealWorld a
v' <- 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 Int
sz
forall a. IORef a -> a -> IO ()
writeIORef IORef Int
szRef Int
sz'
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' }) a
x
mutableVectorOutputStream :: VM.MVector v c =>
IO (OutputStream c, IO (v (PrimState IO) c))
mutableVectorOutputStream :: forall (v :: * -> * -> *) c.
MVector v c =>
IO (OutputStream c, IO (v (PrimState IO) c))
mutableVectorOutputStream = forall (v :: * -> * -> *) c.
MVector v c =>
Int -> IO (OutputStream c, IO (v (PrimState IO) c))
mutableVectorOutputStreamSized Int
dEFAULT_BUFSIZ
mutableVectorOutputStreamSized :: VM.MVector v c =>
Int
-> IO (OutputStream c, IO (v (PrimState IO) c))
mutableVectorOutputStreamSized :: forall (v :: * -> * -> *) c.
MVector v c =>
Int -> IO (OutputStream c, IO (v (PrimState IO) c))
mutableVectorOutputStreamSized Int
initialSize = do
MVar (VectorFillInfo v c)
r <- forall (v :: * -> * -> *) a.
MVector v a =>
Int -> IO (VectorFillInfo v a)
vfNew Int
initialSize forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a. a -> IO (MVar a)
newMVar
OutputStream c
c <- forall r a. Consumer r a -> IO (OutputStream r)
S.fromConsumer forall a b. (a -> b) -> a -> b
$ forall {v :: * -> * -> *} {b}.
MVector v b =>
MVar (VectorFillInfo v b) -> Consumer b ()
consumer MVar (VectorFillInfo v c)
r
forall (m :: * -> *) a. Monad m => a -> m a
return (OutputStream c
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 b) -> Consumer b ()
consumer MVar (VectorFillInfo v b)
r = Consumer b ()
go
where
go :: Consumer b ()
go = forall r. Consumer r (Maybe r)
S.await forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
(forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! ()) forall a b. (a -> b) -> a -> b
$ \b
c -> do
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar (VectorFillInfo v b)
r forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> a -> IO (VectorFillInfo v a)
vfAppend b
c
Consumer b ()
go)
flush :: MVar (VectorFillInfo v a) -> IO (v RealWorld a)
flush MVar (VectorFillInfo v a)
r = forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar (VectorFillInfo v a)
r forall a b. (a -> b) -> a -> b
$ \VectorFillInfo v a
vfi -> do
!v RealWorld a
v <- forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> IO (v (PrimState IO) a)
vfFinish VectorFillInfo v a
vfi
VectorFillInfo v a
vfi' <- forall (v :: * -> * -> *) a.
MVector v a =>
Int -> IO (VectorFillInfo v a)
vfNew Int
initialSize
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! (VectorFillInfo v a
vfi', v RealWorld a
v)
{-# INLINE mutableVectorOutputStreamSized #-}
outputToMutableVector :: MVector v a =>
(OutputStream a -> IO b)
-> IO (v (PrimState IO) a)
outputToMutableVector :: forall (v :: * -> * -> *) a b.
MVector v a =>
(OutputStream a -> IO b) -> IO (v (PrimState IO) a)
outputToMutableVector = forall (v :: * -> * -> *) a b.
MVector v a =>
Int -> (OutputStream a -> IO b) -> IO (v (PrimState IO) a)
outputToMutableVectorSized Int
dEFAULT_BUFSIZ
{-# INLINE outputToMutableVector #-}
outputToMutableVectorSized :: MVector v a =>
Int
-> (OutputStream a -> IO b)
-> IO (v (PrimState IO) a)
outputToMutableVectorSized :: forall (v :: * -> * -> *) a b.
MVector v a =>
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) <- 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 #-}
outputToVector :: Vector v a => (OutputStream a -> IO b) -> IO (v a)
outputToVector :: forall (v :: * -> *) a b.
Vector v a =>
(OutputStream a -> IO b) -> IO (v a)
outputToVector = forall (v :: * -> *) a b.
Vector v a =>
Int -> (OutputStream a -> IO b) -> IO (v a)
outputToVectorSized Int
dEFAULT_BUFSIZ
{-# INLINE outputToVector #-}
outputToVectorSized :: Vector v a =>
Int
-> (OutputStream a -> IO b)
-> IO (v a)
outputToVectorSized :: forall (v :: * -> *) a b.
Vector v a =>
Int -> (OutputStream a -> IO b) -> IO (v a)
outputToVectorSized Int
n = forall (v :: * -> * -> *) a b.
MVector v a =>
Int -> (OutputStream a -> IO b) -> IO (v (PrimState IO) a)
outputToMutableVectorSized Int
n forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> forall (v :: * -> *) a.
Vector v a =>
Mutable v RealWorld a -> IO (v a)
basicUnsafeFreezeCompat
{-# INLINE outputToVectorSized #-}
chunkVector :: Vector v a => Int -> InputStream a -> IO (InputStream (v a))
chunkVector :: forall (v :: * -> *) a.
Vector v a =>
Int -> InputStream a -> IO (InputStream (v a))
chunkVector Int
n InputStream a
input = if Int
n forall a. Ord a => a -> a -> Bool
<= Int
0
then forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"chunkVector: bad size: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
n
else forall (v :: * -> * -> *) a.
MVector v a =>
Int -> IO (VectorFillInfo v a)
vfNew Int
n forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall r a. Generator r a -> IO (InputStream r)
fromGenerator forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> IO (v (PrimState IO) a)
vfFinish VectorFillInfo (Mutable v) a
vfi forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
V.unsafeFreeze) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall r. r -> Generator r ()
yield
!VectorFillInfo (Mutable v) a
vfi' <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ 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 forall a. Ord a => a -> a -> Bool
<= Int
0 = VectorFillInfo (Mutable v) a -> Generator (v a) ()
doneChunk VectorFillInfo (Mutable v) a
vfi
| Bool
otherwise = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a. InputStream a -> IO (Maybe a)
S.read InputStream a
input) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall (v :: * -> * -> *) a.
MVector v a =>
VectorFillInfo v a -> IO (v (PrimState IO) a)
vfFinish VectorFillInfo (Mutable v) a
vfi forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
V.unsafeFreeze)
if forall (v :: * -> *) a. Vector v a => v a -> Bool
V.null v a
v then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! () else forall r. r -> Generator r ()
yield v a
v
chunk :: a -> Generator (v a) ()
chunk a
x = do
!VectorFillInfo (Mutable v) a
vfi' <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ 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 forall a. Num a => a -> a -> a
- Int
1) VectorFillInfo (Mutable v) a
vfi'
{-# INLINE chunkVector #-}
writeVector :: Vector v a => v a -> OutputStream a -> IO ()
writeVector :: forall (v :: * -> *) a.
Vector v a =>
v a -> OutputStream a -> IO ()
writeVector v a
v OutputStream a
out = forall (m :: * -> *) (v :: * -> *) a b.
(Monad m, Vector v a) =>
(a -> m b) -> v a -> m ()
V.mapM_ (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Maybe a -> OutputStream a -> IO ()
S.write OutputStream a
out forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Maybe a
Just) v a
v
{-# INLINE writeVector #-}
dEFAULT_BUFSIZ :: Int
dEFAULT_BUFSIZ :: Int
dEFAULT_BUFSIZ = Int
64