[Utilities] add DoubleDicts.outer_keys (#2052)

* DoubleDict outer_keys and nonempty_outer_keys

Co-authored-by: Oscar Dowson <[email protected]>

Author: joaquimg

docs/src/submodules/Utilities/reference.md

@@ -220,4 +220,6 @@ Utilities.DoubleDicts.DoubleDict
 Utilities.DoubleDicts.DoubleDictInner
 Utilities.DoubleDicts.IndexDoubleDict
 Utilities.DoubleDicts.IndexDoubleDictInner
+Utilities.DoubleDicts.outer_keys
+Utilities.DoubleDicts.nonempty_outer_keys
 ```

src/Utilities/DoubleDicts.jl

@@ -299,6 +299,55 @@ function Base.keys(d::AbstractDoubleDictInner{F,S}) where {F,S}
     return MOI.ConstraintIndex{F,S}.(keys(d.dict))
 end
 
+"""
+    outer_keys(d::AbstractDoubleDict)
+
+Return an iterator over the outer keys of the `AbstractDoubleDict` `d`.
+Each outer key is a `Tuple{Type,Type}` so that a double loop can
+be easily used:
+```julia
+for (F, S) in DoubleDicts.outer_keys(dict)
+    for (k, v) in dict[F, S]
+        # ...
+    end
+end
+```
+
+For performance, it is recommended that the inner loop lies in a separate
+function to gurantee type-stability.
+Some outer keys `(F, S)` might lead to an empty `dict[F, S]`.
+If you want only nonempty `dict[F, S]`, use [`nonempty_outer_keys`](@ref).
+"""
+outer_keys(d::AbstractDoubleDict) = keys(d.dict)
+
+"""
+    nonempty_outer_keys(d::AbstractDoubleDict)
+
+Return a vector of outer keys of the `AbstractDoubleDict` `d`.
+
+Only outer keys that have a nonempty set of inner keys will be returned.
+
+Each outer key is a `Tuple{Type,Type}` so that a double loop can
+be easily used
+```julia
+for (F, S) in DoubleDicts.nonempty_outer_keys(dict)
+    for (k, v) in dict[F, S]
+        # ...
+    end
+end
+For performance, it is recommended that the inner loop lies in a separate
+function to gurantee type-stability.
+
+If you want an iterator of all current outer keys, use [`outer_keys`](@ref).
+```
+"""
+function nonempty_outer_keys(d::AbstractDoubleDict)
+    return Base.Iterators.Filter(
+        fs -> length(d[fs[1], fs[2]]) > 0,
+        outer_keys(d),
+    )
+end
+
 # Base.iterate
 
 function Base.iterate(d::AbstractDoubleDict)

test/Utilities/DoubleDicts.jl

@@ -177,6 +177,83 @@ function test_IndexDoubleDict()
     return
 end
 
+function test_IndexDoubleDict_double_loop()
+    dict = DoubleDicts.IndexDoubleDict()
+    keys = [
+        MOI.ConstraintIndex{MOI.VariableIndex,MOI.Integer}(1),
+        MOI.ConstraintIndex{MOI.VariableIndex,MOI.Integer}(2),
+        MOI.ConstraintIndex{MOI.VariableIndex,MOI.ZeroOne}(1),
+    ]
+    vals = keys
+    for (k, v) in zip(keys, vals)
+        dict[k] = v
+    end
+    # outer_keys
+    number_elements = 0
+    number_outer_pairs = 0
+    for (F, S) in DoubleDicts.outer_keys(dict)
+        number_outer_pairs += 1
+        for (k, v) in dict[F, S]
+            @test dict[k] == v
+            number_elements += 1
+        end
+    end
+    @test number_elements == 3
+    @test number_outer_pairs == 2
+    # nonempty_outer_keys
+    number_elements = 0
+    number_outer_pairs = 0
+    for (F, S) in DoubleDicts.nonempty_outer_keys(dict)
+        number_outer_pairs += 1
+        for (k, v) in dict[F, S]
+            @test dict[k] == v
+            number_elements += 1
+        end
+    end
+    @test number_elements == 3
+    @test number_outer_pairs == 2
+    return
+end
+
+function test_IndexDoubleDict_double_loop_with_delete()
+    dict = DoubleDicts.IndexDoubleDict()
+    keys = [
+        MOI.ConstraintIndex{MOI.VariableIndex,MOI.Integer}(1),
+        MOI.ConstraintIndex{MOI.VariableIndex,MOI.Integer}(2),
+        MOI.ConstraintIndex{MOI.VariableIndex,MOI.ZeroOne}(1),
+    ]
+    vals = keys
+    for (k, v) in zip(keys, vals)
+        dict[k] = v
+    end
+    delete!(dict, MOI.ConstraintIndex{MOI.VariableIndex,MOI.ZeroOne}(1))
+    # outer_keys
+    number_elements = 0
+    number_outer_pairs = 0
+    for (F, S) in DoubleDicts.outer_keys(dict)
+        number_outer_pairs += 1
+        for (k, v) in dict[F, S]
+            @test dict[k] == v
+            number_elements += 1
+        end
+    end
+    @test number_elements == 2
+    @test number_outer_pairs == 2
+    # nonempty_outer_keys
+    number_elements = 0
+    number_outer_pairs = 0
+    for (F, S) in DoubleDicts.nonempty_outer_keys(dict)
+        number_outer_pairs += 1
+        for (k, v) in dict[F, S]
+            @test dict[k] == v
+            number_elements += 1
+        end
+    end
+    @test number_elements == 2
+    @test number_outer_pairs == 1
+    return
+end
+
 end  # module
 
 TestDoubleDicts.runtests()