Specialize fill! (#253)

Author: jishnub

Project.toml

@@ -23,10 +23,11 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 PooledArrays = "2dfb63ee-cc39-5dd5-95bd-886bf059d720"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9"
 WeakRefStrings = "ea10d353-3f73-51f8-a26c-33c1cb351aa5"
 
 [targets]
-test = ["Test", "JLArrays", "StaticArrays", "OffsetArrays", "PooledArrays", "TypedTables", "WeakRefStrings", "Documenter"]
+test = ["Test", "JLArrays", "StaticArrays", "OffsetArrays", "PooledArrays", "TypedTables", "WeakRefStrings", "Documenter", "SparseArrays"]

src/structarray.jl

@@ -425,6 +425,12 @@ function Base.copyto!(I::StructArray, doffs::Integer, J::StructArray, soffs::Int
     return I
 end
 
+function Base.fill!(s::StructArray{T}, x) where {T}
+    xT = maybe_convert_elt(T, x)
+    foreachfield(fill!, s, xT)
+    return s
+end
+
 function Base.resize!(s::StructArray, i::Integer)
     for a in components(s)
         resize!(a, i)

test/runtests.jl

@@ -8,6 +8,7 @@ using DataAPI: refarray, refvalue
 using Adapt: adapt, Adapt
 using JLArrays
 using Test
+using SparseArrays
 
 using Documenter: doctest
 if Base.VERSION >= v"1.6" && Int === Int64
@@ -569,6 +570,54 @@ end
     @test t[5] == p
 end
 
+@testset "fill!" begin
+    @testset "dense array, complex" begin
+        A = zeros(3,3)
+        B = zeros(3,3)
+        S = StructArray{Complex{eltype(A)}}((A, B))
+        fill!(S, 2+3im)
+        @test all(==(2), A)
+        @test all(==(3), B)
+    end
+
+    @testset "offset array, custom struct" begin
+        struct Vec3D{T} <: FieldVector{3, T}
+            x :: T
+            y :: T
+            z :: T
+        end
+
+        A = zeros(3:6, 3:6)
+        B = zeros(3:6, 3:6)
+        C = zeros(3:6, 3:6)
+        S = StructArray{Vec3D{eltype(A)}}((A, B, C))
+        fill!(S, Vec3D(1,2,3))
+        @test all(==(1), A)
+        @test all(==(2), B)
+        @test all(==(3), C)
+    end
+
+    @testset "Tuple" begin
+        S = StructArray{Tuple{Int,Int}}(([1,2], [3,4]))
+        fill!(S, (4,5))
+        @test all(==((4,5)), S)
+
+        S = StructArray{@NamedTuple{a::Int,b::Int}}(([1,2], [3,4]))
+        fill!(S, (a=10.0, b=20.0))
+        @test all(==(10), S.a)
+        @test all(==(20), S.b)
+    end
+
+    @testset "sparse matrix, complex" begin
+        A = spzeros(3)
+        B = spzeros(3)
+        S = StructArray{Complex{eltype(A)}}((A,B))
+        fill!(S, 0)
+        @test all(iszero, A)
+        @test all(iszero, B)
+    end
+end
+
 @testset "concat" begin
     t = StructArray{Pair{Int, String}}(([3, 5], ["a", "b"]))
     push!(t, (2 => "d"))