Merge pull request #5519 from JuliaLang/kms/get!

RFC: Add `get!` for Dicts (addresses #1529)

base/dict.jl

@@ -419,13 +419,34 @@ function empty!{K,V}(h::Dict{K,V})
     return h
 end
 
-function setindex!{K,V}(h::Dict{K,V}, v0, key0)
-    key = convert(K,key0)
-    if !isequal(key,key0)
-        error(key0, " is not a valid key for type ", K)
+# get the index where a key is stored, or -1 if not present
+function ht_keyindex{K,V}(h::Dict{K,V}, key)
+    sz = length(h.keys)
+    iter = 0
+    maxprobe = max(16, sz>>6)
+    index = hashindex(key, sz)
+    keys = h.keys
+
+    while true
+        if isslotempty(h,index)
+            break
+        end
+        if !isslotmissing(h,index) && isequal(key,keys[index])
+            return index
+        end
+
+        index = (index & (sz-1)) + 1
+        iter+=1
+        iter > maxprobe && break
     end
-    v   = convert(V,  v0)
 
+    return -1
+end
+
+# get the index where a key is stored, or -pos if not present
+# and the key would be inserted at pos
+# This version is for use by setindex! and get!
+function ht_keyindex2{K,V}(h::Dict{K,V}, key)
     sz = length(h.keys)
 
     if h.ndel >= ((3*sz)>>2) || h.count*3 > sz*2
@@ -437,71 +458,90 @@ function setindex!{K,V}(h::Dict{K,V}, v0, key0)
     iter = 0
     maxprobe = max(16, sz>>6)
     index = hashindex(key, sz)
-    avail = -1  # an available slot
-    keys = h.keys; vals = h.vals
+    avail = 0
+    keys = h.keys
 
     while true
         if isslotempty(h,index)
-            if avail > 0; index = avail; end
-            h.slots[index] = 0x1
-            h.keys[index] = key
-            h.vals[index] = v
-            h.count += 1
-            return h
+            avail < 0 && return avail
+            return -index
         end
 
         if isslotmissing(h,index)
-            if avail<0
+            if avail == 0
                 # found an available slot, but need to keep scanning
                 # in case "key" already exists in a later collided slot.
-                avail = index
+                avail = -index
             end
         elseif isequal(key, keys[index])
-            vals[index] = v
-            return h
+            return index
         end
 
         index = (index & (sz-1)) + 1
         iter+=1
         iter > maxprobe && break
     end
 
-    if avail>0
-        index = avail
-        h.slots[index] = 0x1
-        h.keys[index] = key
-        h.vals[index] = v
-        h.count += 1
-        return h
-    end
+    avail < 0 && return avail
 
     rehash(h, h.count > 64000 ? sz*2 : sz*4)
 
-    setindex!(h, v, key)
+    return ht_keyindex2(h, key)
 end
 
-# get the index where a key is stored, or -1 if not present
-function ht_keyindex{K,V}(h::Dict{K,V}, key)
-    sz = length(h.keys)
-    iter = 0
-    maxprobe = max(16, sz>>6)
-    index = hashindex(key, sz)
-    keys = h.keys
+function _setindex!(h::Dict, v, key, index)
+    h.slots[index] = 0x1
+    h.keys[index] = key
+    h.vals[index] = v
+    h.count += 1
+end
 
-    while true
-        if isslotempty(h,index)
-            break
-        end
-        if !isslotmissing(h,index) && isequal(key,keys[index])
-            return index
-        end
+function setindex!{K,V}(h::Dict{K,V}, v0, key0)
+    key = convert(K,key0)
+    if !isequal(key,key0)
+        error(key0, " is not a valid key for type ", K)
+    end
+    v = convert(V,  v0)
 
-        index = (index & (sz-1)) + 1
-        iter+=1
-        iter > maxprobe && break
+    index = ht_keyindex2(h, key)
+
+    if index > 0
+        h.vals[index] = v
+    else
+        _setindex!(h, v, key, -index)
     end
 
-    return -1
+    return h
+end
+
+function get!{K,V}(h::Dict{K,V}, key0, default)
+    key = convert(K,key0)
+    if !isequal(key,key0)
+        error(key0, " is not a valid key for type ", K)
+    end
+
+    index = ht_keyindex2(h, key)
+
+    index > 0 && return h.vals[index]
+
+    v = convert(V,  default)
+    _setindex!(h, v, key, -index)
+    return v
+end
+
+function get!{K,V}(default::Function, h::Dict{K,V}, key0)
+    key = convert(K,key0)
+    if !isequal(key,key0)
+        error(key0, " is not a valid key for type ", K)
+    end
+
+    index = ht_keyindex2(h, key)
+
+    index > 0 && return h.vals[index]
+
+    v = convert(V,  default())
+    _setindex!(h, v, key, -index)
+    return v
 end
 
 function getindex{K,V}(h::Dict{K,V}, key)
@@ -514,6 +554,11 @@ function get{K,V}(h::Dict{K,V}, key, deflt)
     return (index<0) ? deflt : h.vals[index]::V
 end
 
+function get{K,V}(deflt::Function, h::Dict{K,V}, key)
+    index = ht_keyindex(h, key)
+    return (index<0) ? deflt() : h.vals[index]::V
+end
+
 haskey(h::Dict, key) = (ht_keyindex(h, key) >= 0)
 in{T<:Dict}(key, v::KeyIterator{T}) = (ht_keyindex(v.dict, key) >= 0)
 
@@ -625,6 +670,9 @@ function getkey{K}(wkh::WeakKeyDict{K}, kk, deflt)
 end
 
 get{K}(wkh::WeakKeyDict{K}, key, def) = get(wkh.ht, key, def)
+get{K}(def::Function, wkh::WeakKeyDict{K}, key) = get(def, wkh.ht, key)
+get!{K}(wkh::WeakKeyDict{K}, key, def) = get!(wkh.ht, key, def)
+get!{K}(def::Function, wkh::WeakKeyDict{K}, key) = get!(def, wkh.ht, key)
 pop!{K}(wkh::WeakKeyDict{K}, key) = pop!(wkh.ht, key)
 pop!{K}(wkh::WeakKeyDict{K}, key, def) = pop!(wkh.ht, key, def)
 delete!{K}(wkh::WeakKeyDict{K}, key) = delete!(wkh.ht, key)

doc/stdlib/base.rst

@@ -727,6 +727,32 @@ Given a dictionary ``D``, the syntax ``D[x]`` returns the value of key ``x`` (if
 
    Return the value stored for the given key, or the given default value if no mapping for the key is present.
 
+.. function:: get(f::Function, collection, key)
+
+   Return the value stored for the given key, or if no mapping for the key is present, return ``f()``.  Use ``get!`` to also store the default value in the dictionary.
+
+   This is intended to be called using ``do`` block syntax::
+
+     get(dict, key) do
+         # default value calculated here
+	 time()
+     end
+
+.. function:: get!(collection, key, default)
+
+   Return the value stored for the given key, or if no mapping for the key is present, store ``key => default``, and return ``default``.
+
+.. function:: get!(f::Function, collection, key)
+
+   Return the value stored for the given key, or if no mapping for the key is present, store ``key => f()``, and return ``f()``.
+
+   This is intended to be called using ``do`` block syntax::
+
+     get!(dict, key) do
+         # default value calculated here
+	 time()
+     end
+
 .. function:: getkey(collection, key, default)
 
    Return the key matching argument ``key`` if one exists in ``collection``, otherwise return ``default``.

test/collections.jl

@@ -147,6 +147,30 @@ d4[1001] = randstring(3)
 @test !isequal({1 => 2}, {"dog" => "bone"})
 @test isequal(Dict{Int, Int}(), Dict{String, String}())
 
+# get! (get with default values assigned to the given location)
+
+let
+    f(x) = x^2
+    d = {8=>19}
+    def = {}
+    @test get!(d, 8, 5) == 19
+    @test get!(d, 19, 2) == 2
+
+    @test get!(d, 42) do  # d is updated with f(2)
+        f(2)
+    end == 4
+
+    @test get!(d, 42) do  # d is not updated
+        f(200)
+    end == 4
+
+    @test get(d, 13) do   # d is not updated
+        f(4)
+    end == 16
+
+    @test d == {8=>19, 19=>2, 42=>4}
+end
+
 # issue #2540
 d = {x => 1
     for x in ['a', 'b', 'c']}