module Data.Repa.Eval.Stream
( streamOfArray
, unstreamToArray
, unstreamToArrayIO)
where
import Data.Repa.Array.Generic.Index as A
import Data.Repa.Array.Internals.Bulk as A
import Data.Repa.Array.Internals.Target as A
import qualified Data.Vector.Fusion.Stream as SS
import qualified Data.Vector.Fusion.Stream.Monadic as S
import qualified Data.Vector.Fusion.Stream.Size as S
import qualified Data.Vector.Fusion.Util as S
import System.IO.Unsafe
#include "repa-array.h"
streamOfArray
:: (Monad m, Bulk l a, Index l ~ Int)
=> A.Array l a
-> S.Stream m a
streamOfArray vec
= S.generate (A.length vec)
(\i -> A.index vec i)
unstreamToArray
:: Target l a
=> Name l -> S.Stream S.Id a -> Array l a
unstreamToArray nDst s
= unsafePerformIO
$ unstreamToArrayIO nDst
$ SS.liftStream s
unstreamToArrayIO
:: Target l a
=> Name l -> S.Stream IO a -> IO (Array l a)
unstreamToArrayIO nDst (S.Stream step s0 sz)
= case sz of
S.Exact i -> unstreamToArrayIO_max i
S.Max i -> unstreamToArrayIO_max i
S.Unknown -> unstreamToArrayIO_unknown 32
where unstreamToArrayIO_max !nMax
= do !vec0 <- unsafeNewBuffer (create nDst zeroDim)
!vec <- unsafeGrowBuffer vec0 nMax
let go_unstreamIO_max !sPEC !i !s
= step s >>= \m
-> case m of
S.Yield e s'
-> do unsafeWriteBuffer vec i e
go_unstreamIO_max sPEC (i + 1) s'
S.Skip s'
-> go_unstreamIO_max sPEC i s'
S.Done
-> do buf' <- unsafeSliceBuffer 0 i vec
arr <- unsafeFreezeBuffer buf'
return arr
go_unstreamIO_max S.SPEC 0 s0
unstreamToArrayIO_unknown !nStart
= do !vec0 <- unsafeNewBuffer (create nDst zeroDim)
!vec1 <- unsafeGrowBuffer vec0 nStart
let go_unstreamIO_unknown !sPEC !vec !i !n !s
= go_unstreamIO_unknown1 vec i n s
(\vec' i' n' s' -> go_unstreamIO_unknown sPEC vec' i' n' s')
(\result -> return result)
go_unstreamIO_unknown1 !vec !i !n !s cont done
= step s >>= \r
-> case r of
S.Yield e s'
-> do (vec', n')
<- if i >= n
then do vec' <- unsafeGrowBuffer vec n
return (vec', n + n)
else return (vec, n)
unsafeWriteBuffer vec' i e
cont vec' (i + 1) n' s'
S.Skip s'
-> cont vec i n s'
S.Done
-> do
vec' <- unsafeSliceBuffer 0 i vec
arr <- unsafeFreezeBuffer vec'
done arr
go_unstreamIO_unknown S.SPEC vec1 0 nStart s0