From 2a0fa7c1a8c2d9725acea894a7fc270179841870 Mon Sep 17 00:00:00 2001
From: Twan Koolen <koolen.twan@gmail.com>
Date: Fri, 28 Jun 2019 20:05:01 -0400
Subject: [PATCH 1/8] Add CatVector.

---
 src/custom_collections/CatVector.jl          | 54 ++++++++++++++++++++
 src/custom_collections/custom_collections.jl |  4 +-
 test/test_custom_collections.jl              | 28 ++++++++++
 3 files changed, 85 insertions(+), 1 deletion(-)
 create mode 100644 src/custom_collections/CatVector.jl

diff --git a/src/custom_collections/CatVector.jl b/src/custom_collections/CatVector.jl
new file mode 100644
index 00000000..2aaff9a0
--- /dev/null
+++ b/src/custom_collections/CatVector.jl
@@ -0,0 +1,54 @@
+struct CatVector{T, N, I<:Tuple{Vararg{AbstractVector{T}, N}}} <: AbstractVector{T}
+    vecs::I
+end
+
+@inline Base.size(vec::CatVector) = (mapreduce(length, +, vec.vecs; init=0),)
+Base.eltype(vec::CatVector) = eltype(eltype(vec.vecs))
+
+# Note: getindex and setindex are pretty naive.
+Base.@propagate_inbounds function Base.getindex(vec::CatVector, i::Int)
+    @boundscheck checkbounds(vec, i)
+    I = 1
+    @inbounds while true
+        subvec = vec.vecs[I]
+        l = length(subvec)
+        if i <= l
+            return subvec[eachindex(subvec)[i]]
+        else
+            i -= l
+            I += 1
+        end
+    end
+    error()
+end
+
+Base.@propagate_inbounds function Base.setindex!(vec::CatVector, val, i::Int)
+    @boundscheck checkbounds(vec, i)
+    I = 1
+    while true
+        subvec = vec.vecs[I]
+        l = length(subvec)
+        if i <= l
+            subvec[eachindex(subvec)[i]] = val
+            return val
+        else
+            i -= l
+            I += 1
+        end
+    end
+    error()
+end
+
+Base.@propagate_inbounds function Base.copyto!(dest::AbstractVector, src::CatVector)
+    @boundscheck length(dest) == length(src) || throw(DimensionMismatch())
+    dest_indices = eachindex(dest)
+    k = 1
+    @inbounds for i in eachindex(src.vecs)
+        vec = src.vecs[i]
+        for j in eachindex(vec)
+            dest[dest_indices[k]] = vec[j]
+            k += 1
+        end
+    end
+    return dest
+end
diff --git a/src/custom_collections/custom_collections.jl b/src/custom_collections/custom_collections.jl
index 5a53041e..f3021d6d 100644
--- a/src/custom_collections/custom_collections.jl
+++ b/src/custom_collections/custom_collections.jl
@@ -10,7 +10,8 @@ export
     CacheIndexDict,
     SegmentedVector,
     SegmentedBlockDiagonalMatrix,
-    UnorderedPair
+    UnorderedPair,
+    CatVector
 
 export
     foreach_with_extra_args,
@@ -34,5 +35,6 @@ include("IndexDict.jl")
 include("SegmentedVector.jl")
 include("SegmentedBlockDiagonalMatrix.jl")
 include("UnorderedPair.jl")
+include("CatVector.jl")
 
 end # module
diff --git a/test/test_custom_collections.jl b/test/test_custom_collections.jl
index cc2be3ba..3da51bc6 100644
--- a/test/test_custom_collections.jl
+++ b/test/test_custom_collections.jl
@@ -153,4 +153,32 @@ Base.axes(m::NonOneBasedMatrix) = ((1:m.m) .- 2, (1:m.n) .+ 1)
         dict = Dict(p1 => 3)
         @test dict[p2] == 3
     end
