Skip to content

Commit

Permalink
Add enumerate(::IndexStyle, A) for index/value iteration
Browse files Browse the repository at this point in the history
`enumerate(A)` doesn't guarantee that the counter corresponds to the index; so when you need an index, call this method.
  • Loading branch information
timholy committed Feb 16, 2017
1 parent 5f20475 commit b9fab65
Show file tree
Hide file tree
Showing 3 changed files with 83 additions and 3 deletions.
2 changes: 1 addition & 1 deletion base/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1642,7 +1642,7 @@ _sub2ind_vec(inds, out, i) = (@_inline_meta; sub2ind(inds, out...))
function ind2sub{N}(inds::Union{DimsInteger{N},Indices{N}}, ind::AbstractVector{<:Integer})
M = length(ind)
t = ntuple(n->similar(ind),Val{N})
for (i,idx) in enumerate(ind) # FIXME: change to eachindexvalue
for (i,idx) in enumerate(IndexLinear(), ind)
sub = ind2sub(inds, idx)
for j = 1:N
t[j][i] = sub[j]
Expand Down
75 changes: 73 additions & 2 deletions base/iterators.jl
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ module Iterators

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

using Base: tuple_type_cons, SizeUnknown, HasLength, HasShape, IsInfinite, EltypeUnknown, HasEltype, OneTo
using Base: tuple_type_cons, SizeUnknown, HasLength, HasShape, IsInfinite, EltypeUnknown, HasEltype, OneTo, @propagate_inbounds

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

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

struct IndexValue{I,A<:AbstractArray}
data::A
itr::I
end

"""
enumerate(IndexLinear(), A)
enumerate(IndexCartesian(), A)
enumerate(IndexStyle(A), A)
An iterator that accesses each element of the array `A`, returning
`(i, x)`, where `i` is the index for the element and `x = A[i]`. This
is similar to `enumerate(A)`, except `i` will always be a valid index
for `A`.
Specifying `IndexLinear()` ensures that `i` will be an integer;
specifying `IndexCartesian()` ensures that `i` will be a
`CartesianIndex`; specifying `IndexStyle(A)` chooses whichever has
been defined as the native indexing style for array `A`.
```jldoctest
julia> A = ["a" "d"; "b" "e"; "c" "f"];
julia> for (index, value) in enumerate(IndexStyle(A), A)
println("\$index \$value")
end
1 a
2 b
3 c
4 d
5 e
6 f
julia> S = view(A, 1:2, :);
julia> for (index, value) in enumerate(IndexStyle(S), S)
println("\$index \$value")
end
CartesianIndex{2}((1, 1)) a
CartesianIndex{2}((2, 1)) b
CartesianIndex{2}((1, 2)) d
CartesianIndex{2}((2, 2)) e
```
Note that `enumerate(A)` returns `i` as a *counter* (always starting
at 1), whereas `enumerate(IndexLinear(), A)` returns `i` as an *index*
(starting at the first linear index of `A`, which may or may not be
1).
See also: [`IndexStyle`](@ref), [`indices`](@ref).
"""
enumerate(::IndexLinear, A::AbstractArray) = IndexValue(A, linearindices(A))
enumerate(::IndexCartesian, A::AbstractArray) = IndexValue(A, CartesianRange(indices(A)))

length(v::IndexValue) = length(v.itr)
indices(v::IndexValue) = indices(v.itr)
size(v::IndexValue) = size(v.itr)
@inline start(v::IndexValue) = start(v.itr)
@propagate_inbounds function next(v::IndexValue, state)
indx, n = next(v.itr, state)
item = v.data[indx]
(indx, item), n
end
@inline done(v::IndexValue, state) = done(v.itr, state)

eltype{I,A}(::Type{IndexValue{I,A}}) = Tuple{eltype(I), eltype(A)}

iteratorsize{I}(::Type{IndexValue{I}}) = iteratorsize(I)
iteratoreltype{I}(::Type{IndexValue{I}}) = iteratoreltype(I)

# zip

abstract type AbstractZipIterator end
Expand Down
9 changes: 9 additions & 0 deletions test/arrayops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1121,6 +1121,15 @@ end
end
end

@testset "eachindexvalue" begin
A14 = [11 13; 12 14]
R = CartesianRange(indices(A14))
@test [a for (a,b) in enumerate(IndexLinear(), A14)] == [1,2,3,4]
@test [a for (a,b) in enumerate(IndexCartesian(), A14)] == vec(collect(R))
@test [b for (a,b) in enumerate(IndexLinear(), A14)] == [11,12,13,14]
@test [b for (a,b) in enumerate(IndexCartesian(), A14)] == [11,12,13,14]
end

@testset "reverse" begin
@test reverse([2,3,1]) == [1,3,2]
@test reverse([1:10;],1,4) == [4,3,2,1,5,6,7,8,9,10]
Expand Down

0 comments on commit b9fab65

Please sign in to comment.