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

Simplify scalar indexing #19958

Merged
merged 1 commit into from
Jan 28, 2017
Merged
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
135 changes: 37 additions & 98 deletions base/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -806,65 +806,45 @@ error_if_canonical_indexing(::LinearIndexing, ::AbstractArray, ::Any...) = nothi
_getindex(::LinearIndexing, A::AbstractArray, I...) = error("indexing $(typeof(A)) with types $(typeof(I)) is not supported")

## LinearFast Scalar indexing: canonical method is one Int
_getindex(::LinearFast, A::AbstractVector, i::Int) = (@_propagate_inbounds_meta; getindex(A, i))
_getindex(::LinearFast, A::AbstractArray, i::Int) = (@_propagate_inbounds_meta; getindex(A, i))
_getindex{T}(::LinearFast, A::AbstractArray{T,0}) = A[1]
function _getindex{T,N}(::LinearFast, A::AbstractArray{T,N}, I::Vararg{Int,N})
# We must check bounds for sub2ind; so we can then use @inbounds
_getindex(::LinearFast, A::AbstractArray, i::Int) = (@_propagate_inbounds_meta; getindex(A, i))
_getindex(::LinearFast, A::AbstractArray) = (@_propagate_inbounds_meta; getindex(A, _to_linear_index(A)))
function _getindex(::LinearFast, A::AbstractArray, I::Int...)
@_inline_meta
@boundscheck checkbounds(A, I...)
@inbounds r = getindex(A, sub2ind(A, I...))
r
end
function _getindex(::LinearFast, A::AbstractVector, I1::Int, I::Int...)
@_inline_meta
@boundscheck checkbounds(A, I1, I...)
@inbounds r = getindex(A, I1)
r
end
function _getindex(::LinearFast, A::AbstractArray, I::Int...) # TODO: DEPRECATE FOR #14770
@_inline_meta
@boundscheck checkbounds(A, I...)
@inbounds r = getindex(A, sub2ind(A, I...))
@boundscheck checkbounds(A, I...) # generally _to_linear_index requires bounds checking
@inbounds r = getindex(A, _to_linear_index(A, I...))
r
end

_to_linear_index(A::AbstractArray, i::Int) = i
_to_linear_index(A::AbstractVector, i::Int, I::Int...) = i # TODO: DEPRECATE FOR #14770
_to_linear_index{T,N}(A::AbstractArray{T,N}, I::Vararg{Int,N}) = (@_inline_meta; sub2ind(A, I...))
_to_linear_index(A::AbstractArray) = 1 # TODO: DEPRECATE FOR #14770
_to_linear_index(A::AbstractArray, I::Int...) = (@_inline_meta; sub2ind(A, I...)) # TODO: DEPRECATE FOR #14770

## LinearSlow Scalar indexing: Canonical method is full dimensionality of Ints
_getindex{T,N}(::LinearSlow, A::AbstractArray{T,N}, I::Vararg{Int, N}) = (@_propagate_inbounds_meta; getindex(A, I...))
function _getindex(::LinearSlow, A::AbstractArray, i::Int)
# ind2sub requires all dimensions to be > 0; may as well just check bounds
_getindex(::LinearSlow, A::AbstractArray) = (@_propagate_inbounds_meta; getindex(A, _to_subscript_indices(A)...))
function _getindex(::LinearSlow, A::AbstractArray, I::Int...)
@_inline_meta
@boundscheck checkbounds(A, i)
@inbounds r = getindex(A, ind2sub(A, i)...)
@boundscheck checkbounds(A, I...) # generally _to_subscript_indices requires bounds checking
@inbounds r = getindex(A, _to_subscript_indices(A, I...)...)
r
end
@generated function _getindex{T,AN}(::LinearSlow, A::AbstractArray{T,AN}, I::Int...) # TODO: DEPRECATE FOR #14770
N = length(I)
if N > AN
# Drop trailing ones
Isplat = Expr[:(I[$d]) for d = 1:AN]
Osplat = Expr[:(I[$d] == 1) for d = AN+1:N]
quote
@_propagate_inbounds_meta
@boundscheck (&)($(Osplat...)) || throw_boundserror(A, I)
getindex(A, $(Isplat...))
end
else
# Expand the last index into the appropriate number of indices
Isplat = Expr[:(I[$d]) for d = 1:N-1]
sz = Expr(:tuple)
sz.args = Expr[:(size(A, $d)) for d=max(N,1):AN]
szcheck = Expr[:(size(A, $d) > 0) for d=max(N,1):AN]
last_idx = N > 0 ? :(I[$N]) : 1
quote
# ind2sub requires all dimensions to be > 0:
@_propagate_inbounds_meta
@boundscheck (&)($(szcheck...)) || throw_boundserror(A, I)
getindex(A, $(Isplat...), ind2sub($sz, $last_idx)...)
end
end
_getindex{T,N}(::LinearSlow, A::AbstractArray{T,N}, I::Vararg{Int, N}) = (@_propagate_inbounds_meta; getindex(A, I...))
_to_subscript_indices(A::AbstractArray, i::Int) = (@_inline_meta; _unsafe_ind2sub(A, i))
_to_subscript_indices{T,N}(A::AbstractArray{T,N}) = (@_inline_meta; fill_to_length((), 1, Val{N})) # TODO: DEPRECATE FOR #14770
_to_subscript_indices{T}(A::AbstractArray{T,0}) = () # TODO: REMOVE FOR #14770
_to_subscript_indices{T}(A::AbstractArray{T,0}, i::Int) = () # TODO: REMOVE FOR #14770
_to_subscript_indices{T}(A::AbstractArray{T,0}, I::Int...) = () # TODO: DEPRECATE FOR #14770
function _to_subscript_indices{T,N}(A::AbstractArray{T,N}, I::Int...) # TODO: DEPRECATE FOR #14770
@_inline_meta
J, _ = IteratorsMD.split(I, Val{N}) # (maybe) drop any trailing indices
sz = _remaining_size(J, size(A)) # compute trailing size (overlapping the final index)
(front(J)..., _unsafe_ind2sub(sz, last(J))...) # (maybe) extend the last index
end
_to_subscript_indices{T,N}(A::AbstractArray{T,N}, I::Vararg{Int,N}) = I
_remaining_size(::Tuple{Any}, t::Tuple) = t
_remaining_size(h::Tuple, t::Tuple) = (@_inline_meta; _remaining_size(tail(h), tail(t)))
_unsafe_ind2sub(::Tuple{}, i) = () # ind2sub may throw(BoundsError()) in this case
_unsafe_ind2sub(sz, i) = (@_inline_meta; ind2sub(sz, i))

