Add copyto, map for CatVector.

Author: tkoolen

src/custom_collections/CatVector.jl

@@ -1,5 +1,5 @@
-struct CatVector{T, N, I<:Tuple{Vararg{AbstractVector{T}, N}}} <: AbstractVector{T}
-    vecs::I
+struct CatVector{T, N, V<:AbstractVector{T}} <: AbstractVector{T}
+    vecs::NTuple{N, V}
 end
 
 @inline Base.size(vec::CatVector) = (mapreduce(length, +, vec.vecs; init=0),)
@@ -39,7 +39,7 @@ Base.@propagate_inbounds function Base.setindex!(vec::CatVector, val, i::Int)
     error()
 end
 
-Base.@propagate_inbounds function Base.copyto!(dest::AbstractVector, src::CatVector)
+Base.@propagate_inbounds function Base.copyto!(dest::AbstractVector{T}, src::CatVector{T}) where {T}
     @boundscheck length(dest) == length(src) || throw(DimensionMismatch())
     dest_indices = eachindex(dest)
     k = 1
@@ -55,3 +55,30 @@ end
 
 Base.similar(vec::CatVector) = CatVector(map(similar, vec.vecs))
 Base.similar(vec::CatVector, ::Type{T}) where {T} = CatVector(map(x -> similar(x, T), vec.vecs))
+
+function check_cat_vectors_line_up(x::CatVector, ys::CatVector...)
+    for j in eachindex(ys)
+        y = ys[j]
+        length(x.vecs) == length(y.vecs) || throw(ArgumentError("Subvectors must line up"))
+        for i in eachindex(x.vecs)
+            length(x.vecs[i]) == length(y.vecs[i]) || throw(ArgumentError("Subvectors must line up"))
+        end
+    end
+end
+
+@inline function Base.copyto!(dest::CatVector, src::CatVector)
+    @boundscheck check_cat_vectors_line_up(dest, src)
+    @inbounds for i in eachindex(dest.vecs)
+        copyto!(dest.vecs[i], src.vecs[i])
+    end
+    return dest
+end
+
+@inline function Base.map!(f::F, dest::CatVector, args::CatVector...) where F
+    @boundscheck check_cat_vectors_line_up(dest, args...)
+    @inbounds for i in eachindex(dest.vecs)
+        map!(f, dest.vecs[i], map(arg -> arg.vecs[i], args)...)
+    end
+    return dest
+end
+

test/test_custom_collections.jl

@@ -200,5 +200,13 @@ Base.axes(m::NonOneBasedMatrix) = ((1:m.m) .- 2, (1:m.n) .+ 1)
             @test length(y3.vecs[i]) == length(y.vecs[i])
             @test y3.vecs[i] !== y.vecs[i]
         end
+
+        y4 = similar(y)
+        copyto!(y4, y)
+        @test y4 == y
+
+        y5 = similar(y)
+        map!(+, y5, y, y)
+        @test Vector(y5) == Vector(y) + Vector(y)
     end
 end