{-# LANGUAGE MagicHash #-}
module Data.Array.Repa.Repr.Delayed
        ( D, Array(..)
        , fromFunction, toFunction
        , delay)
where
import Data.Array.Repa.Eval.Load
import Data.Array.Repa.Eval.Target
import Data.Array.Repa.Eval.Chunked
import Data.Array.Repa.Eval.Cursored
import Data.Array.Repa.Eval.Elt
import Data.Array.Repa.Index
import Data.Array.Repa.Shape
import Data.Array.Repa.Base
import Debug.Trace
import GHC.Exts

-- | Delayed arrays are represented as functions from the index to element value.
--
--   Every time you index into a delayed array the element at that position 
--   is recomputed.
data D

-- | Compute elements of a delayed array.
instance Source D a where
 data Array D sh a
        = ADelayed  
                !sh 
                (sh -> a) 

 index :: Array D sh a -> sh -> a
index       (ADelayed _  f) sh
ix  = sh -> a
f sh
ix
 {-# INLINE index #-}

 linearIndex :: Array D sh a -> Int -> a
linearIndex (ADelayed sh f) Int
ix  = sh -> a
f (sh -> Int -> sh
forall sh. Shape sh => sh -> Int -> sh
fromIndex sh
sh Int
ix)
 {-# INLINE linearIndex #-}

 extent :: Array D sh a -> sh
extent (ADelayed sh _)
        = sh
sh
 {-# INLINE extent #-}

 deepSeqArray :: Array D sh a -> b -> b
deepSeqArray (ADelayed sh f) b
y
        = sh
sh sh -> b -> b
forall sh a. Shape sh => sh -> a -> a
`deepSeq` sh -> a
f (sh -> a) -> b -> b
`seq` b
y
 {-# INLINE deepSeqArray #-}


-- Load -----------------------------------------------------------------------
-- | Compute all elements in an array.
instance Shape sh => Load D sh e where
 loadP :: Array D sh e -> MVec r2 e -> IO ()
loadP (ADelayed sh getElem) MVec r2 e
mvec
  = MVec r2 e
mvec 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[Delayed]: start"
        Int -> (Int -> e -> IO ()) -> (Int -> e) -> IO ()
forall a. Int -> (Int -> a -> IO ()) -> (Int -> a) -> IO ()
fillChunkedP (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
mvec) (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
mvec
        String -> IO ()
traceEventIO String
"Repa.loadP[Delayed]: end"
 {-# INLINE [4] loadP #-}

 loadS :: Array D sh e -> MVec r2 e -> IO ()
loadS (ADelayed sh getElem) MVec r2 e
mvec
  = MVec r2 e
mvec MVec r2 e -> IO () -> IO ()
forall r e a. Target r e => MVec r e -> a -> a
`deepSeqMVec` 
    do  String -> IO ()
traceEventIO String
"Repa.loadS[Delayed]: start"
        Int -> (Int -> e -> IO ()) -> (Int -> e) -> IO ()
forall a. Int -> (Int -> a -> IO ()) -> (Int -> a) -> IO ()
fillLinearS (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
mvec) (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
mvec
        String -> IO ()
traceEventIO String
"Repa.loadS[Delayed]: end"
 {-# INLINE [4] loadS #-}


-- | Compute a range of elements in a rank-2 array.
instance Elt e => LoadRange D DIM2 e where
 loadRangeP :: Array D DIM2 e -> MVec r2 e -> DIM2 -> DIM2 -> IO ()
loadRangeP  (ADelayed (Z :. _h :. (I# w)) getElem) MVec r2 e
mvec
             (DIM0
Z :. (I# Int#
y0) :. (I# Int#
x0)) (DIM0
Z :. (I# Int#
h0) :. (I# Int#
w0))
  = MVec r2 e
mvec MVec r2 e -> IO () -> IO ()
forall r e a. Target r e => MVec r e -> a -> a
`deepSeqMVec` 
    do  String -> IO ()
traceEventIO String
"Repa.loadRangeP[Delayed]: start"
        (Int -> e -> IO ())
-> (DIM2 -> e) -> Int# -> Int# -> Int# -> Int# -> Int# -> IO ()
forall a.
Elt a =>
(Int -> a -> IO ())
-> (DIM2 -> a) -> Int# -> Int# -> Int# -> Int# -> Int# -> IO ()
fillBlock2P (MVec r2 e -> Int -> e -> IO ()
forall r e. Target r e => MVec r e -> Int -> e -> IO ()
unsafeWriteMVec MVec r2 e
mvec) 
                        DIM2 -> e
getElem
                        Int#
w Int#
x0 Int#
y0 Int#
w0 Int#
h0
        MVec r2 e -> IO ()
forall r e. Target r e => MVec r e -> IO ()
touchMVec MVec r2 e
mvec
        String -> IO ()
traceEventIO String
"Repa.loadRangeP[Delayed]: end"
 {-# INLINE [1] loadRangeP #-}

 loadRangeS :: Array D DIM2 e -> MVec r2 e -> DIM2 -> DIM2 -> IO ()
loadRangeS  (ADelayed (Z :. _h :. (I# w)) getElem) MVec r2 e
mvec
             (DIM0
Z :. (I# Int#
y0) :. (I# Int#
x0)) (DIM0
Z :. (I# Int#
h0) :. (I# Int#
w0))
  = MVec r2 e
mvec MVec r2 e -> IO () -> IO ()
forall r e a. Target r e => MVec r e -> a -> a
`deepSeqMVec`
    do  String -> IO ()
traceEventIO String
"Repa.loadRangeS[Delayed]: start"
        (Int -> e -> IO ())
-> (DIM2 -> e) -> Int# -> Int# -> Int# -> Int# -> Int# -> IO ()
forall a.
(Int -> a -> IO ())
-> (DIM2 -> a) -> Int# -> Int# -> Int# -> Int# -> Int# -> IO ()
fillBlock2S (MVec r2 e -> Int -> e -> IO ()
forall r e. Target r e => MVec r e -> Int -> e -> IO ()
unsafeWriteMVec MVec r2 e
mvec) 
                DIM2 -> e
getElem
                Int#
w Int#
x0 Int#
y0 Int#
w0 Int#
h0
        MVec r2 e -> IO ()
forall r e. Target r e => MVec r e -> IO ()
touchMVec MVec r2 e
mvec
        String -> IO ()
traceEventIO String
"Repa.loadRangeS[Delayed]: end"
 {-# INLINE [1] loadRangeS #-}


-- Conversions ----------------------------------------------------------------
-- | O(1). Wrap a function as a delayed array.
fromFunction :: sh -> (sh -> a) -> Array D sh a
fromFunction :: sh -> (sh -> a) -> Array D sh a
fromFunction sh
sh sh -> a
f 
        = sh -> (sh -> a) -> Array D sh a
forall sh a. sh -> (sh -> a) -> Array D sh a
ADelayed sh
sh sh -> a
f 
{-# INLINE fromFunction #-}


-- | O(1). Produce the extent of an array, and a function to retrieve an
--   arbitrary element.
toFunction 
        :: (Shape sh, Source r1 a)
        => Array r1 sh a -> (sh, sh -> a)
toFunction :: Array r1 sh a -> (sh, sh -> a)
toFunction Array r1 sh a
arr
 = case Array r1 sh a -> Array D sh a
forall sh r e.
(Shape sh, Source r e) =>
Array r sh e -> Array D sh e
delay Array r1 sh a
arr of
        ADelayed sh f -> (sh
sh, sh -> a
f)
{-# INLINE toFunction #-}


-- | O(1). Delay an array.
--   This wraps the internal representation to be a function from
--   indices to elements, so consumers don't need to worry about
--   what the previous representation was.
--
delay   :: Shape sh => Source r e
        => Array r sh e -> Array D sh e
delay :: Array r sh e -> Array D sh e
delay Array r sh e
arr = sh -> (sh -> e) -> Array D sh e
forall sh a. sh -> (sh -> a) -> Array D sh a
ADelayed (Array r sh e -> sh
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r sh e
arr) (Array r sh e -> sh -> e
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
unsafeIndex Array r sh e
arr)
{-# INLINE delay #-}