content: {-# LANGUAGE BangPatterns          #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE RankNTypes            #-}
  {-# LANGUAGE CPP                   #-}
  
  -- | 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 (..), 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 = V.basicUnsafeFreeze
  #endif
  
  ------------------------------------------------------------------------------
  -- | 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 = fromGenerator . V.mapM_ 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 = toVectorSized 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 n = toMutableVectorSized n >=> basicUnsafeFreezeCompat
  {-# 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 = toMutableVectorSized 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 initialSize input = vfNew initialSize >>= go
    where
      go vfi = S.read input >>= maybe (vfFinish vfi) (vfAppend vfi >=> 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 = vectorOutputStreamSized 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 n = do
      (os, flush) <- mutableVectorOutputStreamSized n
      return $! (os, flush >>= basicUnsafeFreezeCompat)
  
  
  ------------------------------------------------------------------------------
  data VectorFillInfo v c = VectorFillInfo {
        _vec :: !(v (PrimState IO) c)
      , _idx :: {-# UNPACK #-} !(IORef Int)
  
      -- TODO: vector contains its own size
      , _sz  :: {-# UNPACK #-} !(IORef Int)
      }
  
  
  ------------------------------------------------------------------------------
  vfNew :: MVector v a => Int -> IO (VectorFillInfo v a)
  vfNew initialSize = do
      v  <- VM.unsafeNew initialSize
      i  <- newIORef 0
      sz <- newIORef initialSize
      return $! VectorFillInfo v i sz
  
  
  ------------------------------------------------------------------------------
  vfFinish :: MVector v a =>
              VectorFillInfo v a
           -> IO (v (PrimState IO) a)
  vfFinish vfi = liftM (flip VM.unsafeTake v) $ readIORef i
    where
      v = _vec vfi
      i = _idx vfi
  
  
  ------------------------------------------------------------------------------
  vfAppend :: MVector v a =>
           VectorFillInfo v a
        -> a
        -> IO (VectorFillInfo v a)
  vfAppend vfi !x = do
      i  <- readIORef iRef
      sz <- readIORef szRef
      if i < sz then add i else grow sz
    where
      v     = _vec vfi
      iRef  = _idx vfi
      szRef = _sz vfi
  
      add i = do
          VM.unsafeWrite v i x
          writeIORef iRef $! i + 1
          return vfi
  
      grow sz = do
          let !sz' = sz * 2
          v' <- VM.unsafeGrow v sz
          writeIORef szRef sz'
          vfAppend (vfi { _vec = v' }) 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 = mutableVectorOutputStreamSized 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 initialSize = do
      r <- vfNew initialSize >>= newMVar
      c <- S.fromConsumer $ consumer r
      return (c, flush r)
  
    where
      consumer r = go
        where
          go = S.await >>=
               (maybe (return $! ()) $ \c -> do
                   liftIO $ modifyMVar_ r $ flip vfAppend c
                   go)
  
      flush r = modifyMVar r $ \vfi -> do
                                  !v   <- vfFinish vfi
                                  vfi' <- vfNew initialSize
                                  return $! (vfi', 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 = outputToMutableVectorSized 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 n f = do
      (os, getVec) <- mutableVectorOutputStreamSized n
      _ <- f os
      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 = outputToVectorSized 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 n = outputToMutableVectorSized n >=> basicUnsafeFreezeCompat
  {-# 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 n input = if n <= 0
                          then error $ "chunkVector: bad size: " ++ show n
                          else vfNew n >>= fromGenerator . go n
    where
      doneChunk !vfi = do
          liftIO (vfFinish vfi >>= V.unsafeFreeze) >>= yield
          !vfi' <- liftIO $ vfNew n
          go n vfi'
  
      go !k !vfi | k <= 0    = doneChunk vfi
                 | otherwise = liftIO (S.read input) >>= maybe finish chunk
        where
          finish = do
              v <- liftIO (vfFinish vfi >>= V.unsafeFreeze)
              if V.null v then return $! () else yield v
  
          chunk x = do
              !vfi' <- liftIO $ vfAppend vfi x
              go (k - 1) 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 out = V.mapM_ (flip S.write out . Just) v
  {-# INLINE writeVector #-}
  
  
  ------------------------------------------------------------------------------
  dEFAULT_BUFSIZ :: Int
  dEFAULT_BUFSIZ = 64
  