From 68e5b47bfbcd0d4a9eb75d49aef54ccad93f254c Mon Sep 17 00:00:00 2001
From: Jameson Nash <vtjnash@gmail.com>
Date: Sun, 6 Dec 2020 18:09:06 -0500
Subject: [PATCH] fixup! fix numerous issues with WeakKeyDict

---
 base/weakkeydict.jl | 66 ++++++++++++++++++++++++++++++++++++++-------
 test/dict.jl        | 21 ++++++++++++---
 2 files changed, 73 insertions(+), 14 deletions(-)

diff --git a/base/weakkeydict.jl b/base/weakkeydict.jl
index 41e11e9033f67..47e655e4b78a9 100644
--- a/base/weakkeydict.jl
+++ b/base/weakkeydict.jl
@@ -81,12 +81,16 @@ end
 sizehint!(d::WeakKeyDict, newsz) = sizehint!(d.ht, newsz)
 empty(d::WeakKeyDict, ::Type{K}, ::Type{V}) where {K, V} = WeakKeyDict{K, V}()
 
+IteratorSize(::Type{<:WeakKeyDict}) = SizeUnknown()
+
 islocked(wkh::WeakKeyDict) = islocked(wkh.lock)
 lock(f, wkh::WeakKeyDict) = lock(f, wkh.lock)
 trylock(f, wkh::WeakKeyDict) = trylock(f, wkh.lock)
 
 function setindex!(wkh::WeakKeyDict{K}, v, key) where K
     !isa(key, K) && throw(ArgumentError("$(limitrepr(key)) is not a valid key for type $K"))
+    # 'nothing' is not valid both because 'finalizer' will reject it,
+    # and because we therefore use it as a sentinel value
     key === nothing && throw(ArgumentError("`nothing` is not a valid key for type WeakKeyDict"))
     lock(wkh) do
         _cleanup_locked(wkh)
@@ -103,7 +107,7 @@ function setindex!(wkh::WeakKeyDict{K}, v, key) where K
 end
 function get!(wkh::WeakKeyDict{K}, key, default) where {K}
     v = lock(wkh) do
-        if haskey(wkh.ht, key)
+        if key !== nothing && haskey(wkh.ht, key)
             wkh.ht[key]
         else
             wkh[key] = default
@@ -113,7 +117,7 @@ function get!(wkh::WeakKeyDict{K}, key, default) where {K}
 end
 function get!(default::Callable, wkh::WeakKeyDict{K}, key) where {K}
     v = lock(wkh) do
-        if haskey(wkh.ht, key)
+        if key !== nothing && haskey(wkh.ht, key)
             wkh.ht[key]
         else
             wkh[key] = default()
@@ -132,14 +136,56 @@ function getkey(wkh::WeakKeyDict{K}, kk, default) where K
 end
 
 map!(f, iter::ValueIterator{<:WeakKeyDict})= map!(f, values(iter.dict.ht))
-get(wkh::WeakKeyDict{K}, key, default) where {K} = lock(() -> get(wkh.ht, key, default), wkh)
-get(default::Callable, wkh::WeakKeyDict{K}, key) where {K} = lock(() -> get(default, wkh.ht, key), wkh)
-pop!(wkh::WeakKeyDict{K}, key) where {K} = lock(() -> pop!(wkh.ht, key), wkh)
-pop!(wkh::WeakKeyDict{K}, key, default) where {K} = lock(() -> pop!(wkh.ht, key, default), wkh)
-delete!(wkh::WeakKeyDict, key) = (lock(() -> delete!(wkh.ht, key), wkh); wkh)
-empty!(wkh::WeakKeyDict) = (lock(() -> empty!(wkh.ht), wkh); wkh)
-haskey(wkh::WeakKeyDict{K}, key) where {K} = lock(() -> haskey(wkh.ht, key), wkh)
-getindex(wkh::WeakKeyDict{K}, key) where {K} = lock(() -> getindex(wkh.ht, key), wkh)
+
+function get(wkh::WeakKeyDict{K}, key, default) where {K}
+    key === nothing && throw(KeyError(nothing))
+    lock(wkh) do
+        return get(wkh.ht, key, default)
+    end
+end
+function get(default::Callable, wkh::WeakKeyDict{K}, key) where {K}
+    key === nothing && throw(KeyError(nothing))
+    lock(wkh) do
+        return get(default, wkh.ht, key)
+    end
+end
+function pop!(wkh::WeakKeyDict{K}, key) where {K}
+    key === nothing && throw(KeyError(nothing))
+    lock(wkh) do
+        return pop!(wkh.ht, key)
+    end
+end
+function pop!(wkh::WeakKeyDict{K}, key, default) where {K}
+    key === nothing && return default
+    lock(wkh) do
+        return pop!(wkh.ht, key, default)
+    end
+end
+function delete!(wkh::WeakKeyDict, key)
+    key === nothing && return wkh
+    lock(wkh) do
+        delete!(wkh.ht, key)
+    end
+    return wkh
+end
+function empty!(wkh::WeakKeyDict)
+    lock(wkh) do
+        empty!(wkh.ht)
+    end
+    return wkh
+end
+function haskey(wkh::WeakKeyDict{K}, key) where {K}
+    key === nothing && return false
+    lock(wkh) do
+        return haskey(wkh.ht, key)
+    end
+end
+function getindex(wkh::WeakKeyDict{K}, key) where {K}
+    key === nothing && throw(KeyError(nothing))
+    lock(wkh) do
+        return getindex(wkh.ht, key)
+    end
+end
 isempty(wkh::WeakKeyDict) = isempty(wkh.ht)
 length(t::WeakKeyDict) = length(t.ht)
 
diff --git a/test/dict.jl b/test/dict.jl
index 0ae0b885684b5..7de84f382f142 100644
--- a/test/dict.jl
+++ b/test/dict.jl
@@ -900,13 +900,26 @@ Dict(1 => rand(2,3), 'c' => "asdf") # just make sure this does not trigger a dep
     @test !isempty(d26939)
     @test count(d26939) == 0
     empty!(d26939)
-    (@noinline d -> d[big(12345)] = 1)(d26939)
-    (@noinline d -> d[big(54321)] = 1)(d26939)
+    for i in 1:8
+        (@noinline (d, i) -> d[big(i + 12345)] = 1)(d26939, i)
+    end
     lock(GC.gc, d26939)
-    @test length(d26939.ht) == 2
-    @test length(d26939) == 2
+    @test length(d26939.ht) == 8
+    @test length(d26939) == 8
     @test !isempty(d26939)
     @test count(d26939) == 0
+    @test !haskey(d26939, nothing)
+    @test_throws KeyError(nothing) d26939[nothing]
+    @test_throws KeyError(nothing) get(d26939, nothing, 1)
+    @test_throws KeyError(nothing) get(() -> 1, d26939, nothing)
+    @test_throws KeyError(nothing) pop!(d26939, nothing)
+    @test getkey(d26939, nothing, 321) === 321
+    @test pop!(d26939, nothing, 321) === 321
+    @test delete!(d26939, nothing) === d26939
+    @test length(d26939.ht) == 8
+    @test_throws ArgumentError d26939[nothing] = 1
+    @test_throws ArgumentError get!(d26939, nothing, 1)
+    @test_throws ArgumentError get!(() -> 1, d26939, nothing)
     Base._cleanup_locked(d26939)
     @test length(d26939.ht) == 0