Cleanup and remove generated functions as much as possible (JuliaArrays#23)

Author: Pietro Vertechi

src/StructArrays.jl

@@ -3,6 +3,7 @@ module StructArrays
 import Requires
 export StructArray
 
+include("interface.jl")
 include("structarray.jl")
 include("utils.jl")
 include("collect.jl")

src/interface.jl

@@ -0,0 +1,20 @@
+@generated function staticschema(::Type{T}) where {T}
+    name_tuple = Expr(:tuple, [QuoteNode(f) for f in fieldnames(T)]...)
+    type_tuple = Expr(:curly, :Tuple, [Expr(:call, :fieldtype, :T, i) for i in 1:fieldcount(T)]...)
+    Expr(:curly, :NamedTuple, name_tuple, type_tuple)
+end
+
+@generated function staticschema(::Type{T}) where {T<:Tuple}
+    name_tuple = Expr(:tuple, [QuoteNode(Symbol("x$f")) for f in fieldnames(T)]...)
+    type_tuple = Expr(:curly, :Tuple, [Expr(:call, :fieldtype, :T, i) for i in 1:fieldcount(T)]...)
+    Expr(:curly, :NamedTuple, name_tuple, type_tuple)
+end
+
+staticschema(::Type{T}) where {T<:NamedTuple} = T
+
+getnames(::Type{NamedTuple{names, types}}) where {names, types} = names
+gettypes(::Type{NamedTuple{names, types}}) where {names, types} = types
+
+function fields(::Type{T}) where {T}
+    getnames(staticschema(T))
+end

src/structarray.jl

@@ -28,29 +28,25 @@ StructArray(; kwargs...) = StructArray(values(kwargs))
 StructArray{T}(args...) where {T} = StructArray{T}(NamedTuple{fields(T)}(args))
 
 _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
 
+_similar(v::AbstractArray, ::Type{Z}; unwrap = t -> false) where {Z} =
+    unwrap(Z) ? buildfromschema(typ -> _similar(v, typ; unwrap = unwrap), Z) : similar(v, Z)
+
+function StructArray{T}(::Base.UndefInitializer, sz::Dims; unwrap = t -> false) where {T}
+    buildfromschema(typ -> _undef_array(typ, sz; unwrap = unwrap), T)
+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
-        StructArray{T}(NamedTuple{fields(T)}($ex))
-    end
+
+function similar_structarray(v::AbstractArray, ::Type{Z}; unwrap = t -> false) where {Z}
+    buildfromschema(typ -> _similar(v, typ; unwrap = unwrap), Z)
 end
 
-@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)); unwrap = unwrap)) for i in 1:fieldcount(T)]...)
-    push = Expr(:block, [:($(syms[i])[j] = getfield(f, $i)) for i in 1:fieldcount(T)]...)
-    quote
-        $init
-        for (j, f) in enumerate(v)
-            @inbounds $push
-        end
-        return StructArray{T}($(syms...))
+function StructArray(v::AbstractArray{T}; unwrap = t -> false) where {T}
+    s = similar_structarray(v, T; unwrap = unwrap)
+    for i in eachindex(v)
+        @inbounds s[i] = v[i]
     end
+    s
 end
 StructArray(s::StructArray) = copy(s)
 
@@ -63,7 +59,15 @@ Base.propertynames(s::StructArray) = fieldnames(typeof(columns(s)))
 
 Base.size(s::StructArray) = size(columns(s)[1])
 
-Base.@propagate_inbounds Base.getindex(s::StructArray, I::Int...) = get_ith(s, I...)
+@generated function Base.getindex(x::StructArray{T, N, NamedTuple{names, types}}, I::Int...) where {T, N, names, types}
+    args = [:(getfield(cols, $i)[I...]) for i in 1:length(names)]
+    return quote
+        cols = columns(x)
+        @boundscheck checkbounds(x, I...)
+        @inbounds $(Expr(:call, :createinstance, :T, args...))
+    end
+end
+
 function Base.getindex(s::StructArray{T, N, C}, I::Union{Int, AbstractArray, Colon}...) where {T, N, C}
     StructArray{T}(map(v -> getindex(v, I...), columns(s)))
 end
