Rework replace and replace!

base/set.jl

@@ -564,6 +564,11 @@ convert(::Type{T}, s::AbstractSet) where {T<:AbstractSet} = T(s)
 
 ## replace/replace! ##
 
+function check_count(count::Integer)
+ count < 0 && throw(DomainError(count, "`count` must not be negative (got $count)"))
+ return min(count, typemax(Int)) % Int
+end
+
 # TODO: use copy!, which is currently unavailable from here since it is defined in Future
 _copy_oftype(x, ::Type{T}) where {T} = copyto!(similar(x, T), x)
 # TODO: use similar() once deprecation is removed and it preserves keys
@@ -574,8 +579,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}) = 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::Int)
 # has to be implemented
 
 """
@@ -600,16 +611,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::Integer=typemax(Int)) =
+ replace_pairs!(A, A, check_count(count), old_new)
 
-function _replace!(A, count::Integer, old_new::Tuple{Vararg{Pair}})
+function replace_pairs!(res, A, count::Int, 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
 
 """
@@ -631,7 +643,7 @@ julia> replace!(isodd, A, 0, count=2)
 ```
 """
 replace!(pred::Callable, A, new; count::Integer=typemax(Int)) =
- replace!(x -> ifelse(pred(x), new, x), A, count=count)
+ replace!(x -> ifelse(pred(x), new, x), A, count=check_count(count))
 
 """
  replace!(new::Function, A; [count::Integer])
@@ -660,12 +672,8 @@ julia> replace!(x->2x, Set([3, 6]))
 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)
- A
-end
+replace!(new::Callable, A; count::Integer=typemax(Int)) =
+ _replace!(new, A, A, check_count(count))
 
 """
  replace(A, old_new::Pair...; [count::Integer])
@@ -674,6 +682,14 @@ Return a copy of collection `A` where, for each pair `old=>new` in `old_new`,
 all occurrences of `old` are replaced by `new`.
 Equality is determined using [`isequal`](@ref).
 If `count` is specified, then replace at most `count` occurrences in total.
+
+The element type of the result is chosen using promotion (see [`promote_type`](@ref))
+based on the element type of `A` and on the types of the `new` values in pairs.
+If `count` is omitted and the element type of `A` is a `Union`, the element type
+of the result will not include singleton types which are replaced with values of
+a different type: for example, `Union{T,Missing}` will become `T` if `missing` is
+replaced.
+
 See also [`replace!`](@ref).
 
 # Examples
@@ -684,18 +700,41 @@ julia> replace([1, 2, 1, 3], 1=>0, 2=>4, count=2)
  4
  1
  3
+
+julia> replace([1, missing], missing=>0)
+2-element Array{Int64,1}:
+ 1
+ 0
 ```
 """
-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_pairs!(_similar_or_copy(A, T), A, typemax(Int), old_new)
+ else
+ U = promote_type(eltype(A), V)
+ replace_pairs!(_similar_or_copy(A, U), A, check_count(count), 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(T::DataType) = isdefined(T, :instance)
+issingletontype(::Type) = false
+function subtract_singletontype(::Type{T}, x::Pair{K}) where {T, K}
+ if issingletontype(K)
+ 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])
 
@@ -715,7 +754,7 @@ julia> replace(isodd, [1, 2, 3, 1], 0, count=2)
 """
 function replace(pred::Callable, A, new; count::Integer=typemax(Int))
  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, check_count(count))
 end
 
 """
@@ -742,7 +781,8 @@ 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::Integer=typemax(Int)) =
+ _replace!(new, _similar_or_copy(A), A, check_count(count))
 
 # Handle ambiguities
 replace!(a::Callable, b::Pair; count::Integer=-1) = throw(MethodError(replace!, (a, b)))
@@ -751,42 +791,65 @@ replace(a::Callable, b::Pair; count::Integer=-1) = throw(MethodError(replace, (a
 replace(a::Callable, b::Pair, c::Pair; count::Integer=-1) = throw(MethodError(replace, (a, b, c)))
 replace(a::AbstractString, b::Pair, c::Pair) = throw(MethodError(replace, (a, b, c)))
 
-
 ### replace! for AbstractDict/AbstractSet
 
 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::Int) 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 # cannot replace elements while iterating over 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::Int)
  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 >= length(A) # simpler loop allows for SIMD
+ for i in eachindex(A)
+ @inbounds Ai = A[i]
+ y = new(Ai)
+ @inbounds res[i] = y
+ end
+ else
+ for i in eachindex(A)
+ @inbounds Ai = A[i]
+ if c < count
+ y = new(Ai)
+ @inbounds res[i] = y
+ c += (Ai !== y)
+ else
+ @inbounds res[i] = Ai
+ end
  end
- c == count && break
  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]