{-# LANGUAGE MagicHash #-}
module Data.Array.Repa.Repr.Cursored
        ( C, Array (..)
        , makeCursored)
where
import Data.Array.Repa.Base
import Data.Array.Repa.Shape
import Data.Array.Repa.Index
import Data.Array.Repa.Repr.Delayed
import Data.Array.Repa.Repr.Undefined
import Data.Array.Repa.Eval.Load
import Data.Array.Repa.Eval.Elt
import Data.Array.Repa.Eval.Cursored
import Data.Array.Repa.Eval.Target
import GHC.Exts
import Debug.Trace

-- | Cursored Arrays.
--   These are produced by Repa's stencil functions, and help the fusion
--   framework to share index compuations between array elements.
--
--   The basic idea is described in ``Efficient Parallel Stencil Convolution'',
--   Ben Lippmeier and Gabriele Keller, Haskell 2011 -- though the underlying
--   array representation has changed since this paper was published.
data C


-- | Compute elements of a cursored array.
instance Source C a where

 data Array C sh a
        = forall cursor. ACursored
        { forall sh a. Array C sh a -> sh
cursoredExtent :: !sh 
                
          -- | Make a cursor to a particular element.
        , ()
makeCursor     :: sh -> cursor

          -- | Shift the cursor by an offset, to get to another element.
        , ()
shiftCursor    :: sh -> cursor -> cursor

          -- | Load\/compute the element at the given cursor.
        , ()
loadCursor     :: cursor -> a }


 index :: forall sh. Shape sh => Array C sh a -> sh -> a
index (ACursored sh
_ sh -> cursor
makec sh -> cursor -> cursor
_ cursor -> a
loadc)
        = cursor -> a
loadc (cursor -> a) -> (sh -> cursor) -> sh -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. sh -> cursor
makec
 {-# INLINE index #-}

 unsafeIndex :: forall sh. Shape sh => Array C sh a -> sh -> a
unsafeIndex    = Array C sh a -> sh -> a
forall sh. Shape sh => Array C sh a -> sh -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
index
 {-# INLINE unsafeIndex #-}
 
 linearIndex :: forall sh. Shape sh => Array C sh a -> Int -> a
linearIndex (ACursored sh
sh sh -> cursor
makec sh -> cursor -> cursor
_ cursor -> a
loadc)
        = cursor -> a
loadc (cursor -> a) -> (Int -> cursor) -> Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. sh -> cursor
makec (sh -> cursor) -> (Int -> sh) -> Int -> cursor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. sh -> Int -> sh
forall sh. Shape sh => sh -> Int -> sh
fromIndex sh
sh
 {-# INLINE linearIndex #-}

 extent :: forall sh. Shape sh => Array C sh a -> sh
extent (ACursored sh
sh sh -> cursor
_ sh -> cursor -> cursor
_ cursor -> a
_)
        = sh
sh
 {-# INLINE extent #-}
        
 deepSeqArray :: forall sh b. Shape sh => Array C sh a -> b -> b
deepSeqArray (ACursored sh
sh sh -> cursor
makec sh -> cursor -> cursor
shiftc cursor -> a
loadc) b
y
  = sh
sh sh -> b -> b
forall a. sh -> a -> a
forall sh a. Shape sh => sh -> a -> a
`deepSeq` sh -> cursor
makec  (sh -> cursor) -> b -> b
forall a b. a -> b -> b
`seq` sh -> cursor -> cursor
shiftc (sh -> cursor -> cursor) -> b -> b
forall a b. a -> b -> b
`seq` cursor -> a
loadc (cursor -> a) -> b -> b
forall a b. a -> b -> b
`seq` b
y
 {-# INLINE deepSeqArray #-}


-- Fill -----------------------------------------------------------------------
-- | Compute all elements in an rank-2 array. 
instance Elt e => Load C DIM2 e where
 loadP :: forall r2. Target r2 e => Array C DIM2 e -> MVec r2 e -> IO ()
loadP (ACursored (DIM0
Z :. (I# Int#
h) :. (I# Int#
w)) DIM2 -> cursor
makec DIM2 -> cursor -> cursor
shiftc cursor -> e
loadc) MVec r2 e
marr
  = do  String -> IO ()
traceEventIO String
"Repa.loadP[Cursored]: start"
        (Int -> e -> IO ())
-> (DIM2 -> cursor)
-> (DIM2 -> cursor -> cursor)
-> (cursor -> e)
-> Int#
-> Int#
-> Int#
-> Int#
-> Int#
-> IO ()
forall a cursor.
Elt a =>
(Int -> a -> IO ())
-> (DIM2 -> cursor)
-> (DIM2 -> cursor -> cursor)
-> (cursor -> a)
-> Int#
-> Int#
-> Int#
-> Int#
-> Int#
-> IO ()
fillCursoredBlock2P 
                (MVec r2 e -> Int -> e -> IO ()
forall r e. Target r e => MVec r e -> Int -> e -> IO ()
unsafeWriteMVec MVec r2 e
marr) 
                DIM2 -> cursor
makec DIM2 -> cursor -> cursor
shiftc cursor -> e
loadc
                Int#
w Int#
0# Int#
0# Int#
w Int#
h
        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[Cursored]: end"
 {-# INLINE loadP #-}
        
 loadS :: forall r2. Target r2 e => Array C DIM2 e -> MVec r2 e -> IO ()
loadS (ACursored (DIM0
Z :. (I# Int#
h) :. (I# Int#
w)) DIM2 -> cursor
makec DIM2 -> cursor -> cursor
shiftc cursor -> e
loadc) MVec r2 e
marr
  = do  String -> IO ()
traceEventIO String
"Repa.loadS[Cursored]: start"
        (Int -> e -> IO ())
-> (DIM2 -> cursor)
-> (DIM2 -> cursor -> cursor)
-> (cursor -> e)
-> Int#
-> Int#
-> Int#
-> Int#
-> Int#
-> IO ()
forall a cursor.
Elt a =>
(Int -> a -> IO ())
-> (DIM2 -> cursor)
-> (DIM2 -> cursor -> cursor)
-> (cursor -> a)
-> Int#
-> Int#
-> Int#
-> Int#
-> Int#
-> IO ()
fillCursoredBlock2S 
                (MVec r2 e -> Int -> e -> IO ()
forall r e. Target r e => MVec r e -> Int -> e -> IO ()
unsafeWriteMVec MVec r2 e
marr) 
                DIM2 -> cursor
makec DIM2 -> cursor -> cursor
shiftc cursor -> e
loadc
                Int#
w Int#
0# Int#
0# Int#
w Int#
h
        MVec r2 e -> IO ()
forall r e. Target r e => MVec r e -> IO ()
touchMVec MVec r2 e
marr
        String -> IO ()
traceEventIO String
"Repa.loadS[Cursored]: end"
 {-# INLINE loadS #-}
        

-- | Compute a range of elements in a rank-2 array.
instance Elt e => LoadRange C DIM2 e where
 loadRangeP :: forall r2.
Target r2 e =>
Array C DIM2 e -> MVec r2 e -> DIM2 -> DIM2 -> IO ()
loadRangeP  (ACursored (DIM0
Z :. Int
_h :. (I# Int#
w)) DIM2 -> cursor
makec DIM2 -> cursor -> cursor
shiftc cursor -> e
loadc) MVec r2 e
marr
             (DIM0
Z :. (I# Int#
y0) :. (I# Int#
x0)) (DIM0
Z :. (I# Int#
h0) :. (I# Int#
w0))
  = do  String -> IO ()
traceEventIO String
"Repa.loadRangeP[Cursored]: start"
        (Int -> e -> IO ())
-> (DIM2 -> cursor)
-> (DIM2 -> cursor -> cursor)
-> (cursor -> e)
-> Int#
-> Int#
-> Int#
-> Int#
-> Int#
-> IO ()
forall a cursor.
Elt a =>
(Int -> a -> IO ())
-> (DIM2 -> cursor)
-> (DIM2 -> cursor -> cursor)
-> (cursor -> a)
-> Int#
-> Int#
-> Int#
-> Int#
-> Int#
-> IO ()
fillCursoredBlock2P 
                (MVec r2 e -> Int -> e -> IO ()
forall r e. Target r e => MVec r e -> Int -> e -> IO ()
unsafeWriteMVec MVec r2 e
marr) 
                DIM2 -> cursor
makec DIM2 -> cursor -> cursor
shiftc cursor -> e
loadc
                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
marr
        String -> IO ()
traceEventIO String
"Repa.loadRangeP[Cursored]: end"
 {-# INLINE loadRangeP #-}
        
 loadRangeS :: forall r2.
Target r2 e =>
Array C DIM2 e -> MVec r2 e -> DIM2 -> DIM2 -> IO ()
loadRangeS  (ACursored (DIM0
Z :. Int
_h :. (I# Int#
w)) DIM2 -> cursor
makec DIM2 -> cursor -> cursor
shiftc cursor -> e
loadc) MVec r2 e
marr
             (DIM0
Z :. (I# Int#
y0) :. (I# Int#
x0)) 
             (DIM0
Z :. (I# Int#
h0) :. (I# Int#
w0))
  = do  String -> IO ()
traceEventIO String
"Repa.loadRangeS[Cursored]: start"
        (Int -> e -> IO ())
-> (DIM2 -> cursor)
-> (DIM2 -> cursor -> cursor)
-> (cursor -> e)
-> Int#
-> Int#
-> Int#
-> Int#
-> Int#
-> IO ()
forall a cursor.
Elt a =>
(Int -> a -> IO ())
-> (DIM2 -> cursor)
-> (DIM2 -> cursor -> cursor)
-> (cursor -> a)
-> Int#
-> Int#
-> Int#
-> Int#
-> Int#
-> IO ()
fillCursoredBlock2S
                (MVec r2 e -> Int -> e -> IO ()
forall r e. Target r e => MVec r e -> Int -> e -> IO ()
unsafeWriteMVec MVec r2 e
marr) 
                DIM2 -> cursor
makec DIM2 -> cursor -> cursor
shiftc cursor -> e
loadc
                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
marr
        String -> IO ()
traceEventIO String
"Repa.loadRangeS[Cursored]: end"
 {-# INLINE loadRangeS #-}
        

-- Conversions ----------------------------------------------------------------
-- | Define a new cursored array.
makeCursored 
        :: sh
        -> (sh -> cursor)               -- ^ Create a cursor for an index.
        -> (sh -> cursor -> cursor)     -- ^ Shift a cursor by an offset.
        -> (cursor -> e)                -- ^ Compute the element at the cursor.
        -> Array C sh e

makeCursored :: forall sh cursor e.
sh
-> (sh -> cursor)
-> (sh -> cursor -> cursor)
-> (cursor -> e)
-> Array C sh e
makeCursored = sh
-> (sh -> cursor)
-> (sh -> cursor -> cursor)
-> (cursor -> e)
-> Array C sh e
forall sh a cursor.
sh
-> (sh -> cursor)
-> (sh -> cursor -> cursor)
-> (cursor -> a)
-> Array C sh a
ACursored
{-# INLINE makeCursored #-}