fix replace(f, ::Union{Set,Dict}, count=0) (#35450)

base/set.jl

@@ -590,6 +590,7 @@ askey(k, ::AbstractSet) = k
 
 function _replace!(new::Callable, res::T, A::T,
                    count::Int) where T<:Union{AbstractDict,AbstractSet}
+    count == 0 && return res
     c = 0
     if res === A # cannot replace elements while iterating over A
         repl = Pair{eltype(A),eltype(A)}[]
@@ -598,8 +599,8 @@ function _replace!(new::Callable, res::T, A::T,
             if x !== y
                 push!(repl, x => y)
                 c += 1
+                c == count && break
             end
-            c == count && break
         end
         for oldnew in repl
             pop!(res, askey(first(oldnew), res))
@@ -614,8 +615,8 @@ function _replace!(new::Callable, res::T, A::T,
                 pop!(res, askey(x, res))
                 push!(res, y)
                 c += 1
+                c == count && break
             end
-            c == count && break
         end
     end
     res

test/sets.jl

@@ -591,8 +591,12 @@ end
     @test s == Set([2, 3])
     @test replace!(x->2x, s, count=0x1) in [Set([4, 3]), Set([2, 6])]
 
-    for count = (0, 0x0, big(0))
-        @test replace([1, 2], 1=>0, 2=>0, count=count) == [1, 2] # count=0 --> no replacements
+    for count = (0, 0x0, big(0)) # count == 0 --> no replacements
+        @test replace([1, 2], 1=>0, 2=>0; count) == [1, 2]
+        for dict = (Dict(1=>2, 2=>3), IdDict(1=>2, 2=>3))
+            @test replace(dict, (1=>2) => (1=>3); count) == dict
+        end
+        @test replace(Set([1, 2]), 2=>-1; count) == Set([1, 2])
     end
 
     # test collisions with AbstractSet/AbstractDict