## Setindex! is defined similarly. We first dispatch to an internal _setindex!
# function that allows dispatch on array storage
Expand All @@ -882,65 +862,24 @@ end
_setindex!(::LinearIndexing, A::AbstractArray, v, I...) = error("indexing $(typeof(A)) with types $(typeof(I)) is not supported")

## LinearFast Scalar indexing
_setindex!(::LinearFast, A::AbstractVector, v, i::Int) = (@_propagate_inbounds_meta; setindex!(A, v, i))
_setindex!(::LinearFast, A::AbstractArray, v, i::Int) = (@_propagate_inbounds_meta; setindex!(A, v, i))
_setindex!{T}(::LinearFast, A::AbstractArray{T,0}, v) = (@_propagate_inbounds_meta; setindex!(A, v, 1))
function _setindex!{T,N}(::LinearFast, A::AbstractArray{T,N}, v, I::Vararg{Int,N})
# We must check bounds for sub2ind; so we can then use @inbounds
@_inline_meta
@boundscheck checkbounds(A, I...)
@inbounds r = setindex!(A, v, sub2ind(A, I...))
r
end
function _setindex!(::LinearFast, A::AbstractVector, v, I1::Int, I::Int...)
@_inline_meta
@boundscheck checkbounds(A, I1, I...)
@inbounds r = setindex!(A, v, I1)
r
end
function _setindex!(::LinearFast, A::AbstractArray, v, I::Int...) # TODO: DEPRECATE FOR #14770
_setindex!(::LinearFast, A::AbstractArray, v) = (@_propagate_inbounds_meta; setindex!(A, v, _to_linear_index(A)))
function _setindex!(::LinearFast, A::AbstractArray, v, I::Int...)
@_inline_meta
@boundscheck checkbounds(A, I...)
@inbounds r = setindex!(A, v, sub2ind(A, I...))
@inbounds r = setindex!(A, v, _to_linear_index(A, I...))
r
end

# LinearSlow Scalar indexing
_setindex!{T,N}(::LinearSlow, A::AbstractArray{T,N}, v, I::Vararg{Int, N}) = (@_propagate_inbounds_meta; setindex!(A, v, I...))
function _setindex!(::LinearSlow, A::AbstractArray, v, i::Int)
# ind2sub requires all dimensions to be > 0; may as well just check bounds
_setindex!(::LinearSlow, A::AbstractArray, v) = (@_propagate_inbounds_meta; setindex!(A, v, _to_subscript_indices(A)...))
function _setindex!(::LinearSlow, A::AbstractArray, v, I::Int...)
@_inline_meta
@boundscheck checkbounds(A, i)
@inbounds r = setindex!(A, v, ind2sub(A, i)...)
@boundscheck checkbounds(A, I...)
@inbounds r = setindex!(A, v, _to_subscript_indices(A, I...)...)
r
end
@generated function _setindex!{T,AN}(::LinearSlow, A::AbstractArray{T,AN}, v, I::Int...) # TODO: DEPRECATE FOR #14770
N = length(I)
if N > AN
# Drop trailing ones
Isplat = Expr[:(I[$d]) for d = 1:AN]
Osplat = Expr[:(I[$d] == 1) for d = AN+1:N]
quote
# We only check the trailing ones, so just propagate @inbounds state
@_propagate_inbounds_meta
@boundscheck (&)($(Osplat...)) || throw_boundserror(A, I)
setindex!(A, v, $(Isplat...))
end
else
# Expand the last index into the appropriate number of indices
Isplat = Expr[:(I[$d]) for d = 1:N-1]
sz = Expr(:tuple)
sz.args = Expr[:(size(A, $d)) for d=max(N,1):AN]
szcheck = Expr[:(size(A, $d) > 0) for d=max(N,1):AN]
last_idx = N > 0 ? :(I[$N]) : 1
quote
# ind2sub requires all dimensions to be > 0:
@_propagate_inbounds_meta
@boundscheck (&)($(szcheck...)) || throw_boundserror(A, I)
setindex!(A, v, $(Isplat...), ind2sub($sz, $last_idx)...)
end
end
end

## get (getindex with a default value) ##

Expand Down