[Bridges] fix support for nested final_touch bridges (#3027)

Author: odow

src/Bridges/Constraint/map.jl

@@ -20,14 +20,14 @@ struct Map <: AbstractDict{MOI.ConstraintIndex,AbstractBridge}
     # of creation so we need `OrderedDict` and not `Dict`.
     # For `VariableIndex` constraints: (variable, set type) -> bridge
     single_variable_constraints::OrderedDict{Tuple{Int64,Type},AbstractBridge}
-    needs_final_touch::OrderedDict{Type,OrderedSet}
+    needs_final_touch::Vector{Any}
 
     function Map()
         return new(
             Union{Nothing,AbstractBridge}[],
             Tuple{Type,Type}[],
             OrderedDict{Tuple{Int64,Type},AbstractBridge}(),
-            OrderedDict{Type,OrderedSet}(),
+            Any[],
         )
     end
 end
@@ -78,6 +78,14 @@ function Base.getindex(
     return map.single_variable_constraints[(ci.value, S)]
 end
 
+function _unregister_for_final_touch(b::Map, bridge)
+    if MOI.Bridges.needs_final_touch(bridge)
+        i = findfirst(==(bridge), b.needs_final_touch)
+        deleteat!(b.needs_final_touch, i)
+    end
+    return
+end
+
 function Base.delete!(map::Map, ci::MOI.ConstraintIndex)
     _unregister_for_final_touch(map, map.bridges[_index(ci)]::AbstractBridge)
     map.bridges[_index(ci)] = nothing
@@ -298,7 +306,6 @@ function add_key_for_bridge(
     ::S,
     is_available::Function,
 ) where {F<:MOI.AbstractFunction,S<:MOI.AbstractSet}
-    _register_for_final_touch(map, bridge)
     _ensure_available(map, F, S, is_available)
     push!(map.bridges, bridge)
     push!(map.constraint_types, (F, S))
@@ -312,39 +319,17 @@ function add_key_for_bridge(
     ::S,
     ::Function,
 ) where {S<:MOI.AbstractScalarSet}
-    _register_for_final_touch(map, bridge)
     map.single_variable_constraints[(func.value, S)] = bridge
     return MOI.ConstraintIndex{MOI.VariableIndex,S}(func.value)
 end
 
-function _register_for_final_touch(map::Map, bridge::BT) where {BT}
-    if MOI.Bridges.needs_final_touch(bridge)
-        if !haskey(map.needs_final_touch, BT)
-            map.needs_final_touch[BT] = OrderedSet{BT}()
-        end
-        push!(map.needs_final_touch[BT], bridge)
-    end
-    return
-end
-
-function _unregister_for_final_touch(b::Map, bridge::BT) where {BT}
-    if MOI.Bridges.needs_final_touch(bridge)
-        delete!(b.needs_final_touch[BT], bridge)
-    end
-    return
-end
-
-# Function barrier to iterate over bridges of the same type in an efficient way.
-function _final_touch(bridges, model)
-    for bridge in bridges
-        MOI.Bridges.final_touch(bridge, model)
-    end
-    return
-end
-
 function MOI.Bridges.final_touch(map::Map, model::MOI.ModelLike)
-    for bridges in values(map.needs_final_touch)
-        _final_touch(bridges, model)
+    # A bridge's `final_touch` may add a new bridge that needs `final_touch`,
+    # but that's okay because it will push a new element to the end of
+    # `map.needs_final_touch` and this iterator will keep going until it reaches
+    # the end (because it's a `Base.Vector{Any}`).
+    for bridge in map.needs_final_touch
+        MOI.Bridges.final_touch(bridge, model)
     end
     return
 end

src/Bridges/bridge_optimizer.jl

@@ -1961,14 +1961,29 @@ function MOI.supports_constraint(
 end
 
 function add_bridged_constraint(b, BridgeType, f, s)
+    map = Constraint.bridges(b)::Constraint.Map
+    # Okay, this part is a bit tricky. `Constraint.bridge_constraint` might add
+    # new constraints which also need `final_touch`. But we need the return
+    # value `bridge` to come _before_ them in the `needs_final_touch` vector.
+    # To make this happen, we remember how many bridges are already in the
+    # vector, and then we insert the bridge into the vector. This is an O(N)
+    # operation in the number of new bridges, but it's a pretty rare edge-case,
+    # and the number of new bridges should be small. (And most of the time, N=0,
+    # because most bridges that need final touch don't add any bridges that
+    # themselves need final touch. Good examples are the ToMILP bridges that
+    # need variable bounds.)
+    n_final_touch = length(map.needs_final_touch)
     bridge = Constraint.bridge_constraint(BridgeType, recursive_model(b), f, s)
+    if MOI.Bridges.needs_final_touch(bridge)
+        insert!(map.needs_final_touch, n_final_touch + 1, bridge)
+    end
     # `MOI.VectorOfVariables` constraint indices have negative indices
     # to distinguish between the indices of the inner model.
     # However, they can clash between the indices created by the variable
     # so we use the last argument to inform the constraint bridge mapping about
     # indices already taken by variable bridges.
     ci = Constraint.add_key_for_bridge(
-        Constraint.bridges(b)::Constraint.Map,
+        map,
         bridge,
         f,
         s,