1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
|
-- 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 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')
=> 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)
=> 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)
=> 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 #-}
|
