Support indexing with reshape

Author: jishnub

src/indexing.jl

@@ -230,12 +230,14 @@ end
 # getindex
 
 @propagate_inbounds function getindex(a::StaticArray, inds::Union{Int, StaticArray{<:Tuple, Int}, SOneTo, Colon}...)
-    _getindex(a, index_sizes(Size(a), inds...), inds)
+    ar = reshape(a, Val(length(inds)))
+    _getindex(ar, index_sizes(Size(ar), inds...), inds)
 end
 
 if isdefined(Base, :IdentityUnitRange)
     @propagate_inbounds function getindex(a::StaticArray, inds::Union{Int, StaticArray{<:Tuple, Int}, SOneTo, Colon, Base.IdentityUnitRange{<:SOneTo}}...)
-        _getindex(a, index_sizes(Size(a), inds...), inds)
+        ar = reshape(a, Val(length(inds)))
+        _getindex(ar, index_sizes(Size(ar), inds...), inds)
     end
 end
 

test/indexing.jl

@@ -137,6 +137,12 @@ using StaticArrays, Test
         @test v[2,1] == 2
         @test_throws BoundsError v[1,2]
         @test_throws BoundsError v[3,1]
+
+        # SOneTo
+        @test (@inferred v[axes(v,1), SOneTo(1)]) === SMatrix{2,1}(v)
+        @test v[axes(v,1), SOneTo(1)] == v[Base.OneTo(length(v)), Base.OneTo(1)]
+        @test (@inferred v[axes(v,1), 1, SOneTo(1)]) === SMatrix{2,1}(v)
+        @test v[axes(v,1), 1, SOneTo(1)] == v[Base.OneTo(length(v)), 1, Base.OneTo(1)]
     end
 
     @testset "2D getindex() on SMatrix" begin
@@ -249,6 +255,29 @@ using StaticArrays, Test
         @test (@inferred getindex(a, SVector(1,2), 1, 1, 1)) == [24,48]
     end
 
+    @testset "indexing with reshape for SMatrix/MMatrix" begin
+        sm = @SMatrix [1 3; 2 4]
+        mm = @MMatrix [1 3; 2 4]
+        for m in Any[sm, mm, view(sm, :, :), view(mm, :, :)]
+            sa = @inferred m[:, SOneTo(1), 1, SOneTo(1)]
+            a =  m[:, Base.OneTo(1), 1, Base.OneTo(1)]
+            @test sa == a
+            @test sa == SArray{Tuple{2,1,1}}(a)
+            if m isa SArray
+                @test sa === SArray{Tuple{2,1,1}}(a)
+            end
+
+            if isdefined(Base, :IdentityUnitRange)
+                sa = @inferred m[:, Base.IdentityUnitRange(SOneTo(1)), 1, SOneTo(1)]
+                @test sa == a
+                @test sa == SArray{Tuple{2,1,1}}(a)
+                if m isa SArray
+                    @test sa === SArray{Tuple{2,1,1}}(a)
+                end
+            end
+        end
+    end
+
     @testset "Indexing with empty vectors" begin
         a = [1.0 2.0; 3.0 4.0]
         @test a[SVector{0,Int}()] == SVector{0,Float64}(())