module Data.Repa.Flow.Chunked.Folds
( folds_i
, FoldsDict)
where
import Data.Repa.Flow.Chunked.Base
import Data.Repa.Flow.States
import Data.Repa.Fusion.Unpack
import Data.Repa.Option
import Data.Repa.Array.Generic.Index as A
import Data.Repa.Array.Generic.Target as A
import Data.Repa.Array.Generic as A hiding (FoldsDict)
import Data.Repa.Array.Meta.Window as A
import Data.Repa.Array.Meta.Tuple as A
import qualified Data.Repa.Flow.Generic as G
#include "repa-flow.h"
type FoldsDict i m lSeg tSeg lElt tElt lGrp tGrp lRes tRes n a b
= ( States i m
, Windowable lSeg (n, Int), Windowable lElt a
, BulkI lSeg (n, Int)
, BulkI lElt a
, BulkI lGrp n
, BulkI lRes b
, TargetI lElt a
, TargetI lGrp n
, TargetI lRes b
, Unpack (Buffer lGrp n) tGrp
, Unpack (Buffer lRes b) tRes)
folds_i :: FoldsDict i m lSeg tSeg lElt tElt lGrp tGrp lRes tRes n a b
=> Name lGrp
-> Name lRes
-> (a -> b -> b)
-> b
-> Sources i m lSeg (n, Int)
-> Sources i m lElt a
-> m (Sources i m (T2 lGrp lRes) (n, b))
folds_i _ _ f z sLens@(G.Sources nLens _)
sVals@(G.Sources nVals _)
= do
let nFolds = min nLens nVals
refsState <- newRefs nFolds None3
refsNameLens <- newRefs nFolds Nothing
refsVals <- newRefs nFolds Nothing
refsValsDone <- newRefs nFolds False
let pull_folds i eat eject
= do mNameLens <- folds_loadChunkNameLens sLens refsNameLens i
mVals <- folds_loadChunkVals sVals refsVals refsValsDone i
case (mNameLens, mVals) of
(Nothing, _) -> eject
(_, Nothing) -> eject
(Just cNameLens, Just cVals)
-> cNameLens `seq` cVals `seq`
do
mState <- readRefs refsState i
let (cResults, sFolds)
= A.foldsWith name name f z
(fromOption3 mState) cNameLens cVals
folds_update
refsState refsNameLens refsVals i
cNameLens cVals sFolds
valsDone <- readRefs refsValsDone i
if A.length cResults == 0
&& A.length cNameLens >= 0
&& valsDone
then eject
else eat cResults
return $ G.Sources nFolds pull_folds
folds_loadChunkNameLens
:: States i m
=> Sources i m l1 (n, Int)
-> Refs i m (Maybe (Array l1 (n, Int)))
-> i
-> m (Maybe (Array l1 (n, Int)))
folds_loadChunkNameLens (G.Sources _ pullLens) refsLens i
= do mChunkLens <- readRefs refsLens i
case mChunkLens of
Nothing
-> let eatLens_folds chunk
= writeRefs refsLens i (Just chunk)
ejectLens_folds = return ()
in do
pullLens i eatLens_folds ejectLens_folds
readRefs refsLens i
jc@(Just _)
-> return jc
folds_loadChunkVals
:: (States i m, TargetI l2 a)
=> Sources i m l2 a
-> Refs i m (Maybe (Array l2 a))
-> Refs i m Bool
-> i
-> m (Maybe (Array l2 a))
folds_loadChunkVals (G.Sources _ pullVals) refsVals refsValsDone i
= do mChunkVals <- readRefs refsVals i
case mChunkVals of
Nothing
-> let eatVals_folds chunk
= writeRefs refsVals i (Just chunk)
ejectVals_folds
= do writeRefs refsVals i (Just $ A.fromList name [])
writeRefs refsValsDone i True
in do
pullVals i eatVals_folds ejectVals_folds
readRefs refsVals i
jc@(Just _)
-> return jc
folds_update
:: ( States i m
, Windowable l1 (n, Int), Windowable l2 a
, A.Index l1 ~ Int, A.Index l2 ~ Int)
=> Refs i m (Option3 n Int b)
-> Refs i m (Maybe (Array l1 (n, Int)))
-> Refs i m (Maybe (Array l2 a))
-> i
-> Array l1 (n, Int)
-> Array l2 a
-> Folds Int Int n a b
-> m ()
folds_update refsState refsLens refsVals i cLens cVals sFolds
= do
writeRefs refsState i
$ case _nameSeg sFolds of
Some n -> Some3 n (_lenSeg sFolds) (_valSeg sFolds)
None -> None3
let !posLens = _stateLens sFolds
let !nLensRemain = A.length cLens posLens
writeRefs refsLens i
$ if nLensRemain <= 0
then Nothing
else Just $ A.window posLens nLensRemain cLens
let !posVals = _stateVals sFolds
let !nValsRemain = A.length cVals posVals
writeRefs refsVals i
$ if nValsRemain <= 0
then Nothing
else Just $ A.window posVals nValsRemain cVals