From b9fab65c54eeb7af04ca7420911ea9828e5d0a55 Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Thu, 16 Feb 2017 11:05:18 -0600
Subject: [PATCH] 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.
---
 base/abstractarray.jl |  2 +-
 base/iterators.jl     | 75 +++++++++++++++++++++++++++++++++++++++++--
 test/arrayops.jl      |  9 ++++++
 3 files changed, 83 insertions(+), 3 deletions(-)

diff --git a/base/abstractarray.jl b/base/abstractarray.jl
index 5171452b6e8f4c..7afa37a614e5f2 100644
--- a/base/abstractarray.jl
+++ b/base/abstractarray.jl
@@ -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]
diff --git a/base/iterators.jl b/base/iterators.jl
index 8dbc0ec50a6dac..333255e3a7e405 100644
--- a/base/iterators.jl
+++ b/base/iterators.jl
@@ -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
 
@@ -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"];
@@ -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
diff --git a/test/arrayops.jl b/test/arrayops.jl
index 6970a12bc2e5f3..d9fe6ffcc5f79e 100644
--- a/test/arrayops.jl
+++ b/test/arrayops.jl
@@ -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]