module Data.Repa.Stream.Ratchet
( unsafeRatchetS)
where
import Control.Monad.Primitive
import Data.IORef
import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..))
import qualified Data.Vector.Generic as GV
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Fusion.Stream.Size as S
#include "repa-stream.h"
unsafeRatchetS
:: (GM.MVector vm Int, GV.Vector vv Int)
=> vm (PrimState IO) Int
-> vv Int
-> IORef (vm (PrimState IO) Int)
-> Stream IO Int
unsafeRatchetS !mvStarts !vMax !rmvLens
= Stream ostep (0, Nothing, 0, 0) S.Unknown
where
!iSegMax = GM.length mvStarts 1
ostep (iSeg, mvmLens, oSeg, oLen)
= ostep' iSeg mvmLens oSeg oLen
ostep' !iSeg !mvmLens !oSeg !oLen
| iSeg <= iSegMax
= do !iVal <- GM.unsafeRead mvStarts iSeg
let !iNext = vMax `GV.unsafeIndex` iSeg
if iVal >= iNext
then return $ Skip (iSeg + 1, mvmLens, oSeg, oLen)
else do
GM.unsafeWrite mvStarts iSeg (iVal + 1)
return $ Yield iVal (iSeg + 1, mvmLens, oSeg, oLen + 1)
| oLen > 0
= do
!vmLens <- case mvmLens of
Nothing -> readIORef rmvLens
Just vmLens -> return $ vmLens
let !oSegLen = GM.length vmLens
if oSeg >= oSegLen
then do
!vmLens' <- GM.unsafeGrow vmLens (GM.length vmLens)
writeIORef rmvLens vmLens'
GM.unsafeWrite vmLens' oSeg oLen
return $ Skip (0, Just vmLens', oSeg + 1, 0)
else do
GM.unsafeWrite vmLens oSeg oLen
return $ Skip (0, Just vmLens, oSeg + 1, 0)
| otherwise
= do !vmLens <- case mvmLens of
Nothing -> readIORef rmvLens
Just vmLens -> return $ vmLens
let !vmLens' = GM.unsafeSlice 0 oSeg vmLens
writeIORef rmvLens vmLens'
return Done