+
+    @testset "CatVector" begin
+        CatVector = RigidBodyDynamics.CatVector
+        Random.seed!(52)
+        vecs = ntuple(i -> rand(rand(0 : 5)), Val(10))
+        l = sum(length, vecs)
+        x = zeros(l)
+        y = CatVector(vecs)
+        @test length(y) == l
+        x .= y
+        for i in eachindex(x)
+            @test x[i] == y[i]
+        end
+        x .= 0
+        @test x != y
+        copyto!(x, y)
+        for i in eachindex(x)
+            @test x[i] == y[i]
+        end
+        for i in eachindex(y)
+            x[i] = y[i]
+        end
+        @test x == y
+        allocs = let x=x, vecs=vecs
+            @allocated copyto!(x, RigidBodyDynamics.CatVector(vecs))
+        end
+        @test allocs == 0
+    end
 end

From 350896cb57b35016471f9a8e5246689b4b8d7365 Mon Sep 17 00:00:00 2001
From: Twan Koolen <koolen.twan@gmail.com>
Date: Fri, 28 Jun 2019 21:55:16 -0400
Subject: [PATCH 2/8] Add similar for CatVector

---
 src/custom_collections/CatVector.jl |  3 +++
 test/test_custom_collections.jl     | 20 ++++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/src/custom_collections/CatVector.jl b/src/custom_collections/CatVector.jl
index 2aaff9a0..a9efb9d0 100644
--- a/src/custom_collections/CatVector.jl
+++ b/src/custom_collections/CatVector.jl
@@ -52,3 +52,6 @@ Base.@propagate_inbounds function Base.copyto!(dest::AbstractVector, src::CatVec
     end
     return dest
 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))
diff --git a/test/test_custom_collections.jl b/test/test_custom_collections.jl
index 3da51bc6..d71a692a 100644
--- a/test/test_custom_collections.jl
+++ b/test/test_custom_collections.jl
@@ -161,24 +161,44 @@ Base.axes(m::NonOneBasedMatrix) = ((1:m.m) .- 2, (1:m.n) .+ 1)
         l = sum(length, vecs)
         x = zeros(l)
         y = CatVector(vecs)
+
         @test length(y) == l
+
         x .= y
         for i in eachindex(x)
             @test x[i] == y[i]
         end
+
         x .= 0
         @test x != y
         copyto!(x, y)
         for i in eachindex(x)
             @test x[i] == y[i]
         end
+
+        x .= 0
         for i in eachindex(y)
             x[i] = y[i]
         end
         @test x == y
+
         allocs = let x=x, vecs=vecs
             @allocated copyto!(x, RigidBodyDynamics.CatVector(vecs))
         end
         @test allocs == 0
+
+        y2 = similar(y)
+        @test eltype(y2) == eltype(y)
+        for i in eachindex(y.vecs)
+            @test length(y2.vecs[i]) == length(y.vecs[i])
+            @test y2.vecs[i] !== y.vecs[i]
+        end
+
+        y3 = similar(y, Int)
+        @test eltype(y3) == Int
+        for i in eachindex(y.vecs)
+            @test length(y3.vecs[i]) == length(y.vecs[i])
+            @test y3.vecs[i] !== y.vecs[i]
+        end
     end
 end

From b43949b58e8f96653e32464b0d870dfbef590fbf Mon Sep 17 00:00:00 2001
From: Twan Koolen <koolen.twan@gmail.com>
Date: Fri, 28 Jun 2019 22:20:26 -0400
Subject: [PATCH 3/8] Add copyto, map for CatVector.

---
 src/custom_collections/CatVector.jl | 33 ++++++++++++++++++++++++++---
 test/test_custom_collections.jl     |  8 +++++++
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/src/custom_collections/CatVector.jl b/src/custom_collections/CatVector.jl
index a9efb9d0..7a7dc00b 100644
--- a/src/custom_collections/CatVector.jl
+++ b/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
+
diff --git a/test/test_custom_collections.jl b/test/test_custom_collections.jl
index d71a692a..3dcc16ed 100644
--- a/test/test_custom_collections.jl
+++ b/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

