From 0f5bb4bc009a71991171557791016d503ba43c8f Mon Sep 17 00:00:00 2001
From: Gerhard Aigner <gerhard.aigner@gmail.com>
Date: Sat, 21 Dec 2024 12:29:06 +0100
Subject: [PATCH] make replace/replace! work with count

---
 src/chainedvector.jl  | 42 ++++++++++++++++++++++++++++++++++++++++--
 test/chainedvector.jl |  7 +++++++
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/src/chainedvector.jl b/src/chainedvector.jl
index a5dd69e..6e74529 100644
--- a/src/chainedvector.jl
+++ b/src/chainedvector.jl
@@ -981,7 +981,45 @@ end
 
 Base.replace(f::Base.Callable, a::ChainedVector) = ChainedVector([replace(f, A) for A in a.arrays])
 Base.replace!(f::Base.Callable, a::ChainedVector) = (foreach(A -> replace!(f, A), a.arrays); return a)
-Base.replace(a::ChainedVector, old_new::Pair...; count::Union{Integer,Nothing}=nothing) = ChainedVector([replace(A, old_new...; count=count) for A in a.arrays])
-Base.replace!(a::ChainedVector, old_new::Pair...; count::Integer=typemax(Int)) = (foreach(A -> replace!(A, old_new...; count=count), a.arrays); return a)
+
+function _check_count(count::Integer)
+    count < 0 && throw(DomainError(count, "`count` must not be negative"))
+    return min(count, typemax(Int)) % Int
+end
+
+function Base.replace(A::ChainedVector, old_new::Pair...; count::Integer=typemax(Int))
+    count = _check_count(count)
+    R = similar(A)
+    isempty(A) && return R
+    reps = 0
+    for (ai,ri) in zip(eachindex(A),eachindex(R))
+        R[ri] = A[ai]
+        for (old, new) in old_new
+            if R[ri] == old && reps < count
+                R[ri] = new
+                reps += 1
+                continue
+            end
+        end
+    end
+    return R
+end
+
+function Base.replace!(A::ChainedVector, old_new::Pair...; count::Integer=typemax(Int))
+    count = _check_count(count)
+    isempty(A) && return A
+    reps = 0
+    for i in eachindex(A)
+        for (old, new) in old_new
+            reps == count && return A
+            if A[i] == old
+                A[i] = new
+                reps += 1
+                continue
+            end
+        end
+    end
+    return A
+end
 
 Base.Broadcast.broadcasted(f::F, A::ChainedVector) where {F} = map(f, A)
diff --git a/test/chainedvector.jl b/test/chainedvector.jl
index faf3d93..07f0b65 100644
--- a/test/chainedvector.jl
+++ b/test/chainedvector.jl
@@ -350,6 +350,13 @@
     @test map(x -> x == 1 ? 2.0 : x, x) == replace!(x, 1 => 2)
     @test isempty(x)
 
+    @test replace!(ChainedVector([[1,2], [1,2]]), 2=>20) == [1,20,1,20]
+    @test replace!(ChainedVector([[1,2], [1,2]]), 2=>20, count=1) == [1,20,1,2]
+    @test replace!(ChainedVector([[1,2], [1,2]]), 2=>20, count=2) == [1,20,1,20]
+    x = [1,2]
+    @test replace!(ChainedVector([x,[2,3]]), 2=>99) == [1,99,99,3]
+    @test x == [1,99]
+
     # copyto!
     # ChainedVector dest: doffs, soffs, n
     x = ChainedVector([[1,2,3], [4,5,6], [7,8,9,10]])