Enable linear indexing with multidimensional index (#644)

Author: TheBB

src/indexing.jl

@@ -112,15 +112,15 @@ end
     end
 end
 
-@propagate_inbounds function getindex(a::StaticArray, inds::StaticVector{<:Any, Int})
-    _getindex(a, Length(inds), inds)
+@propagate_inbounds function getindex(a::StaticArray, inds::StaticArray{<:Tuple, Int})
+    _getindex(a, Size(inds), inds)
 end
 
-@generated function _getindex(a::StaticArray, ::Length{L}, inds::StaticVector{<:Any, Int}) where {L}
-    exprs = [:(a[inds[$i]]) for i = 1:L]
+@generated function _getindex(a::StaticArray, s::Size{S}, inds::StaticArray{<:Tuple, Int}) where {S}
+    exprs = [:(a[inds[$i]]) for i = 1:prod(S)]
     return quote
         @_propagate_inbounds_meta
-        similar_type(a, Size(L))(tuple($(exprs...)))
+        similar_type(a, s)(tuple($(exprs...)))
     end
 end
 
@@ -159,36 +159,36 @@ end
     end
 end
 
-@propagate_inbounds function setindex!(a::StaticArray, v, inds::StaticVector{<:Any, Int})
-    _setindex!(a, v, Length(inds), inds)
+@propagate_inbounds function setindex!(a::StaticArray, v, inds::StaticArray{<:Tuple, Int})
+    _setindex!(a, v, Size(inds), inds)
     return v
 end
 
-@generated function _setindex!(a::StaticArray, v, ::Length{L}, inds::StaticVector{<:Any, Int}) where {L}
-    exprs = [:(a[inds[$i]] = v) for i = 1:L]
+@generated function _setindex!(a::StaticArray, v, s::Size{S}, inds::StaticArray{<:Tuple, Int}) where {S}
+    exprs = [:(a[inds[$i]] = v) for i = 1:prod(S)]
     return quote
         @_propagate_inbounds_meta
-        similar_type(a, Size(L))(tuple($(exprs...)))
+        similar_type(a, s)(tuple($(exprs...)))
     end
 end
 
-@generated function _setindex!(a::StaticArray, v::AbstractArray, ::Length{L}, inds::StaticVector{<:Any, Int}) where {L}
-    exprs = [:(a[inds[$i]] = v[$i]) for i = 1:L]
+@generated function _setindex!(a::StaticArray, v::AbstractArray, s::Size{S}, inds::StaticArray{<:Tuple, Int}) where {S}
+    exprs = [:(a[inds[$i]] = v[$i]) for i = 1:prod(S)]
     return quote
         @_propagate_inbounds_meta
-        if length(v) != L
-            throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$L destination"))
+        if length(v) != $(prod(S))
+            throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$(length(inds)) destination"))
         end
         $(Expr(:block, exprs...))
     end
 end
 
-@generated function _setindex!(a::StaticArray, v::StaticArray, ::Length{L}, inds::StaticVector{<:Any, Int}) where {L}
-    exprs = [:(a[inds[$i]] = v[$i]) for i = 1:L]
+@generated function _setindex!(a::StaticArray, v::StaticArray, s::Size{S}, inds::StaticArray{<:Tuple, Int}) where {S}
+    exprs = [:(a[inds[$i]] = v[$i]) for i = 1:prod(S)]
     return quote
         @_propagate_inbounds_meta
-        if Length(typeof(v)) != L
-            throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$L destination"))
+        if Length(typeof(v)) != Length(s)
+            throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$(length(inds)) destination"))
         end
         $(Expr(:block, exprs...))
     end

test/indexing.jl

@@ -9,6 +9,9 @@ using StaticArrays, Test
 
         # Colon
         @test (@inferred getindex(sv,:)) === sv
+
+        # SArray
+        @test (@inferred getindex(sv, SMatrix{2,2}(1,4,2,3))) === SMatrix{2,2}(4,7,5,6)
     end
 
     @testset "Linear getindex() on SMatrix" begin
@@ -20,6 +23,9 @@ using StaticArrays, Test
 
         # Colon
         @test (@inferred getindex(sm,:)) === sv
+
+        # SArray
+        @test (@inferred getindex(sm, SMatrix{2,2}(1,4,2,3))) === SMatrix{2,2}(4,7,5,6)
     end
 
     @testset "Linear getindex()/setindex!() on MVector" begin
@@ -39,6 +45,9 @@ using StaticArrays, Test
         mv = MVector(0,0,0)
         @test (mv[SVector(1,3)] = SVector(4, 5); (@inferred mv == MVector(4,0,5)))
 
+        mv = MVector(0,0,0)
+        @test (mv[SMatrix{2,1}(1,3)] = SMatrix{2,1}(4, 5); (@inferred mv == MVector(4,0,5)))
+
         # Colon
         mv = MVector{4,Int}(undef)
         @test (mv[:] = vec; (@inferred getindex(mv, :))::MVector{4,Int} == MVector((4,5,6,7)))
@@ -61,6 +70,12 @@ using StaticArrays, Test
         # Colon
         mm = MMatrix{2,2,Int}(undef)
         @test (mm[:] = vec; (@inferred getindex(mm, :))::MVector{4,Int} == MVector((4,5,6,7)))
+
+        # SMatrix
+        mm = MMatrix{2,2,Int}(undef)
+        mi = MMatrix{2,2}(4,2,1,3)
+        data = @SMatrix [4 5; 6 7]
+        @test (mm[mi] = data; (@inferred getindex(mm, :))::MVector{4,Int} == MVector((5,6,7,4)))
     end
 
     @testset "Linear getindex()/setindex!() with a SVector on an Array" begin