Skip to content

Commit

Permalink
fixup! fix numerous issues with WeakKeyDict
Browse files Browse the repository at this point in the history
  • Loading branch information
vtjnash committed Dec 6, 2020
1 parent 02a0c51 commit 68e5b47
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 14 deletions.
66 changes: 56 additions & 10 deletions base/weakkeydict.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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()
Expand All @@ -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)

Expand Down
21 changes: 17 additions & 4 deletions test/dict.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down

0 comments on commit 68e5b47

Please sign in to comment.