From 53886a64d1cf5a4c1dafb8488e6509dcc13eaf6b Mon Sep 17 00:00:00 2001
From: Twan Koolen <koolen.twan@gmail.com>
Date: Fri, 28 Jun 2019 23:24:24 -0400
Subject: [PATCH 4/8] Add broadcast for CatVector. Works, but it's pretty slow.

---
 src/custom_collections/CatVector.jl | 35 +++++++++++++++++++++++------
 1 file changed, 28 insertions(+), 7 deletions(-)

diff --git a/src/custom_collections/CatVector.jl b/src/custom_collections/CatVector.jl
index 7a7dc00b..5840830a 100644
--- a/src/custom_collections/CatVector.jl
+++ b/src/custom_collections/CatVector.jl
@@ -56,16 +56,17 @@ 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
+@inline function check_cat_vectors_line_up(x::CatVector, y::CatVector)
+    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
+    nothing
 end
 
+@inline check_cat_vectors_line_up(x::CatVector, y) = nothing
+@inline check_cat_vectors_line_up(x::CatVector, y, tail...) = (check_cat_vectors_line_up(x, y); check_cat_vectors_line_up(x, tail...))
+
 @inline function Base.copyto!(dest::CatVector, src::CatVector)
     @boundscheck check_cat_vectors_line_up(dest, src)
     @inbounds for i in eachindex(dest.vecs)
@@ -82,3 +83,23 @@ end
     return dest
 end
 
+Base.@propagate_inbounds catvec_broadcast_getindex(vec::CatVector, i::Int, j::Int, k::Int) = vec.vecs[i][j]
+Base.@propagate_inbounds catvec_broadcast_getindex(x, i::Int, j::Int, k::Int) = Broadcast._broadcast_getindex(x, i)
+
+@inline function Base.copyto!(dest::CatVector, bc::Broadcast.Broadcasted{Nothing})
+    flat = Broadcast.flatten(bc)
+    index = 1
+    dest_vecs = dest.vecs
+    @boundscheck check_cat_vectors_line_up(dest, bc.args...)
+    @inbounds for i in eachindex(dest_vecs)
+        vec = dest_vecs[i]
+        for j in eachindex(vec)
+            k = axes(flat)[1][index]
+            let f = flat.f, args = flat.args, i = i, j = j, k = k
+                vec[j] = Broadcast._broadcast_getindex_evalf(f, map(arg -> catvec_broadcast_getindex(arg, i, j, k), args)...)
+            end
+            index += 1
+        end
+    end
+    return dest
+end

From 5ad06accf00851d0b55b32f8ba2b9f0fef9d6620 Mon Sep 17 00:00:00 2001
From: Twan Koolen <koolen.twan@gmail.com>
Date: Sat, 29 Jun 2019 01:17:49 -0400
Subject: [PATCH 5/8] Faster broadcast.

---
 src/custom_collections/CatVector.jl | 56 ++++++++++++++---------------
 test/test_custom_collections.jl     | 13 +++++++
 2 files changed, 40 insertions(+), 29 deletions(-)

diff --git a/src/custom_collections/CatVector.jl b/src/custom_collections/CatVector.jl
index 5840830a..d88a9bad 100644
--- a/src/custom_collections/CatVector.jl
+++ b/src/custom_collections/CatVector.jl
@@ -6,34 +6,36 @@ end
 Base.eltype(vec::CatVector) = eltype(eltype(vec.vecs))
 
 # Note: getindex and setindex are pretty naive.
-Base.@propagate_inbounds function Base.getindex(vec::CatVector, i::Int)
-    @boundscheck checkbounds(vec, i)
-    I = 1
+Base.@propagate_inbounds function Base.getindex(vec::CatVector, index::Int)
+    @boundscheck checkbounds(vec, index)
+    i = 1
+    j = index
     @inbounds while true
