From 0cd77ddd05a7c9d9c03a38d8e06bbddecb1ad908 Mon Sep 17 00:00:00 2001
From: Rafael Fourquet <fourquet.rafael@gmail.com>
Date: Wed, 15 Apr 2020 09:33:02 +0200
Subject: [PATCH] fix replace(f, ::Union{Set,Dict}, count=0) (#35450)

---
 base/set.jl  | 5 +++--
 test/sets.jl | 8 ++++++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/base/set.jl b/base/set.jl
index 1f1c3599ac83d..f766d59b57c49 100644
--- a/base/set.jl
+++ b/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
diff --git a/test/sets.jl b/test/sets.jl
index 677f2ee776e95..23e0e8175dd05 100644
--- a/test/sets.jl
+++ b/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