Fix replace/replace! when count is supplied

src/chainedvector.jl

@@ -242,6 +242,7 @@ end
 
 # custom index type used in eachindex
 struct ChainedVectorIndex{A} <: Integer
+    arrays::Vector{A}
     arrays_i::Int
     array::A
     array_i::Int
@@ -258,17 +259,23 @@ Base.hash(x::ChainedVectorIndex, h::UInt) = hash(x.i, h)
 
 @inline Base.getindex(x::ChainedVectorIndex) = @inbounds x.array[x.array_i]
 
+# An index from `eachindex(B)` may be used on another ChainedVector `A`, e.g. `similar(B)`.
+# It addresses the chunk it holds only when used on `B`; elsewhere it's the linear index `x.i`.
+@inline ownsindex(A::ChainedVector, x::ChainedVectorIndex) = x.arrays === A.arrays
+
 function Base.checkbounds(::Type{Bool}, A::ChainedVector, ind::ChainedVectorIndex)
-    @assert ind.array === A.arrays[ind.arrays_i] "indexing ChainedVector with wrong ChainedVectorIndex"
+    ownsindex(A, ind) || return checkbounds(Bool, A, ind.i)
     return 1 <= ind.array_i <= length(ind.array)
 end
 
 Base.@propagate_inbounds function Base.getindex(A::ChainedVector, x::ChainedVectorIndex)
+    ownsindex(A, x) || return A[x.i]
     @boundscheck checkbounds(A, x)
     return @inbounds x.array[x.array_i]
 end
 
 Base.@propagate_inbounds function Base.setindex!(A::ChainedVector, v, x::ChainedVectorIndex)
+    ownsindex(A, x) || return setindex!(A, v, x.i)
     @boundscheck checkbounds(A, x)
     @inbounds x.array[x.array_i] = v
     return v
@@ -278,7 +285,7 @@ function Base.getindex(A::ChainedVector{T}, inds::AbstractVector{<:ChainedVector
     len = length(inds)
     x = Vector{T}(undef, len)
     for i = 1:len
-        x[i] = inds[i][]
+        x[i] = A[inds[i]]
     end
     return x
 end
@@ -298,8 +305,9 @@ function Base.nextind(A::ChainedVector, x::ChainedVectorIndex)
         i += 1
     else
         chunk_i += 1 # make sure this goes out of bounds
+        i += 1
     end
-    return ChainedVectorIndex(chunkidx, chunk, chunk_i, i)
+    return ChainedVectorIndex(x.arrays, chunkidx, chunk, chunk_i, i)
 end
 
 function Base.prevind(A::ChainedVector, x::ChainedVectorIndex)
@@ -317,8 +325,9 @@ function Base.prevind(A::ChainedVector, x::ChainedVectorIndex)
         i -= 1
     else
         chunk_i -= 1 # make sure this goes out of bounds
+        i -= 1
     end
-    return ChainedVectorIndex(chunkidx, chunk, chunk_i, i)
+    return ChainedVectorIndex(x.arrays, chunkidx, chunk, chunk_i, i)
 end
 
 # efficient iteration via eachindex
@@ -342,7 +351,7 @@ end
     chunkidx = chunk_i = 1
     @inbounds chunk = arrays[chunkidx]
     # we already ran cleanup! so chunks are guaranteed non-empty
-    return ChainedVectorIndex(chunkidx, chunk, chunk_i, 1), (arrays, chunkidx, chunk, length(chunk), chunk_i + 1, 2)
+    return ChainedVectorIndex(arrays, chunkidx, chunk, chunk_i, 1), (arrays, chunkidx, chunk, length(chunk), chunk_i + 1, 2)
 end
 
 @inline function Base.iterate(x::IndexIterator, (arrays, chunkidx, chunk, chunklen, chunk_i, i))
@@ -353,7 +362,7 @@ end
         chunklen = length(chunk)
         chunk_i = 1
     end
-    return ChainedVectorIndex(chunkidx, chunk, chunk_i, i), (arrays, chunkidx, chunk, chunklen, chunk_i + 1, i + 1)
+    return ChainedVectorIndex(arrays, chunkidx, chunk, chunk_i, i), (arrays, chunkidx, chunk, chunklen, chunk_i + 1, i + 1)
 end
 
 @inline function Base.iterate(A::ChainedVector)
@@ -955,7 +964,17 @@ function Base.filter!(f, a::ChainedVector)
     return a
 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)
+# Without a `count` limit, chunks are replaced independently. A limit spans chunks,
+# so limited replacements use Base's elementwise methods.
+Base.replace(f::Base.Callable, a::ChainedVector; count::Integer=typemax(Int)) =
+    count < length(a) ? invoke(replace, Tuple{Base.Callable, Any}, f, a; count=count) :
+    ChainedVector([replace(f, A) for A in a.arrays])
+Base.replace!(f::Base.Callable, a::ChainedVector; count::Integer=typemax(Int)) =
+    count < length(a) ? invoke(replace!, Tuple{Base.Callable, Any}, f, a; count=count) :
+    (foreach(A -> replace!(f, A), a.arrays); a)
+Base.replace(a::ChainedVector, old_new::Pair...; count::Union{Integer,Nothing}=nothing) =
+    count === nothing ? ChainedVector([replace(A, old_new...) for A in a.arrays]) :
+    invoke(replace, Tuple{Any, Vararg{Pair}}, a, old_new...; count=count)
+Base.replace!(a::ChainedVector, old_new::Pair...; count::Integer=typemax(Int)) =
+    count < length(a) ? invoke(replace!, Tuple{Any, Vararg{Pair}}, a, old_new...; count=count) :
+    (foreach(A -> replace!(A, old_new...), a.arrays); a)

test/chainedvector.jl

@@ -380,6 +380,17 @@ end
     @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]
