2
2
3
3
# ## Multidimensional iterators
4
4
module IteratorsMD
5
- import Base: eltype, length, size, start, done, next, last, in, getindex,
5
+ import Base: eltype, length, size, start, done, next, first, last, in, getindex,
6
6
setindex!, IndexStyle, min, max, zero, one, isless, eachindex,
7
7
ndims, iteratorsize, convert
8
8
@@ -135,7 +135,6 @@ module IteratorsMD
135
135
136
136
# Iteration
137
137
"""
138
- CartesianRange(Istart::CartesianIndex, Istop::CartesianIndex) -> R
139
138
CartesianRange(sz::Dims) -> R
140
139
CartesianRange(istart:istop, jstart:jstop, ...) -> R
141
140
@@ -166,28 +165,46 @@ module IteratorsMD
166
165
CartesianIndex{3}((2, 2, 2))
167
166
```
168
167
"""
169
- struct CartesianRange{I<: CartesianIndex }
170
- start:: I
171
- stop:: I
168
+ struct CartesianRange{N,R<: NTuple{N,AbstractUnitRange{Int}} }
169
+ indices:: R
172
170
end
173
171
174
- CartesianRange (index:: CartesianIndex ) = CartesianRange (one (index), index)
175
- CartesianRange (:: Tuple{} ) = CartesianRange {CartesianIndex{0}} (CartesianIndex {0} (()),CartesianIndex {0} (()))
176
- CartesianRange (sz:: NTuple{N,Int} ) where {N} = CartesianRange (CartesianIndex (sz))
177
- CartesianRange (rngs:: NTuple{N,Union{Integer,AbstractUnitRange}} ) where {N} =
178
- CartesianRange (CartesianIndex (map (first, rngs)), CartesianIndex (map (last, rngs)))
179
-
180
- convert (:: Type{NTuple{N,UnitRange{Int}}} , R:: CartesianRange{CartesianIndex{N}} ) where {N} =
181
- map ((f,l)-> f: l, first (R). I, last (R). I)
182
- convert (:: Type{NTuple{N,UnitRange}} , R:: CartesianRange ) where {N} =
183
- convert (NTuple{N,UnitRange{Int}}, R)
184
- convert (:: Type {Tuple{Vararg{UnitRange{Int}}}}, R:: CartesianRange{CartesianIndex{N}} ) where {N} =
172
+ CartesianRange (:: Tuple{} ) = CartesianRange {0,typeof(())} (())
173
+ CartesianRange (inds:: NTuple{N,AbstractUnitRange{Int}} ) where {N} =
174
+ CartesianRange {N,typeof(inds)} (inds)
175
+ CartesianRange (inds:: Vararg{AbstractUnitRange{Int},N} ) where {N} =
176
+ CartesianRange (inds)
177
+ CartesianRange (inds:: NTuple{N,AbstractUnitRange{<:Integer}} ) where {N} =
178
+ CartesianRange (map (r-> convert (AbstractUnitRange{Int}, r), inds))
179
+ CartesianRange (inds:: Vararg{AbstractUnitRange{<:Integer},N} ) where {N} =
180
+ CartesianRange (inds)
181
+
182
+ CartesianRange (index:: CartesianIndex ) = CartesianRange (index. I)
183
+ CartesianRange (sz:: NTuple{N,<:Integer} ) where {N} = CartesianRange (map (Base. OneTo, sz))
184
+ CartesianRange (inds:: NTuple{N,Union{<:Integer,AbstractUnitRange{<:Integer}}} ) where {N} =
185
+ CartesianRange (map (i-> first (i): last (i), inds))
186
+
187
+ convert (:: Type{Tuple{}} , R:: CartesianRange{0} ) = ()
188
+ convert (:: Type{NTuple{N,AbstractUnitRange{Int}}} , R:: CartesianRange{N} ) where {N} =
189
+ R. indices
190
+ convert (:: Type{NTuple{N,AbstractUnitRange}} , R:: CartesianRange{N} ) where {N} =
191
+ convert (NTuple{N,AbstractUnitRange{Int}}, R)
192
+ convert (:: Type{NTuple{N,UnitRange{Int}}} , R:: CartesianRange{N} ) where {N} =
193
+ UnitRange {Int} .(convert (NTuple{N,AbstractUnitRange}, R))
194
+ convert (:: Type{NTuple{N,UnitRange}} , R:: CartesianRange{N} ) where {N} =
195
+ UnitRange .(convert (NTuple{N,AbstractUnitRange}, R))
196
+ convert (:: Type {Tuple{Vararg{AbstractUnitRange{Int}}}}, R:: CartesianRange{N} ) where {N} =
197
+ convert (NTuple{N,AbstractUnitRange{Int}}, R)
198
+ convert (:: Type{Tuple{Vararg{AbstractUnitRange}}} , R:: CartesianRange ) =
199
+ convert (Tuple{Vararg{AbstractUnitRange{Int}}}, R)
200
+ convert (:: Type {Tuple{Vararg{UnitRange{Int}}}}, R:: CartesianRange{N} ) where {N} =
185
201
convert (NTuple{N,UnitRange{Int}}, R)
186
202
convert (:: Type{Tuple{Vararg{UnitRange}}} , R:: CartesianRange ) =
187
203
convert (Tuple{Vararg{UnitRange{Int}}}, R)
188
204
189
- ndims (R:: CartesianRange ) = length (R. start)
190
- ndims (:: Type{CartesianRange{I}} ) where {I<: CartesianIndex } = length (I)
205
+ ndims (R:: CartesianRange ) = ndims (typeof (R))
206
+ ndims (:: Type{CartesianRange{N}} ) where {N} = N
207
+ ndims (:: Type{CartesianRange{N,TT}} ) where {N,TT} = N
191
208
192
209
eachindex (:: IndexCartesian , A:: AbstractArray ) = CartesianRange (indices (A))
193
210
@@ -201,17 +218,20 @@ module IteratorsMD
201
218
@inline maxt (a:: Tuple , b:: Tuple{} ) = a
202
219
@inline maxt (a:: Tuple , b:: Tuple ) = (max (a[1 ], b[1 ]), maxt (tail (a), tail (b))... )
203
220
204
- eltype (:: Type{CartesianRange{I}} ) where {I} = I
221
+ eltype (R:: CartesianRange ) = eltype (typeof (R))
222
+ eltype (:: Type{CartesianRange{N}} ) where {N} = CartesianIndex{N}
223
+ eltype (:: Type{CartesianRange{N,TT}} ) where {N,TT} = CartesianIndex{N}
205
224
iteratorsize (:: Type{<:CartesianRange} ) = Base. HasShape ()
206
225
207
- @inline function start (iter:: CartesianRange{<:CartesianIndex} )
208
- if any (map (> , iter. start. I, iter. stop. I))
209
- return iter. stop+ 1
226
+ @inline function start (iter:: CartesianRange )
227
+ iterfirst, iterlast = first (iter), last (iter)
228
+ if any (map (> , iterfirst. I, iterlast. I))
229
+ return iterlast+ 1
210
230
end
211
- iter . start
231
+ iterfirst
212
232
end
213
- @inline function next (iter:: CartesianRange{I} , state) where I <: CartesianIndex
214
- state, I (inc (state. I, iter. start . I, iter. stop . I))
233
+ @inline function next (iter:: CartesianRange , state)
234
+ state, CartesianIndex (inc (state. I, first ( iter) . I, last ( iter) . I))
215
235
end
216
236
# increment & carry
217
237
@inline inc (:: Tuple{} , :: Tuple{} , :: Tuple{} ) = ()
@@ -223,39 +243,38 @@ module IteratorsMD
223
243
newtail = inc (tail (state), tail (start), tail (stop))
224
244
(start[1 ], newtail... )
225
245
end
226
- @inline done (iter:: CartesianRange{<:CartesianIndex} , state) = state. I[end ] > iter. stop . I [end ]
246
+ @inline done (iter:: CartesianRange , state) = state. I[end ] > last ( iter. indices [end ])
227
247
228
248
# 0-d cartesian ranges are special-cased to iterate once and only once
229
- start (iter:: CartesianRange{<:CartesianIndex{0} } ) = false
230
- next (iter:: CartesianRange{<:CartesianIndex{0}} , state) = iter . start , true
231
- done (iter:: CartesianRange{<:CartesianIndex{0} } , state) = state
249
+ start (iter:: CartesianRange{0 } ) = false
250
+ next (iter:: CartesianRange{0} , state) = CartesianIndex () , true
251
+ done (iter:: CartesianRange{0 } , state) = state
232
252
233
- size (iter:: CartesianRange{<:CartesianIndex} ) = map (dimlength, iter. start . I, iter. stop . I)
253
+ size (iter:: CartesianRange ) = map (dimlength, first ( iter) . I, last ( iter) . I)
234
254
dimlength (start, stop) = stop- start+ 1
235
255
236
256
length (iter:: CartesianRange ) = prod (size (iter))
237
257
238
- last (iter:: CartesianRange ) = iter. stop
258
+ first (iter:: CartesianRange ) = CartesianIndex (map (first, iter. indices))
259
+ last (iter:: CartesianRange ) = CartesianIndex (map (last, iter. indices))
239
260
240
- @inline function in (i:: I , r:: CartesianRange{I } ) where I <: CartesianIndex
241
- _in (true , i. I, r . start . I, r . stop . I)
261
+ @inline function in (i:: CartesianIndex{N} , r:: CartesianRange{N } ) where {N}
262
+ _in (true , i. I, first (r) . I, last (r) . I)
242
263
end
243
264
_in (b, :: Tuple{} , :: Tuple{} , :: Tuple{} ) = b
244
265
@inline _in (b, i, start, stop) = _in (b & (start[1 ] <= i[1 ] <= stop[1 ]), tail (i), tail (start), tail (stop))
245
266
246
- simd_outer_range (iter:: CartesianRange{CartesianIndex{0} } ) = iter
267
+ simd_outer_range (iter:: CartesianRange{0 } ) = iter
247
268
function simd_outer_range (iter:: CartesianRange )
248
- start = CartesianIndex (tail (iter. start. I))
249
- stop = CartesianIndex (tail (iter. stop. I))
250
- CartesianRange (start, stop)
269
+ CartesianRange (tail (iter. indices))
251
270
end
252
271
253
- simd_inner_length (iter:: CartesianRange{<:CartesianIndex{0} } , :: CartesianIndex ) = 1
254
- simd_inner_length (iter:: CartesianRange , I:: CartesianIndex ) = iter. stop [1 ]- iter . start[ 1 ] + 1
272
+ simd_inner_length (iter:: CartesianRange{0 } , :: CartesianIndex ) = 1
273
+ simd_inner_length (iter:: CartesianRange , I:: CartesianIndex ) = length ( iter. indices [1 ])
255
274
256
- simd_index (iter:: CartesianRange{<:CartesianIndex{0}} , :: CartesianIndex , I1:: Int ) = iter. start
275
+ simd_index (iter:: CartesianRange{0} , :: CartesianIndex , I1:: Int ) = first ( iter)
257
276
@inline function simd_index (iter:: CartesianRange , Ilast:: CartesianIndex , I1:: Int )
258
- CartesianIndex ((I1+ iter. start [1 ], Ilast. I... ))
277
+ CartesianIndex ((I1+ first ( iter. indices [1 ]) , Ilast. I... ))
259
278
end
260
279
261
280
# Split out the first N elements of a tuple
@@ -859,23 +878,23 @@ function copy!(dest::AbstractArray{T,N}, src::AbstractArray{T,N}) where {T,N}
859
878
end
860
879
861
880
@generated function copy! (dest:: AbstractArray{T1,N} ,
862
- Rdest:: CartesianRange{CartesianIndex{N} } ,
881
+ Rdest:: CartesianRange{N } ,
863
882
src:: AbstractArray{T2,N} ,
864
- Rsrc:: CartesianRange{CartesianIndex{N} } ) where {T1,T2,N}
883
+ Rsrc:: CartesianRange{N } ) where {T1,T2,N}
865
884
quote
866
885
isempty (Rdest) && return dest
867
886
if size (Rdest) != size (Rsrc)
868
887
throw (ArgumentError (" source and destination must have same size (got $(size (Rsrc)) and $(size (Rdest)) )" ))
869
888
end
870
- @boundscheck checkbounds (dest, Rdest. start )
871
- @boundscheck checkbounds (dest, Rdest. stop )
872
- @boundscheck checkbounds (src, Rsrc. start )
873
- @boundscheck checkbounds (src, Rsrc. stop )
874
- ΔI = Rdest. start - Rsrc. start
889
+ @boundscheck checkbounds (dest, first ( Rdest) )
890
+ @boundscheck checkbounds (dest, last ( Rdest) )
891
+ @boundscheck checkbounds (src, first ( Rsrc) )
892
+ @boundscheck checkbounds (src, last ( Rsrc) )
893
+ ΔI = first ( Rdest) - first ( Rsrc)
875
894
# TODO : restore when #9080 is fixed
876
895
# for I in Rsrc
877
896
# @inbounds dest[I+ΔI] = src[I]
878
- @nloops $ N i (n-> Rsrc. start[n] : Rsrc . stop [n]) begin
897
+ @nloops $ N i (n-> Rsrc. indices [n]) begin
879
898
@inbounds @nref ($ N,dest,n-> i_n+ ΔI[n]) = @nref ($ N,src,i)
880
899
end
881
900
dest
0 commit comments