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 n s0 = do
    mv <- liftIO $ VM.new n
    go mv 0 s0
  where
    go !mv !i s
        | i == n    = liftIO $ VG.unsafeFreeze mv
        | otherwise =
            inspect s >>= \case
                Left        ()  -> liftIO $ VG.unsafeFreeze $ VM.take i mv
                Right (a :> s') -> liftIO (VM.write mv i a) >> go mv (i+1) 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 s0 = do
    mv <- liftIO $ VM.new 1024
    go mv 0 s0
  where
    go !mv !i =
        inspect >=> \case
            Left        r  -> liftM (:> r) $ liftIO $ VG.unsafeFreeze $ VM.take i mv
            Right (a :> s) -> do mv' <- if VM.length mv == i then liftIO (VM.grow mv (VM.length mv)) else return mv
                                 liftIO $ VM.write mv' i a
                                 go mv' (i+1) s