+    # results stay ChainedVectors and promote like Base
+    @test replace(ChainedVector([[1,2], [1,2]]), 2=>2.5) isa ChainedVector{Float64}
+    @test replace(ChainedVector([[1,2], [1,2]]), 2=>2.5, count=1) == [1,2.5,1,2]
+    @test replace(ChainedVector([[1,missing], [missing]]), missing=>0) isa ChainedVector{Int}
+
     # copyto!
     # ChainedVector dest: doffs, soffs, n
     x = ChainedVector([[1,2,3], [4,5,6], [7,8,9,10]])
@@ -623,6 +634,48 @@ end
     end
 end
 
+@testset "replace[!] comparison with Vector" begin
+
+    testvecs = (
+        [[1, 2], [3, 2, 5]],
+        [[1, 2]],
+        [[2],[2],[2],[2,3]],
+        [[1,2,missing]],
+        [[missing,1],[missing,2,1]],
+        [[missing]]
+    )
+    function missing_equal(a,b)
+        ismissing(a) && ismissing(b) && return true
+        ismissing(a) ⊻ ismissing(b) && return false
+        return all(skipmissing(a) .== skipmissing(b))
+    end 
+    gen_cv_v(x) = (c = ChainedVector(x); (c, collect(c)))
+    for f in (replace, replace!)
+        for x in testvecs
+            cv, v = gen_cv_v(x)
+            @test missing_equal(f(v, 2 => 22),f(cv, 2 => 22))
+            @test missing_equal(v,cv)
+
+            cv, v = gen_cv_v(x)
+            @test missing_equal(f(x -> x ÷ 2, v), f(x -> x ÷ 2, cv))
+            @test missing_equal(v,cv)
+
+            for c in (0, 1, 2, 3)
+                cv, v = gen_cv_v(x)
+                @test missing_equal(f(x -> x ÷ 2, v, count=c), f(x -> x ÷ 2, cv, count=c))
+                @test missing_equal(v,cv)
+
+                for p in ((2=>2,),(2 => 22,), (2 => 22, 3 => 33))
+                    cv, v = gen_cv_v(x)
+                    @test missing_equal(f(v, p..., count=c), f(cv, p..., count=c))
+                    @test eltype(f(v, p..., count=c)) == eltype(f(cv, p..., count=c))
+                    @test missing_equal(v,cv)
+                end
+            end
+        end
+    end
+end
+
 
 @testset "iteration protocol on ChainedVector" begin
     for len in 0:6
@@ -733,8 +786,29 @@ end
     ind = prevind(x, ind)
     @test_throws BoundsError x[ind]
     # https://github.com/apache/arrow-julia/issues/418
+    # indices from one ChainedVector index others by position, even under @inbounds
     y = ChainedVector([collect(1:i) for i = 10:100])
-    @test_throws AssertionError x[first(eachindex(y))]
+    z = ChainedVector([x[1:7], x[8:end]]) # same values, different chunks
+    for (i, ind) in enumerate(eachindex(y))
+        @test x[ind] == @inbounds(x[ind]) == z[ind] == x[i] == z[i]
+        @inbounds y[ind] = -i
+    end
+    @test y == -(1:length(y))
+    w = similar(y)
+    for ind in eachindex(y)
+        @inbounds w[ind] = y[ind]
+    end
+    @test w == y
+    longer = ChainedVector([x.arrays; [[1]]])
+    @test_throws BoundsError x[collect(eachindex(longer))[end]]
+    p = ChainedVector([[1, 2], [3, 4]])
+    i = first(eachindex(ChainedVector([[10], [20, 30, 40]])))
+    @test p[nextind(p, i)] == 2
+    @test p[prevind(p, nextind(p, nextind(p, i)))] == 2
+    # generic Base code indexes `similar(A)` with `eachindex(A)`
+    a = ChainedVector([[1, 2], [1, 2]])
+    @test replace(x -> 10x, a; count=3) == [10, 20, 10, 2]
+    @test a == [1, 2, 1, 2]
     # https://github.com/JuliaData/SentinelArrays.jl/issues/74
     x = ChainedVector([[true], [false], [true]])
     @test BitVector(x) == [true, false, true]
@@ -782,7 +856,7 @@ end
 end
 
 @testset "getindex with UnitRange" begin
-    x = ChainedVector([collect(1:i) for i = 10:100])
+    x = ChainedVector([collect(1:i) for i = 1:10])
     @test isempty(x[1:0])
     @test x[1:1] == [1]
     @test x[1:end] == x