Adding support for components/broadcasting/views with StructArrays of StaticArrays

Project.toml

@@ -5,11 +5,13 @@ version = "0.5.1"
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [compat]
 Adapt = "1, 2, 3"
 DataAPI = "1"
+StaticArrays = "1"
 Tables = "1"
 julia = "1"
 

src/StructArrays.jl

@@ -13,6 +13,7 @@ include("collect.jl")
 include("sort.jl")
 include("lazy.jl")
 include("tables.jl")
+include("staticarrays_support.jl")
 
 # Implement refarray and refvalue to deal with pooled arrays and weakrefstrings effectively
 import DataAPI: refarray, refvalue

src/staticarrays_support.jl

@@ -0,0 +1,19 @@
+import StaticArrays: StaticArray, tuple_prod
+
+"""
+    StructArrays.staticschema(::Type{<:StaticArray{S, T}}) where {S, T}
+
+The `staticschema` of a `StaticArray` element type is the `staticschema` of the underlying `Tuple`.
+```julia
+julia> StructArrays.staticschema(SVector{2, Float64})
+Tuple{Float64, Float64}
+```
+"""
+@generated function StructArrays.staticschema(::Type{<:StaticArray{S, T}}) where {S, T}
+    return quote
+        Base.@_inline_meta
+        return NTuple{$(tuple_prod(S)), T}
+    end
+end
+StructArrays.createinstance(::Type{T}, args...) where {T<:StaticArray} = T(args)
+StructArrays.component(s::StaticArray, i) = getindex(s, i)

src/structarray.jl

@@ -344,7 +344,7 @@ Base.@propagate_inbounds function Base.getindex(x::StructArray{T, <:Any, <:Any,
     return createinstance(T, get_ith(cols, I)...)
 end
 
-function Base.view(s::StructArray{T, N, C}, I...) where {T, N, C}
+@inline function Base.view(s::StructArray{T, N, C}, I...) where {T, N, C}
     StructArray{T}(map(v -> view(v, I...), components(s)))
 end
 

test/runtests.jl

@@ -1,6 +1,7 @@
 using StructArrays
 using StructArrays: staticschema, iscompatible, _promote_typejoin, append!!
 using OffsetArrays: OffsetArray
+using StaticArrays
 import Tables, PooledArrays, WeakRefStrings
 using TypedTables: Table
 using DataAPI: refarray, refvalue
@@ -809,3 +810,26 @@ Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{MyArray}}, ::Type{El
     s = StructArray{ComplexF64}((MyArray(rand(2,2)), MyArray(rand(2,2))))
     @test_throws MethodError s .+ s
 end
+
+@testset "staticarrays" begin
+
+    # test that staticschema returns the right things
+    for StaticVectorType = [SVector, MVector, SizedVector]    
+        @test StructArrays.staticschema(StaticVectorType{2,Float64}) == Tuple{Float64,Float64}
+    end
+
+    # test broadcast + components for vectors
+    for StaticVectorType = [SVector, MVector, SizedVector]
+        x = @inferred StructArray([StaticVectorType{2}(Float64[i;i+1]) for i = 1:2])
+        y = @inferred StructArray([StaticVectorType{2}(Float64[i+1;i+2]) for i = 1:2])
+        @test StructArrays.components(x) == ([1.0,2.0], [2.0,3.0])
+        @test x .+ y == StructArray([StaticVectorType{2}(Float64[2*i+1;2*i+3]) for i = 1:2])
+    end
+    # test broadcast + components for general arrays
+    for StaticArrayType = [SArray, MArray, SizedArray]
+        x = @inferred StructArray([StaticArrayType{Tuple{1,2}}(ones(1,2) .+ i) for i = 0:1])
+        y = @inferred StructArray([StaticArrayType{Tuple{1,2}}(2*ones(1,2) .+ i) for i = 0:1])
+        @test StructArrays.components(x) == ([1., 2.], [1., 2.])
+        @test x .+ y == StructArray([StaticArrayType{Tuple{1,2}}(3*ones(1,2) .+ 2*i) for i = 0:1])
+    end
+end