Merge pull request #10704 from mbauman/mb/eachindex

Make eachindex more efficient for linear arrays

Author: mbauman

base/abstractarray.jl

@@ -337,12 +337,24 @@ end
 
 zero{T}(x::AbstractArray{T}) = fill!(similar(x), zero(T))
 
-## iteration support for arrays as ranges ##
+## iteration support for arrays by iterating over `eachindex` in the array ##
+# Allows fast iteration by default for both LinearFast and LinearSlow arrays
+
+# While the definitions for LinearFast are all simple enough to inline on their
+# own, LinearSlow's CartesianRange is more complicated and requires explicit
+# inlining. The real @inline macro is not available this early in the bootstrap,
+# so this internal macro splices the meta Expr directly into the function body.
+macro _inline_meta()
+    Expr(:meta, :inline)
+end
+start(A::AbstractArray) = (@_inline_meta(); itr = eachindex(A); (itr, start(itr)))
+next(A::AbstractArray,i) = (@_inline_meta(); (idx, s) = next(i[1], i[2]); (A[idx], (i[1], s)))
+done(A::AbstractArray,i) = done(i[1], i[2])
+
+# eachindex iterates over all indices. LinearSlow definitions are later.
+eachindex(A::AbstractArray) = (@_inline_meta; eachindex(linearindexing(A), A))
+eachindex(::LinearFast, A::AbstractArray) = 1:length(A)
 
-start(A::AbstractArray) = _start(A,linearindexing(A))
-_start(::AbstractArray,::LinearFast) = 1
-next(a::AbstractArray,i) = (a[i],i+1)
-done(a::AbstractArray,i) = (i > length(a))
 isempty(a::AbstractArray) = (length(a) == 0)
 
 ## Conversions ##

base/array.jl

@@ -297,6 +297,11 @@ end
 
 collect(itr) = collect(eltype(itr), itr)
 
+## Iteration ##
+start(A::Array) = 1
+next(a::Array,i) = (a[i],i+1)
+done(a::Array,i) = (i > length(a))
+
 ## Indexing: getindex ##
 
 getindex(a::Array) = arrayref(a,1)

base/dict.jl

@@ -713,7 +713,7 @@ function skip_deleted(h::Dict, i)
 end
 
 start(t::Dict) = skip_deleted(t, 1)
-done(t::Dict, i) = done(t.vals, i)
+done(t::Dict, i) = i > length(t.vals)
 next(t::Dict, i) = ((t.keys[i],t.vals[i]), skip_deleted(t,i+1))
 
 isempty(t::Dict) = (t.count == 0)

base/multidimensional.jl

@@ -1,11 +1,11 @@
 ### Multidimensional iterators
 module IteratorsMD
 
-import Base: eltype, length, start, _start, done, next, last, getindex, setindex!, linearindexing, min, max
+import Base: eltype, length, start, done, next, last, getindex, setindex!, linearindexing, min, max, eachindex
 import Base: simd_outer_range, simd_inner_length, simd_index
 import Base: @nref, @ncall, @nif, @nexprs, LinearFast, LinearSlow, to_index
 
-export CartesianIndex, CartesianRange, eachindex
+export CartesianIndex, CartesianRange
 
 # Traits for linear indexing
 linearindexing(::BitArray) = LinearFast()
@@ -110,60 +110,37 @@ stagedfunction CartesianRange{N}(I::CartesianIndex{N})
 end
 CartesianRange{N}(sz::NTuple{N,Int}) = CartesianRange(CartesianIndex(sz))
 
-stagedfunction eachindex{T,N}(A::AbstractArray{T,N})
+stagedfunction eachindex{T,N}(::LinearSlow, A::AbstractArray{T,N})
     startargs = fill(1, N)
     stopargs = [:(size(A,$i)) for i=1:N]
-    :(CartesianRange(CartesianIndex{$N}($(startargs...)), CartesianIndex{$N}($(stopargs...))))
+    meta = Expr(:meta, :inline)
+    :($meta; CartesianRange(CartesianIndex{$N}($(startargs...)), CartesianIndex{$N}($(stopargs...))))
 end
 
 eltype{I}(::Type{CartesianRange{I}}) = I
 eltype{I}(::CartesianRange{I}) = I
 
