Add CleverDicts as submodule to MOI.Utilities

Project.toml

@@ -5,9 +5,10 @@ version = "0.9.0"
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 
 [compat]
 julia = "1"

src/Utilities/CleverDicts.jl

@@ -0,0 +1,198 @@
+module CleverDicts
+
+import OrderedCollections
+
+"""
+    CleverDict{K, V}
+
+A smart storage type for managing sequential objects with non-decreasing integer
+indices.
+
+Provided no keys are deleted, the backing storage is a `Vector{V}`. Once a key
+has been deleted, the backing storage switches to an `OrderedDict{K, V}`.
+
+The i'th ordered element can be obtained with `c[LinearIndex(i)]`.
+
+Note that querying a `LinearIndex` immediately after deleting a key via
+`delete!` is very slow. (It requires a rebuild of an ordered list of variables.)
+
+Store an item `val` using `new_item(c::CleverDict, val)`. `new_item` returns a
+key corresponding to the stored item.
+
+Overload the functions `index_to_key` and `key_to_index` to enable mappings
+between the integer index of the vector and the dictionary key.
+
+## Example
+
+    struct MyKey
+        x::Int
+    end
+    index_to_key(::Type{MyKey}, i::Int) = MyKey(i)
+    key_to_index(key::MyKey) = key.x
+"""
+mutable struct CleverDict{K, V}
+    last_index::Int
+    vector::Union{Nothing, Vector{V}}
+    dict::Union{Nothing, OrderedCollections.OrderedDict{K, V}}
+    CleverDict{K, V}() where {K, V} = new{K, V}(0, V[], nothing)
+end
+
+"""
+    index_to_key(::Type{K}, index::Int)
+
+Create a new key associated with the integer value `index`.
+"""
+function index_to_key end
+
+"""
+    key_to_index(key::K)
+
+Map `key` to an integer valued index, assuming that there have been no
+deletions.
+"""
+function key_to_index end
+
+"""
+    new_item(c::CleverDict{K, V}, val::Val)::K where {K, V}
+
+Set `val` in the next available key, and return that key.
+"""
+function new_item(c::CleverDict{K, V}, val::V)::K where {K, V}
+    c.last_index += 1
+    key = index_to_key(K, c.last_index)
+    if c.dict === nothing
+        push!(c.vector, val)
+    else
+        c.dict[key] = val
+        # If there is a vector (e.g., because it has been rebuild for
+        # `LinearIndex`), clear it.
+        c.vector = nothing
+    end
+    return key
+end
+
+function Base.empty!(c::CleverDict{K, V})::Nothing where {K, V}
+    c.vector = V[]
+    c.last_index = 0
+    c.dict = nothing
+    return
+end
+
+function Base.getindex(c::CleverDict{K, V}, key::K)::V where {K, V}
+    # Perform this `haskey` check up front to detect getting with keys that are
+    # invalid (i.e., have previously been deleted).
+    if !haskey(c, key)
+        throw(KeyError(key))
+    end
+    # Case  I) no call to `Base.delete!`, so return the element:
+    # Case II) `Base.delete!` must have been called, so return the element
+    #          from the dictionary.
+    return c.dict === nothing ? c.vector[key_to_index(key)] : c.dict[key]
+end
+
+function Base.setindex!(c::CleverDict{K, V}, val::V, key::K)::V where {K, V}
+    # Perform this `haskey` check up front to detect setting with keys that are
+    # invalid (i.e., have already been deleted). You can only call setindex!
+    # with a key obtained from `new_key` that hasn't been deleted.
+    if !haskey(c, key)
+        throw(KeyError(key))
+    elseif c.dict === nothing
+        @assert c.vector !== nothing
+        c.vector[key_to_index(key)] = val
+    else
+        c.dict[key] = val
+    end
+    return val
+end
+
+struct LinearIndex
+    i::Int
+end
+
+function Base.getindex(c::CleverDict{K, V}, index::LinearIndex)::V where {K, V}
+    if !(1 <= index.i <= length(c))
+        throw(KeyError(index))
+    end
+    # Get the `index` linear element. If `c.vector` is currently `nothing`
+    # (i.e., there has been a deletion), rebuild `c.vector`. This is a
+    # trade-off: We could ensure `c.vector` is always updated, but this requires
+    # a `splice!` in `delete!`, making deletions costly. However, it makes this
+    # `getindex` operation trival because we would never have to rebuild the
+    # vector.
+    # The current implemented approach offers quick deletions, but an expensive
+    # rebuild the first time you query a `LinearIndex` after a deletion or a new
+    # key is added. Once the rebuild is done, there are quick queries until the
+    # next deletion or addition. Thus, the worst-case is a user repeatedly
+    # deleting a key and then querying a LinearIndex (e.g., getting the MOI
+    # objective function).
+    if c.vector === nothing
+        c.vector = Vector{V}(undef, length(c))
+        for (i, val) in enumerate(values(c.dict))
+            c.vector[i] = val
+        end
+    end
+    return c.vector[index.i]::V
+end
+
+function Base.delete!(c::CleverDict{K, V}, key::K)::Nothing where {K, V}
+    if c.dict === nothing
+        c.dict = OrderedCollections.OrderedDict{K, Union{Nothing, V}}()
+        for (i, info) in enumerate(c.vector)
+            c.dict[index_to_key(K, i)] = info
+        end
+    end
+    delete!(c.dict, key)
+    c.vector = nothing
+    return
+end
+
+function Base.length(c::CleverDict)::Int
+    return c.dict == nothing ? length(c.vector) : length(c.dict)
+end
+
+Base.haskey(::CleverDict, key) = false
+function Base.haskey(c::CleverDict{K, V}, key::K)::Bool where {K, V}
+    if c.dict === nothing
+        return 1 <= key_to_index(key) <= length(c.vector)
+    else
+        return haskey(c.dict, key)
+    end
+end
+
+function Base.iterate(
+    c::CleverDict{K, V}
+)::Union{Nothing, Tuple{Pair{K, V}, Int}} where {K, V}
+    if length(c) == 0
+        return nothing
+    elseif c.dict !== nothing
+        return iterate(c.dict)
+    else
+        @assert c.vector !== nothing
+        key = index_to_key(K, 1)
+        return key => c.vector[1], 2
+    end
+end
+
+function Base.iterate(
+    c::CleverDict{K, V}, s::Int
+)::Union{Nothing, Tuple{Pair{K, V}, Int}} where {K, V}
+    if s > length(c)
+        return nothing
+    elseif c.dict !== nothing
+        return iterate(c.dict, s)
+    else
+        @assert c.vector !== nothing
+        key = index_to_key(K, s)
+        return key => c.vector[s], s + 1
+    end
+end
+
+function Base.values(c::CleverDict{K, V}) where {K, V}
+    return c.dict !== nothing ? values(c.dict) : c.vector
+end
+
+function Base.keys(c::CleverDict{K, V}) where {K, V}
+    return c.dict !== nothing ? keys(c.dict) : index_to_key.(K, 1:length(c))
+end
+
+end

