module Data.Repa.Chain.Scan
( StepUnfold (..)
, unfoldsC
, scanMaybeC
, groupsByC)
where
import Data.Repa.Chain.Base
import Data.Repa.Scalar.Option
import qualified Data.Vector.Fusion.Stream.Size as S
#include "repa-stream.h"
unfoldsC
:: Monad m
=> (a -> k -> m (StepUnfold k b))
-> k
-> Chain m s a
-> Chain m (s, k, Option a) b
unfoldsC f k0 (Chain _ s0 istep)
= Chain S.Unknown (s0, k0, None) ostep
where
ostep (s1, k1, None)
= istep s1 >>= \rs
-> case rs of
Yield xa s2 -> return $ Skip (s2, k1, Some xa)
Skip s2 -> return $ Skip (s2, k1, None)
Done s2 -> return $ Done (s2, k1, None)
ostep (s1, k1, Some xa)
= f xa k1 >>= \kmb
-> case kmb of
StepUnfoldGive xb k2 -> return $ Yield xb (s1, k2, Some xa)
StepUnfoldNext xb k2 -> return $ Yield xb (s1, k2, None)
StepUnfoldBump k2 -> return $ Skip (s1, k2, Some xa)
StepUnfoldFinish k2 -> return $ Skip (s1, k2, None)
data StepUnfold s a
= StepUnfoldGive a s
| StepUnfoldNext a s
| StepUnfoldBump s
| StepUnfoldFinish s
deriving Show
scanMaybeC
:: Monad m
=> (k -> a -> m (k, Maybe b))
-> k
-> Chain m s a
-> Chain m (s, k) b
scanMaybeC f k0 (Chain sz s0 istep)
= Chain (S.toMax sz) (s0, k0) ostep
where
ostep (s1, k1)
= istep s1 >>= \rs
-> case rs of
Yield x s2
-> f k1 x >>= \rk
-> case rk of
(k2, Nothing) -> return $ Skip (s2, k2)
(k2, Just y) -> return $ Yield y (s2, k2)
Skip s2 -> return $ Skip (s2, k1)
Done s2 -> return $ Done (s2, k1)
groupsByC
:: Monad m
=> (a -> a -> m Bool)
-> Maybe (a, Int)
-> Chain m s a
-> Chain m (s, Maybe (a, Int)) (a, Int)
groupsByC f !s !vec
= scanMaybeC work_groupsByC s vec
where
work_groupsByC !acc !y
= case acc of
Nothing
-> return $ (Just (y, 1), Nothing)
Just (x, n)
-> f x y >>= \rk
-> if rk
then return (Just (x, n + 1), Nothing)
else return (Just (y, 1), Just (x, n))