module Data.Repa.Stream.Compact
( compactS
, compactInS )
where
import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..))
import qualified Data.Vector.Fusion.Stream.Size as S
#include "repa-stream.h"
compactS
:: Monad m
=> (s -> a -> (s, Maybe b))
-> s
-> Stream m a
-> Stream m b
compactS f s0 (Stream istep si0 sz)
= Stream ostep (si0, s0) (S.toMax sz)
where
ostep (si, s)
= istep si >>= \m
-> case m of
Yield x si'
-> case f s x of
(s', Nothing) -> return $ Skip (si', s')
(s', Just y) -> return $ Yield y (si', s')
Skip si' -> return $ Skip (si', s)
Done -> return $ Done
compactInS
:: Monad m
=> (a -> a -> (a, Maybe a))
-> Stream m a
-> Stream m a
compactInS f (Stream istep si0 sz)
= Stream ostep (si0, Nothing) (S.toMax sz)
where
ostep (si, ms@Nothing)
= istep si >>= \m
-> case m of
Yield x si' -> return $ Skip (si', Just x)
Skip si' -> return $ Skip (si', ms)
Done -> return $ Done
ostep (si, ms@(Just s))
= istep si >>= \m
-> case m of
Yield x si'
-> case f s x of
(s', Nothing) -> return $ Skip (si', Just s')
(s', Just y) -> return $ Yield y (si', Just s')
Skip si' -> return $ Skip (si', ms)
Done -> return $ Yield s (si, Nothing)