Skip to content

Commit fb17155

Browse files
committed
Add enumerate(::IndexStyle, A) for index/value iteration
`enumerate(A)` doesn't guarantee that the counter corresponds to the index; so when you need an index, call this method.
1 parent 22cd6b0 commit fb17155

File tree

3 files changed

+83
-3
lines changed

3 files changed

+83
-3
lines changed

base/abstractarray.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1642,7 +1642,7 @@ _sub2ind_vec(inds, out, i) = (@_inline_meta; sub2ind(inds, out...))
16421642
function ind2sub{N}(inds::Union{DimsInteger{N},Indices{N}}, ind::AbstractVector{<:Integer})
16431643
M = length(ind)
16441644
t = ntuple(n->similar(ind),Val{N})
1645-
for (i,idx) in enumerate(ind) # FIXME: change to eachindexvalue
1645+
for (i,idx) in enumerate(IndexLinear(), ind)
16461646
sub = ind2sub(inds, idx)
16471647
for j = 1:N
16481648
t[j][i] = sub[j]

base/iterators.jl

Lines changed: 73 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ module Iterators
44

55
import Base: start, done, next, isempty, length, size, eltype, iteratorsize, iteratoreltype, indices, ndims
66

7-
using Base: tuple_type_cons, SizeUnknown, HasLength, HasShape, IsInfinite, EltypeUnknown, HasEltype, OneTo
7+
using Base: tuple_type_cons, SizeUnknown, HasLength, HasShape, IsInfinite, EltypeUnknown, HasEltype, OneTo, @propagate_inbounds
88

99
export enumerate, zip, rest, countfrom, take, drop, cycle, repeated, product, flatten, partition
1010

@@ -40,7 +40,8 @@ and `x` is the `i`th value from the given iterator. It's useful when
4040
you need not only the values `x` over which you are iterating, but
4141
also the number of iterations so far. Note that `i` may not be valid
4242
for indexing `iter`; it's also possible that `x != iter[i]`, if `iter`
43-
has indices that do not start at 1.
43+
has indices that do not start at 1. See the `enumerate(IndexLinear(),
44+
iter)` method if you want to ensure that `i` is an index.
4445
4546
```jldoctest
4647
julia> a = ["a", "b", "c"];
@@ -69,6 +70,76 @@ eltype{I}(::Type{Enumerate{I}}) = Tuple{Int, eltype(I)}
6970
iteratorsize{I}(::Type{Enumerate{I}}) = iteratorsize(I)
7071
iteratoreltype{I}(::Type{Enumerate{I}}) = iteratoreltype(I)
7172

73+
struct IndexValue{I,A<:AbstractArray}
74+
data::A
75+
itr::I
76+
end
77+
78+
"""
79+
enumerate(IndexLinear(), A)
80+
enumerate(IndexCartesian(), A)
81+
enumerate(IndexStyle(A), A)
82+
83+
An iterator that accesses each element of the array `A`, returning
84+
`(i, x)`, where `i` is the index for the element and `x = A[i]`. This
85+
is similar to `enumerate(A)`, except `i` will always be a valid index
86+
for `A`.
87+
88+
Specifying `IndexLinear()` ensures that `i` will be an integer;
89+
specifying `IndexCartesian()` ensures that `i` will be a
90+
`CartesianIndex`; specifying `IndexStyle(A)` chooses whichever has
91+
been defined as the native indexing style for array `A`.
92+
93+
```jldoctest
94+
julia> A = ["a" "d"; "b" "e"; "c" "f"];
95+
96+
julia> for (index, value) in enumerate(IndexStyle(A), A)
97+
println("\$index \$value")
98+
end
99+
1 a
100+
2 b
101+
3 c
102+
4 d
103+
5 e
104+
6 f
105+
106+
julia> S = view(A, 1:2, :);
107+
108+
julia> for (index, value) in enumerate(IndexStyle(S), S)
109+
println("\$index \$value")
110+
end
111+
CartesianIndex{2}((1, 1)) a
112+
CartesianIndex{2}((2, 1)) b
113+
CartesianIndex{2}((1, 2)) d
114+
CartesianIndex{2}((2, 2)) e
115+
```
116+
117+
Note that `enumerate(A)` returns `i` as a *counter* (always starting
118+
at 1), whereas `enumerate(IndexLinear(), A)` returns `i` as an *index*
119+
(starting at the first linear index of `A`, which may or may not be
120+
1).
121+
122+
See also: [`IndexStyle`](@ref), [`indices`](@ref).
123+
"""
124+
enumerate(::IndexLinear, A::AbstractArray) = IndexValue(A, linearindices(A))
125+
enumerate(::IndexCartesian, A::AbstractArray) = IndexValue(A, CartesianRange(indices(A)))
126+
127+
length(v::IndexValue) = length(v.itr)
128+
indices(v::IndexValue) = indices(v.itr)
129+
size(v::IndexValue) = size(v.itr)
130+
@inline start(v::IndexValue) = start(v.itr)
131+
@propagate_inbounds function next(v::IndexValue, state)
132+
indx, n = next(v.itr, state)
133+
item = v.data[indx]
134+
(indx, item), n
135+
end
136+
@inline done(v::IndexValue, state) = done(v.itr, state)
137+
138+
eltype{I,A}(::Type{IndexValue{I,A}}) = Tuple{eltype(I), eltype(A)}
139+
140+
iteratorsize{I}(::Type{IndexValue{I}}) = iteratorsize(I)
141+
iteratoreltype{I}(::Type{IndexValue{I}}) = iteratoreltype(I)
142+
72143
# zip
73144

74145
abstract type AbstractZipIterator end

test/arrayops.jl

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1121,6 +1121,15 @@ end
11211121
end
11221122
end
11231123

1124+
@testset "eachindexvalue" begin
1125+
A14 = [11 13; 12 14]
1126+
R = CartesianRange(indices(A14))
1127+
@test [a for (a,b) in enumerate(IndexLinear(), A14)] == [1,2,3,4]
1128+
@test [a for (a,b) in enumerate(IndexCartesian(), A14)] == vec(collect(R))
1129+
@test [b for (a,b) in enumerate(IndexLinear(), A14)] == [11,12,13,14]
1130+
@test [b for (a,b) in enumerate(IndexCartesian(), A14)] == [11,12,13,14]
1131+
end
1132+
11241133
@testset "reverse" begin
11251134
@test reverse([2,3,1]) == [1,3,2]
11261135
@test reverse([1:10;],1,4) == [4,3,2,1,5,6,7,8,9,10]

0 commit comments

Comments
 (0)