Merge pull request #425 from jlchan/jc/VoA_StructArray_setindex

Specialize `setindex!` for `VectorOfArray{StructArray}`

Author: ChrisRackauckas

ext/RecursiveArrayToolsStructArraysExt.jl

@@ -3,4 +3,25 @@ module RecursiveArrayToolsStructArraysExt
 import RecursiveArrayTools, StructArrays
 RecursiveArrayTools.rewrap(::StructArrays.StructArray, u) = StructArrays.StructArray(u)
 
-end
+using RecursiveArrayTools: VectorOfArray
+using StructArrays: StructArray
+
+const VectorOfStructArray{T, N} = VectorOfArray{T, N, <:StructArray}
+
+# Since `StructArray` lazily materializes struct entries, the general `setindex!(x, val, I)` 
+# operation `VA.u[I[end]][Base.front(I)...]` will only update a lazily materialized struct 
+# entry of `u`, but will not actually mutate `x::StructArray`. See the StructArray documentation 
+# for more details:
+#
+#   https://juliaarrays.github.io/StructArrays.jl/stable/counterintuitive/#Modifying-a-field-of-a-struct-element 
+# 
+# To avoid this, we can materialize a struct entry, modify it, and then use `setindex!` 
+# with the modified struct entry.
+function Base.setindex!(VA::VectorOfStructArray{T, N}, v,
+        I::Int...) where {T, N}
+    u_I = VA.u[I[end]]
+    u_I[Base.front(I)...] = v
+    return VA.u[I[end]] = u_I
+end
+
+end

test/basic_indexing.jl

@@ -262,3 +262,12 @@ num_allocs = @allocations foo!(u_matrix)
 
 # issue 354
 @test VectorOfArray(ones(1))[:] == ones(1)
+
+# check VectorOfArray indexing for a StructArray of mutable structs
+using StructArrays
+using StaticArrays: MVector
+x = VectorOfArray(StructArray{MVector{1, Float64}}(ntuple(_ -> [1.0, 2.0], 1)))
+
+# check VectorOfArray assignment 
+x[1, 1] = 10
+@test x[1, 1] == 10