    {-# 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 f :: a -> b
    f arr :: 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 f :: a -> b -> c
    f arr1 :: Array r1 sh a
    arr1 arr2 :: Array r2 sh b
    arr2
     = let  get :: sh -> c
    get ix :: 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 f :: 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 f :: c -> a -> b
    f arr1 :: Array r sh c
    arr1 (ACursored sh makec shiftc loadc)
      = let makec' :: sh -> (sh, cursor)
    makec' ix :: sh
    ix               = (sh
    ix, sh -> cursor
    makec sh
    ix)
            {-# INLINE makec' #-}

            shiftc' :: sh -> (sh, cursor) -> (sh, cursor)
    shiftc' off :: sh
    off (ix :: sh
    ix, cur :: 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' (ix :: sh
    ix, cur :: 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 f :: 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 f :: c -> a -> b
    f arr1 :: 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 f :: 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 f :: c -> a -> b
    f arr1 :: 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 f :: 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 f :: c -> a -> b
    f arr1 :: 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     _   (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 _ _ (AUndefined sh) = sh -> Array X sh b
    forall sh e. sh -> Array X sh e
    AUndefined sh
    sh