module Data.Array.Repa.Base ( Array , Repr (..), (!), toList , deepSeqArrays) where import Data.Array.Repa.Shape -- | 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 family Array r sh e -- Repr ----------------------------------------------------------------------- -- | Class of array representations that we can read elements from. -- class Repr r e where -- | 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 arr ix = arr `linearIndex` toIndex (extent arr) ix {-# INLINE unsafeIndex #-} unsafeIndex arr ix = arr `unsafeLinearIndex` toIndex (extent arr) ix -- | O(1). Linear indexing into underlying, row-major, array representation. linearIndex, unsafeLinearIndex :: Shape sh => Array r sh e -> Int -> e {-# INLINE unsafeLinearIndex #-} unsafeLinearIndex = linearIndex -- | Ensure an array's data structure is fully evaluated. deepSeqArray :: Shape sh => Array r sh e -> b -> b -- | O(1). Alias for `index` (!) :: (Repr r e, Shape sh) => Array r sh e -> sh -> e (!) = index -- | O(n). Convert an array to a list. toList :: (Shape sh, Repr r e) => Array r sh e -> [e] {-# INLINE toList #-} toList arr = go 0 where len = size (extent arr) go ix | ix == len = [] | otherwise = unsafeLinearIndex arr ix : go (ix + 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, Repr r e) => [Array r sh e] -> b -> b {-# INLINE deepSeqArrays #-} deepSeqArrays arrs x = case arrs of [] -> x [a1] -> a1 `deepSeqArray` x [a1, a2] -> a1 `deepSeqArray` a2 `deepSeqArray` x [a1, a2, a3] -> a1 `deepSeqArray` a2 `deepSeqArray` a3 `deepSeqArray` x [a1, a2, a3, a4] -> a1 `deepSeqArray` a2 `deepSeqArray` a3 `deepSeqArray` a4 `deepSeqArray` x _ -> error "deepSeqArrays: only works for up to four arrays"