Skip to content

Commit 853e920

Browse files
mbaumanmbauman
committed
Simpler scalar indexing
1 parent ca3b06e commit 853e920

File tree

2 files changed

+42
-98
lines changed

2 files changed

+42
-98
lines changed

‎base/abstractarray.jl

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

808808
## LinearFast Scalar indexing: canonical method is one Int
809-
_getindex(::LinearFast, A::AbstractVector, i::Int) = (@_propagate_inbounds_meta; getindex(A, i))
810-
_getindex(::LinearFast, A::AbstractArray, i::Int) = (@_propagate_inbounds_meta; getindex(A, i))
811-
_getindex{T}(::LinearFast, A::AbstractArray{T,0}) = A[1]
812-
function _getindex{T,N}(::LinearFast, A::AbstractArray{T,N}, I::Vararg{Int,N})
813-
# We must check bounds for sub2ind; so we can then use @inbounds
809+
_getindex(::LinearFast, A::AbstractArray, i::Int) = (@_propagate_inbounds_meta; getindex(A, i))
810+
_getindex(::LinearFast, A::AbstractArray) = (@_propagate_inbounds_meta; getindex(A, _to_linear_index(A)))
811+
function _getindex(::LinearFast, A::AbstractArray, I::Int...)
814812
@_inline_meta
815-
@boundscheck checkbounds(A, I...)
816-
@inbounds r = getindex(A, sub2ind(A, I...))
817-
r
818-
end
819-
function _getindex(::LinearFast, A::AbstractVector, I1::Int, I::Int...)
820-
@_inline_meta
821-
@boundscheck checkbounds(A, I1, I...)
822-
@inbounds r = getindex(A, I1)
823-
r
824-
end
825-
function _getindex(::LinearFast, A::AbstractArray, I::Int...) # TODO: DEPRECATE FOR #14770
826-
@_inline_meta
827-
@boundscheck checkbounds(A, I...)
828-
@inbounds r = getindex(A, sub2ind(A, I...))
813+
@boundscheck checkbounds(A, I...) # generally _to_linear_index requires bounds checking
814+
@inbounds r = getindex(A, _to_linear_index(A, I...))
829815
r
830816
end
831-
817+
_to_linear_index(A::AbstractArray, i::Int) = i
818+
_to_linear_index(A::AbstractVector, i::Int, I::Int...) = i # TODO: DEPRECATE FOR #14770
819+
_to_linear_index{T,N}(A::AbstractArray{T,N}, I::Vararg{Int,N}) = (@_inline_meta; sub2ind(A, I...))
820+
_to_linear_index(A::AbstractArray) = 1 # TODO: DEPRECATE FOR #14770
821+
_to_linear_index(A::AbstractArray, I::Int...) = (@_inline_meta; sub2ind(A, I...)) # TODO: DEPRECATE FOR #14770
832822

