WeakKeyIdDict implementation (keeping WeakKeyDict)

base/abstractdict.jl

@@ -453,7 +453,7 @@ end
 
 function isequal(l::AbstractDict, r::AbstractDict)
  l === r && return true
- if isa(l,IdDict) != isa(r,IdDict)
+ if isa(l,IdDict) != isa(r,IdDict) || isa(l,WeakKeyDict) != isa(r,WeakKeyDict) || isa(l,WeakKeyIdDict) != isa(r,WeakKeyIdDict)
  return false
  end
  if length(l) != length(r) return false end
@@ -467,7 +467,7 @@ end
 
 function ==(l::AbstractDict, r::AbstractDict)
  l === r && return true
- if isa(l,IdDict) != isa(r,IdDict)
+ if isa(l,IdDict) != isa(r,IdDict) || isa(l,WeakKeyDict) != isa(r,WeakKeyDict) || isa(l,WeakKeyIdDict) != isa(r,WeakKeyIdDict)
  return false
  end
  length(l) != length(r) && return false

base/exports.jl

@@ -93,6 +93,7 @@ export
  Vector,
  VersionNumber,
  WeakKeyDict,
+ WeakKeyIdDict,
 
 # Ccall types
  Cchar,

base/weakkeydict.jl

@@ -2,23 +2,79 @@
 
 # weak key dictionaries
 
+# Type to wrap a WeakRef to furbish it with object-id comparison and hashing.
+struct WeakRefForWeakIdDict
+ w::WeakRef
+ WeakRefForWeakIdDict(wr::WeakRef) = new(wr)
+end
+WeakRefForWeakIdDict(val) = WeakRefForWeakIdDict(WeakRef(val))
+==(wr1::WeakRefForWeakIdDict, wr2::WeakRefForWeakIdDict) = wr1.w.value===wr2.w.value
+hash(wr::WeakRefForWeakIdDict, h::UInt) = hash_uint(3h - objectid(wr.w.value))
+
+# Type to wrap a WeakRef to furbish it with == comparison and "normal" hashing of its value
+struct WeakRefForWeakDict
+ w::WeakRef
+ WeakRefForWeakDict(wr::WeakRef) = new(wr)
+end
+WeakRefForWeakDict(val) = WeakRefForWeakDict(WeakRef(val))
+==(wr1::WeakRefForWeakDict, wr2::WeakRefForWeakDict) = wr1.w.value==wr2.w.value
+hash(wr::WeakRefForWeakDict, h::UInt) = hash(wr.w.value, h)
+
+
+abstract type AbstractWeakKeyDict{K,V} <: AbstractDict{K,V} end
+"""
+ WeakKeyIdDict([itr])
+
+`WeakKeyIdDict()` constructs a hash table where the keys are weak
+references to objects, and thus may be garbage collected even when
+referenced in a hash table.
+
+The hashing and comparison are based on object-id and === of the key.
+
+See [`Dict`](@ref) for further help.
+"""
+mutable struct WeakKeyIdDict{K,V} <: AbstractWeakKeyDict{K,V}
+ ht::Dict{WeakRefForWeakIdDict,V}
+ lock::Threads.RecursiveSpinLock
+ finalizer::Function
+
+ # Constructors mirror Dict's
+ function WeakKeyIdDict{K,V}() where V where K
+ t = new(Dict{WeakRefForWeakIdDict,V}(), Threads.RecursiveSpinLock(), identity)
+ t.finalizer = function (k)
+ # when a weak key is finalized, remove from dictionary if it is still there
+ if islocked(t)
+ finalizer(t.finalizer, k)
+ return nothing
+ end
+ delete!(t, k)
+ end
+ return t
+ end
+end
+keytype_ht(::Type{<:WeakKeyIdDict}) = WeakRefForWeakIdDict
+striptp(::Type{<:WeakKeyIdDict}) = WeakKeyIdDict
+
 """
  WeakKeyDict([itr])
 
 `WeakKeyDict()` constructs a hash table where the keys are weak
 references to objects, and thus may be garbage collected even when
 referenced in a hash table.
 
+The hashing and comparison are based on the normal hash and == of the
+key.
+
 See [`Dict`](@ref) for further help.
 """
-mutable struct WeakKeyDict{K,V} <: AbstractDict{K,V}
- ht::Dict{WeakRef,V}
+mutable struct WeakKeyDict{K,V} <: AbstractWeakKeyDict{K,V}
+ ht::Dict{WeakRefForWeakDict,V}
  lock::Threads.RecursiveSpinLock
  finalizer::Function
 
  # Constructors mirror Dict's
  function WeakKeyDict{K,V}() where V where K
- t = new(Dict{Any,V}(), Threads.RecursiveSpinLock(), identity)
+ t = new(Dict{WeakRefForWeakDict,V}(), Threads.RecursiveSpinLock(), identity)
  t.finalizer = function (k)
  # when a weak key is finalized, remove from dictionary if it is still there
  if islocked(t)
