[Bridges] Add MOI.Bridge.FirstBridge attribute (#2211)

* Add FirstBridge

* Add tests

* Update

* Add to docs

* Update src/Utilities/functions.jl

* Update src/Bridges/Constraint/bridges/functionize.jl

---------

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

Author: blegat

docs/src/submodules/Bridges/reference.md

@@ -75,6 +75,7 @@ Bridges.unbridged_variable_function
 Bridges.bridged_function
 Bridges.supports_constraint_bridges
 Bridges.recursive_model
+Bridges.FirstBridge
 ```
 
 ## LazyBridgeOptimizer API

src/Bridges/Constraint/bridges/functionize.jl

@@ -42,6 +42,15 @@ function MOI.get(
     return MOI.Utilities.canonical(f)
 end
 
+# Needed to avoid an ambiguity with the getter for MOI.AbstractConstraintAttribute
+function MOI.get(
+    ::MOI.ModelLike,
+    ::MOI.Bridges.FirstBridge,
+    bridge::AbstractFunctionConversionBridge,
+)
+    return bridge
+end
+
 function MOI.supports(
     model::MOI.ModelLike,
     attr::MOI.AbstractConstraintAttribute,

src/Bridges/bridge.jl

@@ -323,3 +323,26 @@ julia> MOI.Bridges.bridging_cost(
 ```
 """
 bridging_cost(::Type{<:AbstractBridge}) = 1.0
+
+"""
+    struct FirstBridge <: MOI.AbstractConstraintAttribute end
+
+Returns the first bridge used to bridge the constraint.
+
+!!! warning
+    The indices of the bridge correspond to internal indices and may not
+    correspond to indices of the model this attribute is got from.
+"""
+struct FirstBridge <: MOI.AbstractConstraintAttribute end
+
+MOI.is_set_by_optimize(::FirstBridge) = true
+
+MOI.get(::MOI.ModelLike, ::FirstBridge, b::MOI.Bridges.AbstractBridge) = b
+
+function MOI.Utilities.map_indices(
+    ::Function,
+    ::FirstBridge,
+    b::MOI.Bridges.AbstractBridge,
+)
+    return b
+end

src/Utilities/functions.jl

@@ -221,10 +221,18 @@ constraint functions, attribute values and submittable values. If you define a
 new attribute whose values `x::X` contain variable or constraint indices, you
 must also implement this function.
 """
-map_indices(f, ::MOI.AnyAttribute, x) = map_indices(f, x)
+map_indices(f::Function, ::MOI.AnyAttribute, x) = map_indices(f, x)
+
+function map_indices(
+    variable_map::AbstractDict{T,T},
+    attr::MOI.AnyAttribute,
+    x::X,
+)::X where {T<:MOI.Index,X}
+    return map_indices(Base.Fix1(getindex, variable_map), attr, x)
+end
 
 # RawOptimizerAttribute values are passed through un-changed.
-map_indices(::Any, ::MOI.RawOptimizerAttribute, x) = x
+map_indices(::Function, ::MOI.RawOptimizerAttribute, x) = x
 
 """
     map_indices(

test/Bridges/Constraint/functionize.jl

@@ -351,6 +351,16 @@ function test_canonical_constraint_function()
     return
 end
 
+function test_first_bridge()
+    inner = MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}())
+    model = MOI.Bridges.Constraint.ScalarFunctionize{Float64}(inner)
+    x = MOI.add_variable(model)
+    ci = MOI.add_constraint(model, x, MOI.GreaterThan(0.0))
+    b = MOI.get(model, MOI.Bridges.FirstBridge(), ci)
+    @test b isa MOI.Bridges.Constraint.ScalarFunctionizeBridge
+    return
+end
+
 end  # module
 
 TestConstraintFunctionize.runtests()

test/Bridges/bridge_optimizer.jl

@@ -1108,6 +1108,27 @@ function test_list_of_constraints_with_attribute_set()
     return
 end
 
+function test_first_bridge()
+    inner = MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}())
+    bridge = MOI.Bridges.Constraint.ZeroOne{Float64}(inner)
+    model = MOI.Utilities.CachingOptimizer(
+        MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
+        bridge,
+    )
+    x = MOI.add_variable(model)
+    ci = MOI.add_constraint(model, x, MOI.ZeroOne())
+    MOI.Utilities.attach_optimizer(model)
+    b = MOI.get(model, MOI.Bridges.FirstBridge(), ci)
+    @test b isa MOI.Bridges.Constraint.ZeroOneBridge
+    y = MOI.add_variable(model)
+    ci = MOI.add_constraint(model, y, MOI.Integer())
+    @test_throws(
+        MOI.GetAttributeNotAllowed{MOI.Bridges.FirstBridge},
+        MOI.get(model, MOI.Bridges.FirstBridge(), ci),
+    )
+    return
+end
+
 end  # module
 
 TestBridgeOptimizer.runtests()