diff --git a/base/abstractarray.jl b/base/abstractarray.jl index 0624fffd2a42c..6dee41d3d7615 100644 --- a/base/abstractarray.jl +++ b/base/abstractarray.jl @@ -113,10 +113,10 @@ linearindexing{T<:AbstractArray}(::Type{T}) = LinearSlow() linearindexing{T<:Array}(::Type{T}) = LinearFast() linearindexing{T<:Range}(::Type{T}) = LinearFast() -*(::LinearFast, ::LinearFast) = LinearFast() -*(::LinearSlow, ::LinearFast) = LinearSlow() -*(::LinearFast, ::LinearSlow) = LinearSlow() -*(::LinearSlow, ::LinearSlow) = LinearSlow() +linearindexing(A::AbstractArray, B::AbstractArray) = linearindexing(linearindexing(A), linearindexing(B)) +linearindexing(A::AbstractArray, B::AbstractArray...) = linearindexing(linearindexing(A), linearindexing(B...)) +linearindexing(::LinearFast, ::LinearFast) = LinearFast() +linearindexing(::LinearIndexing, ::LinearIndexing) = LinearSlow() # 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. @@ -385,9 +385,14 @@ eachindex(::LinearFast, A::AbstractArray) = 1:length(A) function eachindex(A::AbstractArray, B::AbstractArray) @_inline_meta - eachindex(linearindexing(A)*linearindexing(B), A, B) + eachindex(linearindexing(A,B), A, B) +end +function eachindex(A::AbstractArray, B::AbstractArray...) + @_inline_meta + eachindex(linearindexing(A,B...), A, B...) end eachindex(::LinearFast, A::AbstractArray, B::AbstractArray) = 1:max(length(A),length(B)) +eachindex(::LinearFast, A::AbstractArray, B::AbstractArray...) = 1:max(length(A), map(length, B)...) isempty(a::AbstractArray) = (length(a) == 0) diff --git a/base/multidimensional.jl b/base/multidimensional.jl index 87096095155b7..5b5e9ff9a6402 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -94,10 +94,14 @@ ndims(R::CartesianRange) = length(R.start) :($meta; CartesianRange(CartesianIndex{$N}($(startargs...)), CartesianIndex{$N}($(stopargs...)))) end -@generated function eachindex{S,T,M,N}(::LinearSlow, A::AbstractArray{S,M}, B::AbstractArray{T,N}) - K = max(M,N) +@generated function eachindex(::LinearSlow, A::AbstractArray, B::AbstractArray...) + K = max(ndims(A), map(ndims, B)...) startargs = fill(1, K) - stopargs = [:(max(size(A,$i),size(B,$i))) for i=1:K] + stopargs = Array(Expr, K) + for i = 1:K + Bargs = [:(size(B[$j],$i)) for j = 1:length(B)] + stopargs[i] = :(max(size(A,$i),$(Bargs...))) + end meta = Expr(:meta, :inline) :($meta; CartesianRange(CartesianIndex{$K}($(startargs...)), CartesianIndex{$K}($(stopargs...)))) end diff --git a/doc/stdlib/arrays.rst b/doc/stdlib/arrays.rst index 9479407574983..44f7483a96e61 100644 --- a/doc/stdlib/arrays.rst +++ b/doc/stdlib/arrays.rst @@ -31,7 +31,7 @@ Basic functions Returns the number of elements in A -.. function:: eachindex(A) +.. function:: eachindex(A...) 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:: @@ -59,6 +59,13 @@ Basic functions (iter.I_1,iter.I_2) = (2,3) A[iter] = 0.8090413606455655 +If you supply more than one ``AbstractArray`` argument, ``eachindex`` +will create an iterable object that is fast for all arguments (a +``UnitRange`` if all inputs have fast linear indexing, a +CartesianRange otherwise). If the arrays have different sizes and/or +dimensionalities, ``eachindex`` returns an interable that spans the +largest range along each dimension. + .. function:: Base.linearindexing(A) ``linearindexing`` defines how an AbstractArray most efficiently accesses its elements. If ``Base.linearindexing(A)`` returns ``Base.LinearFast()``, this means that linear indexing with only one index is an efficient operation. If it instead returns ``Base.LinearSlow()`` (by default), this means that the array intrinsically accesses its elements with indices specified for every dimension. Since converting a linear index to multiple indexing subscripts is typically very expensive, this provides a traits-based mechanism to enable efficient generic code for all array types. diff --git a/test/arrayops.jl b/test/arrayops.jl index dcda66cfc9967..887ceaaa06e60 100644 --- a/test/arrayops.jl +++ b/test/arrayops.jl @@ -1105,6 +1105,12 @@ R = CartesianRange((0,3)) R = CartesianRange((3,0)) @test done(R, start(R)) == true +@test eachindex(Base.LinearSlow(),zeros(3),zeros(2,2),zeros(2,2,2),zeros(2,2)) == CartesianRange((3,2,2)) +@test eachindex(Base.LinearFast(),zeros(3),zeros(2,2),zeros(2,2,2),zeros(2,2)) == 1:8 +@test eachindex(zeros(3),sub(zeros(3,3),1:2,1:2),zeros(2,2,2),zeros(2,2)) == CartesianRange((3,2,2)) +@test eachindex(zeros(3),zeros(2,2),zeros(2,2,2),zeros(2,2)) == 1:8 + + #rotates a = [1 0 0; 0 0 0]