module Data.Repa.Flow.Generic.Array.Chunk
(chunk_i)
where
import Data.Repa.Flow.Generic.Base
import Data.Repa.Array.Generic as A
import Data.Repa.Array.Generic.Target as A
import Data.Repa.Array.Generic.Index as A
#include "repa-flow.h"
chunk_i :: (Target lDst a, Index lDst ~ Int, States i IO)
=> Name lDst
-> Int
-> Sources i IO a
-> IO (Sources i IO (Array lDst a))
chunk_i nDst !maxLen (Sources n pullX)
= do
final <- newRefs n Nothing
let pull_chunk i eat eject
= do
chunk <- unsafeNewBuffer (A.create nDst maxLen)
let loop_chunk !ix
| ix >= maxLen
= writeRefs final i (Just ix)
| otherwise
= pullX i eat_chunk eject_chunk
where
eat_chunk x
= do unsafeWriteBuffer chunk ix x
loop_chunk (ix + 1)
eject_chunk
| ix == 0 = writeRefs final i Nothing
| otherwise = writeRefs final i (Just ix)
loop_chunk 0
mlen <- readRefs final i
case mlen of
Nothing -> eject
Just len
-> do chunk' <- unsafeSliceBuffer 0 len chunk
arr <- unsafeFreezeBuffer chunk'
eat arr
return $ Sources n pull_chunk