Rework replace and replace!

Introduce a new replace!(new::Callable, res::T, A::T, count::Union{Nothing,Int}) method
which custom types can implement to support all replace and replace! methods automatically,
instead of the current replace!(new::Callable, A::T, count::Int). This offers several advantages:
- For arrays, instead of copying the input and then replace elements, we can do the copy and replace
operations at the same time, which is quite faster for arrays when count=nothing.
- For dicts and sets, copying up-front is still faster as long as most original elements are preserved,
but for replace(), we can apply replacements directly instead of storing a them in a temporary vector.
- When the LHS of a pair contains a singleton type, we can subtract it from the element type
of the result, e.g. Union{T,Missing} becomes T.

Also simplify the dispatch logic by removing the internal _replace! method in favor of replace!.

Author: nalimilan

base/set.jl

@@ -574,8 +574,14 @@ _copy_oftype(x::AbstractArray{T}, ::Type{T}) where {T} = copy(x)
 _copy_oftype(x::AbstractDict{K,V}, ::Type{Pair{K,V}}) where {K,V} = copy(x)
 _copy_oftype(x::AbstractSet{T}, ::Type{T}) where {T} = copy(x)
 
+_similar_or_copy(x::Any) = similar(x)
+_similar_or_copy(x::Any, ::Type{T}) where {T} = similar(x, T)
+# Make a copy on construction since it is faster than inserting elements separately
+_similar_or_copy(x::Union{AbstractDict,AbstractSet}) where {T} = copy(x)
+_similar_or_copy(x::Union{AbstractDict,AbstractSet}, ::Type{T}) where {T} = _copy_oftype(x, T)
+
 # to make replace/replace! work for a new container type Cont, only
-# replace!(new::Callable, A::Cont; count::Integer=typemax(Int))
+# replace!(new::Callable, res::Cont, A::Cont; count::Integer=typemax(Int))
 # has to be implemented
 
 """
@@ -600,16 +606,17 @@ julia> replace!(Set([1, 2, 3]), 1=>0)
 Set([0, 2, 3])
 ```
 """
-replace!(A, old_new::Pair...; count::Integer=typemax(Int)) = _replace!(A, count, old_new)
+replace!(A, old_new::Pair...; count::Union{Integer,Nothing}=nothing) =
+    replace!(A, A, count, old_new)
 
-function _replace!(A, count::Integer, old_new::Tuple{Vararg{Pair}})
+function replace!(res, A, count::Union{Integer,Nothing}, old_new::Tuple{Vararg{Pair}})
     @inline function new(x)
         for o_n in old_new
             isequal(first(o_n), x) && return last(o_n)
         end
         return x # no replace
     end
-    replace!(new, A, count=count)
+    replace!(new, res, A, count)
 end
 
 """
@@ -630,7 +637,7 @@ julia> replace!(isodd, A, 0, count=2)
  1
 ```
 """
-replace!(pred::Callable, A, new; count::Integer=typemax(Int)) =
+replace!(pred::Callable, A, new; count::Union{Integer,Nothing}=nothing) =
     replace!(x -> ifelse(pred(x), new, x), A, count=count)
 
 """
@@ -661,9 +668,14 @@ Set([6, 12])
 ```
 """
 function replace!(new::Callable, A::Union{AbstractArray,AbstractDict,AbstractSet};
-                  count::Integer=typemax(Int))
-    count < 0 && throw(DomainError(count, "`count` must not be negative"))
-    count != 0 && _replace!(new, A, min(count, typemax(Int)) % Int)
+                  count::Union{Integer,Nothing}=nothing)
+    if count === nothing
+        replace!(new, A, A, nothing)
+    elseif count < 0
+        throw(DomainError(count, "`count` must not be negative"))
+    elseif count != 0
+        replace!(new, A, A, min(count, typemax(Int)) % Int)
+    end
     A
 end
 
