Add setwith!

src/Dictionaries.jl

@@ -10,7 +10,7 @@ export getindices, setindices!
 export AbstractDictionary, AbstractIndices, IndexError, ArrayIndices, Indices, Dictionary, ArrayDictionary, MappedDictionary, ReverseIndices, ReverseDictionary, DictionaryView, FilteredDictionary, FilteredIndices, BroadcastedDictionary, FillDictionary, UnorderedIndices, UnorderedDictionary
 
 export dictionary, index, distinct, disjoint, isdictequal, filterview, sortkeys, sortpairs, sortkeys!, sortpairs!
-export issettable, isinsertable, set!, unset!
+export issettable, isinsertable, set!, setwith!, unset!
 export istokenizable, tokentype, tokens, tokenized, gettoken, gettokenvalue, istokenassigned, settokenvalue!, gettoken!, deletetoken!, sharetokens
 
 include("AbstractDictionary.jl")
@@ -47,7 +47,6 @@ end # module
 # * TODO: have `delete!` return next key, `deletetoken!` return next token.
 #   For these kinds of algorithms, probably need: firstindex, firsttoken, nextind, prevind,
 #   nexttoken, prevtoken, lastindex, lasttoken.
-# * A surface interface for updates like https://github.com/JuliaLang/julia/pull/31367
 # * More operations for "ordered" indices/sets (sort-based dictionaries and
 #   B-trees). We can probably formalize an interface around a trait here. Certain operations
 #   like slicing out an interval or performing a sort-merge co-iteration for `merge` become

src/insertion.jl

@@ -30,6 +30,15 @@ function.
 """
 isinsertable(::AbstractIndices) = false
 
+safe_convert(::Type{I}, i::I) where {I} = i
+function safe_convert(::Type{I}, i) where {I}
+    i2 = convert(I, i)
+    if !isequal(i, i2)
+        throw(ArgumentError("$i is not a valid key for type $I"))
+    end
+    return i2
+end
+
 ### Underlying token interface functions
 
 """
@@ -93,10 +102,7 @@ end
 Insert the new index `i` into `indices`. An error is thrown if `i` already exists.
 """
 @propagate_inbounds function Base.insert!(indices::AbstractIndices{I}, i) where {I}
-    i2 = convert(I, i)
-    if !isequal(i, i2)
-        throw(ArgumentError("$i is not a valid key for type $I"))
-    end
+    i2 = safe_convert(I, i)
     return insert!(indices, i2)
 end
 
@@ -118,10 +124,7 @@ Hint: Use `setindex!` to update an existing value, and `set!` to perform an "ups
 (update-or-insert) operation.
 """
 function Base.insert!(d::AbstractDictionary{I}, i, value) where {I}
-    i2 = convert(I, i)
-    if !isequal(i, i2)
-        throw(ArgumentError("$i is not a valid key for type $I"))
-    end
+    i2 = safe_convert(I, i)
     return insert!(d, i2, value)
 end
 
@@ -161,10 +164,7 @@ Hint: Use `setindex!` to exclusively update an existing value, and `insert!` to
 insert a new value. See also `get!`.
 """
 @propagate_inbounds function set!(d::AbstractDictionary{I}, i, value) where {I}
-    i2 = convert(I, i)
-    if !isequal(i, i2)
-        throw(ArgumentError("$i is not a valid key for type $I"))
-    end
+    i2 = safe_convert(I, i)
     return set!(d, i2, value)
 end
 
@@ -191,6 +191,22 @@ function set!(indices::AbstractIndices{I}, i1::I, i2::I) where {I}
     end
 end
 
+"""
+    setwith!(f, dict::AbstractDictionary, i, value)
+
+Update the value at `i` with the function `f` (`f(dict[i], value)`) or insert `value`.
+
+Hint: Use [`mergewith!`](@ref) to merge `Dictionary`s together.
+"""
+function setwith!(f, d::AbstractDictionary{I}, i, value) where {I}
+    (had_token, token) = gettoken!(d, i)
+    new_value = had_token ? f(@inbounds(gettokenvalue(d, token)), value) : value
+    @inbounds settokenvalue!(d, token, new_value)
+    return d
+end
+
+setwith!(f, ::AbstractIndices, i, value) = error("`setwith!` does not work with `AbstractIndices`")
+
 """
     set!(indices::AbstractIndices, i)
 
