Skip to content

Commit

Permalink
Broadcast improvements.
Browse files Browse the repository at this point in the history
  • Loading branch information
tkoolen committed Jan 12, 2020
1 parent ac27a70 commit a7d3567
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 9 deletions.
12 changes: 8 additions & 4 deletions src/custom_collections/CatVector.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
10 changes: 5 additions & 5 deletions test/test_custom_collections.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit a7d3567

Please sign in to comment.