module Data.Repa.Flow.Chunked.Fold
( foldlS
, foldlAllS )
where
import Data.Repa.Flow.Chunked.Base
import Data.Repa.Flow.States as S
import qualified Data.Repa.Flow.Generic as G
import qualified Data.Repa.Array.Generic.Target as A
import qualified Data.Repa.Array.Generic.Index as A
import qualified Data.Repa.Array.Generic as A
#include "repa-flow.h"
foldlS
:: ( States Int m
, A.Target lDst a, A.Index lDst ~ Int
, A.BulkI lSrc b)
=> A.Name lDst
-> (a -> b -> a)
-> a
-> Sources Int m lSrc b
-> m (A.Array lDst a)
foldlS nDst f z (G.Sources n pull)
= do
refsState <- newRefs n z
let loop_foldlS !ix
= pull ix eat_foldlS eject_foldlS
where
eat_foldlS arr
= do s <- readRefs refsState ix
let !s' = A.foldl f s arr
writeRefs refsState ix s'
loop_foldlS ix
eject_foldlS
= case next ix n of
Nothing -> return ()
Just ix' -> loop_foldlS ix'
loop_foldlS first
ls <- S.toListM refsState
return $ A.fromList nDst ls
foldlAllS
:: ( States () m
, A.BulkI lSrc b)
=> (a -> b -> a)
-> a
-> Sources Int m lSrc b
-> m a
foldlAllS f z (G.Sources n pull)
= do
ref <- newRefs () Nothing
let loop_foldlAllS !s !ix
= pull ix eat_foldlAllS eject_foldlAllS
where
eat_foldlAllS arr
= do let s' = A.foldl f s arr
loop_foldlAllS s' ix
eject_foldlAllS
= case next ix n of
Nothing -> writeRefs ref () (Just s)
Just ix' -> loop_foldlAllS s ix'
loop_foldlAllS z first
Just x <- readRefs ref ()
return x