@@ -218,10 +234,7 @@ set to `default`, which is returned.
 See also `get`, `set!`.
 """
 @propagate_inbounds function Base.get!(d::AbstractDictionary{I}, i, default) where {I}
-    i2 = convert(I, i)
-    if !isequal(i, i2)
-        throw(ArgumentError("$i is not a valid key for type $I"))
-    end
+    i2 = safe_convert(I, i)
     return get!(d, i2, default)
 end
 
@@ -246,10 +259,7 @@ Return the value `dict[i]` if index `i` exists. Otherwise, a new index `i` is in
 set to the value `f()`, which is returned.
 """
 function Base.get!(f::Callable, d::AbstractDictionary{I}, i) where {I}
-    i2 = convert(I, i)
-    if !isequal(i, i2)
-        throw(ArgumentError("$i is not a valid key for type $I"))
-    end
+    i2 = safe_convert(I, i)
     get!(f, d, i2)
 end
 
@@ -278,10 +288,7 @@ Delete the index `i` from `dict`. An error is thrown if `i` does not exist.
 See also `unset!`, `insert!`.
 """
 @propagate_inbounds function Base.delete!(d::AbstractDictionary{I}, i) where {I}
-    i2 = convert(I, i)
-    if !isequal(i, i2)
-        throw(ArgumentError("$i is not a valid key for type $I"))
-    end
+    i2 = safe_convert(I, i)
     return delete!(d, i2)
 end
 
@@ -307,10 +314,7 @@ Delete the index `i` from `dict` if it exists, or do nothing otherwise.
 See also `delete!`, `set!`.
 """
 @propagate_inbounds function unset!(indices::AbstractDictionary{I}, i) where {I}
-    i2 = convert(I, i)
-    if !isequal(i, i2)
-        throw(ArgumentError("$i is not a valid key for type $I"))
-    end
+    i2 = safe_convert(I, i)
     return unset!(indices, i2)
 end
 

test/ArrayDictionary.jl

@@ -40,12 +40,19 @@
     @test !isempty(d)
     @test isequal(copy(d), d)
     @test_throws IndexError insert!(d, 10, 12)
-    @test d[10] == 11
-    set!(d, 10, 12)
+    @test d[10.0] == 11
+    set!(d, 10.0, 12.0)
     @test length(d) == 1
     @test d[10] == 12
-    d[10] = 13
+    setwith!(+, d, 10.0, 2.0)
+    @test length(d) == 1
+    @test d[10] == 14
+    setwith!(+, d, 2.0, 2.0)
+    @test length(d) == 2
+    @test d[2] == 2
+    delete!(d, 2)
     @test length(d) == 1
+    d[10.0] = 13.0
     @test d[10] == 13
     io = IOBuffer(); print(io, d); @test String(take!(io)) == "{10 = 13}"
     if VERSION < v"1.6-"

test/Dictionary.jl

@@ -94,8 +94,15 @@
     set!(d, 10.0, 12.0)
     @test length(d) == 1
     @test d[10] == 12
-    d[10.0] = 13.0
+    setwith!(+, d, 10.0, 2.0)
+    @test length(d) == 1
+    @test d[10] == 14
+    setwith!(+, d, 2.0, 2.0)
+    @test length(d) == 2
+    @test d[2] == 2
+    delete!(d, 2)
     @test length(d) == 1
+    d[10.0] = 13.0
     @test d[10] == 13
     io = IOBuffer(); print(io, d); @test String(take!(io)) == "{10 = 13}"
     if VERSION < v"1.6-"

test/Indices.jl

@@ -54,6 +54,7 @@
     @test length(set!(h, 2, 2)) == 2
     @test length(set!(h, 3.0, 3.0)) == 3
     @test_throws ErrorException set!(h, 4, 5)
+    @test_throws ErrorException setwith!(+, h, 3, 2)
 
     @testset "Comparison" begin
         i1 = Indices([1,2,3])

test/UnorderedDictionary.jl

@@ -72,8 +72,15 @@
     set!(d, 10.0, 12.0)
     @test length(d) == 1
     @test d[10] == 12
-    d[10.0] = 13.0
+    setwith!(+, d, 10.0, 2.0)
+    @test length(d) == 1
+    @test d[10] == 14
+    setwith!(+, d, 2.0, 2.0)
+    @test length(d) == 2
+    @test d[2] == 2
+    delete!(d, 2)
     @test length(d) == 1
+    d[10.0] = 13.0
     @test d[10] == 13
     io = IOBuffer(); print(io, d); @test String(take!(io)) == "{10 = 13}"
     if VERSION < v"1.6-"