Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

RFC: Allow any index type in nonscalar indexing #12567

Closed
wants to merge 2 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 2 additions & 8 deletions base/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -494,13 +494,10 @@ _getindex{T}(::LinearSlow, A::AbstractArray{T,0}) = error("indexing not defined
_getindex(::LinearSlow, A::AbstractVector) = (@_inline_meta; getindex(A, 1))
_getindex(l::LinearSlow, A::AbstractArray) = (@_inline_meta; _getindex(l, A, 1))
_unsafe_getindex(::LinearFast, A::AbstractArray) = (@_inline_meta; unsafe_getindex(A, 1))
_unsafe_getindex{T}(::LinearSlow, A::AbstractArray{T,0}) = error("indexing not defined for ", typeof(A))
_unsafe_getindex{T}(::LinearSlow, A::AbstractArray{T,0}) = (@_inline_meta; getindex(A))
_unsafe_getindex(::LinearSlow, A::AbstractVector) = (@_inline_meta; unsafe_getindex(A, 1))
_unsafe_getindex(l::LinearSlow, A::AbstractArray) = (@_inline_meta; _unsafe_getindex(l, A, 1))

_getindex(::LinearIndexing, A::AbstractArray, I...) = error("indexing $(typeof(A)) with types $(typeof(I)) is not supported")
_unsafe_getindex(::LinearIndexing, A::AbstractArray, I...) = (@_inline_meta; getindex(A, I...))

## LinearFast Scalar indexing
_getindex(::LinearFast, A::AbstractArray, I::Int) = error("indexing not defined for ", typeof(A))
function _getindex(::LinearFast, A::AbstractArray, I::Real...)
Expand Down Expand Up @@ -596,13 +593,10 @@ _setindex!{T}(::LinearSlow, A::AbstractArray{T,0}, v) = error("indexing not defi
_setindex!(::LinearSlow, A::AbstractVector, v) = (@_inline_meta; setindex!(A, v, 1))
_setindex!(l::LinearSlow, A::AbstractArray, v) = (@_inline_meta; _setindex!(l, A, v, 1))
_unsafe_setindex!(::LinearFast, A::AbstractArray, v) = (@_inline_meta; unsafe_setindex!(A, v, 1))
_unsafe_setindex!{T}(::LinearSlow, A::AbstractArray{T,0}, v) = error("indexing not defined for ", typeof(A))
_unsafe_setindex!{T}(::LinearSlow, A::AbstractArray{T,0}, v) = (@_inline_meta; setindex!(A, v))
_unsafe_setindex!(::LinearSlow, A::AbstractVector, v) = (@_inline_meta; unsafe_setindex!(A, v, 1))
_unsafe_setindex!(l::LinearSlow, A::AbstractArray, v) = (@_inline_meta; _unsafe_setindex!(l, A, v, 1))

_setindex!(::LinearIndexing, A::AbstractArray, v, I...) = error("indexing $(typeof(A)) with types $(typeof(I)) is not supported")
_unsafe_setindex!(::LinearIndexing, A::AbstractArray, v, I...) = (@_inline_meta; setindex!(A, v, I...))

## LinearFast Scalar indexing
_setindex!(::LinearFast, A::AbstractArray, v, I::Int) = error("indexed assignment not defined for ", typeof(A))
function _setindex!(::LinearFast, A::AbstractArray, v, I::Real...)
Expand Down
20 changes: 1 addition & 19 deletions base/deprecated.jl
Original file line number Diff line number Diff line change
Expand Up @@ -441,24 +441,6 @@ end
end

# 10458
to_index_nodep(i::Real) = convert(Int,i)::Int

@noinline function to_index(i::Real)
depwarn("indexing with non Integer Reals is deprecated", :to_index)
to_index_nodep(i)
end

to_index{T<:Integer}(A::AbstractArray{T}) = A
@noinline function to_index{T<:Real}(A::AbstractArray{T})
depwarn("indexing with non Integer AbstractArrays is deprecated", :to_index)
Int[to_index_nodep(x) for x in A]
end

@noinline function to_index(I::Tuple)
depwarn("to_index(I::Tuple) is deprecated, use to_indexes(I...) instead.", :to_index)
to_indexes(I...)
end

@deprecate getindex(c::Char, I::Real...) getindex(c, map(Int, I)...)
@deprecate getindex(s::AbstractString, x::Real) getindex(s, Int(x))
@deprecate checkbounds(s::AbstractString, i::Real) checkbounds(s, Int(i))
Expand Down Expand Up @@ -841,4 +823,4 @@ for f in (:remotecall, :remotecall_fetch, :remotecall_wait)
@deprecate ($f)(w::Worker, f::Function, args...) ($f)(f, w::Worker, args...)
@deprecate ($f)(id::Integer, f::Function, args...) ($f)(f, id::Integer, args...)
end
end
end
177 changes: 112 additions & 65 deletions base/multidimensional.jl
Original file line number Diff line number Diff line change
Expand Up @@ -162,39 +162,94 @@ index_lengths(A::AbstractArray, I::AbstractArray) = (length(I),)
index_lengths_dim(A, dim) = ()
index_lengths_dim(A, dim, ::Colon) = (trailingsize(A, dim),)
@inline index_lengths_dim(A, dim, ::Colon, i, I...) = (size(A, dim), index_lengths_dim(A, dim+1, i, I...)...)
@inline index_lengths_dim(A, dim, ::Real, I...) = (1, index_lengths_dim(A, dim+1, I...)...)
@inline index_lengths_dim(A, dim, ::Any, I...) = (1, index_lengths_dim(A, dim+1, I...)...)
@inline index_lengths_dim(A, dim, i::AbstractArray{Bool}, I...) = (sum(i), index_lengths_dim(A, dim+1, I...)...)
@inline index_lengths_dim(A, dim, i::AbstractArray, I...) = (length(i), index_lengths_dim(A, dim+1, I...)...)

# shape of array to create for getindex() with indexes I, dropping trailing scalars
index_shape(A::AbstractArray, I::AbstractArray) = size(I) # Linear index reshape
index_shape(A::AbstractArray, I::AbstractArray{Bool}) = (sum(I),) # Logical index
index_shape(A::AbstractArray, I::Colon) = (length(A),)
@inline index_shape(A::AbstractArray, I...) = index_shape_dim(A, 1, I...)
index_shape_dim(A, dim, I::Real...) = ()
index_shape_dim(A, dim, ::Colon) = (trailingsize(A, dim),)
@inline index_shape_dim(A, dim, ::Colon, i, I...) = (size(A, dim), index_shape_dim(A, dim+1, i, I...)...)
@inline index_shape_dim(A, dim, ::Real, I...) = (1, index_shape_dim(A, dim+1, I...)...)
@inline index_shape_dim(A, dim, i::AbstractVector{Bool}, I...) = (sum(i), index_shape_dim(A, dim+1, I...)...)
@inline index_shape_dim(A, dim, i::AbstractVector, I...) = (length(i), index_shape_dim(A, dim+1, I...)...)
@generated function index_shape(A::AbstractArray, I...)
N = length(I)
sz = Expr(:tuple)
for d=1:N
if !any(i->i<:Union{AbstractArray,Colon}, I[d:end])
break
elseif I[d] <: Colon
push!(sz.args, d < N ? :(size(A, $d)) : :(trailingsize(A, Val{$d})))
elseif I[d] <: AbstractArray{Bool}
push!(sz.args, :(sum(I[$d])))
elseif I[d] <: AbstractVector
push!(sz.args, :(length(I[$d])))
else
push!(sz.args, 1)
end
end
quote
$(Expr(:meta, :inline))
$sz
end
end

to_nonscalar_index(I::AbstractArray{Bool}) = find(I)
to_nonscalar_index(A::AbstractArray) = A
to_nonscalar_index(i) = i

to_nonscalar_indexes() = ()
to_nonscalar_indexes(i1) = (to_nonscalar_index(i1),)
to_nonscalar_indexes(i1, I...) = (to_nonscalar_index(i1), to_nonscalar_indexes(I...)...)

# Cartesian indexing code generation
function cartindex_exprs(indexes, syms)
exprs = Any[]
for (i,ind) in enumerate(indexes)
if ind <: CartesianIndex
for j = 1:length(ind)
push!(exprs, :($syms[$i][$j]))
end
else
push!(exprs, :($syms[$i]))
end
end
if isempty(exprs)
push!(exprs, 1) # Handle the zero-dimensional case
end
exprs
end

### From abstractarray.jl: Internal multidimensional indexing definitions ###
# These are not defined on directly ongetindex and unsafe_getindex to avoid
# ambiguities for AbstractArray subtypes. See the note in abstractarray.jl

# Note that it's most efficient to call checkbounds first, and then to_index
@inline function _getindex(l::LinearIndexing, A::AbstractArray, I::Union{Real, AbstractArray, Colon}...)
checkbounds(A, I...)
_unsafe_getindex(l, A, I...)
@generated function _getindex(l::LinearIndexing, A::AbstractArray, I...)
if any(i->i<:CartesianIndex, I)
:($(Expr(:meta, :inline)); getindex(A, $(cartindex_exprs(I, :I)...)))
elseif !any(i->i<:Union{AbstractArray,Colon}, I)
:(error("indexing $(typeof(A)) with types $(typeof(I)) is not supported"))
else
quote
$(Expr(:meta, :inline))
checkbounds(A, I...)
_unsafe_getindex(l, A, I...)
end
end
end
@generated function _unsafe_getindex(l::LinearIndexing, A::AbstractArray, I::Union{Real, AbstractArray, Colon}...)
@generated function _unsafe_getindex(l::LinearIndexing, A::AbstractArray, I...)
N = length(I)
quote
# This is specifically *not* inlined.
@nexprs $N d->(I_d = to_index(I[d]))
dest = similar(A, @ncall $N index_shape A I)
@ncall $N checksize dest I
@ncall $N _unsafe_getindex! dest l A I
if any(i->i<:CartesianIndex, I)
:($(Expr(:meta, :inline)); unsafe_getindex(A, $(cartindex_exprs(I, :I)...)))
elseif !any(i->i<:Union{AbstractArray,Colon}, I)
:($(Expr(:meta, :inline)); getindex(A, I...))
else
quote
# This is specifically *not* inlined, so use Cartesian to avoid splats
@nexprs $N d->(I_d = to_nonscalar_index(I[d]))
dest = similar(A, @ncall $N index_shape A I)
@ncall $N checksize dest I
@ncall $N _unsafe_getindex! dest l A I
end
end
end

Expand All @@ -219,9 +274,14 @@ function _unsafe_getindex(::LinearIndexing, src::AbstractArray, I::AbstractArray
dest
end

indexref(idx, i::Int) = idx
indexref(::Colon, i::Int) = i
@inline indexref(A::AbstractArray, i::Int) = unsafe_getindex(A, i)

# Indexing with an array of indices is inherently linear in the source, but
# might be able to be optimized with fast dividing integers
@inline function _unsafe_getindex!(dest::AbstractArray, ::LinearIndexing, src::AbstractArray, I::AbstractArray)
@inline _unsafe_getindex!(dest::AbstractArray, ::LinearSlow, src::AbstractArray, I::AbstractArray) = _unsafe_getindex!(dest, LinearFast(), src, I)
@inline function _unsafe_getindex!(dest::AbstractArray, ::LinearFast, src::AbstractArray, I::AbstractArray)
D = eachindex(dest)
Ds = start(D)
for idx in I
Expand All @@ -232,7 +292,7 @@ end
end

# Fast source - compute the linear index
@generated function _unsafe_getindex!(dest::AbstractArray, ::LinearFast, src::AbstractArray, I::Union{Real, AbstractVector, Colon}...)
@generated function _unsafe_getindex!(dest::AbstractArray, ::LinearFast, src::AbstractArray, I...)
N = length(I)
quote
$(Expr(:meta, :inline))
Expand All @@ -241,7 +301,7 @@ end
$(symbol(:offset_, N)) = 1
D = eachindex(dest)
Ds = start(D)
@nloops $N i dest d->(offset_{d-1} = offset_d + (unsafe_getindex(I[d], i_d)-1)*stride_d) begin
@nloops $N i dest d->(offset_{d-1} = offset_d + (indexref(I[d], i_d)-1)*stride_d) begin
d, Ds = next(D, Ds)
unsafe_setindex!(dest, unsafe_getindex(src, offset_0), d)
end
Expand All @@ -250,13 +310,13 @@ end
end
# Slow source - index with the indices provided.
# TODO: this may not be the full dimensionality; that case could be optimized
@generated function _unsafe_getindex!(dest::AbstractArray, ::LinearSlow, src::AbstractArray, I::Union{Real, AbstractVector, Colon}...)
@generated function _unsafe_getindex!(dest::AbstractArray, ::LinearSlow, src::AbstractArray, I...)
N = length(I)
quote
$(Expr(:meta, :inline))
D = eachindex(dest)
Ds = start(D)
@nloops $N i dest d->(j_d = unsafe_getindex(I[d], i_d)) begin
@nloops $N i dest d->(j_d = indexref(I[d], i_d)) begin
d, Ds = next(D, Ds)
v = @ncall $N unsafe_getindex src j
unsafe_setindex!(dest, v, d)
Expand All @@ -274,10 +334,10 @@ checksize(A::AbstractArray, I::AbstractArray{Bool}) = length(A) == sum(I) || thr
@nexprs $N d->(_checksize(A, d, I[d]) || throw(DimensionMismatch("index $d selects $(length(I[d])) elements, but size(A, $d) = $(size(A,d))")))
end
end
_checksize(A::AbstractArray, dim, I) = size(A, dim) == length(I)
_checksize(A::AbstractArray, dim, I::AbstractVector) = size(A, dim) == length(I)
_checksize(A::AbstractArray, dim, I::AbstractVector{Bool}) = size(A, dim) == sum(I)
_checksize(A::AbstractArray, dim, ::Colon) = true
_checksize(A::AbstractArray, dim, ::Real) = size(A, dim) == 1
_checksize(A::AbstractArray, dim, ::Any) = size(A, dim) == 1

@inline unsafe_setindex!(v::BitArray, x::Bool, ind::Int) = (Base.unsafe_bitsetindex!(v.chunks, x, ind); v)
@inline unsafe_setindex!(v::BitArray, x, ind::Real) = (Base.unsafe_bitsetindex!(v.chunks, convert(Bool, x), to_index(ind)); v)
Expand All @@ -288,12 +348,30 @@ _checksize(A::AbstractArray, dim, ::Real) = size(A, dim) == 1
# before redispatching to the _unsafe_batchsetindex!
_iterable(v::AbstractArray) = v
_iterable(v) = repeated(v)
@inline function _setindex!(l::LinearIndexing, A::AbstractArray, x, J::Union{Real,AbstractArray,Colon}...)
checkbounds(A, J...)
_unsafe_setindex!(l, A, x, J...)
@generated function _setindex!(l::LinearIndexing, A::AbstractArray, x, I...)
if any(i->i<:CartesianIndex, I)
:($(Expr(:meta, :inline)); setindex!(A, x, $(cartindex_exprs(I, :I)...)))
elseif !any(i->i<:Union{AbstractArray,Colon}, I)
:(error("indexed assignment of $(typeof(A)) with types $(typeof(I)) is not supported"))
else
quote
$(Expr(:meta, :inline))
checkbounds(A, I...)
_unsafe_setindex!(l, A, x, I...)
end
end
end
@inline function _unsafe_setindex!(l::LinearIndexing, A::AbstractArray, x, J::Union{Real,AbstractArray,Colon}...)
_unsafe_batchsetindex!(l, A, _iterable(x), to_indexes(J...)...)
@generated function _unsafe_setindex!(l::LinearIndexing, A::AbstractArray, x, I...)
if any(i->i<:CartesianIndex, I)
:($(Expr(:meta, :inline)); unsafe_setindex!(A, x, $(cartindex_exprs(I, :I)...)))
elseif !any(i->i<:Union{AbstractArray,Colon}, I)
:($(Expr(:meta, :inline)); setindex!(A, x, I...))
else
quote
$(Expr(:meta, :inline))
_unsafe_batchsetindex!(l, A, _iterable(x), to_nonscalar_indexes(I...)...)
end
end
end

# 1-d logical indexing: override the above to avoid calling find (in to_index)
Expand All @@ -316,7 +394,7 @@ function _unsafe_setindex!(::LinearIndexing, A::AbstractArray, x, I::AbstractArr
end

# Use iteration over X so we don't need to worry about its storage
@generated function _unsafe_batchsetindex!(::LinearFast, A::AbstractArray, X, I::Union{Real,AbstractArray,Colon}...)
@generated function _unsafe_batchsetindex!(::LinearFast, A::AbstractArray, X, I...)
N = length(I)
quote
@nexprs $N d->(I_d = I[d])
Expand All @@ -326,59 +404,28 @@ end
stride_1 = 1
@nexprs $N d->(stride_{d+1} = stride_d*size(A,d))
$(symbol(:offset_, N)) = 1
@nloops $N i d->(1:idxlens[d]) d->(offset_{d-1} = offset_d + (unsafe_getindex(I_d, i_d)-1)*stride_d) begin
@nloops $N i d->(1:idxlens[d]) d->(offset_{d-1} = offset_d + (indexref(I_d, i_d)-1)*stride_d) begin
v, Xs = next(X, Xs)
unsafe_setindex!(A, v, offset_0)
end
A
end
end
@generated function _unsafe_batchsetindex!(::LinearSlow, A::AbstractArray, X, I::Union{Real,AbstractArray,Colon}...)
@generated function _unsafe_batchsetindex!(::LinearSlow, A::AbstractArray, X, I...)
N = length(I)
quote
@nexprs $N d->(I_d = I[d])
idxlens = @ncall $N index_lengths A I
@ncall $N setindex_shape_check X (d->idxlens[d])
Xs = start(X)
@nloops $N i d->(1:idxlens[d]) d->(j_d = unsafe_getindex(I_d, i_d)) begin
@nloops $N i d->(1:idxlens[d]) d->(j_d = indexref(I_d, i_d)) begin
v, Xs = next(X, Xs)
@ncall $N unsafe_setindex! A v j
end
A
end
end

# Cartesian indexing
function cartindex_exprs(indexes, syms)
exprs = Any[]
for (i,ind) in enumerate(indexes)
if ind <: CartesianIndex
for j = 1:length(ind)
push!(exprs, :($syms[$i][$j]))
end
else
push!(exprs, :($syms[$i]))
end
end
if isempty(exprs)
push!(exprs, 1) # Handle the zero-dimensional case
end
exprs
end
@generated function _getindex{T,N}(l::LinearIndexing, A::AbstractArray{T,N}, I::Union{Real,AbstractArray,Colon,CartesianIndex}...)
:($(Expr(:meta, :inline)); getindex(A, $(cartindex_exprs(I, :I)...)))
end
@generated function _unsafe_getindex{T,N}(l::LinearIndexing, A::AbstractArray{T,N}, I::Union{Real,AbstractArray,Colon,CartesianIndex}...)
:($(Expr(:meta, :inline)); unsafe_getindex(A, $(cartindex_exprs(I, :I)...)))
end
@generated function _setindex!{T,N}(l::LinearIndexing, A::AbstractArray{T,N}, v, I::Union{Real,AbstractArray,Colon,CartesianIndex}...)
:($(Expr(:meta, :inline)); setindex!(A, v, $(cartindex_exprs(I, :I)...)))
end
@generated function _unsafe_setindex!{T,N}(l::LinearIndexing, A::AbstractArray{T,N}, v, I::Union{Real,AbstractArray,Colon,CartesianIndex}...)
:($(Expr(:meta, :inline)); unsafe_setindex!(A, v, $(cartindex_exprs(I, :I)...)))
end


##

@generated function findn{T,N}(A::AbstractArray{T,N})
Expand Down
7 changes: 1 addition & 6 deletions base/operators.jl
Original file line number Diff line number Diff line change
Expand Up @@ -309,14 +309,9 @@ function setindex_shape_check{T}(X::AbstractArray{T,2}, i, j)
end
setindex_shape_check(X, I...) = nothing # Non-arrays broadcast to all idxs

# convert to a supported index type (Array, Colon, or Int)
# convert to a supported scalar index type (Int)
to_index(i::Int) = i
to_index(i::Integer) = convert(Int,i)::Int
to_index(c::Colon) = c
to_index(I::AbstractArray{Bool}) = find(I)
to_index(A::AbstractArray) = A
to_index{T<:AbstractArray}(A::AbstractArray{T}) = throw(ArgumentError("invalid index: $A"))
to_index(A::AbstractArray{Colon}) = throw(ArgumentError("invalid index: $A"))
to_index(i) = throw(ArgumentError("invalid index: $i"))

to_indexes() = ()
Expand Down
Loading