Indexing performance + lazy dense dicts

src/Utilities/cachingoptimizer.jl

@@ -184,28 +184,6 @@ function _standardize(d::AbstractDict{MOI.Index, MOI.Index})
     return map
 end
 function _standardize(d::IndexMap)
-    # return d
-    # if we return d as is, its not possible to add variables
-    # if there was a dense dict inside
-    # the solution would be to allow automatically swtiching
-    # a ClevelDenseDict...
-    return IndexMap(_standard_dict(d.varmap), d.conmap)
-end
-function _standard_dict(
-    d::D
-)::D where {D<:Dict{MOI.VariableIndex, MOI.VariableIndex}}
-    return d
-end
-function _standard_dict(
-    d::D
-)::Dict{MOI.VariableIndex, MOI.VariableIndex} where {
-    D<:AbstractDict{MOI.VariableIndex, MOI.VariableIndex}
-}
-    ret = Dict{MOI.VariableIndex, MOI.VariableIndex}()
-    sizehint!(ret, length(d))
-    for (k,v) in d
-        ret[k] = v
-    end
     return d
 end
 

src/Utilities/copy.jl

@@ -78,12 +78,14 @@ function dense_variable_dict(::Type{V}, n) where V
 end
 
 struct IndexMap <: AbstractDict{MOI.Index, MOI.Index}
+    # cannot remove the standard dict without breaking
     varmap::Union{DenseVariableDict{MOI.VariableIndex},
-                  Dict{MOI.VariableIndex, MOI.VariableIndex}}
+        Dict{MOI.VariableIndex, MOI.VariableIndex}}
     conmap::DoubleDicts.MainIndexDoubleDict
 end
-IndexMap() = IndexMap(Dict{MOI.VariableIndex, MOI.VariableIndex}(),
-                      DoubleDicts.IndexDoubleDict())
+IndexMap() = IndexMap(
+    Dict{MOI.VariableIndex, MOI.VariableIndex}(),
+    DoubleDicts.IndexDoubleDict())
 function IndexMap(n)
     IndexMap(dense_variable_dict(MOI.VariableIndex, n),
              DoubleDicts.IndexDoubleDict())

src/Utilities/dense_dict.jl

@@ -4,39 +4,112 @@
         inverse_hash::I
         set::BitSet
         map::Vector{V}
+        dict::Dict{K,V}
     end
 
-Same as `Dict{K, V}` but `hash(key)` is assumed to belong to `eachindex(map)`.
+Same as `Dict{K, V}` but very performant if `hash(key)` belongs to
+`eachindex(map)`. If `hash(key) == length(map) + 1` the dictionary
+grows continuosly in aperformant fashion. Otherwise, a regular `Dict{K, V}`
+is used.
 """
 struct DenseDict{K, V, F, I} <: AbstractDict{K, V}
     hash::F
     inverse_hash::I
     set::BitSet
     map::Vector{V}
-    function DenseDict{K, V}(hash, inverse_hash, n) where {K, V}
+    dict::Dict{K,V}
+    function DenseDict{K, V}(hash, inverse_hash, n = 0) where {K, V}
         set = BitSet()
         sizehint!(set, n)
-        return new{K, V, typeof(hash), typeof(inverse_hash)}(hash, inverse_hash, set, Vector{K}(undef, n))
+        return new{K, V, typeof(hash), typeof(inverse_hash)}(
+            hash, inverse_hash, set, Vector{K}(undef, n), Dict{K,V}()
+            )
     end
 end
 
+_is_dense(d::DenseDict) = !isempty(d.map)
+
 # Implementation of the `AbstractDict` API.
-# Base.empty(::DenseDict, ::Type{K}, ::Type{V}) not implemented
+function Base.empty!(d::DenseDict)
+    if _is_dense(d)
+        empty!(d.set)
+        empty!(d.map)
+    else
+        empty!(d.dict)
+    end
+end
 function Base.iterate(d::DenseDict{K,V}, args...) where {K,V}
-    itr = iterate(d.set, args...)
-    if itr === nothing
-        return nothing
+    if _is_dense(d)
+        itr = iterate(d.set, args...)
+        if itr === nothing
+            return nothing
+        else
+            el, i = itr
+            return d.inverse_hash(el)::K => d.map[el]::V, i
+        end
+    else
+        return Base.iterate(d.dict, args...)
+    end
+end
+function Base.length(d::DenseDict)
+    if _is_dense(d)
+        return length(d.set)
+    else
+        return length(d.dict)
+    end
+end
+function Base.haskey(d::DenseDict, key)
+    if _is_dense(d)
+        return d.hash(key) in d.set
     else
-        el, i = itr
-        return d.inverse_hash(el)::K => d.map[el]::V, i
+        return Base.haskey(d.dict, key)
+    end
+end
+function Base.getindex(d::DenseDict, key)
+    if _is_dense(d)
+        return d.map[d.hash(key)]
+    else
+        return d.dict[key]
+    end
+end
+function Base.setindex!(d::DenseDict, value, key)
+    h = d.hash(key)
+    if h <= length(d.map) && _is_dense(d)
+        push!(d.set, h)
+        d.map[h] = value
+    elseif h == length(d.map) + 1 && _is_dense(d)
+        push!(d.set, h)
+        push!(d.map, value)
+    else
+        if _is_dense(d)
+            _rehash(d)
+        end
+        d.dict[key] = value
+    end
+end
+function Base.sizehint!(d::DenseDict, n)
+    if _is_dense(d)
+        sizehint!(d.set, n)
+        sizehint!(d.map, n)
+    else
+        sizehint!(d.dict, n)
     end
 end
 
-Base.length(d::DenseDict) = length(d.set)
-Base.haskey(dict::DenseDict, key) = dict.hash(key) in dict.set
-Base.getindex(dict::DenseDict, key) = dict.map[dict.hash(key)]
-function Base.setindex!(dict::DenseDict, value, key)
-    h = dict.hash(key)
-    push!(dict.set, h)
-    dict.map[h] = value
+function _rehash(d::DenseDict{K}) where K
+    sizehint!(d.dict, length(d.set))
+    # assumes dict is currently dense
+    # iterator protocol from DenseDict is used
+    for (k,v) in d
+        d.dict[k] = v
+    end
+    empty!(d.set)
+    empty!(d.map)
+end
+
+function Base.delete!(d::DenseDict{K}, k::K) where K
+    if _is_dense(d)
+        _rehash(d)
+    end
+    delete!(d.dict, k)
 end

test/Utilities/dense_dict.jl

@@ -32,3 +32,28 @@ d[6] = 0.75
 @test d[2] == 1.5
 @test d[4] == 0.25
 @test d[6] == 0.75
+
+d[8] = 1.75
+
+sizehint!(d, 4)
+@test Base.length(d) == 4
+
+delete!(d, 6)
+@test !haskey(d, 6)
+@test d[2] == 1.5
+@test d[4] == 0.25
+@test d[8] == 1.75
+@test Base.length(d) == 3
+
+sizehint!(d, 4)
+d[24] = 2.5
+@test d[24] == 2.5
+
+@test sort(collect(d)) == sort([2 => 1.5, 4 => 0.25, 8 => 1.75, 24 => 2.5])
+
+empty!(d)
+@test length(d) == 0
+
+d = MOI.Utilities.DenseDict{Int, Float64}(div2, mul2, 3)
+empty!(d)
+@test length(d) == 0