Fix reshaping of a mutable StaticArray (#933)

* Fix reshaping of a mutable StaticArray

* more reshape changes

Author: mateuszbaran

Project.toml

@@ -1,6 +1,6 @@
 name = "StaticArrays"
 uuid = "90137ffa-7385-5640-81b9-e52037218182"
-version = "1.2.5"
+version = "1.2.6"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/abstractarray.jl

@@ -183,26 +183,20 @@ homogenize_shape(shape::Tuple{Vararg{SOneTo}}) = Size(map(last, shape))
 homogenize_shape(shape::Tuple{Vararg{HeterogeneousShape}}) = map(last, shape)
 
 
-@inline reshape(a::StaticArray, s::Size) = similar_type(a, s)(Tuple(a))
-@inline reshape(a::AbstractArray, s::Size) = _reshape(a, IndexStyle(a), s)
-@inline reshape(a::StaticArray, s::Tuple{SOneTo,Vararg{SOneTo}}) = reshape(a, homogenize_shape(s))
-@generated function _reshape(a::AbstractArray, indexstyle, s::Size{S}) where {S}
-    if indexstyle == IndexLinear
-        exprs = [:(a[$i]) for i = 1:prod(S)]
-    else
-        exprs = [:(a[$(inds)]) for inds ∈ CartesianIndices(S)]
-    end
-
-    return quote
-        @_inline_meta
-        if length(a) != prod(s)
-            throw(DimensionMismatch("Tried to resize dynamic object of size $(size(a)) to $s"))
-        end
-        return similar_type(a, s)(tuple($(exprs...)))
-    end
+@inline reshape(a::SArray, s::Size) = similar_type(a, s)(Tuple(a))
+@inline reshape(a::AbstractArray, s::Size) = __reshape(a, ((typeof(s).parameters...)...,), s)
+@inline reshape(a::SArray, s::Tuple{SOneTo,Vararg{SOneTo}}) = reshape(a, homogenize_shape(s))
+@inline function reshape(a::StaticArray, s::Tuple{SOneTo,Vararg{SOneTo}})
+    return __reshape(a, map(u -> last(u), s), homogenize_shape(s))
+end
+@inline function __reshape(a, shape, ::Size{S}) where {S}
+    return SizedArray{Tuple{S...}}(Base._reshape(a, shape))
+end
+@inline function __reshape(a::SizedArray, shape, ::Size{S}) where {S}
+    return SizedArray{Tuple{S...}}(Base._reshape(a.data, shape))
 end
 
-reshape(a::Array, ::Size{S}) where {S} = SizedArray{Tuple{S...}}(a)
+reshape(a::Vector, ::Size{S}) where {S} = SizedArray{Tuple{S...}}(a)
 
 Base.rdims(out::Val{N}, inds::Tuple{SOneTo, Vararg{SOneTo}}) where {N} = Base.rdims(ntuple(i -> SOneTo(1), Val(N)), inds)
 Base.rdims(out::Tuple{Any}, inds::Tuple{SOneTo, Vararg{SOneTo}}) = (SOneTo(Base.rdims_trailing(inds...)),)

test/abstractarray.jl

@@ -109,7 +109,7 @@ using StaticArrays, Test, LinearAlgebra
         @test @inferred(reshape(SVector(1,2,3,4), axes(SMatrix{2,2}(1,2,3,4)))) === SMatrix{2,2}(1,2,3,4)
         @test @inferred(reshape(SVector(1,2,3,4), Size(2,2))) === SMatrix{2,2}(1,2,3,4)
         @test @inferred(reshape([1,2,3,4], Size(2,2)))::SizedArray{Tuple{2,2},Int,2,1} == [1 3; 2 4]
-        @test_throws DimensionMismatch reshape([1 2; 3 4], Size(2,1,2))
+        @test_throws DimensionMismatch reshape([1 2; 3 4], Size(2,2,2))
 
         @test @inferred(vec(SMatrix{2, 2}([1 2; 3 4])))::SVector{4,Int} == [1, 3, 2, 4]
 
@@ -119,6 +119,21 @@ using StaticArrays, Test, LinearAlgebra
         # IndexLinear
         @test reshape(view(ones(4, 4), 1, 1:4), Size(4, 1)) == SMatrix{4,1}(ones(4, 1))
         @test_throws DimensionMismatch reshape(view(ones(4,4), 1:4, 1:2), Size(5, 2))
+
+        # mutation
+        m = @MMatrix [1 2; 3 4]
+        mr = reshape(m, SOneTo(4))
+        mr[2] = 10
+        @test m == SA[1 2; 10 4]
+
+        mrs = reshape(m, Size(4))
+        mrs[2] = 10
+        @test m == SA[1 2; 10 4]
+
+        ms = SizedMatrix{2,2}([1 2; 3 4])
+        msr = reshape(ms, SOneTo(4))
+        msr[2] = 10
+        @test ms == SA[1 2; 10 4]
     end
 
     @testset "copy" begin

test/broadcast.jl

@@ -270,4 +270,15 @@ end
             @test DS2 == DS .* DS
         end
     end
+
+    @testset "broadcast! with reshaping" begin
+        m = @MMatrix [1 2; 3 4]
+        m[1:2] .= 10
+        @test m == SA[10 2; 10 4]
+
+        ms = SizedMatrix{2,2}([1 2; 3 4])
+        ms[1:2] .= 10
+        @test ms == SA[10 2; 10 4]
+    end
+
 end