module Data.Array.Repa.Base
        ( Source (..), (!), toList
        , deepSeqArrays)
where
import Data.Array.Repa.Shape


-- Source -----------------------------------------------------------------------
-- | Class of array representations that we can read elements from.
class Source r e where
 -- Arrays with a representation tag, shape, and element type.
 --   Use one of the type tags like `D`, `U` and so on for @r@, 
 --   one of `DIM1`, `DIM2` ... for @sh@.
 data Array r sh e

 -- | O(1). Take the extent (size) of an array.
 extent :: Shape sh => Array r sh e -> sh

 -- | O(1). Shape polymorphic indexing.
 index, unsafeIndex 
        :: Shape sh => Array r sh e -> sh -> e

 {-# INLINE index #-}
 index Array r sh e
arr sh
ix           = Array r sh e
arr Array r sh e -> Int -> e
forall r e sh. (Source r e, Shape sh) => Array r sh e -> Int -> e
`linearIndex`       sh -> sh -> Int
forall sh. Shape sh => sh -> sh -> Int
toIndex (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) sh
ix

 {-# INLINE unsafeIndex #-}
 unsafeIndex Array r sh e
arr sh
ix     = Array r sh e
arr Array r sh e -> Int -> e
forall r e sh. (Source r e, Shape sh) => Array r sh e -> Int -> e
`unsafeLinearIndex` sh -> sh -> Int
forall sh. Shape sh => sh -> sh -> Int
toIndex (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) sh
ix

 -- | O(1). Linear indexing into underlying, row-major, array representation.
 linearIndex, unsafeLinearIndex
        :: Shape sh => Array r sh e -> Int -> e

 {-# INLINE unsafeLinearIndex #-}
 unsafeLinearIndex      = Array r sh e -> Int -> e
forall r e sh. (Source r e, Shape sh) => Array r sh e -> Int -> e
linearIndex

 -- | Ensure an array's data structure is fully evaluated.
 deepSeqArray 
        :: Shape sh =>Array r sh e -> b -> b


-- | O(1). Alias for `index`
(!) :: Shape sh => Source r e => Array r sh e -> sh -> e
(!) = Array r sh e -> sh -> e
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
index


-- | O(n). Convert an array to a list.
toList  :: Shape sh => Source r e
        => Array r sh e -> [e]
{-# INLINE toList #-}
toList :: Array r sh e -> [e]
toList Array r sh e
arr 
 = Int -> [e]
go Int
0 
 where  len :: Int
len     = sh -> Int
forall sh. Shape sh => sh -> Int
size (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)
        go :: Int -> [e]
go Int
ix
         | Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len    = []
         | Bool
otherwise    = Array r sh e -> Int -> e
forall r e sh. (Source r e, Shape sh) => Array r sh e -> Int -> e
unsafeLinearIndex Array r sh e
arr Int
ix e -> [e] -> [e]
forall a. a -> [a] -> [a]
: Int -> [e]
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)


-- | Apply `deepSeqArray` to up to four arrays. 
---
--   NOTE: this shouldn't be needed anymore, as we've made all the shape fields strict.
--      
--   The implementation of this function has been hand-unwound to work for up to
--   four arrays. Putting more in the list yields `error`.
-- 
--   For functions that are /not/ marked as INLINE, you should apply `deepSeqArrays`
--   to argument arrays before using them in a @compute@ or @copy@ expression.
--   For example:
--
-- @  processArrays 
--     :: Monad m 
--     => Array U DIM2 Int -> Array U DIM2 Int 
--     -> m (Array U DIM2 Int)
--  processArrays arr1 arr2
--   = [arr1, arr2] \`deepSeqArrays\` 
--     do arr3 <- computeP $ map f arr1
--        arr4 <- computeP $ zipWith g arr3 arr2
--        return arr4
--  @
--
--  Applying `deepSeqArrays` tells the GHC simplifier that it's ok to unbox 
--  size fields and the pointers to the underlying array data at the start
--  of the function. Without this, they may be unboxed repeatedly when
--  computing elements in the result arrays, which will make your program slow.
--
--  If you INLINE @processArrays@ into the function that computes @arr1@ and @arr2@,
--  then you don't need to apply `deepSeqArrays`. This is because a pointer
--  to the underlying data will be passed directly to the consumers and never boxed.
--
--  If you're not sure, then just follow the example code above.
--   
deepSeqArrays 
        :: Shape sh => Source r e
        => [Array r sh e] -> b -> b
{-# INLINE deepSeqArrays #-}
deepSeqArrays :: [Array r sh e] -> b -> b
deepSeqArrays [Array r sh e]
arrs b
x
 = case [Array r sh e]
arrs of
        []              -> b
x

        [Array r sh e
a1]
         -> Array r sh e
a1 Array r sh e -> b -> b
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray` b
x

        [Array r sh e
a1, Array r sh e
a2]
         -> Array r sh e
a1 Array r sh e -> Array r sh e -> Array r sh e
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray` Array r sh e
a2 Array r sh e -> b -> b
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray` b
x

        [Array r sh e
a1, Array r sh e
a2, Array r sh e
a3]
         -> Array r sh e
a1 Array r sh e -> Array r sh e -> Array r sh e
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray` Array r sh e
a2 Array r sh e -> Array r sh e -> Array r sh e
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray` Array r sh e
a3 Array r sh e -> b -> b
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray` b
x

        [Array r sh e
a1, Array r sh e
a2, Array r sh e
a3, Array r sh e
a4]
         -> Array r sh e
a1 Array r sh e -> Array r sh e -> Array r sh e
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray` Array r sh e
a2 Array r sh e -> Array r sh e -> Array r sh e
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray` Array r sh e
a3 Array r sh e -> Array r sh e -> Array r sh e
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray` Array r sh e
a4 Array r sh e -> b -> b
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray` b
x

        [Array r sh e]
_ -> [Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"deepSeqArrays: only works for up to four arrays"