module Data.Array.Repa.Repr.HintInterleave
        (I, Array (..), hintInterleave)
where
import Data.Array.Repa.Eval.Load
import Data.Array.Repa.Eval.Target
import Data.Array.Repa.Eval.Interleaved
import Data.Array.Repa.Repr.Delayed
import Data.Array.Repa.Shape
import Data.Array.Repa.Base
import Debug.Trace


-- | Hints that computing this array will be an unbalanced workload
--   and evaluation should be interleaved between the processors.
data I r1

instance Source r1 a => Source (I r1) a where
 data Array (I r1) sh a
        = AInterleave !(Array r1 sh a)

 extent :: Array (I r1) sh a -> sh
extent (AInterleave arr) 
        = Array r1 sh a -> sh
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r1 sh a
arr
 {-# INLINE extent #-}

 index :: Array (I r1) sh a -> sh -> a
index  (AInterleave arr) sh
ix
        = Array r1 sh a -> sh -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
index Array r1 sh a
arr sh
ix
 {-# INLINE index #-}

 unsafeIndex :: Array (I r1) sh a -> sh -> a
unsafeIndex (AInterleave arr) sh
ix
        = Array r1 sh a -> sh -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
unsafeIndex Array r1 sh a
arr sh
ix
 {-# INLINE unsafeIndex #-}

 linearIndex :: Array (I r1) sh a -> Int -> a
linearIndex (AInterleave arr) Int
ix
        = Array r1 sh a -> Int -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> Int -> e
linearIndex Array r1 sh a
arr Int
ix
 {-# INLINE linearIndex #-}

 unsafeLinearIndex :: Array (I r1) sh a -> Int -> a
unsafeLinearIndex (AInterleave arr) Int
ix
        = Array r1 sh a -> Int -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> Int -> e
unsafeLinearIndex Array r1 sh a
arr Int
ix
 {-# INLINE unsafeLinearIndex #-}

 deepSeqArray :: Array (I r1) sh a -> b -> b
deepSeqArray (AInterleave arr) b
x
        = Array r1 sh a -> b -> b
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
deepSeqArray Array r1 sh a
arr b
x
 {-# INLINE deepSeqArray #-}


deriving instance Show (Array r1 sh e) 
        => Show (Array (I r1) sh e)

deriving instance Read (Array r1 sh e) 
        => Read (Array (I r1) sh e)


-- | Wrap an array with a unbalanced-ness hint.
hintInterleave :: Array r1 sh e -> Array (I r1) sh e
hintInterleave :: Array r1 sh e -> Array (I r1) sh e
hintInterleave = Array r1 sh e -> Array (I r1) sh e
forall r1 sh a. Array r1 sh a -> Array (I r1) sh a
AInterleave


-- Load -----------------------------------------------------------------------
instance (Shape sh, Load D sh e) 
        => Load (I D) sh e where
 loadP :: Array (I D) sh e -> MVec r2 e -> IO ()
loadP (AInterleave (ADelayed sh getElem)) MVec r2 e
marr
  = MVec r2 e
marr MVec r2 e -> IO () -> IO ()
forall r e a. Target r e => MVec r e -> a -> a
`deepSeqMVec`
    do  String -> IO ()
traceEventIO String
"Repa.loadP[Interleaved]: start"
        Int -> (Int -> e -> IO ()) -> (Int -> e) -> IO ()
forall a. Int -> (Int -> a -> IO ()) -> (Int -> a) -> IO ()
fillInterleavedP (sh -> Int
forall sh. Shape sh => sh -> Int
size sh
sh) (MVec r2 e -> Int -> e -> IO ()
forall r e. Target r e => MVec r e -> Int -> e -> IO ()
unsafeWriteMVec MVec r2 e
marr) (sh -> e
getElem (sh -> e) -> (Int -> sh) -> Int -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. sh -> Int -> sh
forall sh. Shape sh => sh -> Int -> sh
fromIndex sh
sh) 
        MVec r2 e -> IO ()
forall r e. Target r e => MVec r e -> IO ()
touchMVec MVec r2 e
marr
        String -> IO ()
traceEventIO String
"Repa.loadP[Interleaved]: end"
 {-# INLINE [4] loadP #-}

 -- The fact that the workload is unbalanced doesn't affect us when the
 -- program is run sequentially, so just use the filling method of the inner
 -- representation
 loadS :: Array (I D) sh e -> MVec r2 e -> IO ()
loadS (AInterleave arr) MVec r2 e
marr
  = Array D sh e -> MVec r2 e -> IO ()
forall r1 sh e r2.
(Load r1 sh e, Target r2 e) =>
Array r1 sh e -> MVec r2 e -> IO ()
loadS Array D sh e
arr MVec r2 e
marr
 {-# INLINE loadS #-}