module Data.Repa.Flow.Generic.Map
( map_i, map_o
, smap_i, smap_o
, szipWith_ii, szipWith_io, szipWith_oi)
where
import Data.Repa.Flow.Generic.Base
import Control.Monad
import Prelude as P
#include "repa-flow.h"
map_i :: Monad m
=> (a -> b) -> Sources i m a -> m (Sources i m b)
map_i f s = smap_i (\_ x -> f x) s
smap_i :: Monad m
=> (i -> a -> b) -> Sources i m a -> m (Sources i m b)
smap_i f (Sources n pullsA)
= return $ Sources n pullsB_map
where
pullsB_map i eat eject
= pullsA i eat_a eject_a
where
eat_a v = eat (f i v)
eject_a = eject
map_o :: Monad m
=> (a -> b) -> Sinks i m b -> m (Sinks i m a)
map_o f k = smap_o (\_ x -> f x) k
smap_o :: Monad m
=> (i -> a -> b) -> Sinks i m b -> m (Sinks i m a)
smap_o f (Sinks n pushB ejectB)
= return $ Sinks n pushA_map ejectA_map
where
pushA_map i a = pushB i (f i a)
ejectA_map i = ejectB i
szipWith_ii
:: (Ord i, Monad m)
=> (i -> a -> b -> c)
-> Sources i m a -> Sources i m b
-> m (Sources i m c)
szipWith_ii f (Sources nA pullA) (Sources nB pullB)
= return $ Sources (min nA nB) pull_szipWith
where
pull_szipWith i eat eject
= pullA i eatA eject
where
eatA xA = pullB i eatB eject
where
eatB xB = eat (f i xA xB)
szipWith_io
:: (Ord i, Monad m)
=> (i -> a -> b -> c)
-> Sinks i m c -> Sources i m a
-> m (Sinks i m b)
szipWith_io f (Sinks nC pushC ejectC) (Sources nA pullA)
= return $ Sinks nB pushB ejectC
where
!nB = min nC nA
pushB i xB
| i > nB = return ()
| otherwise = pullA i eatA (ejectC i)
where
eatA xA = pushC i (f i xA xB)
szipWith_oi
:: (Ord i, Monad m)
=> (i -> a -> b -> c)
-> Sinks i m c -> Sources i m b
-> m (Sinks i m a)
szipWith_oi f (Sinks nC pushC ejectC) (Sources nB pullB)
= return $ Sinks nA pushA ejectC
where
!nA = min nC nB
pushA i xA
| i > nA = return ()
| otherwise = pullB i eatB (ejectC i)
where
eatB xB = pushC i (f i xA xB)