-        subvec = vec.vecs[I]
+        subvec = vec.vecs[i]
         l = length(subvec)
-        if i <= l
-            return subvec[eachindex(subvec)[i]]
+        if j <= l
+            return subvec[eachindex(subvec)[j]]
         else
-            i -= l
-            I += 1
+            j -= l
+            i += 1
         end
     end
     error()
 end
 
-Base.@propagate_inbounds function Base.setindex!(vec::CatVector, val, i::Int)
-    @boundscheck checkbounds(vec, i)
-    I = 1
+Base.@propagate_inbounds function Base.setindex!(vec::CatVector, val, index::Int)
+    @boundscheck checkbounds(vec, index)
+    i = 1
+    j = index
     while true
-        subvec = vec.vecs[I]
+        subvec = vec.vecs[i]
         l = length(subvec)
-        if i <= l
-            subvec[eachindex(subvec)[i]] = val
+        if j <= l
+            subvec[eachindex(subvec)[j]] = val
             return val
         else
-            i -= l
-            I += 1
+            j -= l
+            i += 1
         end
     end
     error()
@@ -83,22 +85,18 @@ end
     return dest
 end
 
-Base.@propagate_inbounds catvec_broadcast_getindex(vec::CatVector, i::Int, j::Int, k::Int) = vec.vecs[i][j]
-Base.@propagate_inbounds catvec_broadcast_getindex(x, i::Int, j::Int, k::Int) = Broadcast._broadcast_getindex(x, i)
+Base.@propagate_inbounds catvec_broadcast_vec(x::CatVector, k::Int) = x.vecs[k]
+Base.@propagate_inbounds catvec_broadcast_vec(x::Number, k::Int) = x
 
 @inline function Base.copyto!(dest::CatVector, bc::Broadcast.Broadcasted{Nothing})
     flat = Broadcast.flatten(bc)
-    index = 1
-    dest_vecs = dest.vecs
-    @boundscheck check_cat_vectors_line_up(dest, bc.args...)
-    @inbounds for i in eachindex(dest_vecs)
-        vec = dest_vecs[i]
-        for j in eachindex(vec)
-            k = axes(flat)[1][index]
-            let f = flat.f, args = flat.args, i = i, j = j, k = k
-                vec[j] = Broadcast._broadcast_getindex_evalf(f, map(arg -> catvec_broadcast_getindex(arg, i, j, k), args)...)
-            end
-            index += 1
+    @boundscheck check_cat_vectors_line_up(dest, flat.args...)
+    @inbounds for i in eachindex(dest.vecs)
+        let i = i, f = flat.f, args = flat.args
+            dest′ = catvec_broadcast_vec(dest, i)
+            args′ = map(arg -> catvec_broadcast_vec(arg, i), args)
+            axes′ = (eachindex(dest′),)
+            copyto!(dest′, Broadcast.Broadcasted{Nothing}(f, args′, axes′))
         end
     end
     return dest
diff --git a/test/test_custom_collections.jl b/test/test_custom_collections.jl
index 3dcc16ed..35994e5c 100644
--- a/test/test_custom_collections.jl
+++ b/test/test_custom_collections.jl
@@ -208,5 +208,18 @@ Base.axes(m::NonOneBasedMatrix) = ((1:m.m) .- 2, (1:m.n) .+ 1)
         y5 = similar(y)
         map!(+, y5, y, y)
         @test Vector(y5) == Vector(y) + Vector(y)
+
+        z = similar(y)
+        rand!(z)
+        yvec = Vector(y)
+        zvec = Vector(z)
+
+        z .= muladd.(1e-3, y, z)
+        zvec .= muladd.(1e-3, yvec, zvec)
+        @test zvec == z
+        allocs = let y=y, z=z
+            @allocated z .= muladd.(1e-3, y, z)
+        end
+        @test allocs == 0
     end
 end

