-- | Helper functions for turnings streams into vectors.
--
-- Mostly very similar to bundle conversion functions from the @vector@
-- package.

module Data.Vector.Generic.Unstream where

import           Control.Monad.ST
import           GHC.Conc (pseq)
import qualified Data.Vector.Fusion.Stream.Monadic as SM
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM
import           System.IO.Unsafe (unsafePerformIO)

-- for testing

import qualified Data.Vector.Unboxed as VU



-- | Turns a stream into a vector.
--
-- TODO insert index checks? Generalized @flag devel@

streamToVectorM :: forall m v a . (Monad m, VG.Vector v a) => SM.Stream m a -> m (v a)
streamToVectorM :: Stream m a -> m (v a)
streamToVectorM Stream m a
s = do
  let mv' :: Mutable v RealWorld a
mv' = IO (Mutable v RealWorld a) -> Mutable v RealWorld a
forall a. IO a -> a
unsafePerformIO (IO (Mutable v RealWorld a) -> Mutable v RealWorld a)
-> IO (Mutable v RealWorld a) -> 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)
VGM.unsafeNew Int
1
  let put :: (v RealWorld a, Int) -> a -> m (v RealWorld a, Int)
put (v RealWorld a
v',Int
i) a
x =
        do let v :: v RealWorld a
v = IO (v RealWorld a) -> v RealWorld a
forall a. IO a -> a
unsafePerformIO (IO (v RealWorld a) -> v RealWorld a)
-> IO (v RealWorld a) -> v RealWorld a
forall a b. (a -> b) -> a -> b
$ if (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< v RealWorld a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
VGM.length v RealWorld a
v') then v RealWorld a -> IO (v RealWorld a)
forall (m :: * -> *) a. Monad m => a -> m a
return v RealWorld a
v' else 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)
VGM.unsafeGrow v RealWorld a
v (PrimState IO) a
v' (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ v RealWorld a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
VGM.length v RealWorld a
v')
           () -> m (v RealWorld a, Int) -> m (v RealWorld a, Int)
seq (IO () -> ()
forall a. IO a -> a
unsafePerformIO (IO () -> ()) -> IO () -> ()
forall a b. (a -> b) -> a -> b
$ v (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite v RealWorld a
v (PrimState IO) a
v Int
i a
x) ((v RealWorld a, Int) -> m (v RealWorld a, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (v RealWorld a
v,Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))
      {-# Inline [0] put #-}
  (Mutable v RealWorld a
mv,Int
written) <- ((Mutable v RealWorld a, Int)
 -> a -> m (Mutable v RealWorld a, Int))
-> (Mutable v RealWorld a, Int)
-> Stream m a
-> m (Mutable v RealWorld a, Int)
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Stream m b -> m a
SM.foldlM' (Mutable v RealWorld a, Int) -> a -> m (Mutable v RealWorld a, Int)
forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, Monad m) =>
(v RealWorld a, Int) -> a -> m (v RealWorld a, Int)
put (Mutable v RealWorld a
mv',Int
0) Stream m a
s
  Mutable v RealWorld a
mv Mutable v RealWorld a -> m (v a) -> m (v a)
forall a b. a -> b -> b
`pseq` v a -> m (v a)
forall (m :: * -> *) a. Monad m => a -> m a
return (v a -> m (v a))
-> (Mutable v RealWorld a -> v a)
-> Mutable v RealWorld a
-> m (v a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (v a) -> v a
forall a. IO a -> a
unsafePerformIO (IO (v a) -> v a)
-> (Mutable v RealWorld a -> IO (v a))
-> Mutable v RealWorld a
-> v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mutable v RealWorld a -> IO (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
VG.freeze (Mutable v RealWorld a -> m (v a))
-> Mutable v RealWorld a -> m (v a)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Mutable v RealWorld a -> Mutable v RealWorld a
forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
VGM.unsafeSlice Int
0 Int
written Mutable v RealWorld a
mv
{-# Inline streamToVectorM #-}