repa-flow-4.2.2.1: Data/Repa/Flow/Chunked/Folds.hs
module Data.Repa.Flow.Chunked.Folds
( folds_i
, FoldsDict)
where
import Data.Repa.Flow.Chunked.Base
import Data.Repa.Flow.States
import Data.Repa.Scalar.Option
import Data.Repa.Array.Generic.Index 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"
-- | Dictionaries needed to perform a segmented fold.
type FoldsDict i m lSeg lElt lGrp lRes 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)
-- Folds ----------------------------------------------------------------------
-- | Segmented fold over vectors of segment lengths and input values.
folds_i :: FoldsDict i m lSeg lElt lGrp lRes n a b
=> Name lGrp -- ^ Layout for group names.
-> Name lRes -- ^ Layout for fold results.
-> (a -> b -> b) -- ^ Worker function.
-> b -- ^ Initial state when folding each segment.
-> Sources i m lSeg (n, Int) -- ^ Segment lengths.
-> Sources i m lElt a -- ^ Input elements to fold.
-> m (Sources i m (T2 lGrp lRes) (n, b)) -- ^ Result elements.
folds_i _ _ f z sLens@(G.Sources nLens _)
sVals@(G.Sources nVals _)
= do
-- Arity of the result bundle is the minimum of the two inputs.
let nFolds = min nLens nVals
-- Refs to hold partial fold states between chunks.
refsState <- newRefs nFolds None3
-- Refs to hold the current chunk of lengths data for each stream.
refsNameLens <- newRefs nFolds Nothing
-- Refs to hold the current chunk of vals data for each stream.
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
-- If we couldn't get a chunk for both sides then we can't
-- produce anymore results, and the merge is done.
(Nothing, _) -> eject
(_, Nothing) -> eject
-- We've got a chunk for both sides, time to do some work.
(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 we're not producing output while we still
-- have segment lengths then we're done.
if A.length cResults == 0
&& A.length cNameLens >= 0
&& valsDone
then eject
else eat cResults
{-# INLINE pull_folds #-}
return $ G.Sources nFolds pull_folds
{-# INLINE_FLOW folds_i #-}
-- Load the current chunk of lengths data.
-- If we already have one in the state then use that,
-- otherwise try to pull a new chunk from the source.
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)
{-# INLINE eatLens_folds #-}
ejectLens_folds = return ()
{-# INLINE ejectLens_folds #-}
in do
pullLens i eatLens_folds ejectLens_folds
readRefs refsLens i
jc@(Just _)
-> return jc
{-# NOINLINE folds_loadChunkNameLens #-}
-- NOINLINE as this doesn't need to be specialized,
--- and we want to hide the case from the simplifier.
-- Grab the current chunk of values data.
-- If we already have one in the state then use that,
-- otherwise try to pull a new chunk from the source.
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)
{-# INLINE eatVals_folds #-}
-- When there are no more values then shim in an
-- empty chunk so that we can keep calling A.folds
-- this is needed when there are zero lengthed
-- segments on the end of the stream
ejectVals_folds
= do writeRefs refsVals i (Just $ A.fromList name [])
writeRefs refsValsDone i True
{-# INLINE ejectVals_folds #-}
in do
pullVals i eatVals_folds ejectVals_folds
readRefs refsVals i
jc@(Just _)
-> return jc
{-# NOINLINE folds_loadChunkVals #-}
-- NOINLINE as this doesn't need to be specialized,
-- and we want to hide the case from the simplifier.
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
-- Remember state for the final segment.
writeRefs refsState i
$ case _nameSeg sFolds of
Some n -> Some3 n (_lenSeg sFolds) (_valSeg sFolds)
None -> None3
-- Slice down the lengths chunk to just the elements
-- that we haven't already consumed. If we've consumed
-- them all then clear the chunk reference so a new one
-- will be loaded the next time around.
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
-- Likewise for the values chunk.
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
{-# NOINLINE folds_update #-}
-- NOINLINE because it is only called once per chunk
-- and does not need to be specialised.