-stagedfunction start{I}(iter::CartesianRange{I})
-    N=length(I)
-    finishedex = Expr(:(||), [:(iter.stop[$i] < iter.start[$i]) for i = 1:N]...)
-    :(return $finishedex, iter.start)
-end
-
-stagedfunction _start{T,N}(A::AbstractArray{T,N}, ::LinearSlow)
-    args = fill(1, N)
-    :(return isempty(A), CartesianIndex{$N}($(args...)))
-end
-
-# Prevent an ambiguity warning
-next(R::StepRange, state::(Bool, CartesianIndex{1})) = (index=state[2]; return R[index], (index[1]==length(R), CartesianIndex{1}(index[1]+1)))
-next{T}(R::UnitRange{T}, state::(Bool, CartesianIndex{1})) = (index=state[2]; return R[index], (index[1]==length(R), CartesianIndex{1}(index[1]+1)))
-done(R::StepRange, state::(Bool, CartesianIndex{1})) = state[1]
-done(R::UnitRange, state::(Bool, CartesianIndex{1})) = state[1]
-
-stagedfunction next{T,N}(A::AbstractArray{T,N}, state::(Bool, CartesianIndex{N}))
-    I = state[2]
-    finishedex = (N==0 ? true : :(newindex[$N] > size(A, $N)))
+start(iter::CartesianRange) = iter.start
+stagedfunction next{I<:CartesianIndex}(iter::CartesianRange{I}, state)
+    N = length(I)
     meta = Expr(:meta, :inline)
     quote
         $meta
-        index=state[2]
-        @inbounds v = A[index]
-        newindex=@nif $N d->(index[d] < size(A, d)) d->@ncall($N, $I, k->(k>d ? index[k] : k==d ? index[k]+1 : 1))
-        finished=$finishedex
-        v, (finished,newindex)
+        index=state
+        @nif $N d->(index[d] < iter.stop[d]) d->(@nexprs($N, k->(ind_k = ifelse(k>=d, index[k] + (k==d), iter.start[k]))))
+        newindex = @ncall $N $I ind
+        index, newindex
     end
 end
-stagedfunction next{I<:CartesianIndex}(iter::CartesianRange{I}, state::(Bool, I))
+stagedfunction done{I<:CartesianIndex}(iter::CartesianRange{I}, state)
     N = length(I)
-    finishedex = (N==0 ? true : :(newindex[$N] > iter.stop[$N]))
-    meta = Expr(:meta, :inline)
-    quote
-        $meta
-        index=state[2]
-        newindex=@nif $N d->(index[d] < iter.stop[d]) d->@ncall($N, $I, k->(k>d ? index[k] : k==d ? index[k]+1 : iter.start[k]))
-        finished=$finishedex
-        index, (finished,newindex)
-    end
+    :(state[$N] > iter.stop[$N])
 end
 
-done{T,N}(A::AbstractArray{T,N}, state::(Bool, CartesianIndex{N})) = state[1]
-done{I<:CartesianIndex}(iter::CartesianRange{I}, state::(Bool, I)) = state[1]
+# 0-d cartesian ranges are special-cased to iterate once and only once
+start{I<:CartesianIndex{0}}(iter::CartesianRange{I}) = false
+next{I<:CartesianIndex{0}}(iter::CartesianRange{I}, state) = iter.start, true
+done{I<:CartesianIndex{0}}(iter::CartesianRange{I}, state) = state
 
 stagedfunction length{I<:CartesianIndex}(iter::CartesianRange{I})
     N = length(I)
@@ -502,7 +479,7 @@ function copy!{T,N}(dest::AbstractArray{T,N}, src::AbstractArray{T,N})
             break
         end
     end
-    if samesize
+    if samesize && linearindexing(dest) == linearindexing(src)
         for I in eachindex(dest)
             @inbounds dest[I] = src[I]
         end

base/reducedim.jl

@@ -209,7 +209,7 @@ function _mapreducedim!{T,N}(f, op, R::AbstractArray, A::AbstractArray{T,N})
         end
     else
         sizeR = CartesianIndex{N}(size(R))
