module Bio.Streaming.Vector ( stream2vector, stream2vectorN ) where

import Bio.Prelude
import Streaming

import qualified Data.Vector.Generic            as VG
import qualified Data.Vector.Generic.Mutable    as VM

-- | Equivalent to @stream2vector . Streaming.Prelude.take n@, but
-- terminates early and is thereby more efficient.
stream2vectorN :: (MonadIO m, VG.Vector v a) => Int -> Stream (Of a) m () -> m (v a)
stream2vectorN :: Int -> Stream (Of a) m () -> m (v a)
stream2vectorN n :: Int
n s0 :: Stream (Of a) m ()
s0 = do
    Mutable v RealWorld a
mv <- IO (Mutable v RealWorld a) -> m (Mutable v RealWorld a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Mutable v RealWorld a) -> m (Mutable v RealWorld a))
-> IO (Mutable v RealWorld a) -> m (Mutable v RealWorld a)
forall a b. (a -> b) -> a -> b
$ Int -> IO (Mutable v (PrimState IO) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
VM.new Int
n
    Mutable v RealWorld a -> Int -> Stream (Of a) m () -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(MonadIO m, Vector v a) =>
Mutable v RealWorld a -> Int -> Stream (Of a) m () -> m (v a)
go Mutable v RealWorld a
mv 0 Stream (Of a) m ()
s0
  where
    go :: Mutable v RealWorld a -> Int -> Stream (Of a) m () -> m (v a)
go !Mutable v RealWorld a
mv !Int
i s :: Stream (Of a) m ()
s
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n    = IO (v a) -> m (v a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (v a) -> m (v a)) -> IO (v a) -> m (v a)
forall a b. (a -> b) -> a -> b
$ Mutable v (PrimState IO) a -> IO (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
VG.unsafeFreeze Mutable v RealWorld a
Mutable v (PrimState IO) a
mv
        | Bool
otherwise =
            Stream (Of a) m () -> m (Either () (Of a (Stream (Of a) m ())))
forall (m :: * -> *) (f :: * -> *) r.
Monad m =>
Stream f m r -> m (Either r (f (Stream f m r)))
inspect Stream (Of a) m ()
s m (Either () (Of a (Stream (Of a) m ())))
-> (Either () (Of a (Stream (Of a) m ())) -> m (v a)) -> m (v a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                Left        ()  -> IO (v a) -> m (v a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (v a) -> m (v a)) -> IO (v a) -> m (v a)
forall a b. (a -> b) -> a -> b
$ Mutable v (PrimState IO) a -> IO (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
VG.unsafeFreeze (Mutable v (PrimState IO) a -> IO (v a))
-> Mutable v (PrimState IO) a -> IO (v a)
forall a b. (a -> b) -> a -> b
$ Int -> Mutable v RealWorld a -> Mutable v RealWorld a
forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
VM.take Int
i Mutable v RealWorld a
mv
                Right (a :: a
a :> s' :: Stream (Of a) m ()
s') -> IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Mutable v (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VM.write Mutable v RealWorld a
Mutable v (PrimState IO) a
mv Int
i a
a) m () -> m (v a) -> m (v a)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Mutable v RealWorld a -> Int -> Stream (Of a) m () -> m (v a)
go Mutable v RealWorld a
mv (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Stream (Of a) m ()
s'

-- | Reads the whole stream into a 'VG.Vector'.
stream2vector :: (MonadIO m, VG.Vector v a) => Stream (Of a) m r -> m (Of (v a) r)
stream2vector :: Stream (Of a) m r -> m (Of (v a) r)
stream2vector s0 :: Stream (Of a) m r
s0 = do
    Mutable v RealWorld a
mv <- IO (Mutable v RealWorld a) -> m (Mutable v RealWorld a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Mutable v RealWorld a) -> m (Mutable v RealWorld a))
-> IO (Mutable v RealWorld a) -> m (Mutable v RealWorld a)
forall a b. (a -> b) -> a -> b
$ Int -> IO (Mutable v (PrimState IO) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
VM.new 1024
    Mutable v RealWorld a -> Int -> Stream (Of a) m r -> m (Of (v a) r)
forall (m :: * -> *) (v :: * -> *) a b.
(MonadIO m, Vector v a) =>
Mutable v RealWorld a -> Int -> Stream (Of a) m b -> m (Of (v a) b)
go Mutable v RealWorld a
mv 0 Stream (Of a) m r
s0
  where
    go :: Mutable v RealWorld a -> Int -> Stream (Of a) m b -> m (Of (v a) b)
go !Mutable v RealWorld a
mv !Int
i =
        Stream (Of a) m b -> m (Either b (Of a (Stream (Of a) m b)))
forall (m :: * -> *) (f :: * -> *) r.
Monad m =>
Stream f m r -> m (Either r (f (Stream f m r)))
inspect (Stream (Of a) m b -> m (Either b (Of a (Stream (Of a) m b))))
-> (Either b (Of a (Stream (Of a) m b)) -> m (Of (v a) b))
-> Stream (Of a) m b
-> m (Of (v a) b)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> \case
            Left        r :: b
r  -> (v a -> Of (v a) b) -> m (v a) -> m (Of (v a) b)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM (v a -> b -> Of (v a) b
forall a b. a -> b -> Of a b
:> b
r) (m (v a) -> m (Of (v a) b)) -> m (v a) -> m (Of (v a) b)
forall a b. (a -> b) -> a -> b
$ IO (v a) -> m (v a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (v a) -> m (v a)) -> IO (v a) -> m (v a)
forall a b. (a -> b) -> a -> b
$ Mutable v (PrimState IO) a -> IO (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
VG.unsafeFreeze (Mutable v (PrimState IO) a -> IO (v a))
-> Mutable v (PrimState IO) a -> IO (v a)
forall a b. (a -> b) -> a -> b
$ Int -> Mutable v RealWorld a -> Mutable v RealWorld a
forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
VM.take Int
i Mutable v RealWorld a
mv
            Right (a :: a
a :> s :: Stream (Of a) m b
s) -> do Mutable v RealWorld a
mv' <- if Mutable v RealWorld a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
VM.length Mutable v RealWorld a
mv Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i then IO (Mutable v RealWorld a) -> m (Mutable v RealWorld a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Mutable v (PrimState IO) a
-> Int -> IO (Mutable v (PrimState IO) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (v (PrimState m) a)
VM.grow Mutable v RealWorld a
Mutable v (PrimState IO) a
mv (Mutable v RealWorld a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
VM.length Mutable v RealWorld a
mv)) else Mutable v RealWorld a -> m (Mutable v RealWorld a)
forall (m :: * -> *) a. Monad m => a -> m a
return Mutable v RealWorld a
mv
                                 IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Mutable v (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VM.write Mutable v RealWorld a
Mutable v (PrimState IO) a
mv' Int
i a
a
                                 Mutable v RealWorld a -> Int -> Stream (Of a) m b -> m (Of (v a) b)
go Mutable v RealWorld a
mv' (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Stream (Of a) m b
s