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)
import qualified Data.Vector.Unboxed as VU
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 #-}