@@ -72,29 +76,21 @@ function Base.view(s::StructArray{T, N, C}, I...) where {T, N, C}
     StructArray{T}(map(v -> view(v, I...), columns(s)))
 end
 
-Base.@propagate_inbounds Base.setindex!(s::StructArray, val, I::Int...) = set_ith!(s, val, I...)
-
-fields(::Type{<:NamedTuple{K}}) where {K} = K
-@generated function fields(t::Type{T}) where {T}
-   return :($(Expr(:tuple, [QuoteNode(f) for f in fieldnames(T)]...)))
-end
-@generated function fields(t::Type{T}) where {T<:Tuple}
-    return :($(Expr(:tuple, [QuoteNode(Symbol("x$f")) for f in fieldnames(T)]...)))
+function Base.setindex!(s::StructArray, vals, I::Int...)
+    @boundscheck checkbounds(s, I...)
+    @inbounds foreachcolumn((col, val) -> (col[I...] = val), s, vals)
+    s
 end
 
 @inline getfieldindex(v::Tuple, field::Symbol, index::Integer) = getfield(v, index)
 @inline getfieldindex(v, field::Symbol, index::Integer) = getproperty(v, field)
 
-@generated function Base.push!(s::StructArray{T, 1}, vals) where {T}
-    exprs = foreach_expr((args...) -> Expr(:call, :push!, args...), T, :s, :vals)
-    push!(exprs, :s)
-    Expr(:block, exprs...)
+function Base.push!(s::StructArray, vals)
+    foreachcolumn(push!, s, vals)
 end
 
-@generated function Base.append!(s::StructArray{T, 1}, vals) where {T}
-    exprs = foreach_expr((args...) -> Expr(:call, :append!, args...), T, :s, :vals)
-    push!(exprs, :s)
-    Expr(:block, exprs...)
+function Base.append!(s::StructArray, vals)
+    foreachcolumn(append!, s, vals)
 end
 
 function Base.cat(args::StructArray...; dims)