@@ -30,81 +86,95 @@ mutable struct WeakKeyDict{K,V} <: AbstractDict{K,V}
  return t
  end
 end
-function WeakKeyDict{K,V}(kv) where V where K
- h = WeakKeyDict{K,V}()
+keytype_ht(::Type{<:WeakKeyDict}) = WeakRefForWeakDict
+striptp(::Type{<:WeakKeyDict}) = WeakKeyDict
+
+# Constructors as for Dict
+function (::Type{W})(kv) where W<:AbstractWeakKeyDict{K,V} where {K, V}
+ h = W()
  for (k,v) in kv
  h[k] = v
  end
  return h
 end
-WeakKeyDict{K,V}(p::Pair) where V where K = setindex!(WeakKeyDict{K,V}(), p.second, p.first)
-function WeakKeyDict{K,V}(ps::Pair...) where V where K
- h = WeakKeyDict{K,V}()
+(::Type{W})(p::Pair) where W<:AbstractWeakKeyDict{K,V} where {K, V} = setindex!(W(), p.second, p.first)
+function (::Type{W})(ps::Pair...) where W<:AbstractWeakKeyDict{K,V} where {K, V}
+ h = W()
  sizehint!(h, length(ps))
  for p in ps
  h[p.first] = p.second
  end
  return h
 end
-WeakKeyDict() = WeakKeyDict{Any,Any}()
+(::Type{W})() where W<:AbstractWeakKeyDict = W{Any,Any}()
 
-WeakKeyDict(kv::Tuple{}) = WeakKeyDict()
-copy(d::WeakKeyDict) = WeakKeyDict(d)
+(::Type{W})(kv::Tuple{}) where W<:AbstractWeakKeyDict = W()
+copy(d::W) where W<:AbstractWeakKeyDict = W(d)
 
-WeakKeyDict(ps::Pair{K,V}...) where {K,V} = WeakKeyDict{K,V}(ps)
-WeakKeyDict(ps::Pair{K}...) where {K} = WeakKeyDict{K,Any}(ps)
-WeakKeyDict(ps::(Pair{K,V} where K)...) where {V} = WeakKeyDict{Any,V}(ps)
-WeakKeyDict(ps::Pair...) = WeakKeyDict{Any,Any}(ps)
+(::Type{W})(ps::Pair{K,V}...) where W<:AbstractWeakKeyDict where {K,V} = W{K,V}(ps)
+(::Type{W})(ps::Pair{K}...) where W<:AbstractWeakKeyDict where {K} = W{K,Any}(ps)
+(::Type{W})(ps::(Pair{K,V} where K)...) where W<:AbstractWeakKeyDict where {V} = W{Any,V}(ps)
+(::Type{W})(ps::Pair...) where W<:AbstractWeakKeyDict = W{Any,Any}(ps)
+(::Type{W})(ps::Pair) where W<:AbstractWeakKeyDict = W{Any,Any}(ps)
 
-function WeakKeyDict(kv)
+function (::Type{W})(kv) where W<:AbstractWeakKeyDict
  try
- Base.dict_with_eltype((K, V) -> WeakKeyDict{K, V}, kv, eltype(kv))
+ Base.dict_with_eltype((K, V) -> W{K, V}, kv, eltype(kv))
  catch e
  if !isiterable(typeof(kv)) || !all(x->isa(x,Union{Tuple,Pair}),kv)
- throw(ArgumentError("WeakKeyDict(kv): kv needs to be an iterator of tuples or pairs"))
+ throw(ArgumentError("$W(kv): kv needs to be an iterator of tuples or pairs"))
  else
  rethrow(e)
  end
  end
 end
 
-empty(d::WeakKeyDict, ::Type{K}, ::Type{V}) where {K, V} = WeakKeyDict{K, V}()
+empty(d::W, ::Type{K}, ::Type{V}) where W<:AbstractWeakKeyDict{K,V} where {K, V} = striptp(W){K, V}()
 
-islocked(wkh::WeakKeyDict) = islocked(wkh.lock)
-lock(f, wkh::WeakKeyDict) = lock(f, wkh.lock)
-trylock(f, wkh::WeakKeyDict) = trylock(f, wkh.lock)
+islocked(wkh::AbstractWeakKeyDict) = islocked(wkh.lock)
+lock(f, wkh::AbstractWeakKeyDict) = lock(f, wkh.lock)
+trylock(f, wkh::AbstractWeakKeyDict) = trylock(f, wkh.lock)
 