src/Utilities/Utilities.jl

@@ -29,4 +29,6 @@ include("mockoptimizer.jl")
 include("cachingoptimizer.jl")
 include("universalfallback.jl")
 
+include("CleverDicts.jl")
+
 end # module

test/Utilities/CleverDicts.jl

@@ -0,0 +1,163 @@
+using MathOptInterface, Test
+
+const CleverDicts = MathOptInterface.Utilities.CleverDicts
+
+struct MyKey
+    x::Int
+end
+
+CleverDicts.key_to_index(key::MyKey) = key.x
+CleverDicts.index_to_key(::Type{MyKey}, index::Int) = MyKey(index)
+
+@testset "CleverDict" begin
+    @testset "get/set" begin
+        d = CleverDicts.CleverDict{MyKey, String}()
+        key = CleverDicts.new_item(d, "first")
+        @test key == MyKey(1)
+        @test d[key] == "first"
+        @test haskey(d, key) == true
+        @test_throws KeyError d[MyKey(2)]
+        delete!(d, key)
+        @test_throws KeyError d[key]
+        @test_throws KeyError d[key] = "key"
+        @test haskey(d, key) == false
+        key2 = CleverDicts.new_item(d, "second")
+        @test key2 == MyKey(2)
+        @test d[key2] == "second"
+        @test d.vector === nothing
+        @test d.dict !== nothing
+        d[key2] = "third"
+        @test d[key2] == "third"
+
+        empty!(d)
+
+        key = CleverDicts.new_item(d, "first")
+        @test key == MyKey(1)
+        @test d[key] == "first"
+        d[key] = "zeroth"
+        @test d[key] == "zeroth"
+        @test haskey(d, key) == true
+        @test_throws KeyError d[MyKey(2)]
+        delete!(d, key)
+        @test_throws KeyError d[key]
+        @test_throws KeyError d[key] = "key"
+        @test haskey(d, key) == false
+        key2 = CleverDicts.new_item(d, "second")
+        @test key2 == MyKey(2)
+        @test d[key2] == "second"
+        @test d.vector === nothing
+        @test d.dict !== nothing
+    end
+
+    @testset "LinearIndex" begin
+        d = CleverDicts.CleverDict{MyKey, String}()
+        key = CleverDicts.new_item(d, "first")
+        @test d[CleverDicts.LinearIndex(1)] == "first"
+        key2 = CleverDicts.new_item(d, "second")
+        @test d[CleverDicts.LinearIndex(2)] == "second"
+        @test length(d) == 2
+        delete!(d, key)
+        @test d.vector === nothing
+        @test d[CleverDicts.LinearIndex(1)] == "second"
+        @test_throws KeyError d[CleverDicts.LinearIndex(2)]
+        @test length(d) == 1
+        @test d.vector !== nothing
+    end
+
+    @testset "keys/values" begin
+        d = CleverDicts.CleverDict{MyKey, String}()
+        key = CleverDicts.new_item(d, "first")
+        key2 = CleverDicts.new_item(d, "second")
+        @test collect(keys(d)) == [MyKey(1), MyKey(2)]
+        @test collect(values(d)) == ["first", "second"]
+        delete!(d, key)
+        key3 = CleverDicts.new_item(d, "third")
+        @test collect(keys(d)) == [MyKey(2), MyKey(3)]
+        @test collect(values(d)) == ["second", "third"]
+    end
+
+    @testset "iterate" begin
+        d = CleverDicts.CleverDict{MyKey, String}()
+        key = CleverDicts.new_item(d, "first")
+        key2 = CleverDicts.new_item(d, "second")
+        my_keys = MyKey[]
+        my_values = String[]
+        for (k, v) in d
+           push!(my_keys, k)
+           push!(my_values, v)
+        end
+        @test my_keys == [MyKey(1), MyKey(2)]
+        @test my_values == ["first", "second"]
+        delete!(d, key)
+        key3 = CleverDicts.new_item(d, "third")
+        my_keys = MyKey[]
+        my_values = String[]
+        for (k, v) in d
+           push!(my_keys, k)
+           push!(my_values, v)
+        end
+        @test my_keys == [MyKey(2), MyKey(3)]
+        @test my_values == ["second", "third"]
+    end
+
+    @testset "iterate ii" begin
+        d = CleverDicts.CleverDict{MyKey, String}()
+        key = CleverDicts.new_item(d, "first")
+        key2 = CleverDicts.new_item(d, "second")
+        my_keys = MyKey[]
+        my_values = String[]
+        for (k, v) in d
+            push!(my_keys, k)
+            push!(my_values, v)
+        end
+        @test my_keys == [MyKey(1), MyKey(2)]
+        @test my_values == ["first", "second"]
+        delete!(d, key)
+        @test d[CleverDicts.LinearIndex(1)] == "second"
+        key3 = CleverDicts.new_item(d, "third")
+        my_keys = MyKey[]
+        my_values = String[]
+        for (k, v) in d
+            push!(my_keys, k)
+            push!(my_values, v)
+        end
+        @test my_keys == [MyKey(2), MyKey(3)]
+        @test my_values == ["second", "third"]
+    end
+
+    @testset "iterate iii" begin
+        d = CleverDicts.CleverDict{MyKey, String}()
+        y = 0
+        for (k, v) in d
+            y += 1
+        end
+        @test y == 0
+    end
+
+    @testset "haskey" begin
+        d = CleverDicts.CleverDict{MyKey, String}()
+        @test !haskey(d, 1)
+        k = CleverDicts.new_item(d, "a")
+        @test haskey(d, k)
+        j = CleverDicts.new_item(d, "b")
+        @test haskey(d, j)
+        delete!(d, k)
+        @test !haskey(d, k)
+        @test haskey(d, j)
+    end
+
+    @testset "delete!" begin
+        d = CleverDicts.CleverDict{MyKey, String}()
+        @test length(d) == 0
+        @test delete!(d, MyKey(0)) == nothing
+        k1 = CleverDicts.new_item(d, "a")
+        k2 = CleverDicts.new_item(d, "b")
+        d[CleverDicts.LinearIndex(2)] == "b"
+        delete!(d, k1)
+        d[CleverDicts.LinearIndex(1)] == "b"
+        k3 = CleverDicts.new_item(d, "c")
+        @test d[k3] == "c"
+        @test d[CleverDicts.LinearIndex(1)] == "b"
+        @test d[CleverDicts.LinearIndex(2)] == "c"
+    end
+end

test/Utilities/Utilities.jl

@@ -27,3 +27,7 @@ end
 @testset "Copy" begin
     include("copy.jl")
 end
+
+@testset "CleverDicts" begin
+    include("CleverDicts.jl")
+end