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