module Data.Repa.Stream.Extract
(extractS)
where
import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..))
import qualified Data.Vector.Fusion.Stream.Size as S
#include "repa-stream.h"
extractS
:: Monad m
=> (Int -> a)
-> Stream m (Int, Int)
-> Stream m a
extractS get (Stream istep si0 _)
= Stream ostep (si0, Nothing) S.Unknown
where
ostep (si, Nothing)
= istep si >>= \m
-> case m of
Yield (iStart, iLen) si'
-> return $ Skip (si', Just (iStart, iStart + iLen))
Skip si' -> return $ Skip (si', Nothing)
Done -> return $ Done
ostep (si, Just (iPos, iTop))
| iPos >= iTop = return $ Skip (si, Nothing)
| otherwise = return $ Yield (get iPos)
(si, Just (iPos + 1, iTop))