module Data.Repa.Stream.Pad
(padForwardS)
where
import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..))
import Data.Repa.Scalar.Option
import qualified Data.Vector.Fusion.Stream.Size as S
#include "repa-stream.h"
padForwardS
:: (Monad m, Ord k)
=> (k -> k)
-> Stream m (k, v)
-> Stream m (k, v)
padForwardS ksucc (Stream istep si0 _)
= Stream ostep (si0, None2, None2) S.Unknown
where
ostep (si, sPrev@None2, sBound)
= istep si >>= \m
-> case m of
Yield (k0, v0) si'
-> return $ Yield (k0, v0) (si', Some2 k0 v0, sBound)
Skip si'
-> return $ Skip (si', sPrev, sBound)
Done
-> return $ Done
ostep (si, sPrev@(Some2 kPrev _vPrev), sTarget@None2)
= istep si >>= \m
-> case m of
Yield (kStep, vStep) si'
| kExpect <- ksucc kPrev
, kStep > kExpect
-> return $ Skip (si', sPrev, Some2 kStep vStep)
| otherwise
-> return $ Yield (kStep, vStep) (si', Some2 kStep vStep, None2)
Skip si' -> return $ Skip (si', sPrev, sTarget)
Done -> return $ Done
ostep (si, _sPrev@(Some2 kPrev vPrev), sTarget@(Some2 kTarget vTarget))
= let kNext = ksucc kPrev
in if kNext >= kTarget
then return $ Yield (kTarget, vTarget) (si, Some2 kTarget vTarget, None2)
else return $ Yield (kNext, vPrev) (si, Some2 kNext vPrev, sTarget)