module Data.Repa.Flow.Chunked.Map
( smap_i, smap_o
, szipWith_ii)
where
import Data.Repa.Flow.Chunked.Base
import Data.Repa.Flow.States
import Data.Repa.Array.Generic as A
import Data.Repa.Array.Generic.Index as A
import Data.Repa.Array.Meta.Window as A
import Data.Repa.Array.Meta.Delayed as A
import qualified Data.Repa.Flow.Generic as G
#include "repa-flow.h"
smap_i :: (Flow i m l1 a, A.TargetI l2 b)
=> (i -> a -> b) -> Sources i m l1 a -> m (Sources i m l2 b)
smap_i f s0 = G.smap_i (\i c -> A.computeS name $ A.map (f i) c) s0
smap_o :: (Flow i m l1 a, A.TargetI l2 b)
=> (i -> a -> b) -> Sinks i m l2 b -> m (Sinks i m l1 a)
smap_o f s0 = G.smap_o (\i c -> A.computeS name $ A.map (f i) c) s0
szipWith_ii
:: ( Ord i, States i m
, BulkI lSrc1 a, BulkI lSrc2 b
, TargetI lDst c
, Windowable lSrc1 a, Windowable lSrc2 b)
=> Name lDst
-> (i -> a -> b -> c)
-> Sources i m lSrc1 a -> Sources i m lSrc2 b
-> m (Sources i m lDst c)
szipWith_ii nDst f (G.Sources nA pullA) (G.Sources nB pullB)
= do
let nC = min nA nB
bitsA <- newRefs nC Nothing
bitsB <- newRefs nC Nothing
let pullC i eatC ejectC
| not $ check i nC = ejectC
| otherwise = getA
where
getA
= do mA <- readRefs bitsA i
case mA of
Just chunkA -> getB chunkA
Nothing -> pullA i getB ejectC
getB chunkA
= do mB <- readRefs bitsB i
case mB of
Just chunkB -> zipAB chunkA chunkB
Nothing -> pullB i (zipAB chunkA) ejectC
zipAB chunkA chunkB
= do
let !lenA = A.length chunkA
let !lenB = A.length chunkB
let !lenC = min lenA lenB
let !hereA = A.window 0 lenC chunkA
let !restA = A.window lenC (lenA lenC) chunkA
let !hereB = A.window 0 lenC chunkB
let !restB = A.window lenC (lenB lenC) chunkB
let Just !hereC = A.map2S nDst (f i) hereA hereB
(if A.length restA > 0
then writeRefs bitsA i (Just restA)
else writeRefs bitsA i Nothing)
(if A.length restB > 0
then writeRefs bitsB i (Just restB)
else writeRefs bitsB i Nothing)
eatC hereC
return $ G.Sources nC pullC