Skip to content

Commit 0f5bb4b

Browse files
committed
make replace/replace! work with count
1 parent 245d224 commit 0f5bb4b

File tree

2 files changed

+47
-2
lines changed

2 files changed

+47
-2
lines changed

src/chainedvector.jl

+40-2
Original file line numberDiff line numberDiff line change
@@ -981,7 +981,45 @@ end
981981

982982
Base.replace(f::Base.Callable, a::ChainedVector) = ChainedVector([replace(f, A) for A in a.arrays])
983983
Base.replace!(f::Base.Callable, a::ChainedVector) = (foreach(A -> replace!(f, A), a.arrays); return a)
984-
Base.replace(a::ChainedVector, old_new::Pair...; count::Union{Integer,Nothing}=nothing) = ChainedVector([replace(A, old_new...; count=count) for A in a.arrays])
985-
Base.replace!(a::ChainedVector, old_new::Pair...; count::Integer=typemax(Int)) = (foreach(A -> replace!(A, old_new...; count=count), a.arrays); return a)
984+
985+
function _check_count(count::Integer)
986+
count < 0 && throw(DomainError(count, "`count` must not be negative"))
987+
return min(count, typemax(Int)) % Int
988+
end
989+
990+
function Base.replace(A::ChainedVector, old_new::Pair...; count::Integer=typemax(Int))
991+
count = _check_count(count)
992+
R = similar(A)
993+
isempty(A) && return R
994+
reps = 0
995+
for (ai,ri) in zip(eachindex(A),eachindex(R))
996+
R[ri] = A[ai]
997+
for (old, new) in old_new
998+
if R[ri] == old && reps < count
999+
R[ri] = new
1000+
reps += 1
1001+
continue
1002+
end
1003+
end
1004+
end
1005+
return R
1006+
end
1007+
1008+
function Base.replace!(A::ChainedVector, old_new::Pair...; count::Integer=typemax(Int))
1009+
count = _check_count(count)
1010+
isempty(A) && return A
1011+
reps = 0
1012+
for i in eachindex(A)
1013+
for (old, new) in old_new
1014+
reps == count && return A
1015+
if A[i] == old
1016+
A[i] = new
1017+
reps += 1
1018+
continue
1019+
end
1020+
end
1021+
end
1022+
return A
1023+
end
9861024

9871025
Base.Broadcast.broadcasted(f::F, A::ChainedVector) where {F} = map(f, A)

test/chainedvector.jl

+7
Original file line numberDiff line numberDiff line change
@@ -350,6 +350,13 @@
350350
@test map(x -> x == 1 ? 2.0 : x, x) == replace!(x, 1 => 2)
351351
@test isempty(x)
352352

353+
@test replace!(ChainedVector([[1,2], [1,2]]), 2=>20) == [1,20,1,20]
354+
@test replace!(ChainedVector([[1,2], [1,2]]), 2=>20, count=1) == [1,20,1,2]
355+
@test replace!(ChainedVector([[1,2], [1,2]]), 2=>20, count=2) == [1,20,1,20]
356+
x = [1,2]
357+
@test replace!(ChainedVector([x,[2,3]]), 2=>99) == [1,99,99,3]
358+
@test x == [1,99]
359+
353360
# copyto!
354361
# ChainedVector dest: doffs, soffs, n
355362
x = ChainedVector([[1,2,3], [4,5,6], [7,8,9,10]])

0 commit comments

Comments
 (0)