module Data.Repa.Stream.Segment
( findSegmentsS
, startLengthsOfSegsS)
where
import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..))
import qualified Data.Vector.Fusion.Stream.Size as S
#include "repa-stream.h"
findSegmentsS
:: Monad m
=> (a -> Bool)
-> (a -> Bool)
-> i
-> Stream m (i, a)
-> Stream m (i, i)
findSegmentsS pStart pEnd iEnd (Stream istep s sz)
= Stream ostep (s, True, Nothing) (S.toMax sz)
where
ostep (_, False, _)
= return Done
ostep (si, f, n@Nothing)
= do m <- istep si
case m of
Yield (i, x) si'
| pStart x
-> if pEnd x
then return $ Yield (i, i) (si', f, Nothing)
else return $ Skip (si', f, Just i)
| otherwise -> return $ Skip (si', f, n)
Skip si' -> return $ Skip (si', f, n)
Done -> return $ Done
ostep (si, f, j@(Just iStart))
= do m <- istep si
case m of
Yield (i, x) si'
| pEnd x -> return $ Yield (iStart, i)
(si', f, Nothing)
| otherwise -> return $ Skip (si', f, j)
Skip si' -> return $ Skip (si', f, j)
Done -> return $ Yield (iStart, iEnd)
(si, False, Nothing)
startLengthsOfSegsS
:: Monad m
=> Stream m (Int, Int)
-> Stream m (Int, Int)
startLengthsOfSegsS (Stream istep s sz)
= Stream ostep (s, True, Nothing) sz
where
ostep (_, False, _)
= return Done
ostep (si, f, n@Nothing)
= do m <- istep si
case m of
Yield x si' -> return $ Skip (si', f, Just x)
Skip si' -> return $ Skip (si', f, n)
Done -> return $ Done
ostep (si, f, j@(Just (iStart, iEnd)))
= do m <- istep si
case m of
Yield x si' -> return $ Yield (iStart, iEnd iStart + 1)
(si', f, Just x)
Skip si' -> return $ Skip (si', f, j)
Done -> return $ Yield (iStart, iEnd iStart + 1)
(si, False, Nothing)