-function setindex!(wkh::WeakKeyDict{K}, v, key) where K
- k = convert(K, key)
- finalizer(wkh.finalizer, k)
+function setindex!(wkh::W, v, key) where W<:AbstractWeakKeyDict{K} where K
+ !isa(key, K) && throw(ArgumentError("$key is not a valid key for type $K"))
+ finalizer(wkh.finalizer, key)
  lock(wkh) do
- wkh.ht[WeakRef(k)] = v
+ wkh.ht[keytype_ht(W)(key)] = v
  end
  return wkh
 end
 
-function getkey(wkh::WeakKeyDict{K}, kk, default) where K
+function getkey(wkh::W, kk, default) where W<:AbstractWeakKeyDict{K} where K
  return lock(wkh) do
- k = getkey(wkh.ht, kk, secret_table_token)
+ k = getkey(wkh.ht, keytype_ht(W)(kk), secret_table_token)
  k === secret_table_token && return default
- return k.value::K
+ return k.w.value::K
  end
 end
 
-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)
-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)
-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)
-isempty(wkh::WeakKeyDict) = isempty(wkh.ht)
-length(t::WeakKeyDict) = length(t.ht)
-
-function iterate(t::WeakKeyDict{K,V}) where V where K
+get(wkh::W, key, default) where W<:AbstractWeakKeyDict{K} where {K} =
+ lock(() -> get(wkh.ht, keytype_ht(W)(key), default), wkh)
+get(default::Callable, wkh::W, key) where W<:AbstractWeakKeyDict{K} where {K} =
+ lock(() -> get(default, wkh.ht, keytype_ht(W)(key)), wkh)
+get!(wkh::W, key, default) where W<:AbstractWeakKeyDict{K} where {K} =
+ lock(() -> get!(wkh.ht, keytype_ht(W)(key), default), wkh)
+get!(default::Callable, wkh::W, key) where W<:AbstractWeakKeyDict{K} where {K} =
+ lock(() -> get!(default, wkh.ht, keytype_ht(W)(key)), wkh)
+pop!(wkh::W, key) where W<:AbstractWeakKeyDict{K} where {K} =
+ lock(() -> pop!(wkh.ht, keytype_ht(W)(key)), wkh)
+pop!(wkh::W, key, default) where W<:AbstractWeakKeyDict{K} where {K} =
+ lock(() -> pop!(wkh.ht, keytype_ht(W)(key), default), wkh)
+delete!(wkh::W, key) where W<:AbstractWeakKeyDict =
+ lock(() -> delete!(wkh.ht, keytype_ht(W)(key)), wkh)
+empty!(wkh::AbstractWeakKeyDict) = (lock(() -> empty!(wkh.ht), wkh); wkh)
+haskey(wkh::W, key) where W<:AbstractWeakKeyDict{K} where {K} =
+ lock(() -> haskey(wkh.ht, keytype_ht(W)(key)), wkh)
+getindex(wkh::W, key) where W<:AbstractWeakKeyDict{K} where {K} =
+ lock(() -> getindex(wkh.ht, keytype_ht(W)(key)), wkh)
+isempty(wkh::AbstractWeakKeyDict) = isempty(wkh.ht)
+length(t::AbstractWeakKeyDict) = length(t.ht)
+
+function iterate(t::AbstractWeakKeyDict)
  gc_token = Ref{Bool}(false) # no keys will be deleted via finalizers until this token is gc'd
  finalizer(gc_token) do r
  if r[]
@@ -115,13 +185,13 @@ function iterate(t::WeakKeyDict{K,V}) where V where K
  s = lock(t.lock)
  iterate(t, (gc_token,))
 end
-function iterate(t::WeakKeyDict{K,V}, state) where V where K
+function iterate(t::W, state) where W<:AbstractWeakKeyDict{K,V} where {K,V}
  gc_token = first(state)
  y = iterate(t.ht, tail(state)...)
  y === nothing && return nothing
  wkv, i = y
- kv = Pair{K,V}(wkv[1].value::K, wkv[2])
+ kv = Pair{K,V}(wkv[1].w.value::K, wkv[2])
  return (kv, (gc_token, i))
 end
 
-filter!(f, d::WeakKeyDict) = filter_in_one_pass!(f, d)
+filter!(f, d::AbstractWeakKeyDict) = filter_in_one_pass!(f, d)

doc/src/base/collections.md

@@ -43,6 +43,7 @@ Fully implemented by:
  * [`IdDict`](@ref)
  * [`Dict`](@ref)
  * [`WeakKeyDict`](@ref)
+ * [`WeakKeyIdDict`](@ref)
  * `EachLine`
  * `AbstractString`
  * [`Set`](@ref)
@@ -79,6 +80,7 @@ Fully implemented by:
  * [`IdDict`](@ref)
  * [`Dict`](@ref)
  * [`WeakKeyDict`](@ref)
+ * [`WeakKeyIdDict`](@ref)
  * `AbstractString`
  * [`Set`](@ref)
  * [`NamedTuple`](@ref)