Make setindex! return the array rather than the value (#832)

Author: c42f

src/FieldArray.jl

@@ -118,7 +118,7 @@ abstract type FieldVector{N, T} <: FieldArray{Tuple{N}, T, 1} end
 end
 
 @propagate_inbounds getindex(a::FieldArray, i::Int) = getfield(a, i)
-@propagate_inbounds setindex!(a::FieldArray, x, i::Int) = setfield!(a, i, x)
+@propagate_inbounds setindex!(a::FieldArray, x, i::Int) = (setfield!(a, i, x); a)
 
 Base.cconvert(::Type{<:Ptr}, a::FieldArray) = Base.RefValue(a)
 Base.unsafe_convert(::Type{Ptr{T}}, m::Base.RefValue{FA}) where {N,T,D,FA<:FieldArray{N,T,D}} =

src/MArray.jl

@@ -99,7 +99,7 @@ end
         error("setindex!() with non-isbitstype eltype is not supported by StaticArrays. Consider using SizedArray.")
     end
 
-    return val
+    return v
 end
 
 @inline Tuple(v::MArray) = v.data

src/indexing.jl

@@ -40,6 +40,7 @@ end
 @propagate_inbounds function setindex!(a::StaticArray, value, inds::Int...)
     @boundscheck checkbounds(a, inds...)
     _setindex!_scalar(Size(a), a, value, inds...)
+    return a
 end
 
 @generated function _setindex!_scalar(::Size{S}, a::StaticArray, value, inds::Int...) where S
@@ -126,14 +127,15 @@ end
 
 @inline function setindex!(a::StaticArray, v, ::Colon)
     _setindex!(a::StaticArray, v, Length(a), :)
-    return v
+    return a
 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...))
+        return a
     end
 end
 
@@ -145,6 +147,7 @@ end
             throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$L destination"))
         end
         @inbounds $(Expr(:block, exprs...))
+        return a
     end
 end
 
@@ -156,19 +159,21 @@ end
             throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$L destination"))
         end
         $(Expr(:block, exprs...))
+        return a
     end
 end
 
 @propagate_inbounds function setindex!(a::StaticArray, v, inds::StaticArray{<:Tuple, Int})
     _setindex!(a, v, Size(inds), inds)
-    return v
+    return a
 end
 
 @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, s)(tuple($(exprs...)))
+        return a
     end
 end
 
@@ -180,6 +185,7 @@ end
             throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$(length(inds)) destination"))
         end
         $(Expr(:block, exprs...))
+        return a
     end
 end
 
@@ -191,6 +197,7 @@ end
             throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$(length(inds)) destination"))
         end
         $(Expr(:block, exprs...))
+        return a
     end
 end
 
@@ -305,7 +312,7 @@ end
     quote
         @_propagate_inbounds_meta
         $(exprs...)
-        return value
+        return a
     end
 end
 
@@ -348,7 +355,7 @@ end
         quote
             @_propagate_inbounds_meta
             $(exprs...)
-            return v
+            return a
         end
     else
         quote
@@ -358,7 +365,7 @@ end
                 throw(DimensionMismatch("tried to assign $(length(v))-element array to $newsize destination"))
             end
             $(exprs...)
-            return v
+            return a
         end
     end
 end

src/util.jl

@@ -77,7 +77,7 @@ end
 
 size(a::TrivialView) = size(a.a)
 getindex(a::TrivialView, inds...) = getindex(a.a, inds...)
-setindex!(a::TrivialView, inds...) = setindex!(a.a, inds...)
+setindex!(a::TrivialView, inds...) = (setindex!(a.a, inds...); a)
 Base.IndexStyle(::Type{<:TrivialView{A}}) where {A} = IndexStyle(A)
 
 TrivialView(a::AbstractArray{T,N}) where {T,N} = TrivialView{typeof(a),T,N}(a)

test/FieldMatrix.jl

@@ -64,6 +64,7 @@
 
         p = Tensor2x2(0.0, 0.0, 0.0, 0.0)
         p[1,1] = 1.0
+        @test setindex!(p, 1.0, 1,1) === p
         p[2,1] = 2.0
 
         @test size(p) === (2,2)

test/MArray.jl

@@ -152,13 +152,15 @@
         v[2] = 12
         v[3] = 13
         @test v.data === (11, 12, 13)
+        @test setindex!(v, 11, 1) === v
 
         m = @MArray [0 0; 0 0]
         m[1] = 11
         m[2] = 12
         m[3] = 13
         m[4] = 14
         @test m.data === (11, 12, 13, 14)
+        @test setindex!(m, 11, 1, 1) === m
 
         @test_throws BoundsError setindex!(v, 4, -1)
         mm = @MArray zeros(3,3,3,3)

test/indexing.jl

@@ -34,6 +34,7 @@ using StaticArrays, Test
         # SVector
         mv = MVector{4,Int}(undef)
         @test (mv[SVector(1,2,3,4)] = vec; (@inferred getindex(mv, SVector(4,3,2,1)))::MVector{4,Int} == MVector((7,6,5,4)))
+        @test setindex!(mv, vec, SVector(1,2,3,4)) === mv
 
         mv = MVector{4,Int}(undef)
         @test (mv[SVector(1,2,3,4)] = [4, 5, 6, 7]; (@inferred getindex(mv, SVector(4,3,2,1)))::MVector{4,Int} == MVector((7,6,5,4)))
@@ -53,6 +54,7 @@ using StaticArrays, Test
         @test (mv[:] = vec; (@inferred getindex(mv, :))::MVector{4,Int} == MVector((4,5,6,7)))
         @test (mv[:] = [4, 5, 6, 7]; (@inferred getindex(mv, :))::MVector{4,Int} == MVector((4,5,6,7)))
         @test (mv[:] = 2; (@inferred getindex(mv, :))::MVector{4,Int} == MVector((2,2,2,2)))
+        @test setindex!(mv, 2, :) === mv
 
         @test_throws DimensionMismatch setindex!(mv, SVector(1,2,3), SVector(1,2,3,4))
         @test_throws DimensionMismatch setindex!(mv, SVector(1,2,3), :)