Skip to content

Commit 1445962

Browse files
mbaumantkelman
mbauman
authored and
tkelman
committed
Simpler scalar indexing (#19958)
1 parent 1635d98 commit 1445962

File tree

1 file changed

+37
-98
lines changed

1 file changed

+37
-98
lines changed

‎base/abstractarray.jl

Lines changed: 37 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -823,65 +823,45 @@ error_if_canonical_indexing(::LinearIndexing, ::AbstractArray, ::Any...) = nothi
823823
_getindex(::LinearIndexing, A::AbstractArray, I...) = error("indexing $(typeof(A)) with types $(typeof(I)) is not supported")
824824

825825
## LinearFast Scalar indexing: canonical method is one Int
826-
_getindex(::LinearFast, A::AbstractVector, i::Int) = (@_propagate_inbounds_meta; getindex(A, i))
827-
_getindex(::LinearFast, A::AbstractArray, i::Int) = (@_propagate_inbounds_meta; getindex(A, i))
828-
_getindex{T}(::LinearFast, A::AbstractArray{T,0}) = A[1]
829-
function _getindex{T,N}(::LinearFast, A::AbstractArray{T,N}, I::Vararg{Int,N})
830-
# We must check bounds for sub2ind; so we can then use @inbounds
826+
_getindex(::LinearFast, A::AbstractArray, i::Int) = (@_propagate_inbounds_meta; getindex(A, i))
827+
_getindex(::LinearFast, A::AbstractArray) = (@_propagate_inbounds_meta; getindex(A, _to_linear_index(A)))
828+
function _getindex(::LinearFast, A::AbstractArray, I::Int...)
831829
@_inline_meta
832-
@boundscheck checkbounds(A, I...)
833-
@inbounds r = getindex(A, sub2ind(A, I...))
834-
r
835-
end
836-
function _getindex(::LinearFast, A::AbstractVector, I1::Int, I::Int...)
837-
@_inline_meta
838-
@boundscheck checkbounds(A, I1, I...)
839-
@inbounds r = getindex(A, I1)
840-
r
841-
end
842-
function _getindex(::LinearFast, A::AbstractArray, I::Int...) # TODO: DEPRECATE FOR #14770
843-
@_inline_meta
844-
@boundscheck checkbounds(A, I...)
845-
@inbounds r = getindex(A, sub2ind(A, I...))
830+
@boundscheck checkbounds(A, I...) # generally _to_linear_index requires bounds checking
831+
@inbounds r = getindex(A, _to_linear_index(A, I...))
846832
r
847833
end
848-
834+
_to_linear_index(A::AbstractArray, i::Int) = i
835+
_to_linear_index(A::AbstractVector, i::Int, I::Int...) = i # TODO: DEPRECATE FOR #14770
836+
_to_linear_index{T,N}(A::AbstractArray{T,N}, I::Vararg{Int,N}) = (@_inline_meta; sub2ind(A, I...))
837+
_to_linear_index(A::AbstractArray) = 1 # TODO: DEPRECATE FOR #14770
838+
_to_linear_index(A::AbstractArray, I::Int...) = (@_inline_meta; sub2ind(A, I...)) # TODO: DEPRECATE FOR #14770
849839

850840
## LinearSlow Scalar indexing: Canonical method is full dimensionality of Ints
851-
_getindex{T,N}(::LinearSlow, A::AbstractArray{T,N}, I::Vararg{Int, N}) = (@_propagate_inbounds_meta; getindex(A, I...))
852-
function _getindex(::LinearSlow, A::AbstractArray, i::Int)
853-
# ind2sub requires all dimensions to be > 0; may as well just check bounds
841+
_getindex(::LinearSlow, A::AbstractArray) = (@_propagate_inbounds_meta; getindex(A, _to_subscript_indices(A)...))
842+
function _getindex(::LinearSlow, A::AbstractArray, I::Int...)
854843
@_inline_meta
855-
@boundscheck checkbounds(A, i)
856-
@inbounds r = getindex(A, ind2sub(A, i)...)
844+
@boundscheck checkbounds(A, I...) # generally _to_subscript_indices requires bounds checking
845+
@inbounds r = getindex(A, _to_subscript_indices(A, I...)...)
857846
r
858847
end
859-
@generated function _getindex{T,AN}(::LinearSlow, A::AbstractArray{T,AN}, I::Int...) # TODO: DEPRECATE FOR #14770
860-
N = length(I)
861-
if N > AN
862-
# Drop trailing ones
863-
Isplat = Expr[:(I[$d]) for d = 1:AN]
864-
Osplat = Expr[:(I[$d] == 1) for d = AN+1:N]
865-
quote
866-
@_propagate_inbounds_meta
867-
@boundscheck (&)($(Osplat...)) || throw_boundserror(A, I)
868-
getindex(A, $(Isplat...))
869-
end
870-
else
871-
# Expand the last index into the appropriate number of indices
872-
Isplat = Expr[:(I[$d]) for d = 1:N-1]
873-
sz = Expr(:tuple)
874-
sz.args = Expr[:(size(A, $d)) for d=max(N,1):AN]
875-
szcheck = Expr[:(size(A, $d) > 0) for d=max(N,1):AN]
876-
last_idx = N > 0 ? :(I[$N]) : 1
877-
quote
878-
# ind2sub requires all dimensions to be > 0:
879-
@_propagate_inbounds_meta
880-
@boundscheck (&)($(szcheck...)) || throw_boundserror(A, I)
881-
getindex(A, $(Isplat...), ind2sub($sz, $last_idx)...)
882-
end
883-
end
848+
_getindex{T,N}(::LinearSlow, A::AbstractArray{T,N}, I::Vararg{Int, N}) = (@_propagate_inbounds_meta; getindex(A, I...))
849+
_to_subscript_indices(A::AbstractArray, i::Int) = (@_inline_meta; _unsafe_ind2sub(A, i))
850+
_to_subscript_indices{T,N}(A::AbstractArray{T,N}) = (@_inline_meta; fill_to_length((), 1, Val{N})) # TODO: DEPRECATE FOR #14770
851+
_to_subscript_indices{T}(A::AbstractArray{T,0}) = () # TODO: REMOVE FOR #14770
852+
_to_subscript_indices{T}(A::AbstractArray{T,0}, i::Int) = () # TODO: REMOVE FOR #14770
853+
_to_subscript_indices{T}(A::AbstractArray{T,0}, I::Int...) = () # TODO: DEPRECATE FOR #14770
854+
function _to_subscript_indices{T,N}(A::AbstractArray{T,N}, I::Int...) # TODO: DEPRECATE FOR #14770
855+
@_inline_meta
856+
J, _ = IteratorsMD.split(I, Val{N}) # (maybe) drop any trailing indices
857+
sz = _remaining_size(J, size(A)) # compute trailing size (overlapping the final index)
858+
(front(J)..., _unsafe_ind2sub(sz, last(J))...) # (maybe) extend the last index
884859
end
860+
_to_subscript_indices{T,N}(A::AbstractArray{T,N}, I::Vararg{Int,N}) = I
861+
_remaining_size(::Tuple{Any}, t::Tuple) = t
862+
_remaining_size(h::Tuple, t::Tuple) = (@_inline_meta; _remaining_size(tail(h), tail(t)))
863+
_unsafe_ind2sub(::Tuple{}, i) = () # ind2sub may throw(BoundsError()) in this case
864+
_unsafe_ind2sub(sz, i) = (@_inline_meta; ind2sub(sz, i))
885865

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

901881
## LinearFast Scalar indexing
902-
_setindex!(::LinearFast, A::AbstractVector, v, i::Int) = (@_propagate_inbounds_meta; setindex!(A, v, i))
903882
_setindex!(::LinearFast, A::AbstractArray, v, i::Int) = (@_propagate_inbounds_meta; setindex!(A, v, i))
904-
_setindex!{T}(::LinearFast, A::AbstractArray{T,0}, v) = (@_propagate_inbounds_meta; setindex!(A, v, 1))
905-
function _setindex!{T,N}(::LinearFast, A::AbstractArray{T,N}, v, I::Vararg{Int,N})
906-
# We must check bounds for sub2ind; so we can then use @inbounds
907-
@_inline_meta
908-
@boundscheck checkbounds(A, I...)
909-
@inbounds r = setindex!(A, v, sub2ind(A, I...))
910-
r
911-
end
912-
function _setindex!(::LinearFast, A::AbstractVector, v, I1::Int, I::Int...)
913-
@_inline_meta
914-
@boundscheck checkbounds(A, I1, I...)
915-
@inbounds r = setindex!(A, v, I1)
916-
r
917-
end
918-
function _setindex!(::LinearFast, A::AbstractArray, v, I::Int...) # TODO: DEPRECATE FOR #14770
883+
_setindex!(::LinearFast, A::AbstractArray, v) = (@_propagate_inbounds_meta; setindex!(A, v, _to_linear_index(A)))
884+
function _setindex!(::LinearFast, A::AbstractArray, v, I::Int...)
919885
@_inline_meta
920886
@boundscheck checkbounds(A, I...)
921-
@inbounds r = setindex!(A, v, sub2ind(A, I...))
887+
@inbounds r = setindex!(A, v, _to_linear_index(A, I...))
922888
r
923889
end
924890

925891
# LinearSlow Scalar indexing
926892
_setindex!{T,N}(::LinearSlow, A::AbstractArray{T,N}, v, I::Vararg{Int, N}) = (@_propagate_inbounds_meta; setindex!(A, v, I...))
927-
function _setindex!(::LinearSlow, A::AbstractArray, v, i::Int)
928-
# ind2sub requires all dimensions to be > 0; may as well just check bounds
893+
_setindex!(::LinearSlow, A::AbstractArray, v) = (@_propagate_inbounds_meta; setindex!(A, v, _to_subscript_indices(A)...))
894+
function _setindex!(::LinearSlow, A::AbstractArray, v, I::Int...)
929895
@_inline_meta
930-
@boundscheck checkbounds(A, i)
931-
@inbounds r = setindex!(A, v, ind2sub(A, i)...)
896+
@boundscheck checkbounds(A, I...)
897+
@inbounds r = setindex!(A, v, _to_subscript_indices(A, I...)...)
932898
r
933899
end
934-
@generated function _setindex!{T,AN}(::LinearSlow, A::AbstractArray{T,AN}, v, I::Int...) # TODO: DEPRECATE FOR #14770
935-
N = length(I)
936-
if N > AN
937-
# Drop trailing ones
938-
Isplat = Expr[:(I[$d]) for d = 1:AN]
939-
Osplat = Expr[:(I[$d] == 1) for d = AN+1:N]
940-
quote
941-
# We only check the trailing ones, so just propagate @inbounds state
942-
@_propagate_inbounds_meta
943-
@boundscheck (&)($(Osplat...)) || throw_boundserror(A, I)
944-
setindex!(A, v, $(Isplat...))
945-
end
946-
else
947-
# Expand the last index into the appropriate number of indices
948-
Isplat = Expr[:(I[$d]) for d = 1:N-1]
949-
sz = Expr(:tuple)
950-
sz.args = Expr[:(size(A, $d)) for d=max(N,1):AN]
951-
szcheck = Expr[:(size(A, $d) > 0) for d=max(N,1):AN]
952-
last_idx = N > 0 ? :(I[$N]) : 1
953-
quote
954-
# ind2sub requires all dimensions to be > 0:
955-
@_propagate_inbounds_meta
956-
@boundscheck (&)($(szcheck...)) || throw_boundserror(A, I)
957-
setindex!(A, v, $(Isplat...), ind2sub($sz, $last_idx)...)
958-
end
959-
end
960-
end
961900

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

0 commit comments

Comments
 (0)