{-# LANGUAGE CPP, NoMonomorphismRestriction #-}
#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


-- VSegd Streamers ------------------------------------------------------------
-- | Take a stream of virtual segment and segment element indices,
--   and convert it to a stream of physical segment and segment element indices.
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
        {-# INLINE_INNER mkStep' #-}
        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
{-# INLINE_STREAM streamSrcIxsThroughVSegids #-}


-- SSegd Streamers ------------------------------------------------------------
-- | Take a stream of segment and segment element indices,
--   and convert it to a stream of chunk and chunk element indices.
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
   
        {-# INLINE_INNER mkStep' #-}
        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
{-# INLINE_STREAM streamSrcIxsThroughUSSegd #-}