-- Generic Traversal
module Data.Array.Repa.Operators.Traversal
        ( traverse,  unsafeTraverse
        , traverse2, unsafeTraverse2
        , traverse3, unsafeTraverse3
        , traverse4, unsafeTraverse4)
where
import Data.Array.Repa.Base
import Data.Array.Repa.Shape
import Data.Array.Repa.Repr.Delayed
import Prelude hiding (traverse)


-- | Unstructured traversal.
traverse, unsafeTraverse
        :: forall r sh sh' a b
        .  ( Source r a
           , Shape  sh)
        => Array r sh a                 -- ^ Source array.
        -> (sh  -> sh')                 -- ^ Function to produce the extent of the result.
        -> ((sh -> a) -> sh' -> b)      -- ^ Function to produce elements of the result.
                                        --   It is passed a lookup function to get elements of the source.
        -> Array D sh' b

traverse :: Array r sh a
-> (sh -> sh') -> ((sh -> a) -> sh' -> b) -> Array D sh' b
traverse Array r sh a
arr sh -> sh'
transExtent (sh -> a) -> sh' -> b
newElem
 = sh' -> (sh' -> b) -> Array D sh' b
forall sh a. sh -> (sh -> a) -> Array D sh a
fromFunction (sh -> sh'
transExtent (Array r sh a -> sh
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r sh a
arr)) ((sh -> a) -> sh' -> b
newElem (Array r sh a -> sh -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
index Array r sh a
arr))
{-# INLINE [3] traverse #-}

unsafeTraverse :: Array r sh a
-> (sh -> sh') -> ((sh -> a) -> sh' -> b) -> Array D sh' b
unsafeTraverse Array r sh a
arr sh -> sh'
transExtent (sh -> a) -> sh' -> b
newElem
 = sh' -> (sh' -> b) -> Array D sh' b
forall sh a. sh -> (sh -> a) -> Array D sh a
fromFunction (sh -> sh'
transExtent (Array r sh a -> sh
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r sh a
arr)) ((sh -> a) -> sh' -> b
newElem (Array r sh a -> sh -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
unsafeIndex Array r sh a
arr))
{-# INLINE [3] unsafeTraverse #-}


-- | Unstructured traversal over two arrays at once.
traverse2, unsafeTraverse2
        :: forall r1 r2 sh sh' sh'' a b c
        .  ( Source r1 a, Source r2 b
           , Shape sh, Shape sh')
        => Array r1 sh  a               -- ^ First source array.
        -> Array r2 sh' b               -- ^ Second source array.
        -> (sh -> sh' -> sh'')          -- ^ Function to produce the extent of the result.
        -> ((sh -> a) -> (sh' -> b)
                      -> (sh'' -> c))   -- ^ Function to produce elements of the result.
                                        --   It is passed lookup functions to get elements of the
                                        --   source arrays.
        -> Array D sh'' c

traverse2 :: Array r1 sh a
-> Array r2 sh' b
-> (sh -> sh' -> sh'')
-> ((sh -> a) -> (sh' -> b) -> sh'' -> c)
-> Array D sh'' c
traverse2 Array r1 sh a
arrA Array r2 sh' b
arrB sh -> sh' -> sh''
transExtent (sh -> a) -> (sh' -> b) -> sh'' -> c
newElem
 = sh'' -> (sh'' -> c) -> Array D sh'' c
forall sh a. sh -> (sh -> a) -> Array D sh a
fromFunction  (sh -> sh' -> sh''
transExtent (Array r1 sh a -> sh
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r1 sh a
arrA) (Array r2 sh' b -> sh'
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r2 sh' b
arrB))
                 ((sh -> a) -> (sh' -> b) -> sh'' -> c
newElem     (Array r1 sh a -> sh -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
index  Array r1 sh a
arrA) (Array r2 sh' b -> sh' -> b
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
index  Array r2 sh' b
arrB))
{-# INLINE [3] traverse2 #-}

unsafeTraverse2 :: Array r1 sh a
-> Array r2 sh' b
-> (sh -> sh' -> sh'')
-> ((sh -> a) -> (sh' -> b) -> sh'' -> c)
-> Array D sh'' c
unsafeTraverse2 Array r1 sh a
arrA Array r2 sh' b
arrB sh -> sh' -> sh''
transExtent (sh -> a) -> (sh' -> b) -> sh'' -> c
newElem
 = sh'' -> (sh'' -> c) -> Array D sh'' c
forall sh a. sh -> (sh -> a) -> Array D sh a
fromFunction  (sh -> sh' -> sh''
transExtent (Array r1 sh a -> sh
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r1 sh a
arrA) (Array r2 sh' b -> sh'
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r2 sh' b
arrB))
                 ((sh -> a) -> (sh' -> b) -> sh'' -> c
newElem     (Array r1 sh a -> sh -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
unsafeIndex Array r1 sh a
arrA) (Array r2 sh' b -> sh' -> b
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
unsafeIndex Array r2 sh' b
arrB))
{-# INLINE [3] unsafeTraverse2 #-}


-- | Unstructured traversal over three arrays at once.
traverse3, unsafeTraverse3
        :: forall r1  r2  r3
                  sh1 sh2 sh3 sh4
                  a   b   c   d
        .  ( Source r1 a, Source r2 b, Source r3 c
           , Shape sh1,   Shape sh2,   Shape sh3)
        => Array r1 sh1 a
        -> Array r2 sh2 b
        -> Array r3 sh3 c
        -> (sh1 -> sh2 -> sh3 -> sh4)
        -> (  (sh1 -> a) -> (sh2 -> b)
           -> (sh3 -> c)
           ->  sh4 -> d )
        -> Array D sh4 d

traverse3 :: Array r1 sh1 a
-> Array r2 sh2 b
-> Array r3 sh3 c
-> (sh1 -> sh2 -> sh3 -> sh4)
-> ((sh1 -> a) -> (sh2 -> b) -> (sh3 -> c) -> sh4 -> d)
-> Array D sh4 d
traverse3 Array r1 sh1 a
arrA Array r2 sh2 b
arrB Array r3 sh3 c
arrC sh1 -> sh2 -> sh3 -> sh4
transExtent (sh1 -> a) -> (sh2 -> b) -> (sh3 -> c) -> sh4 -> d
newElem
 = sh4 -> (sh4 -> d) -> Array D sh4 d
forall sh a. sh -> (sh -> a) -> Array D sh a
fromFunction (sh1 -> sh2 -> sh3 -> sh4
transExtent (Array r1 sh1 a -> sh1
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r1 sh1 a
arrA) (Array r2 sh2 b -> sh2
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r2 sh2 b
arrB) (Array r3 sh3 c -> sh3
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r3 sh3 c
arrC))
                ((sh1 -> a) -> (sh2 -> b) -> (sh3 -> c) -> sh4 -> d
newElem     (Array r1 sh1 a -> sh1 -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
index Array r1 sh1 a
arrA)  (Array r2 sh2 b -> sh2 -> b
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
index Array r2 sh2 b
arrB)  (Array r3 sh3 c -> sh3 -> c
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
index  Array r3 sh3 c
arrC))
{-# INLINE [3] traverse3 #-}

unsafeTraverse3 :: Array r1 sh1 a
-> Array r2 sh2 b
-> Array r3 sh3 c
-> (sh1 -> sh2 -> sh3 -> sh4)
-> ((sh1 -> a) -> (sh2 -> b) -> (sh3 -> c) -> sh4 -> d)
-> Array D sh4 d
unsafeTraverse3 Array r1 sh1 a
arrA Array r2 sh2 b
arrB Array r3 sh3 c
arrC sh1 -> sh2 -> sh3 -> sh4
transExtent (sh1 -> a) -> (sh2 -> b) -> (sh3 -> c) -> sh4 -> d
newElem
 = sh4 -> (sh4 -> d) -> Array D sh4 d
forall sh a. sh -> (sh -> a) -> Array D sh a
fromFunction (sh1 -> sh2 -> sh3 -> sh4
transExtent (Array r1 sh1 a -> sh1
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r1 sh1 a
arrA) (Array r2 sh2 b -> sh2
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r2 sh2 b
arrB) (Array r3 sh3 c -> sh3
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r3 sh3 c
arrC))
                ((sh1 -> a) -> (sh2 -> b) -> (sh3 -> c) -> sh4 -> d
newElem     (Array r1 sh1 a -> sh1 -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
unsafeIndex Array r1 sh1 a
arrA) (Array r2 sh2 b -> sh2 -> b
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
unsafeIndex Array r2 sh2 b
arrB) (Array r3 sh3 c -> sh3 -> c
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
unsafeIndex Array r3 sh3 c
arrC))
{-# INLINE [3] unsafeTraverse3 #-}


-- | Unstructured traversal over four arrays at once.
traverse4, unsafeTraverse4
        :: forall r1  r2  r3  r4
                  sh1 sh2 sh3 sh4 sh5
                  a   b   c   d   e
        .  ( Source r1 a, Source r2 b, Source r3 c, Source r4 d
           , Shape sh1, Shape sh2, Shape sh3, Shape sh4)
        => Array r1 sh1 a
        -> Array r2 sh2 b
        -> Array r3 sh3 c
        -> Array r4 sh4 d
        -> (sh1 -> sh2 -> sh3 -> sh4 -> sh5 )
        -> (  (sh1 -> a) -> (sh2 -> b)
           -> (sh3 -> c) -> (sh4 -> d)
           ->  sh5 -> e )
        -> Array D sh5 e

traverse4 :: Array r1 sh1 a
-> Array r2 sh2 b
-> Array r3 sh3 c
-> Array r4 sh4 d
-> (sh1 -> sh2 -> sh3 -> sh4 -> sh5)
-> ((sh1 -> a)
    -> (sh2 -> b) -> (sh3 -> c) -> (sh4 -> d) -> sh5 -> e)
-> Array D sh5 e
traverse4 Array r1 sh1 a
arrA Array r2 sh2 b
arrB Array r3 sh3 c
arrC Array r4 sh4 d
arrD sh1 -> sh2 -> sh3 -> sh4 -> sh5
transExtent (sh1 -> a) -> (sh2 -> b) -> (sh3 -> c) -> (sh4 -> d) -> sh5 -> e
newElem
 = sh5 -> (sh5 -> e) -> Array D sh5 e
forall sh a. sh -> (sh -> a) -> Array D sh a
fromFunction (sh1 -> sh2 -> sh3 -> sh4 -> sh5
transExtent (Array r1 sh1 a -> sh1
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r1 sh1 a
arrA) (Array r2 sh2 b -> sh2
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r2 sh2 b
arrB) (Array r3 sh3 c -> sh3
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r3 sh3 c
arrC) (Array r4 sh4 d -> sh4
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r4 sh4 d
arrD))
                ((sh1 -> a) -> (sh2 -> b) -> (sh3 -> c) -> (sh4 -> d) -> sh5 -> e
newElem     (Array r1 sh1 a -> sh1 -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
index  Array r1 sh1 a
arrA) (Array r2 sh2 b -> sh2 -> b
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
index  Array r2 sh2 b
arrB) (Array r3 sh3 c -> sh3 -> c
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
index  Array r3 sh3 c
arrC) (Array r4 sh4 d -> sh4 -> d
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
index  Array r4 sh4 d
arrD))
{-# INLINE [3] traverse4 #-}


unsafeTraverse4 :: Array r1 sh1 a
-> Array r2 sh2 b
-> Array r3 sh3 c
-> Array r4 sh4 d
-> (sh1 -> sh2 -> sh3 -> sh4 -> sh5)
-> ((sh1 -> a)
    -> (sh2 -> b) -> (sh3 -> c) -> (sh4 -> d) -> sh5 -> e)
-> Array D sh5 e
unsafeTraverse4 Array r1 sh1 a
arrA Array r2 sh2 b
arrB Array r3 sh3 c
arrC Array r4 sh4 d
arrD sh1 -> sh2 -> sh3 -> sh4 -> sh5
transExtent (sh1 -> a) -> (sh2 -> b) -> (sh3 -> c) -> (sh4 -> d) -> sh5 -> e
newElem
 = sh5 -> (sh5 -> e) -> Array D sh5 e
forall sh a. sh -> (sh -> a) -> Array D sh a
fromFunction (sh1 -> sh2 -> sh3 -> sh4 -> sh5
transExtent (Array r1 sh1 a -> sh1
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r1 sh1 a
arrA) (Array r2 sh2 b -> sh2
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r2 sh2 b
arrB) (Array r3 sh3 c -> sh3
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r3 sh3 c
arrC) (Array r4 sh4 d -> sh4
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r4 sh4 d
arrD))
                ((sh1 -> a) -> (sh2 -> b) -> (sh3 -> c) -> (sh4 -> d) -> sh5 -> e
newElem     (Array r1 sh1 a -> sh1 -> a
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
unsafeIndex Array r1 sh1 a
arrA) (Array r2 sh2 b -> sh2 -> b
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
unsafeIndex Array r2 sh2 b
arrB) (Array r3 sh3 c -> sh3 -> c
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
unsafeIndex Array r3 sh3 c
arrC) (Array r4 sh4 d -> sh4 -> d
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
unsafeIndex Array r4 sh4 d
arrD))
{-# INLINE [3] unsafeTraverse4 #-}