add setindex! specialization for VectorOfArray{StructArray}

Author: jlchan

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