Allow recursive unwrapping (JuliaArrays#19)

* wip allow recursive unwrapping

* added unwrap keyword argument

* Add tests

* Use keyword argument anyway and define _similar

* spacing

* rename to _undef_array

Author: Pietro Vertechi

src/structarray.jl

@@ -27,17 +27,23 @@ StructArray{T}(; kwargs...) where {T} = StructArray{T}(values(kwargs))
 StructArray(; kwargs...) = StructArray(values(kwargs))
 
 StructArray{T}(args...) where {T} = StructArray{T}(NamedTuple{fields(T)}(args))
-@generated function StructArray{T}(::Base.UndefInitializer, d::Integer...) where {T}
-    ex = Expr(:tuple, [:(Array{$(fieldtype(T, i))}(undef, sz)) for i in 1:fieldcount(T)]...)
+
+_undef_array(::Type{T}, sz; unwrap = t -> false) where {T} = unwrap(T) ? StructArray{T}(undef, sz; unwrap = unwrap) : Array{T}(undef, sz)
+function _similar(v::S, ::Type{Z}; unwrap = t -> false) where {S <: AbstractArray{T, N}, Z} where {T, N}
+    unwrap(Z) ? StructArray{Z}(map(t -> _similar(v, fieldtype(Z, t); unwrap = unwrap), fields(Z))) : similar(v, Z)
+end
+
+StructArray{T}(u::Base.UndefInitializer, d::Integer...; unwrap = t -> false) where {T} = StructArray{T}(u, convert(Dims, d); unwrap = unwrap)
+@generated function StructArray{T}(::Base.UndefInitializer, sz::Dims; unwrap = t -> false) where {T}
+    ex = Expr(:tuple, [:(_undef_array($(fieldtype(T, i)), sz; unwrap = unwrap)) for i in 1:fieldcount(T)]...)
     return quote
-        sz = convert(Tuple{Vararg{Int}}, d)
         StructArray{T}(NamedTuple{fields(T)}($ex))
     end
 end
 
-@generated function StructArray(v::AbstractArray{T, N}) where {T, N}
+@generated function StructArray(v::AbstractArray{T, N}; unwrap = t -> false) where {T, N}
     syms = [gensym() for i in 1:fieldcount(T)]
-    init = Expr(:block, [:($(syms[i]) = similar(v, $(fieldtype(T, i)))) for i in 1:fieldcount(T)]...)
+    init = Expr(:block, [:($(syms[i]) = _similar(v, $(fieldtype(T, i)); unwrap = unwrap)) for i in 1:fieldcount(T)]...)
     push = Expr(:block, [:($(syms[i])[j] = f.$(fieldname(T, i))) for i in 1:fieldcount(T)]...)
     quote
         $init

test/runtests.jl

@@ -96,8 +96,15 @@ end
 end
 
 f_infer() = StructArray{ComplexF64}(rand(2,2), rand(2,2))
+unwrap(::Type) = false
+unwrap(::Type{<:NamedTuple}) = true
+
+g_infer() = StructArray([(a=(b=1,), c=2)], unwrap = unwrap)
+
 @testset "inferrability" begin
     @inferred f_infer()
+    @inferred g_infer()
+    @test g_infer().a.b == [1]
 end
 
 @testset "propertynames" begin