-- 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


-- | Unstructured traversal.
traverse, unsafeTraverse
	:: forall r sh sh' a b
	.  ( Source r a
           , Shape sh, 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 arr transExtent newElem
 = fromFunction (transExtent (extent arr)) (newElem (index arr))
{-# INLINE [3] traverse #-}

unsafeTraverse arr transExtent newElem
 = fromFunction (transExtent (extent arr)) (newElem (unsafeIndex 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', 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 arrA arrB transExtent newElem
 = fromFunction  (transExtent (extent arrA) (extent arrB))
 	         (newElem     (index  arrA) (index  arrB))
{-# INLINE [3] traverse2 #-}

unsafeTraverse2 arrA arrB transExtent newElem
 = fromFunction  (transExtent (extent arrA) (extent arrB))
                 (newElem     (unsafeIndex arrA) (unsafeIndex 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,   Shape sh4)
        => 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 arrA arrB arrC transExtent newElem
 = fromFunction (transExtent (extent arrA) (extent arrB) (extent arrC))
 	        (newElem     (index arrA)  (index arrB)  (index  arrC))
{-# INLINE [3] traverse3 #-}

unsafeTraverse3 arrA arrB arrC transExtent newElem
 = fromFunction	(transExtent (extent arrA) (extent arrB) (extent arrC))
	        (newElem     (unsafeIndex arrA) (unsafeIndex arrB) (unsafeIndex 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, Shape sh5)
        => 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 arrA arrB arrC arrD transExtent newElem
 = fromFunction	(transExtent (extent arrA) (extent arrB) (extent arrC) (extent arrD))
		(newElem     (index  arrA) (index  arrB) (index  arrC) (index  arrD))
{-# INLINE [3] traverse4 #-}


unsafeTraverse4 arrA arrB arrC arrD transExtent newElem
 = fromFunction (transExtent (extent arrA) (extent arrB) (extent arrC) (extent arrD))
		(newElem     (unsafeIndex arrA) (unsafeIndex arrB) (unsafeIndex arrC) (unsafeIndex arrD))
{-# INLINE [3] unsafeTraverse4 #-}