module Data.Array.Repa.Repr.HintSmall
        (S, Array (..), hintSmall)
where
import Data.Array.Repa.Eval.Load
import Data.Array.Repa.Base
import Data.Array.Repa.Shape


-- | Hints that evaluating this array is only a small amount of work.
--   It will be evaluated sequentially in the main thread, instead of
--   in parallel on the gang. This avoids the associated scheduling overhead.
data S r1

instance Source r1 a => Source (S r1) a where
 data Array (S r1) sh a
        = ASmall !(Array r1 sh a)

 extent :: Array (S r1) sh a -> sh
extent (ASmall 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 (S r1) sh a -> sh -> a
index  (ASmall 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 (S r1) sh a -> sh -> a
unsafeIndex (ASmall 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 (S r1) sh a -> Int -> a
linearIndex (ASmall 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 (S r1) sh a -> Int -> a
unsafeLinearIndex (ASmall 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 (S r1) sh a -> b -> b
deepSeqArray (ASmall 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 #-}


-- | Wrap an array with a smallness hint.
hintSmall :: Array r1 sh e -> Array (S r1) sh e
hintSmall :: Array r1 sh e -> Array (S r1) sh e
hintSmall = Array r1 sh e -> Array (S r1) sh e
forall r1 sh a. Array r1 sh a -> Array (S r1) sh a
ASmall


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

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


-- Load ----------------------------------------------------------------------
instance ( Shape sh, Load r1 sh e) 
        => Load (S r1) sh e where
 loadP :: Array (S r1) sh e -> MVec r2 e -> IO ()
loadP (ASmall arr) MVec r2 e
marr
  = Array r1 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 r1 sh e
arr MVec r2 e
marr
 {-# INLINE loadP #-}

 loadS :: Array (S r1) sh e -> MVec r2 e -> IO ()
loadS (ASmall arr) MVec r2 e
marr
  = Array r1 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 r1 sh e
arr MVec r2 e
marr
 {-# INLINE loadS #-}


-- LoadRange ------------------------------------------------------------------
instance ( Shape sh, LoadRange r1 sh e)
        => LoadRange (S r1) sh e where
 loadRangeP :: Array (S r1) sh e -> MVec r2 e -> sh -> sh -> IO ()
loadRangeP (ASmall arr) MVec r2 e
marr sh
ix1 sh
ix2
  = Array r1 sh e -> MVec r2 e -> sh -> sh -> IO ()
forall r1 sh e r2.
(LoadRange r1 sh e, Target r2 e) =>
Array r1 sh e -> MVec r2 e -> sh -> sh -> IO ()
loadRangeS Array r1 sh e
arr MVec r2 e
marr sh
ix1 sh
ix2
 {-# INLINE loadRangeP #-}

 loadRangeS :: Array (S r1) sh e -> MVec r2 e -> sh -> sh -> IO ()
loadRangeS (ASmall arr) MVec r2 e
marr sh
ix1 sh
ix2
  = Array r1 sh e -> MVec r2 e -> sh -> sh -> IO ()
forall r1 sh e r2.
(LoadRange r1 sh e, Target r2 e) =>
Array r1 sh e -> MVec r2 e -> sh -> sh -> IO ()
loadRangeS Array r1 sh e
arr MVec r2 e
marr sh
ix1 sh
ix2
 {-# INLINE loadRangeS #-}