|
| 1 | +module StructArraysStaticArraysExt |
| 2 | + |
| 3 | +using StructArrays |
| 4 | +using StaticArrays: StaticArray, FieldArray, tuple_prod |
| 5 | + |
| 6 | +""" |
| 7 | + StructArrays.staticschema(::Type{<:StaticArray{S, T}}) where {S, T} |
| 8 | +
|
| 9 | +The `staticschema` of a `StaticArray` element type is the `staticschema` of the underlying `Tuple`. |
| 10 | +```julia |
| 11 | +julia> StructArrays.staticschema(SVector{2, Float64}) |
| 12 | +Tuple{Float64, Float64} |
| 13 | +``` |
| 14 | +The one exception to this rule is `<:StaticArrays.FieldArray`, since `FieldArray` is based on a |
| 15 | +struct. In this case, `staticschema(<:FieldArray)` returns the `staticschema` for the struct |
| 16 | +which subtypes `FieldArray`. |
| 17 | +""" |
| 18 | +@generated function StructArrays.staticschema(::Type{<:StaticArray{S, T}}) where {S, T} |
| 19 | + return quote |
| 20 | + Base.@_inline_meta |
| 21 | + return NTuple{$(tuple_prod(S)), T} |
| 22 | + end |
| 23 | +end |
| 24 | +StructArrays.createinstance(::Type{T}, args...) where {T<:StaticArray} = T(args) |
| 25 | +StructArrays.component(s::StaticArray, i) = getindex(s, i) |
| 26 | + |
| 27 | +# invoke general fallbacks for a `FieldArray` type. |
| 28 | +@inline function StructArrays.staticschema(T::Type{<:FieldArray}) |
| 29 | + invoke(StructArrays.staticschema, Tuple{Type{<:Any}}, T) |
| 30 | +end |
| 31 | +StructArrays.component(s::FieldArray, i) = invoke(StructArrays.component, Tuple{Any, Any}, s, i) |
| 32 | +StructArrays.createinstance(T::Type{<:FieldArray}, args...) = invoke(StructArrays.createinstance, Tuple{Type{<:Any}, Vararg}, T, args...) |
| 33 | + |
| 34 | +# Broadcast overload |
| 35 | +using StaticArrays: StaticArrayStyle, similar_type, Size, SOneTo |
| 36 | +using StaticArrays: broadcast_flatten, broadcast_sizes, first_statictype, __broadcast |
| 37 | +using StructArrays: isnonemptystructtype |
| 38 | +using Base.Broadcast: Broadcasted |
| 39 | + |
| 40 | +# StaticArrayStyle has no similar defined. |
| 41 | +# Overload `try_struct_copy` instead. |
| 42 | +@inline function StructArrays.try_struct_copy(bc::Broadcasted{StaticArrayStyle{M}}) where {M} |
| 43 | + flat = broadcast_flatten(bc); as = flat.args; f = flat.f |
| 44 | + argsizes = broadcast_sizes(as...) |
| 45 | + ax = axes(bc) |
| 46 | + ax isa Tuple{Vararg{SOneTo}} || error("Dimension is not static. Please file a bug at `StaticArrays.jl`.") |
| 47 | + return _broadcast(f, Size(map(length, ax)), argsizes, as...) |
| 48 | +end |
| 49 | + |
| 50 | +@inline function _broadcast(f, sz::Size{newsize}, s::Tuple{Vararg{Size}}, a...) where {newsize} |
| 51 | + first_staticarray = first_statictype(a...) |
| 52 | + elements, ET = if prod(newsize) == 0 |
| 53 | + # Use inference to get eltype in empty case (see also comments in _map) |
| 54 | + eltys = Tuple{map(eltype, a)...} |
| 55 | + (), Core.Compiler.return_type(f, eltys) |
| 56 | + else |
| 57 | + temp = __broadcast(f, sz, s, a...) |
| 58 | + temp, eltype(temp) |
| 59 | + end |
| 60 | + if isnonemptystructtype(ET) |
| 61 | + @static if VERSION >= v"1.7" |
| 62 | + arrs = ntuple(Val(fieldcount(ET))) do i |
| 63 | + @inbounds similar_type(first_staticarray, fieldtype(ET, i), sz)(_getfields(elements, i)) |
| 64 | + end |
| 65 | + else |
| 66 | + similarET(::Type{SA}, ::Type{T}) where {SA, T} = i -> @inbounds similar_type(SA, fieldtype(T, i), sz)(_getfields(elements, i)) |
| 67 | + arrs = ntuple(similarET(first_staticarray, ET), Val(fieldcount(ET))) |
| 68 | + end |
| 69 | + return StructArray{ET}(arrs) |
| 70 | + end |
| 71 | + @inbounds return similar_type(first_staticarray, ET, sz)(elements) |
| 72 | +end |
| 73 | + |
| 74 | +@inline function _getfields(x::Tuple, i::Int) |
| 75 | + if @generated |
| 76 | + return Expr(:tuple, (:(getfield(x[$j], i)) for j in 1:fieldcount(x))...) |
| 77 | + else |
| 78 | + return map(Base.Fix2(getfield, i), x) |
| 79 | + end |
| 80 | +end |
| 81 | + |
| 82 | +end |
0 commit comments