From e46a068891787021a2be70a759a09ec4fb38bf36 Mon Sep 17 00:00:00 2001
From: Twan Koolen <koolen.twan@gmail.com>
Date: Tue, 31 Dec 2019 17:28:59 -0500
Subject: [PATCH 6/8] Minor tweaks.

---
 src/custom_collections/CatVector.jl | 28 ++++++++++++++++------------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/src/custom_collections/CatVector.jl b/src/custom_collections/CatVector.jl
index d88a9bad..af97c8a2 100644
--- a/src/custom_collections/CatVector.jl
+++ b/src/custom_collections/CatVector.jl
@@ -3,11 +3,11 @@ struct CatVector{T, N, V<:AbstractVector{T}} <: AbstractVector{T}
 end
 
 @inline Base.size(vec::CatVector) = (mapreduce(length, +, vec.vecs; init=0),)
-Base.eltype(vec::CatVector) = eltype(eltype(vec.vecs))
 
-# Note: getindex and setindex are pretty naive.
+# Note: getindex and setindex are pretty naive. Consider precomputing map from
+# index to vector upon CatVector construction.
 Base.@propagate_inbounds function Base.getindex(vec::CatVector, index::Int)
-    @boundscheck checkbounds(vec, index)
+    @boundscheck index >= 1 || throw(BoundsError(vec, index))
     i = 1
     j = index
     @inbounds while true
@@ -20,11 +20,11 @@ Base.@propagate_inbounds function Base.getindex(vec::CatVector, index::Int)
             i += 1
         end
     end
-    error()
+    throw(BoundsError(vec, index))
 end
 
 Base.@propagate_inbounds function Base.setindex!(vec::CatVector, val, index::Int)
-    @boundscheck checkbounds(vec, index)
+    @boundscheck index >= 1 || throw(BoundsError(vec, index))
     i = 1
     j = index
     while true
@@ -38,7 +38,7 @@ Base.@propagate_inbounds function Base.setindex!(vec::CatVector, val, index::Int
             i += 1
         end
     end
-    error()
+    throw(BoundsError(vec, index))
 end
 
 Base.@propagate_inbounds function Base.copyto!(dest::AbstractVector{T}, src::CatVector{T}) where {T}
@@ -58,20 +58,24 @@ 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))
 
+@noinline cat_vectors_line_up_error() = throw(ArgumentError("Subvectors must line up"))
+
 @inline function check_cat_vectors_line_up(x::CatVector, y::CatVector)
-    length(x.vecs) == length(y.vecs) || throw(ArgumentError("Subvectors must line up"))
+    length(x.vecs) == length(y.vecs) || cat_vectors_line_up_error()
     for i in eachindex(x.vecs)
-        length(x.vecs[i]) == length(y.vecs[i]) || throw(ArgumentError("Subvectors must line up"))
+        length(x.vecs[i]) == length(y.vecs[i]) || cat_vectors_line_up_error()
     end
     nothing
 end
 
 @inline check_cat_vectors_line_up(x::CatVector, y) = nothing
-@inline check_cat_vectors_line_up(x::CatVector, y, tail...) = (check_cat_vectors_line_up(x, y); check_cat_vectors_line_up(x, tail...))
+@inline function check_cat_vectors_line_up(x::CatVector, y, tail...)
+    check_cat_vectors_line_up(x, y)
+    check_cat_vectors_line_up(x, tail...)
+end
 
-@inline function Base.copyto!(dest::CatVector, src::CatVector)
-    @boundscheck check_cat_vectors_line_up(dest, src)
-    @inbounds for i in eachindex(dest.vecs)
+@propagate_inbounds function Base.copyto!(dest::CatVector, src::CatVector)
+    for i in eachindex(dest.vecs)
         copyto!(dest.vecs[i], src.vecs[i])
     end
     return dest

