Add specialized Base.typed_vcat and improve append! for VectorOfArrays

Author: oschulz

src/vector_of_arrays.jl

@@ -263,26 +263,38 @@ Base.length(A::VectorOfArrays) = length(A.kernel_size)
 end
 
 
-function Base.append!(A::VectorOfArrays{T,N}, B::VectorOfArrays{U,N}) where {T,N,U}
-    if !isempty(B)
-        # Implementation supports A === B
-
-        A_ep = A.elem_ptr
-        B_ep = B.elem_ptr
-        idxs_B = firstindex(B_ep):(lastindex(B_ep) - 1)
-        delta_ep_idx = lastindex(A_ep) + 1 - firstindex(B_ep)
-        delta_ep = last(A_ep) - first(B_ep)
-        resize!(A_ep, length(eachindex(A_ep)) + length(idxs_B))
-        @assert checkbounds(Bool, B_ep, idxs_B)
-        @assert checkbounds(Bool, A_ep, broadcast(+, idxs_B, delta_ep_idx))
-        @inbounds @simd for i_B in idxs_B
-            A_ep[i_B + delta_ep_idx] = B_ep[i_B + 1] + delta_ep
+function append_elemptr!(A::AbstractVector{<:Integer}, B::AbstractVector{<:Integer})
+    idxs_A = LinearIndices(A)
+    idxs_B = LinearIndices(B)
+    length_B = length(idxs_B)
+
+    A_from = last(idxs_A) + 1
+    A_to = A_from - 1 + length_B - 1
+    resize!(A, length(first(idxs_A):A_to))
+
+    B_from = first(idxs_B) + 1
+    B_to = last(idxs_B)
+
+    A_idxs_offs = A_from - B_from
+
+    checkindex(Bool, eachindex(B), B_from:B_to)
+    checkindex(Bool, eachindex(A), (B_from:B_to) .+ A_idxs_offs)
+    @inbounds begin
+        value_offset = A[A_from - 1] - B[B_from - 1]
+        @simd for i in B_from:B_to
+            A[i + A_idxs_offs] = B[i] + value_offset
         end
+    end
+
+    A
+end
+
 
+function Base.append!(A::VectorOfArrays{T,N}, B::VectorOfArrays{U,N}) where {T,N,U}
+    if !isempty(B)
         append!(A.data, B.data)
+        append_elemptr!(A.elem_ptr, B.elem_ptr)
         append!(A.kernel_size, B.kernel_size)
-
-        simple_consistency_checks(A)
     end
     A
 end
@@ -310,6 +322,23 @@ function Base.append!(A::VectorOfArrays{T,N}, B::AbstractVector{<:AbstractArray{
 end
 
 
+Base.vcat(V::VectorOfArrays) = V
+
+function Base.vcat(Vs::(VectorOfArrays{U,N} where U)...) where {N}
+    data = vcat(map(x -> x.data, Vs)...)
+
+    elem_ptr = similar(Vs[1].elem_ptr, 1)
+    elem_ptr[1] = firstindex(data)
+    @inbounds for i in eachindex(Vs)
+        append_elemptr!(elem_ptr, Vs[i].elem_ptr)
+    end
+
+    kernel_size = vcat(map(x -> x.kernel_size, Vs)...)
+
+    VectorOfArrays(data, elem_ptr, kernel_size, no_consistency_checks)
+end
+
+
 Base.@propagate_inbounds function Base.unsafe_view(A::VectorOfArrays, idxs::AbstractUnitRange{<:Integer})
     from = first(idxs)
     to = last(idxs)

test/vector_of_arrays.jl

@@ -5,12 +5,20 @@ using Test
 
 using UnsafeArrays
 
+using ArraysOfArrays: full_consistency_checks, append_elemptr!
+
 
 @testset "vector_of_arrays" begin
     ref_AoA1(T::Type, n::Integer) = n == 0 ? [Array{T}(undef, 5)][1:0] : [rand(T, rand(1:9)) for i in 1:n]
     ref_AoA2(T::Type, n::Integer) = n == 0 ? [Array{T}(undef, 4, 2)][1:0] : [rand(T, rand(1:4), rand(1:4)) for i in 1:n]
     ref_AoA3(T::Type, n::Integer) = n == 0 ? [Array{T}(undef, 3, 2, 4)][1:0] : [rand(T, rand(1:3), rand(1:3), rand(1:3)) for i in 1:n]
 
+    @testset "element pointer handling" begin
+        A = [2, 4, 5, 9]; B = [3, 6, 8]
+        @test @inferred(append_elemptr!(deepcopy(A), B)) == [2, 4, 5, 9, 12, 14]
+        @test @inferred(append_elemptr!([2], [5])) == [2]
+    end
+
 
     @testset "ctors" begin
         A1 = ref_AoA1(Float32, 5)
@@ -65,6 +73,33 @@ using UnsafeArrays
     end
 
 
+    @testset "append! and vcat" begin
+        A1 = ref_AoA3(Float32, 3); A2 = ref_AoA3(Float32, 0)
+        A3 = ref_AoA3(Float32, 4); A4 = ref_AoA3(Float64, 2)
+
+        B1 = VectorOfArrays(A1); B2 = VectorOfArrays(A2);
+        B3 = VectorOfArrays(A3); B4 = VectorOfArrays(A4);
+
+        @test @inferred(vcat(B1)) === B1
+        full_consistency_checks(vcat(B1))
+
+        @test @inferred(vcat(B1, B2)) isa VectorOfArrays
+        @test vcat(B1, B2) == vcat(A1, A2)
+        @test eltype(vcat(B1, B2)) == Array{Float32,3}
+        full_consistency_checks(vcat(B1, B2))
+
+        @test @inferred(vcat(B1, B3)) isa VectorOfArrays
+        @test vcat(B1, B3) == vcat(A1, A3)
+        @test eltype(vcat(B1, B3)) == Array{Float32,3}
+        full_consistency_checks(vcat(B1, B3))
+
+        @test @inferred(vcat(B1, B2, B3, B4)) isa VectorOfArrays
+        @test vcat(B1, B2, B3, B4) == vcat(A1, A2, A3, A4)
+        @test eltype(vcat(B1, B2, B3, B4)) == Array{Float64,3}
+        full_consistency_checks(vcat(B1, B2, B3, B4))
+    end
+
+
     @testset "examples" begin
         VA = VectorOfArrays{Float64, 2}()