Add identity_index_map and fix DoubleDict bug (#1329)

* Add identity_index_map and fix DoubleDict bug

* Move model to first argument

* Document IndexMap and identity_index_map

docs/src/submodules/Utilities/reference.md

@@ -58,6 +58,8 @@ Utilities.latex_formulation
 ```@docs
 Utilities.automatic_copy_to
 Utilities.default_copy_to
+Utilities.IndexMap
+Utilities.identity_index_map
 ```
 
 ### [Allocate-Load API](@id allocate_load_api_ref)

src/Utilities/DoubleDicts.jl

@@ -417,7 +417,9 @@ end
 end
 
 function initialize_inner!(d::AbstractWithType{F,S,V,D}) where {F,S,V,D}
-    return d.inner = D()
+    d.inner = D()
+    d.dict.dict[(F, S)] = d.inner
+    return
 end
 
 inner_is_empty(d::AbstractWithType)::Bool = d.inner === nothing

src/Utilities/copy.jl

@@ -58,6 +58,16 @@ function dense_variable_dict(::Type{V}, n) where {V}
     )
 end
 
+"""
+    struct IndexMap <: AbstractDict{MOI.Index,MOI.Index}
+        varmap::DenseVariableDict{MOI.VariableIndex}
+        conmap::DoubleDicts.MainIndexDoubleDict
+    end
+
+Dictionary-like object returned by [`MathOptInterface.copy_to`](@ref) that
+contains the mapping between variable indices in `varmap` and between
+constraint indices in `conmap`.
+"""
 struct IndexMap <: AbstractDict{MOI.Index,MOI.Index}
     varmap::DenseVariableDict{MOI.VariableIndex}
     conmap::DoubleDicts.MainIndexDoubleDict
@@ -70,6 +80,33 @@ function IndexMap(n = 0)
     )
 end
 
+function _identity_constraints_map(
+    model,
+    index_map::MOIU.DoubleDicts.IndexWithType{F,S},
+) where {F,S}
+    for ci in MOI.get(model, MOI.ListOfConstraintIndices{F,S}())
+        index_map[ci] = ci
+    end
+end
+
+"""
+    identity_index_map(model::MOI.ModelLike)
+
+Return an [`IndexMap`](@ref) that maps all variable and constraint indices of
+`model` to themselves.
+"""
+function identity_index_map(model::MOI.ModelLike)
+    vis = MOI.get(model, MOI.ListOfVariableIndices())
+    index_map = IndexMap(length(vis))
+    for vi in vis
+        index_map[vi] = vi
+    end
+    for (F, S) in MOI.get(model, MOI.ListOfConstraints())
+        _identity_constraints_map(model, index_map.conmap[F, S])
+    end
+    return index_map
+end
+
 """
     index_map_for_variable_indices(variables)
 

test/Utilities/copy.jl

@@ -13,6 +13,22 @@ function compare_without_moi(x::String, y::String)
     @test remove_moi(x) == remove_moi(y)
 end
 
+function identity_index_map_test(::Type{T}) where {T}
+    model = MOI.Utilities.Model{T}()
+    x, y = MOI.add_variables(model, 2)
+    fx = MOI.SingleVariable(x)
+    fy = MOI.SingleVariable(y)
+    cx = MOI.add_constraint(model, fx, MOI.EqualTo(one(T)))
+    cy = MOI.add_constraint(model, fy, MOI.EqualTo(zero(T)))
+    c = MOI.add_constraint(model, one(T) * fx + fy, MOI.LessThan(zero(T)))
+    index_map = MOI.Utilities.identity_index_map(model)
+    @test x == index_map[x]
+    @test y == index_map[y]
+    @test cx == index_map[cx]
+    @test cy == index_map[cy]
+    @test c == index_map[c]
+end
+
 @testset "IndexMap" begin
     map = MOIU.IndexMap()
     @test length(map) == 0
@@ -695,3 +711,8 @@ end
     test_pass_copy(Int)
     test_pass_copy(Float64)
 end
+
+@testset "identity_index_map" begin
+    identity_index_map_test(Int)
+    identity_index_map_test(Float64)
+end