From ac27a70c88c7e47b7baf81b7a43f4bab1c08fd4d Mon Sep 17 00:00:00 2001
From: Twan Koolen <koolen.twan@gmail.com>
Date: Tue, 31 Dec 2019 17:32:05 -0500
Subject: [PATCH 7/8] Add docstring.

---
 src/custom_collections/CatVector.jl | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/custom_collections/CatVector.jl b/src/custom_collections/CatVector.jl
index af97c8a2..d673f260 100644
--- a/src/custom_collections/CatVector.jl
+++ b/src/custom_collections/CatVector.jl
@@ -1,3 +1,9 @@
+"""
+$(TYPEDEF)
+
+An `AbstractVector` subtype that acts as a lazy concatenation of a number
+of subvectors.
+"""
 struct CatVector{T, N, V<:AbstractVector{T}} <: AbstractVector{T}
     vecs::NTuple{N, V}
 end

From a7d356752ddf9a384c0da1aac1419db12bf6e35f Mon Sep 17 00:00:00 2001
From: Twan Koolen <koolen.twan@gmail.com>
Date: Tue, 31 Dec 2019 17:53:45 -0500
Subject: [PATCH 8/8] Broadcast improvements.

---
 src/custom_collections/CatVector.jl | 12 ++++++++----
 test/test_custom_collections.jl     | 10 +++++-----
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/custom_collections/CatVector.jl b/src/custom_collections/CatVector.jl
index d673f260..05a8dc0e 100644
--- a/src/custom_collections/CatVector.jl
+++ b/src/custom_collections/CatVector.jl
@@ -95,18 +95,22 @@ end
     return dest
 end
 
-Base.@propagate_inbounds catvec_broadcast_vec(x::CatVector, k::Int) = x.vecs[k]
-Base.@propagate_inbounds catvec_broadcast_vec(x::Number, k::Int) = x
+Base.@propagate_inbounds catvec_broadcast_vec(arg::CatVector, range::UnitRange, k::Int) = arg.vecs[k]
+Base.@propagate_inbounds catvec_broadcast_vec(arg::AbstractVector, range::UnitRange, k::Int) = view(arg, range)
+Base.@propagate_inbounds catvec_broadcast_vec(arg::Number, range::UnitRange, k::Int) = arg
 
 @inline function Base.copyto!(dest::CatVector, bc::Broadcast.Broadcasted{Nothing})
     flat = Broadcast.flatten(bc)
     @boundscheck check_cat_vectors_line_up(dest, flat.args...)
+    offset = 1
     @inbounds for i in eachindex(dest.vecs)
         let i = i, f = flat.f, args = flat.args
-            dest′ = catvec_broadcast_vec(dest, i)
-            args′ = map(arg -> catvec_broadcast_vec(arg, i), args)
+            dest′ = dest.vecs[i]
+            range = offset : offset + length(dest′) - 1
+            args′ = map(arg -> catvec_broadcast_vec(arg, range, i), args)
             axes′ = (eachindex(dest′),)
             copyto!(dest′, Broadcast.Broadcasted{Nothing}(f, args′, axes′))
+            offset = last(range) + 1
         end
     end
     return dest
diff --git a/test/test_custom_collections.jl b/test/test_custom_collections.jl
index 35994e5c..d45859e3 100644
--- a/test/test_custom_collections.jl
+++ b/test/test_custom_collections.jl
@@ -175,13 +175,13 @@ Base.axes(m::NonOneBasedMatrix) = ((1:m.m) .- 2, (1:m.n) .+ 1)
         for i in eachindex(x)
             @test x[i] == y[i]
         end
-
-        x .= 0
-        for i in eachindex(y)
-            x[i] = y[i]
-        end
         @test x == y
 
+        y .= 0
+        rand!(x)
+        y .= x .+ y .+ 1
+        @test x .+ 1 == y
+
         allocs = let x=x, vecs=vecs
             @allocated copyto!(x, RigidBodyDynamics.CatVector(vecs))
         end