#include "fusion-phases.h"
module Data.Array.Parallel.Unlifted.Stream.Ixs
( streamSrcIxsThroughVSegids
, streamSrcIxsThroughUSSegd)
where
import Data.Vector.Fusion.Stream.Monadic
import Data.Array.Parallel.Unlifted.Sequential.USSegd (USSegd(..))
import qualified Data.Array.Parallel.Unlifted.Sequential.USSegd as USSegd
import qualified Data.Vector.Unboxed as U
streamSrcIxsThroughVSegids
:: Monad m
=> U.Vector Int -> Stream m (Int, Int) -> Stream m (Int, Int)
streamSrcIxsThroughVSegids vsegids (Stream mkStep s0 size0)
= vsegids `seq` Stream mkStep' s0 size0
where
mkStep' s
= do step <- mkStep s
case step of
Yield (ix1, ix2) s'
-> let !pseg = U.unsafeIndex vsegids ix1
in return $ Yield (pseg, ix2) s'
Skip s' -> return $ Skip s'
Done -> return Done
streamSrcIxsThroughUSSegd
:: Monad m
=> USSegd -> Stream m (Int, Int) -> Stream m (Int, Int)
streamSrcIxsThroughUSSegd ussegd (Stream mkStep s0 size0)
= ussegd `seq` Stream mkStep' s0 size0
where
!sources = USSegd.takeSources ussegd
!starts = USSegd.takeStarts ussegd
mkStep' s
= do step <- mkStep s
case step of
Yield (ix1, ix2) s'
-> let !src = U.unsafeIndex sources ix1
!start = U.unsafeIndex starts ix1
in return $ Yield (src, start + ix2) s'
Skip s' -> return $ Skip s'
Done -> return Done