module Data.Repa.Stream.Ratchet
( unsafeRatchetS)
where
import Data.IORef
import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..))
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import qualified Data.Vector.Fusion.Stream.Size as S
#include "repa-stream.h"
unsafeRatchetS
:: UM.IOVector Int
-> U.Vector Int
-> IORef (UM.IOVector 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 `G.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 = UM.length vmLens
if oSeg >= oSegLen
then do
!vmLens' <- UM.unsafeGrow vmLens (UM.length vmLens)
writeIORef rmvLens vmLens'
UM.unsafeWrite vmLens' oSeg oLen
return $ Skip (0, Just vmLens', oSeg + 1, 0)
else do
UM.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' = UM.unsafeSlice 0 oSeg vmLens
writeIORef rmvLens vmLens'
return Done