-        @inbounds @simd for IA in eachindex(A)
+        @inbounds @simd for IA in CartesianRange(size(A))
             IR = min(IA, sizeR)
             R[IR] = op(R[IR], f(A[IA]))
         end
@@ -306,7 +306,7 @@ function findminmax!{T,N}(f, Rval, Rind, A::AbstractArray{T,N})
         end
     else
         sizeR = CartesianIndex(size(Rval))
-        @inbounds for IA in eachindex(A)
+        @inbounds for IA in CartesianRange(size(A))
             IR = min(sizeR, IA)
             k += 1
             tmpAv = A[IA]

doc/stdlib/arrays.rst

@@ -33,29 +33,31 @@ Basic functions
 
 .. function:: eachindex(A)
 
-   Creates an iterable object for visiting each multi-dimensional index of the AbstractArray ``A``.  Example for a 2-d array::
+   Creates an iterable object for visiting each index of an AbstractArray ``A`` in an efficient manner. For array types that have opted into fast linear indexing (like ``Array``), this is simply the range ``1:length(A)``. For other array types, this returns a specialized Cartesian range to efficiently index into the array with indices specified for every dimension. Example for a sparse 2-d array::
 
-    julia> A = rand(2,3)
-    2x3 Array{Float64,2}:
-     0.960084  0.629326  0.625155
-     0.432588  0.955903  0.991614
+    julia> A = sprand(2, 3, 0.5)
+    2x3 sparse matrix with 4 Float64 entries:
+        [1, 1]  =  0.598888
+        [1, 2]  =  0.0230247
+        [1, 3]  =  0.486499
+        [2, 3]  =  0.809041
 
     julia> for iter in eachindex(A)
-	       @show iter.I_1, iter.I_2
-	       @show A[iter]
-	   end
+               @show iter.I_1, iter.I_2
+               @show A[iter]
+           end
     (iter.I_1,iter.I_2) = (1,1)
-    A[iter] = 0.9600836263003063
+    A[iter] = 0.5988881393454597
     (iter.I_1,iter.I_2) = (2,1)
-    A[iter] = 0.4325878255452178
+    A[iter] = 0.0
     (iter.I_1,iter.I_2) = (1,2)
-    A[iter] = 0.6293256402775211
+    A[iter] = 0.02302469881746183
     (iter.I_1,iter.I_2) = (2,2)
-    A[iter] = 0.9559027084099654
+    A[iter] = 0.0
     (iter.I_1,iter.I_2) = (1,3)
-    A[iter] = 0.6251548453735303
+    A[iter] = 0.4864987874354343
     (iter.I_1,iter.I_2) = (2,3)
-    A[iter] = 0.9916142534546522
+    A[iter] = 0.8090413606455655
 
 .. function:: countnz(A)
 

test/arrayops.jl

@@ -993,7 +993,7 @@ I2 = CartesianIndex((-1,5,2))
 
 @test length(I1) == 3
 
-a = zeros(2,3)
+a = spzeros(2,3)
 @test CartesianRange(size(a)) == eachindex(a)
 a[CartesianIndex{2}(2,3)] = 5
 @test a[2,3] == 5
@@ -1015,22 +1015,24 @@ indexes = collect(R)
 @test length(R) == 12
 
 r = 2:3
-state = start(eachindex(r))
-@test !done(r, state)
-_, state = next(r, state)
-@test !done(r, state)
-val, state = next(r, state)
-@test done(r, state)
-@test val == 3
-r = 2:3:8
-state = start(eachindex(r))
-@test !done(r, state)
-_, state = next(r, state)
-_, state = next(r, state)
-@test !done(r, state)
-val, state = next(r, state)
-@test val == 8
-@test done(r, state)
+itr = eachindex(r)
+state = start(itr)
+@test !done(itr, state)
+_, state = next(itr, state)
+@test !done(itr, state)
+val, state = next(itr, state)
+@test done(itr, state)
+@test r[val] == 3
+r = sparse(collect(2:3:8))
+itr = eachindex(r)
+state = start(itr)
+@test !done(itr, state)
+_, state = next(itr, state)
+_, state = next(itr, state)
+@test !done(itr, state)
+val, state = next(itr, state)
+@test r[val] == 8
+@test done(itr, state)
 
 
 #rotates