833823
## LinearSlow Scalar indexing: Canonical method is full dimensionality of Ints
834-
_getindex{T,N}(::LinearSlow, A::AbstractArray{T,N}, I::Vararg{Int, N}) = (@_propagate_inbounds_meta; getindex(A, I...))
835-
function _getindex(::LinearSlow, A::AbstractArray, i::Int)
836-
# ind2sub requires all dimensions to be > 0; may as well just check bounds
824+
_getindex(::LinearSlow, A::AbstractArray) = (@_propagate_inbounds_meta; getindex(A, _to_subscript_indices(A)...))
825+
function _getindex(::LinearSlow, A::AbstractArray, I::Int...)
837826
@_inline_meta
838-
@boundscheck checkbounds(A, i)
839-
@inbounds r = getindex(A, ind2sub(A, i)...)
827+
@boundscheck checkbounds(A, I...) # generally _to_subscript_indices requires bounds checking
828+
@inbounds r = getindex(A, _to_subscript_indices(A, I...)...)
840829
r
841830
end
842-
@generated function _getindex{T,AN}(::LinearSlow, A::AbstractArray{T,AN}, I::Int...) # TODO: DEPRECATE FOR #14770
843-
N = length(I)
844-
if N > AN
845-
# Drop trailing ones
846-
Isplat = Expr[:(I[$d]) for d = 1:AN]
847-
Osplat = Expr[:(I[$d] == 1) for d = AN+1:N]
848-
quote
849-
@_propagate_inbounds_meta
850-
@boundscheck (&)($(Osplat...)) || throw_boundserror(A, I)
851-
getindex(A, $(Isplat...))
852-
end
853-
else
854-
# Expand the last index into the appropriate number of indices
855-
Isplat = Expr[:(I[$d]) for d = 1:N-1]
856-
sz = Expr(:tuple)
857-
sz.args = Expr[:(size(A, $d)) for d=max(N,1):AN]
858-
szcheck = Expr[:(size(A, $d) > 0) for d=max(N,1):AN]
859-
last_idx = N > 0 ? :(I[$N]) : 1
860-
quote
861-
# ind2sub requires all dimensions to be > 0:
862-
@_propagate_inbounds_meta
863-
@boundscheck (&)($(szcheck...)) || throw_boundserror(A, I)
864-
getindex(A, $(Isplat...), ind2sub($sz, $last_idx)...)
865-
end
866-
end
831+
_getindex{T,N}(::LinearSlow, A::AbstractArray{T,N}, I::Vararg{Int, N}) = (@_propagate_inbounds_meta; getindex(A, I...))
832+
_to_subscript_indices(A::AbstractArray, i::Int) = (@_inline_meta; _unsafe_ind2sub(A, i))
833+
_to_subscript_indices{T,N}(A::AbstractArray{T,N}) = (@_inline_meta; fill_to_length((), 1, Val{N})) # TODO: DEPRECATE FOR #14770
834+
_to_subscript_indices{T}(A::AbstractArray{T,0}) = () # TODO: REMOVE FOR #14770
835+
_to_subscript_indices{T}(A::AbstractArray{T,0}, i::Int) = () # TODO: REMOVE FOR #14770
836+
_to_subscript_indices{T}(A::AbstractArray{T,0}, I::Int...) = () # TODO: DEPRECATE FOR #14770
837+
function _to_subscript_indices{T,N}(A::AbstractArray{T,N}, I::Int...) # TODO: DEPRECATE FOR #14770
838+
@_inline_meta
839+
J, _ = IteratorsMD.split(I, Val{N}) # (maybe) drop any trailing indices
840+
sz = _remaining_size(J, size(A)) # compute trailing size (overlapping the final index)
841+
(front(J)..., _unsafe_ind2sub(sz, last(J))...) # (maybe) extend the last index
867842
end
843+
_to_subscript_indices{T,N}(A::AbstractArray{T,N}, I::Vararg{Int,N}) = I
844+
_remaining_size(::Tuple{Any}, t::Tuple) = t
845+
_remaining_size(h::Tuple, t::Tuple) = (@_inline_meta; _remaining_size(tail(h), tail(t)))
846+
_unsafe_ind2sub(::Tuple{}, i) = () # ind2sub may throw(BoundsError()) in this case
847+
_unsafe_ind2sub(sz, i) = (@_inline_meta; ind2sub(sz, i))
868848

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

884864
## LinearFast Scalar indexing
885-
_setindex!(::LinearFast, A::AbstractVector, v, i::Int) = (@_propagate_inbounds_meta; setindex!(A, v, i))
886865
_setindex!(::LinearFast, A::AbstractArray, v, i::Int) = (@_propagate_inbounds_meta; setindex!(A, v, i))
887-
_setindex!{T}(::LinearFast, A::AbstractArray{T,0}, v) = (@_propagate_inbounds_meta; setindex!(A, v, 1))
888-
function _setindex!{T,N}(::LinearFast, A::AbstractArray{T,N}, v, I::Vararg{Int,N})
889-
# We must check bounds for sub2ind; so we can then use @inbounds
890-
@_inline_meta
891-
@boundscheck checkbounds(A, I...)
892-
@inbounds r = setindex!(A, v, sub2ind(A, I...))
893-
r
894-
end
895-
function _setindex!(::LinearFast, A::AbstractVector, v, I1::Int, I::Int...)
896-
@_inline_meta
897-
@boundscheck checkbounds(A, I1, I...)
898-
@inbounds r = setindex!(A, v, I1)
899-
r
900-
end
901-
function _setindex!(::LinearFast, A::AbstractArray, v, I::Int...) # TODO: DEPRECATE FOR #14770
866+
_setindex!(::LinearFast, A::AbstractArray, v) = (@_propagate_inbounds_meta; setindex!(A, v, _to_linear_index(A)))
867+
function _setindex!(::LinearFast, A::AbstractArray, v, I::Int...)
902868
@_inline_meta
903869
@boundscheck checkbounds(A, I...)
904-
@inbounds r = setindex!(A, v, sub2ind(A, I...))
870+
@inbounds r = setindex!(A, v, _to_linear_index(A, I...))
905871
r
906872
end
907873

