Make the persistent dict API less strict. (#51336)

JuliaLang · Sep 15, 2023 · ee34a68 · ee34a68
1 parent 3099af9
commit ee34a68
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/base/dict.jl b/base/dict.jl
@@ -916,9 +916,13 @@ Base.PersistentDict{Symbol, Int64} with 1 entry:
 PersistentDict
 
 PersistentDict{K,V}() where {K,V} = PersistentDict(HAMT.HAMT{K,V}())
-PersistentDict(KV::Pair{K,V}) where {K,V} = PersistentDict(HAMT.HAMT(KV...))
+PersistentDict{K,V}(KV::Pair) where {K,V} = PersistentDict(HAMT.HAMT{K,V}(KV...))
+PersistentDict(KV::Pair{K,V}) where {K,V} = PersistentDict(HAMT.HAMT{K,V}(KV...))
 PersistentDict(dict::PersistentDict, pair::Pair) = PersistentDict(dict, pair...)
-function PersistentDict(dict::PersistentDict{K,V}, key::K, val::V) where {K,V}
+PersistentDict{K,V}(dict::PersistentDict{K,V}, pair::Pair) where {K,V} = PersistentDict(dict, pair...)
+function PersistentDict(dict::PersistentDict{K,V}, key, val) where {K,V}
+    key = convert(K, key)
+    val = convert(V, val)
     trie = dict.trie
     h = hash(key)
     found, present, trie, i, bi, top, hs = HAMT.path(trie, key, h, #=persistent=# true)
@@ -975,7 +979,7 @@ function getindex(dict::PersistentDict{K,V}, key::K) where {K,V}
     throw(KeyError(key))
 end
 
-function get(dict::PersistentDict{K,V}, key::K, default::V) where {K,V}
+function get(dict::PersistentDict{K,V}, key::K, default) where {K,V}
     trie = dict.trie
     if HAMT.islevel_empty(trie)
         return default

diff --git a/base/hamt.jl b/base/hamt.jl
@@ -64,14 +64,16 @@ mutable struct HAMT{K, V}
     bitmap::BITMAP
 end
 HAMT{K, V}() where {K, V} = HAMT(Vector{Union{Leaf{K, V}, HAMT{K, V}}}(undef, 0), zero(BITMAP))
-function HAMT(k::K, v::V) where {K, V}
+function HAMT{K,V}(k::K, v) where {K, V}
+    v = convert(V, v)
     # For a single element we can't have a hash-collision
     trie = HAMT(Vector{Union{Leaf{K, V}, HAMT{K, V}}}(undef, 1), zero(BITMAP))
     trie.data[1] = Leaf{K,V}(k,v)
     bi = BitmapIndex(HashState(k))
     set!(trie, bi)
     return trie
 end
+HAMT(k::K, v::V) where {K, V} = HAMT{K,V}(K, V)
 
 struct HashState{K}
     key::K

diff --git a/test/dict.jl b/test/dict.jl
@@ -1101,10 +1101,13 @@ import Base.PersistentDict
         @test length(dict) == 0
         @test isempty(dict)
 
+        dict = PersistentDict{Int, Int}(1=>2.0)
+        @test dict[1] == 2
+
         dict = PersistentDict(1=>2)
         @test dict[1] == 2
 
-        dict = PersistentDict(dict, 1=>3)
+        dict = PersistentDict(dict, 1=>3.0)
         @test dict[1] == 3
 
         dict = PersistentDict(dict, 1, 1)
@@ -1119,7 +1122,7 @@ import Base.PersistentDict
         @test haskey(dict, 1)
         @test !haskey(dict, 2)
 
-        dict2 = PersistentDict(dict, 1, 2)
+        dict2 = PersistentDict{Int, Int}(dict, 1=>2)
         @test dict[1] == 1
         @test dict2[1] == 2