@@ -125,4 +121,3 @@ Base.copy(s::StructArray{T,N,C}) where {T,N,C} = StructArray{T,N,C}(C(copy(x) fo
 function Base.reshape(s::StructArray{T}, d::Dims) where {T}
     StructArray{T}(map(x -> reshape(x, d), columns(s)))
 end
-

src/tables.jl

@@ -5,8 +5,9 @@ Tables.columnaccess(::Type{<:StructArray}) = true
 Tables.rows(s::StructArray) = s
 Tables.columns(s::StructArray) = columns(s)
 
-@generated function Tables.schema(s::StructArray{T}) where {T}
-    names = fieldnames(T)
-    types = map(sym -> fieldtype(T, sym), names)
-    :(Tables.Schema($names, $types))
+function Tables.schema(s::StructArray{T}) where {T}
+    NT = staticschema(T)
+    names = getnames(NT)
+    types = gettypes(NT).parameters
+    Tables.Schema(names, types)
 end

src/utils.jl

@@ -1,36 +1,39 @@
 import Base: tuple_type_cons, tuple_type_head, tuple_type_tail, tail
 
-eltypes(::Type{Tuple{}}) = Tuple{}
-eltypes(::Type{T}) where {T<:Tuple} =
-    tuple_type_cons(eltype(tuple_type_head(T)), eltypes(tuple_type_tail(T)))
-eltypes(::Type{NamedTuple{K, V}}) where {K, V} = NamedTuple{K, eltypes(V)}
+eltypes(::Type{T}) where {T} = map_types(eltype, T)
 
-Base.@pure SkipConstructor(::Type) = false
+map_types(f, ::Type{Tuple{}}) = Tuple{}
+function map_types(f, ::Type{T}) where {T<:Tuple}
+    tuple_type_cons(f(tuple_type_head(T)), map_types(f, tuple_type_tail(T)))
+end
+map_types(f, ::Type{NamedTuple{names, types}}) where {names, types} =
+    NamedTuple{names, map_types(f, types)}
 
-function foreach_expr(f, T, args...)
-    exprs = []
-    for (ind, key) in enumerate(fields(T))
-        new_args = (Expr(:call, :getfieldindex, arg, Expr(:quote, key), ind) for arg in args)
-        push!(exprs, f(new_args...))
-    end
-    exprs
+map_params(f, ::Type{Tuple{}}) = ()
+function map_params(f, ::Type{T}) where {T<:Tuple}
+    (f(tuple_type_head(T)), map_params(f, tuple_type_tail(T))...)
 end
+map_params(f, ::Type{NamedTuple{names, types}}) where {names, types} =
+    NamedTuple{names}(map_params(f, types))
 
-@generated function get_ith(s::StructArray{T}, I...) where {T}
-    exprs = foreach_expr(field -> :($field[I...]), T, :s)
-    return quote
-        @boundscheck checkbounds(s, I...)
-        @inbounds $(Expr(:call, :createinstance, :T, exprs...))
-    end
+buildfromschema(initializer, ::Type{T}) where {T} = buildfromschema(initializer, T, staticschema(T))
+
+function buildfromschema(initializer, ::Type{T}, ::Type{NT}) where {T, NT<:NamedTuple}
+    nt = map_params(initializer, NT)
+    StructArray{T}(nt)
 end
 
-@generated function set_ith!(s::StructArray{T}, vals, I...) where {T}
-    exprs = foreach_expr((field, val) -> :($field[I...] = $val), T, :s, :vals)
-    push!(exprs, :s)
-    return quote
-        @boundscheck checkbounds(s, I...)
-        @inbounds $(Expr(:block, exprs...))
+Base.@pure SkipConstructor(::Type) = false
+
+@generated function foreachcolumn(f, x::StructArray{T, N, NamedTuple{names, types}}, xs...) where {T, N, names, types}
+    exprs = Expr[]
+    for (i, field) in enumerate(names)
+        sym = QuoteNode(field)
+        args = [Expr(:call, :getfieldindex, :(getfield(xs, $j)), sym, i) for j in 1:length(xs)]
+        push!(exprs, Expr(:call, :f, Expr(:., :x, sym), args...))
     end
+    push!(exprs, :(return nothing))
+    Expr(:block, exprs...)
 end
 
 function createinstance(::Type{T}, args...) where {T}
@@ -46,7 +49,7 @@ createinstance(::Type{T}, args...) where {T<:Union{Tuple, NamedTuple}} = T(args)
     Expr(:block, new_tup, construct)
 end
 
-createtype(::Type{T}, ::Type{NamedTuple{names, types}}) where {T, names, types} = createtype(T, names, eltypes(types)) 
+createtype(::Type{T}, ::Type{NamedTuple{names, types}}) where {T, names, types} = createtype(T, names, eltypes(types))
 
 createtype(::Type{T}, names, types) where {T} = T
 createtype(::Type{T}, names, types) where {T<:Tuple} = types

test/runtests.jl

@@ -102,16 +102,14 @@ 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)
+g_infer() = StructArray([(a=(b="1",), c=2)], unwrap = t -> t <: NamedTuple)
 tup_infer() = StructArray([(1, 2), (3, 4)])
 
 @testset "inferrability" begin
     @inferred f_infer()
     @inferred g_infer()
-    @test g_infer().a.b == [1]
+    @test g_infer().a.b == ["1"]
     s = @inferred tup_infer()
     @test Tables.columns(s) == (x1 = [1, 3], x2 = [2, 4])
     @test s[1] == (1, 2)
@@ -150,7 +148,7 @@ StructArrays.SkipConstructor(::Type{<:S}) = true
 end
 
 const initializer = StructArrays.ArrayInitializer(t -> t <: Union{Tuple, NamedTuple, Pair})
-collect_columns(t) = StructArrays.collect_columns(t, initializer = initializer) 
+collect_columns(t) = StructArrays.collect_columns(t, initializer = initializer)
 
 @testset "collectnamedtuples" begin
     v = [(a = 1, b = 2), (a = 1, b = 3)]