908874
# LinearSlow Scalar indexing
909875
_setindex!{T,N}(::LinearSlow, A::AbstractArray{T,N}, v, I::Vararg{Int, N}) = (@_propagate_inbounds_meta; setindex!(A, v, I...))
910-
function _setindex!(::LinearSlow, A::AbstractArray, v, i::Int)
911-
# ind2sub requires all dimensions to be > 0; may as well just check bounds
876+
_setindex!(::LinearSlow, A::AbstractArray, v) = (@_propagate_inbounds_meta; setindex!(A, v, _to_subscript_indices(A)...))
877+
function _setindex!(::LinearSlow, A::AbstractArray, v, I::Int...)
912878
@_inline_meta
913-
@boundscheck checkbounds(A, i)
914-
@inbounds r = setindex!(A, v, ind2sub(A, i)...)
879+
@boundscheck checkbounds(A, I...)
880+
@inbounds r = setindex!(A, v, _to_subscript_indices(A, I...)...)
915881
r
916882
end
917-
@generated function _setindex!{T,AN}(::LinearSlow, A::AbstractArray{T,AN}, v, I::Int...) # TODO: DEPRECATE FOR #14770
918-
N = length(I)
919-
if N > AN
920-
# Drop trailing ones
921-
Isplat = Expr[:(I[$d]) for d = 1:AN]
922-
Osplat = Expr[:(I[$d] == 1) for d = AN+1:N]
923-
quote
924-
# We only check the trailing ones, so just propagate @inbounds state
925-
@_propagate_inbounds_meta
926-
@boundscheck (&)($(Osplat...)) || throw_boundserror(A, I)
927-
setindex!(A, v, $(Isplat...))
928-
end
929-
else
930-
# Expand the last index into the appropriate number of indices
931-
Isplat = Expr[:(I[$d]) for d = 1:N-1]
932-
sz = Expr(:tuple)
933-
sz.args = Expr[:(size(A, $d)) for d=max(N,1):AN]
934-
szcheck = Expr[:(size(A, $d) > 0) for d=max(N,1):AN]
935-
last_idx = N > 0 ? :(I[$N]) : 1
936-
quote
937-
# ind2sub requires all dimensions to be > 0:
938-
@_propagate_inbounds_meta
939-
@boundscheck (&)($(szcheck...)) || throw_boundserror(A, I)
940-
setindex!(A, v, $(Isplat...), ind2sub($sz, $last_idx)...)
941-
end
942-
end
943-
end
944883

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

‎test/TestHelpers.jl

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,11 @@ end
132132
val
133133
end
134134

135+
# Add support for indexing into OffsetVectors with trailing 1s
136+
# These are required because OffsetVectors don't support sub2ind
137+
Base._to_linear_index(A::OffsetVector, I::Int...) = I[1] # TODO: REMOVE FOR #14770
138+
Base._to_linear_index(A::OffsetVector) = error("linear indexing is not defined for one-dimensional arrays")
139+
135140
# Computing a shifted index (subtracting the offset)
136141
offset{N}(offsets::NTuple{N,Int}, inds::NTuple{N,Int}) = _offset((), offsets, inds)
137142
_offset(out, ::Tuple{}, ::Tuple{}) = out

0 commit comments

Comments
 (0)