module Data.Repa.Stream.Insert
(insertS)
where
import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..))
import qualified Data.Vector.Fusion.Stream.Size as S
#include "repa-stream.h"
insertS :: Monad m
=> (Int -> Maybe a)
-> Stream m a
-> Stream m a
insertS fNew (Stream istep si0 _)
= Stream ostep (si0, 0, True, False) S.Unknown
where
ostep (si, ix, tryNew, srcDone)
| tryNew
= case fNew ix of
Just x -> return $ Yield x (si, ix + 1, True, srcDone)
Nothing
| srcDone -> return Done
| otherwise -> ostep_next si ix
| otherwise
= ostep_next si ix
ostep_next !si !ix
= istep si >>= \m
-> case m of
Yield x si' -> return $ Yield x (si', ix + 1, True, False)
Skip si' -> return $ Skip (si', ix, False, False)
Done -> return $ Skip (si, ix, True, True)