{-# LANGUAGE FunctionalDependencies, UndecidableInstances #-}

module Data.Array.Repa.Operators.Mapping
        ( -- * Generic maps
          map
        , zipWith
        , (+^), (-^), (*^), (/^)

          -- * Structured maps
        , Structured(..))
where
import Data.Array.Repa.Shape
import Data.Array.Repa.Base
import Data.Array.Repa.Repr.ByteString
import Data.Array.Repa.Repr.Cursored
import Data.Array.Repa.Repr.Delayed
import Data.Array.Repa.Repr.ForeignPtr
import Data.Array.Repa.Repr.HintSmall
import Data.Array.Repa.Repr.HintInterleave
import Data.Array.Repa.Repr.Partitioned
import Data.Array.Repa.Repr.Unboxed
import Data.Array.Repa.Repr.Undefined
import Prelude hiding (map, zipWith)
import Foreign.Storable
import Data.Word

-- | Apply a worker function to each element of an array, 
--   yielding a new array with the same extent.
--
map     :: (Shape sh, Source r a)
        => (a -> b) -> Array r sh a -> Array D sh b
map :: (a -> b) -> Array r sh a -> Array D sh b
map a -> b
f Array r sh a
arr
 = case Array r 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 r sh a
arr of
        ADelayed sh g -> sh -> (sh -> b) -> Array D sh b
forall sh a. sh -> (sh -> a) -> Array D sh a
ADelayed sh
sh (a -> b
f (a -> b) -> (sh -> a) -> sh -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. sh -> a
g)
{-# INLINE [3] map #-}


-- ZipWith --------------------------------------------------------------------
-- | Combine two arrays, element-wise, with a binary operator.
--      If the extent of the two array arguments differ,
--      then the resulting array's extent is their intersection.
--
zipWith :: (Shape sh, Source r1 a, Source r2 b)
        => (a -> b -> c)
        -> Array r1 sh a -> Array r2 sh b
        -> Array D sh c
zipWith :: (a -> b -> c) -> Array r1 sh a -> Array r2 sh b -> Array D sh c
zipWith a -> b -> c
f Array r1 sh a
arr1 Array r2 sh b
arr2
 = let  get :: sh -> c
get sh
ix  = a -> b -> c
f (Array r1 sh a
arr1 Array r1 sh a -> sh -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
`unsafeIndex` sh
ix) (Array r2 sh b
arr2 Array r2 sh b -> sh -> b
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
`unsafeIndex` sh
ix)
        {-# INLINE get #-}
        
   in   sh -> (sh -> c) -> Array D sh c
forall sh a. sh -> (sh -> a) -> Array D sh a
fromFunction 
                (sh -> sh -> sh
forall sh. Shape sh => sh -> sh -> sh
intersectDim (Array r1 sh a -> sh
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r1 sh a
arr1) (Array r2 sh b -> sh
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r2 sh b
arr2)) 
                sh -> c
get
{-# INLINE [2] zipWith #-}

infixl 7  *^, /^
infixl 6  +^, -^

+^ :: Array r1 sh c -> Array r2 sh c -> Array D sh c
(+^)    = (c -> c -> c) -> Array r1 sh c -> Array r2 sh c -> Array D sh c
forall sh r1 a r2 b c.
(Shape sh, Source r1 a, Source r2 b) =>
(a -> b -> c) -> Array r1 sh a -> Array r2 sh b -> Array D sh c
zipWith c -> c -> c
forall a. Num a => a -> a -> a
(+)
{-# INLINE (+^) #-}

-^ :: Array r1 sh c -> Array r2 sh c -> Array D sh c
(-^)    = (c -> c -> c) -> Array r1 sh c -> Array r2 sh c -> Array D sh c
forall sh r1 a r2 b c.
(Shape sh, Source r1 a, Source r2 b) =>
(a -> b -> c) -> Array r1 sh a -> Array r2 sh b -> Array D sh c
zipWith (-)
{-# INLINE (-^) #-}

*^ :: Array r1 sh c -> Array r2 sh c -> Array D sh c
(*^)    = (c -> c -> c) -> Array r1 sh c -> Array r2 sh c -> Array D sh c
forall sh r1 a r2 b c.
(Shape sh, Source r1 a, Source r2 b) =>
(a -> b -> c) -> Array r1 sh a -> Array r2 sh b -> Array D sh c
zipWith c -> c -> c
forall a. Num a => a -> a -> a
(*)
{-# INLINE (*^) #-}

/^ :: Array r1 sh c -> Array r2 sh c -> Array D sh c
(/^)    = (c -> c -> c) -> Array r1 sh c -> Array r2 sh c -> Array D sh c
forall sh r1 a r2 b c.
(Shape sh, Source r1 a, Source r2 b) =>
(a -> b -> c) -> Array r1 sh a -> Array r2 sh b -> Array D sh c
zipWith c -> c -> c
forall a. Fractional a => a -> a -> a
(/)
{-# INLINE (/^) #-}


-- Structured -------------------------------------------------------------------
-- | Structured versions of @map@ and @zipWith@ that preserve the representation
--   of cursored and partitioned arrays. 
--
--   For cursored (@C@) arrays, the cursoring of the source array is preserved.
--
--   For partitioned (@P@) arrays, the worker function is fused with each array
--   partition separately, instead of treating the whole array as a single
--   bulk object. 
--
--   Preserving the cursored and\/or paritioned representation of an array 
--   is will make follow-on computation more efficient than if the array was
--   converted to a vanilla Delayed (@D@) array as with plain `map` and `zipWith`.
--
--   If the source array is not cursored or partitioned then `smap` and 
--   `szipWith` are identical to the plain functions.
--
class Structured r1 a b where
 -- | The target result representation.
 type TR r1

 -- | Structured @map@.
 smap   :: Shape sh 
        => (a -> b) 
        -> Array r1     sh a 
        -> Array (TR r1) sh b

 -- | Structured @zipWith@.
 --   If you have a cursored or partitioned source array then use that as
 --   the third argument (corresponding to @r1@ here)
 szipWith
        :: (Shape sh, Source r c)
        => (c -> a -> b)
        -> Array r      sh c
        -> Array r1     sh a
        -> Array (TR r1) sh b


-- ByteString -------------------------
instance Structured B Word8 b where
 type TR B = D
 smap :: (Word8 -> b) -> Array B sh Word8 -> Array (TR B) sh b
smap           = (Word8 -> b) -> Array B sh Word8 -> Array (TR B) sh b
forall sh r a b.
(Shape sh, Source r a) =>
(a -> b) -> Array r sh a -> Array D sh b
map
 szipWith :: (c -> Word8 -> b)
-> Array r sh c -> Array B sh Word8 -> Array (TR B) sh b
szipWith       = (c -> Word8 -> b)
-> Array r sh c -> Array B sh Word8 -> Array (TR B) sh b
forall sh r1 a r2 b c.
(Shape sh, Source r1 a, Source r2 b) =>
(a -> b -> c) -> Array r1 sh a -> Array r2 sh b -> Array D sh c
zipWith


-- Cursored ---------------------------
instance Structured C a b where
 type TR C = C

 smap :: (a -> b) -> Array C sh a -> Array (TR C) sh b
smap a -> b
f (ACursored sh makec shiftc loadc)
        = sh
-> (sh -> cursor)
-> (sh -> cursor -> cursor)
-> (cursor -> b)
-> Array C sh b
forall sh a cursor.
sh
-> (sh -> cursor)
-> (sh -> cursor -> cursor)
-> (cursor -> a)
-> Array C sh a
ACursored sh
sh sh -> cursor
makec sh -> cursor -> cursor
shiftc (a -> b
f (a -> b) -> (cursor -> a) -> cursor -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. cursor -> a
loadc)
 {-# INLINE [3] smap #-}

 szipWith :: (c -> a -> b) -> Array r sh c -> Array C sh a -> Array (TR C) sh b
szipWith c -> a -> b
f Array r sh c
arr1 (ACursored sh makec shiftc loadc)
  = let makec' :: sh -> (sh, cursor)
makec' sh
ix               = (sh
ix, sh -> cursor
makec sh
ix)
        {-# INLINE makec' #-}
        
        shiftc' :: sh -> (sh, cursor) -> (sh, cursor)
shiftc' sh
off (sh
ix, cursor
cur)   = (sh -> sh -> sh
forall sh. Shape sh => sh -> sh -> sh
addDim sh
off sh
ix, sh -> cursor -> cursor
shiftc sh
off cursor
cur)
        {-# INLINE shiftc' #-}

        load' :: (sh, cursor) -> b
load' (sh
ix, cursor
cur)         = c -> a -> b
f (Array r sh c
arr1 Array r sh c -> sh -> c
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
`unsafeIndex` sh
ix) (cursor -> a
loadc cursor
cur)
        {-# INLINE load' #-}

    in  sh
-> (sh -> (sh, cursor))
-> (sh -> (sh, cursor) -> (sh, cursor))
-> ((sh, cursor) -> b)
-> Array C sh b
forall sh a cursor.
sh
-> (sh -> cursor)
-> (sh -> cursor -> cursor)
-> (cursor -> a)
-> Array C sh a
ACursored 
                (sh -> sh -> sh
forall sh. Shape sh => sh -> sh -> sh
intersectDim (Array r sh c -> sh
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r sh c
arr1) sh
sh)
                sh -> (sh, cursor)
makec' sh -> (sh, cursor) -> (sh, cursor)
shiftc' (sh, cursor) -> b
load'
 {-# INLINE [2] szipWith #-}


-- Delayed ----------------------------
instance Structured D a b where
 type TR D = D
 smap :: (a -> b) -> Array D sh a -> Array (TR D) sh b
smap           = (a -> b) -> Array D sh a -> Array (TR D) sh b
forall sh r a b.
(Shape sh, Source r a) =>
(a -> b) -> Array r sh a -> Array D sh b
map
 szipWith :: (c -> a -> b) -> Array r sh c -> Array D sh a -> Array (TR D) sh b
szipWith       = (c -> a -> b) -> Array r sh c -> Array D sh a -> Array (TR D) sh b
forall sh r1 a r2 b c.
(Shape sh, Source r1 a, Source r2 b) =>
(a -> b -> c) -> Array r1 sh a -> Array r2 sh b -> Array D sh c
zipWith


-- ForeignPtr -------------------------
instance Storable a => Structured F a b where
 type TR F = D
 smap :: (a -> b) -> Array F sh a -> Array (TR F) sh b
smap           = (a -> b) -> Array F sh a -> Array (TR F) sh b
forall sh r a b.
(Shape sh, Source r a) =>
(a -> b) -> Array r sh a -> Array D sh b
map
 szipWith :: (c -> a -> b) -> Array r sh c -> Array F sh a -> Array (TR F) sh b
szipWith       = (c -> a -> b) -> Array r sh c -> Array F sh a -> Array (TR F) sh b
forall sh r1 a r2 b c.
(Shape sh, Source r1 a, Source r2 b) =>
(a -> b -> c) -> Array r1 sh a -> Array r2 sh b -> Array D sh c
zipWith


-- Partitioned ------------------------
instance (Structured r1 a b
        , Structured r2 a b)
       => Structured (P r1 r2) a b where
 type TR (P r1 r2) = P (TR r1) (TR r2)

 smap :: (a -> b) -> Array (P r1 r2) sh a -> Array (TR (P r1 r2)) sh b
smap a -> b
f (APart sh range arr1 arr2)
        = sh
-> Range sh
-> Array (TR r1) sh b
-> Array (TR r2) sh b
-> Array (P (TR r1) (TR r2)) sh b
forall r1 r2 sh e.
sh
-> Range sh
-> Array r1 sh e
-> Array r2 sh e
-> Array (P r1 r2) sh e
APart sh
sh Range sh
range ((a -> b) -> Array r1 sh a -> Array (TR r1) sh b
forall r1 a b sh.
(Structured r1 a b, Shape sh) =>
(a -> b) -> Array r1 sh a -> Array (TR r1) sh b
smap a -> b
f Array r1 sh a
arr1) ((a -> b) -> Array r2 sh a -> Array (TR r2) sh b
forall r1 a b sh.
(Structured r1 a b, Shape sh) =>
(a -> b) -> Array r1 sh a -> Array (TR r1) sh b
smap a -> b
f Array r2 sh a
arr2)
 {-# INLINE [3] smap #-}

 szipWith :: (c -> a -> b)
-> Array r sh c
-> Array (P r1 r2) sh a
-> Array (TR (P r1 r2)) sh b
szipWith c -> a -> b
f Array r sh c
arr1 (APart sh range arr21 arr22)
        = sh
-> Range sh
-> Array (TR r1) sh b
-> Array (TR r2) sh b
-> Array (P (TR r1) (TR r2)) sh b
forall r1 r2 sh e.
sh
-> Range sh
-> Array r1 sh e
-> Array r2 sh e
-> Array (P r1 r2) sh e
APart sh
sh Range sh
range ((c -> a -> b)
-> Array r sh c -> Array r1 sh a -> Array (TR r1) sh b
forall r1 a b sh r c.
(Structured r1 a b, Shape sh, Source r c) =>
(c -> a -> b)
-> Array r sh c -> Array r1 sh a -> Array (TR r1) sh b
szipWith c -> a -> b
f Array r sh c
arr1 Array r1 sh a
arr21)
                         ((c -> a -> b)
-> Array r sh c -> Array r2 sh a -> Array (TR r2) sh b
forall r1 a b sh r c.
(Structured r1 a b, Shape sh, Source r c) =>
(c -> a -> b)
-> Array r sh c -> Array r1 sh a -> Array (TR r1) sh b
szipWith c -> a -> b
f Array r sh c
arr1 Array r2 sh a
arr22)
 {-# INLINE [2] szipWith #-}


-- Small ------------------------------
instance   Structured r1 a b
        => Structured (S r1) a b where
 type TR (S r1) = S (TR r1)

 smap :: (a -> b) -> Array (S r1) sh a -> Array (TR (S r1)) sh b
smap a -> b
f (ASmall arr1)
        = Array (TR r1) sh b -> Array (S (TR r1)) sh b
forall r1 sh a. Array r1 sh a -> Array (S r1) sh a
ASmall ((a -> b) -> Array r1 sh a -> Array (TR r1) sh b
forall r1 a b sh.
(Structured r1 a b, Shape sh) =>
(a -> b) -> Array r1 sh a -> Array (TR r1) sh b
smap a -> b
f Array r1 sh a
arr1)
 {-# INLINE [3] smap #-}

 szipWith :: (c -> a -> b)
-> Array r sh c -> Array (S r1) sh a -> Array (TR (S r1)) sh b
szipWith c -> a -> b
f Array r sh c
arr1 (ASmall arr2)
        = Array (TR r1) sh b -> Array (S (TR r1)) sh b
forall r1 sh a. Array r1 sh a -> Array (S r1) sh a
ASmall ((c -> a -> b)
-> Array r sh c -> Array r1 sh a -> Array (TR r1) sh b
forall r1 a b sh r c.
(Structured r1 a b, Shape sh, Source r c) =>
(c -> a -> b)
-> Array r sh c -> Array r1 sh a -> Array (TR r1) sh b
szipWith c -> a -> b
f Array r sh c
arr1 Array r1 sh a
arr2)
 {-# INLINE [3] szipWith #-}


-- Interleaved ------------------------
instance   Structured r1 a b
        => Structured (I r1) a b where
 type TR (I r1) = I (TR r1)

 smap :: (a -> b) -> Array (I r1) sh a -> Array (TR (I r1)) sh b
smap a -> b
f (AInterleave arr1)
        = Array (TR r1) sh b -> Array (I (TR r1)) sh b
forall r1 sh a. Array r1 sh a -> Array (I r1) sh a
AInterleave ((a -> b) -> Array r1 sh a -> Array (TR r1) sh b
forall r1 a b sh.
(Structured r1 a b, Shape sh) =>
(a -> b) -> Array r1 sh a -> Array (TR r1) sh b
smap a -> b
f Array r1 sh a
arr1)
 {-# INLINE [3] smap #-}

 szipWith :: (c -> a -> b)
-> Array r sh c -> Array (I r1) sh a -> Array (TR (I r1)) sh b
szipWith c -> a -> b
f Array r sh c
arr1 (AInterleave arr2)
        = Array (TR r1) sh b -> Array (I (TR r1)) sh b
forall r1 sh a. Array r1 sh a -> Array (I r1) sh a
AInterleave ((c -> a -> b)
-> Array r sh c -> Array r1 sh a -> Array (TR r1) sh b
forall r1 a b sh r c.
(Structured r1 a b, Shape sh, Source r c) =>
(c -> a -> b)
-> Array r sh c -> Array r1 sh a -> Array (TR r1) sh b
szipWith c -> a -> b
f Array r sh c
arr1 Array r1 sh a
arr2)
 {-# INLINE [3] szipWith #-}


-- Unboxed ----------------------------
instance Unbox a => Structured U a b where
 type TR U = D
 smap :: (a -> b) -> Array U sh a -> Array (TR U) sh b
smap           = (a -> b) -> Array U sh a -> Array (TR U) sh b
forall sh r a b.
(Shape sh, Source r a) =>
(a -> b) -> Array r sh a -> Array D sh b
map
 szipWith :: (c -> a -> b) -> Array r sh c -> Array U sh a -> Array (TR U) sh b
szipWith       = (c -> a -> b) -> Array r sh c -> Array U sh a -> Array (TR U) sh b
forall sh r1 a r2 b c.
(Shape sh, Source r1 a, Source r2 b) =>
(a -> b -> c) -> Array r1 sh a -> Array r2 sh b -> Array D sh c
zipWith


-- Undefined --------------------------
instance Structured X a b where
 type TR X = X
 smap :: (a -> b) -> Array X sh a -> Array (TR X) sh b
smap     a -> b
_   (AUndefined sh) = sh -> Array X sh b
forall sh e. sh -> Array X sh e
AUndefined sh
sh
 szipWith :: (c -> a -> b) -> Array r sh c -> Array X sh a -> Array (TR X) sh b
szipWith c -> a -> b
_ Array r sh c
_ (AUndefined sh) = sh -> Array X sh b
forall sh e. sh -> Array X sh e
AUndefined sh
sh