@@ -686,16 +698,33 @@ julia> replace([1, 2, 1, 3], 1=>0, 2=>4, count=2)
  3
 ```
 """
-function replace(A, old_new::Pair...; count::Integer=typemax(Int))
+function replace(A, old_new::Pair...; count::Union{Integer,Nothing}=nothing)
     V = promote_valuetype(old_new...)
-    T = promote_type(eltype(A), V)
-    _replace!(_copy_oftype(A, T), count, old_new)
+    if count isa Nothing
+        T = promote_type(subtract_singletontype(eltype(A), old_new...), V)
+        replace!(_similar_or_copy(A, T), A, nothing, old_new)
+    else
+        U = promote_type(eltype(A), V)
+        replace!(_similar_or_copy(A, U), A, min(count, typemax(Int)) % Int, old_new)
+    end
 end
 
 promote_valuetype(x::Pair{K, V}) where {K, V} = V
 promote_valuetype(x::Pair{K, V}, y::Pair...) where {K, V} =
     promote_type(V, promote_valuetype(y...))
 
+# Subtract singleton types which are going to be replaced
+@pure issingletontype(::Type{T}) where {T} = isdefined(T, :instance)
+function subtract_singletontype(::Type{T}, x::Pair{K}) where {T, K}
+    if issingletontype(K) # singleton type
+        Core.Compiler.typesubtract(T, K)
+    else
+        T
+    end
+end
+subtract_singletontype(::Type{T}, x::Pair{K}, y::Pair...) where {T, K} =
+    subtract_singletontype(subtract_singletontype(T, y...), x)
+
 """
     replace(pred::Function, A, new; [count::Integer])
 
@@ -713,9 +742,10 @@ julia> replace(isodd, [1, 2, 3, 1], 0, count=2)
  1
 ```
 """
-function replace(pred::Callable, A, new; count::Integer=typemax(Int))
+function replace(pred::Callable, A, new; count::Union{Integer,Nothing}=nothing)
     T = promote_type(eltype(A), typeof(new))
-    replace!(pred, _copy_oftype(A, T), new, count=count)
+    replace!(x -> ifelse(pred(x), new, x), _similar_or_copy(A, T), A,
+             count === nothing ? nothing : min(count, typemax(Int)) % Int)
 end
 
 """
@@ -742,7 +772,9 @@ Dict{Int64,Int64} with 2 entries:
   1 => 3
 ```
 """
-replace(new::Callable, A; count::Integer=typemax(Int)) = replace!(new, copy(A), count=count)
+replace(new::Callable, A; count::Union{Integer,Nothing}=nothing) =
+    replace!(new, _similar_or_copy(A), A,
+             count === nothing ? nothing : min(count, typemax(Int)) % Int)
 
 # Handle ambiguities
 replace!(a::Callable, b::Pair; count::Integer=-1) = throw(MethodError(replace!, (a, b)))
@@ -757,36 +789,53 @@ replace(a::AbstractString, b::Pair, c::Pair) = throw(MethodError(replace, (a, b,
 askey(k, ::AbstractDict) = k.first
 askey(k, ::AbstractSet) = k
 
-function _replace!(new::Callable, A::Union{AbstractDict,AbstractSet}, count::Int)
-    repl = Pair{eltype(A),eltype(A)}[]
+function replace!(new::Callable, res::T, A::T,
+                  count::Union{Int,Nothing}) where T<:Union{AbstractDict,AbstractSet}
     c = 0
-    for x in A
-        y = new(x)
-        if x !== y
-            push!(repl, x => y)
-            c += 1
+    if res === A
+        repl = Pair{eltype(A),eltype(A)}[]
+        for x in A
+            y = new(x)
+            if x !== y
+                push!(repl, x => y)
+                c += 1
+            end
+            c == count && break
+        end
+        for oldnew in repl
+            pop!(res, askey(first(oldnew), res))
+        end
+        for oldnew in repl
+            push!(res, last(oldnew))
+        end
+    else
+        for x in A
+            y = new(x)
+            if x !== y
+                pop!(res, askey(x, res))
+                push!(res, y)
+                c += 1
+            end
+            c == count && break
         end
-        c == count && break
-    end
-    for oldnew in repl
-        pop!(A, askey(first(oldnew), A))
-    end
-    for oldnew in repl
-        push!(A, last(oldnew))
     end
+    res
 end
 
-### AbstractArray
+### replace! for AbstractArray
 
-function _replace!(new::Callable, A::AbstractArray, count::Int)
+function replace!(new::Callable, res::AbstractArray, A::AbstractArray,
+                  count::Union{Int,Nothing})
     c = 0
     for i in eachindex(A)
         @inbounds Ai = A[i]
-        y = new(Ai)
-        if Ai !== y
-            @inbounds A[i] = y
-            c += 1
+        if count === nothing || c < count
+            y = new(Ai)
+            @inbounds res[i] = y
+            c += (Ai !== y)
+        else
+            @inbounds res[i] = Ai
         end
-        c == count && break
     end
-end
+    res
+end

test/sets.jl

@@ -548,6 +548,15 @@ end
     x = @inferred replace(x -> x > 1, [1, 2], missing)
     @test isequal(x, [1, missing]) && x isa Vector{Union{Int, Missing}}
 
+    x = @inferred replace([1, missing], missing=>2)
+    @test x == [1, 2] && x isa Vector{Int}
+    x = @inferred replace([1, missing], missing=>2, count=1)
+    @test x == [1, 2] && x isa Vector{Union{Int, Missing}}
+    x = @inferred replace([1, missing], missing=>missing)
+    @test isequal(x, [1, missing]) && x isa Vector{Union{Int, Missing}}
+    x = @inferred replace([1, missing], missing=>2, 1=>missing)
+    @test isequal(x, [missing, 2]) && x isa Vector{Union{Int, Missing}}
+
     # test that isequal is used
     @test replace([NaN, 1.0], NaN=>0.0) == [0.0, 1.0]
     @test replace([1, missing], missing=>0) == [1, 0]