diff --git a/src/indexing.jl b/src/indexing.jl index eeedd79eb..2b35270b4 100644 --- a/src/indexing.jl +++ b/src/indexing.jl @@ -2,9 +2,9 @@ getindex(a::StaticArray, i::Int) = error("getindex(::$typeof(a), ::Int) is not defined.") setindex!(a::StaticArray, value, i::Int) = error("setindex!(::$(typeof(a)), value, ::Int) is not defined.") -###################### -## Scalar Indexing ## -###################### +####################################### +## Multidimensional scalar indexing ## +####################################### # Note: all indexing behavior defaults to dense, linear indexing @@ -48,8 +48,6 @@ end end end - - ######################### ## Indexing utilities ## ######################### @@ -79,11 +77,109 @@ _ind(i::Int, ::Int, ::Type{Int}) = :(inds[$i]) _ind(i::Int, j::Int, ::Type{<:StaticArray}) = :(inds[$i][$j]) _ind(i::Int, j::Int, ::Type{Colon}) = j +################################ +## Non-scalar linear indexing ## +################################ + +@inline function getindex(a::StaticArray, ::Colon) + _getindex(a::StaticArray, Length(a), :) +end + +@generated function _getindex(a::StaticArray, ::Length{L}, ::Colon) where {L} + exprs = [:(a[$i]) for i = 1:L] + return quote + @_inline_meta + @inbounds return similar_type(a, Size(L))(tuple($(exprs...))) + end +end + +@propagate_inbounds function getindex(a::StaticArray, inds::StaticArray{Int}) + _getindex(a, Length(inds), inds) +end + +@generated function _getindex(a::StaticArray, ::Length{L}, inds::StaticArray{Int}) where {L} + exprs = [:(a[inds[$i]]) for i = 1:L] + return quote + @_propagate_inbounds_meta + similar_type(a, Size(L))(tuple($(exprs...))) + end +end + +@inline function setindex!(a::StaticArray, v, ::Colon) + _setindex!(a::StaticArray, v, Length(a), :) + return v +end +@generated function _setindex!(a::StaticArray, v, ::Length{L}, ::Colon) where {L} + exprs = [:(a[$i] = v) for i = 1:L] + return quote + @_inline_meta + @inbounds $(Expr(:block, exprs...)) + end +end -##################### -## Array Indexing ## -##################### +@generated function _setindex!(a::StaticArray, v::AbstractArray, ::Length{L}, ::Colon) where {L} + exprs = [:(a[$i] = v[$i]) for i = 1:L] + return quote + @_propagate_inbounds_meta + if length(v) != L + throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$L destination")) + end + @inbounds $(Expr(:block, exprs...)) + end +end + +@generated function _setindex!(a::StaticArray, v::StaticArray, ::Length{L}, ::Colon) where {L} + exprs = [:(a[$i] = v[$i]) for i = 1:L] + return quote + @_propagate_inbounds_meta + if Length(typeof(v)) != L + throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$L destination")) + end + $(Expr(:block, exprs...)) + end +end + +@propagate_inbounds function setindex!(a::StaticArray, v, inds::StaticArray{Int}) + _setindex!(a, v, Length(inds), inds) + return v +end + +@generated function _setindex!(a::StaticArray, v, ::Length{L}, inds::StaticArray{Int}) where {L} + exprs = [:(a[inds[$i]] = v) for i = 1:L] + return quote + @_propagate_inbounds_meta + similar_type(a, Size(L))(tuple($(exprs...))) + end +end + +@generated function _setindex!(a::StaticArray, v::AbstractArray, ::Length{L}, inds::StaticArray{Int}) where {L} + exprs = [:(a[$i] = v[$i]) for i = 1:L] + return quote + @_propagate_inbounds_meta + if length(v) != L + throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$L destination")) + end + $(Expr(:block, exprs...)) + end +end + +@generated function _setindex!(a::StaticArray, v::StaticArray, ::Length{L}, inds::StaticArray{Int}) where {L} + exprs = [:(a[$i] = v[$i]) for i = 1:L] + return quote + @_propagate_inbounds_meta + if Length(typeof(v)) != L + throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$L destination")) + end + $(Expr(:block, exprs...)) + end +end + +########################################### +## Multidimensional non-scalar indexing ## +########################################### + +# getindex @propagate_inbounds function getindex(a::StaticArray, inds::Union{Int, StaticArray{Int}, Colon}...) _getindex(a, index_sizes(Size(a), inds...), inds) @@ -124,7 +220,6 @@ end _getindex(a, index_sizes(i1, i2, i3, i4, inds...), (i1, i2, i3, i4, inds...)) end - @generated function _getindex(a::AbstractArray, ind_sizes::Tuple{Vararg{Size}}, inds) newsize = out_index_size(ind_sizes.parameters...) linearsizes = linear_index_size(ind_sizes.parameters...) @@ -161,7 +256,7 @@ end end end - +# setindex! @propagate_inbounds function setindex!(a::StaticArray, value, inds::Union{Int, StaticArray{Int}, Colon}...) _setindex!(a, value, index_sizes(Size(a), inds...), inds) @@ -264,7 +359,6 @@ end end end - # setindex! from an array @generated function _setindex!(a::AbstractArray, v::AbstractArray, ind_sizes::Tuple{Vararg{Size}}, inds) linearsizes = linear_index_size(ind_sizes.parameters...) diff --git a/test/indexing.jl b/test/indexing.jl index 6cb363844..b4b3c5a71 100644 --- a/test/indexing.jl +++ b/test/indexing.jl @@ -1,6 +1,7 @@ @testset "Indexing" begin @testset "Linear getindex() on SVector" begin sv = SVector{4}(4,5,6,7) + sm = SMatrix{2,2}(4,5,6,7) # SVector @test (@inferred getindex(sv, SVector(4,3,2,1))) === SVector((7,6,5,4)) @@ -9,6 +10,17 @@ @test (@inferred getindex(sv,:)) === sv end + @testset "Linear getindex() on SMatrix" begin + sv = SVector{4}(4,5,6,7) + sm = SMatrix{2,2}(4,5,6,7) + + # SVector + @test (@inferred getindex(sm, SVector(4,3,2,1))) === SVector((7,6,5,4)) + + # Colon + @test (@inferred getindex(sm,:)) === sv + end + @testset "Linear getindex()/setindex!() on MVector" begin vec = @SVector [4,5,6,7]