[MOF] Specialize objective writing for MOF.Model

src/FileFormats/MOF/write.jl

@@ -132,13 +132,98 @@ function _write_objective(
     if sense == MOI.FEASIBILITY_SENSE
         return (; :sense => moi_to_object(sense)), false
     end
+    return _write_objective(model, name_map, sense)
+end
+
+function _write_objective(model, name_map, sense)
     F = MOI.get(model, MOI.ObjectiveFunctionType())
-    objective_function = MOI.get(model, MOI.ObjectiveFunction{F}())
+    return _write_objective(
+        MOI.get(model, MOI.ObjectiveFunction{F}()),
+        name_map,
+        sense,
+    )
+end
+
+function _write_objective(
+    model::MOI.Utilities.UniversalFallback,
+    name_map,
+    sense,
+)
+    if model.objective === nothing
+        return _write_objective(model.model, name_map, sense)
+    end
+    return _write_objective(something(model.objective), name_map, sense)
+end
+
+function _write_objective(model::MOI.Utilities.GenericModel, name_map, sense)
+    return _write_objective(model.objective, name_map, sense)
+end
+
+function _write_objective(
+    objective::MOI.Utilities.ObjectiveContainer{T},
+    name_map,
+    sense,
+) where {T}
+    # Build the complete objective inside each branch, while the function type
+    # is concrete. Returning the function first would merge its types again.
+    if objective.scalar_affine !== nothing
+        return _write_objective(
+            something(objective.scalar_affine),
+            name_map,
+            sense,
+        )
+    elseif objective.single_variable !== nothing
+        return _write_objective(
+            something(objective.single_variable),
+            name_map,
+            sense,
+        )
+    elseif objective.scalar_quadratic !== nothing
+        return _write_objective(
+            something(objective.scalar_quadratic),
+            name_map,
+            sense,
+        )
+    elseif objective.scalar_nonlinear !== nothing
+        return _write_objective(
+            something(objective.scalar_nonlinear),
+            name_map,
+            sense,
+        )
+    elseif objective.vector_variables !== nothing
+        return _write_objective(
+            something(objective.vector_variables),
+            name_map,
+            sense,
+        )
+    elseif objective.vector_affine !== nothing
+        return _write_objective(
+            something(objective.vector_affine),
+            name_map,
+            sense,
+        )
+    elseif objective.vector_quadratic !== nothing
+        return _write_objective(
+            something(objective.vector_quadratic),
+            name_map,
+            sense,
+        )
+    elseif objective.vector_nonlinear !== nothing
+        return _write_objective(
+            something(objective.vector_nonlinear),
+            name_map,
+            sense,
+        )
+    end
+    return _write_objective(zero(MOI.ScalarAffineFunction{T}), name_map, sense)
+end
+
+function _write_objective(f::MOI.AbstractFunction, name_map, sense)
     object = (;
         :sense => moi_to_object(sense),
-        :function => moi_to_object(objective_function, name_map),
+        :function => moi_to_object(f, name_map),
     )
-    return object, (F == MOI.ScalarNonlinearFunction)
+    return object, f isa MOI.ScalarNonlinearFunction
 end
 
 function _write_constraints(

test/FileFormats/MOF/test_MOF.jl

@@ -60,6 +60,17 @@ struct UnsupportedSet <: MOI.AbstractSet end
 
 struct UnsupportedFunction <: MOI.AbstractFunction end
 
+struct ObjectiveStorage{T}
+    inner::T
+end
+
+function MOI.get(
+    storage::ObjectiveStorage,
+    attr::Union{MOI.ObjectiveFunctionType,MOI.ObjectiveFunction},
+)
+    return MOI.get(storage.inner, attr)
+end
+
 function _test_model_equality(
     model_string,
     variables,
@@ -460,6 +471,115 @@ function test_max_objective()
     )
 end
 
+function test_write_objective_storage_traversal()
+    model = MOF.Model()
+    x, y = MOI.add_variables(model, 2)
+    name_map = Dict(x => "x", y => "y")
+    affine = 2.0 * x + y + 3.0
+    quadratic = 2.0 * x * y + affine
+    nonlinear = MOI.ScalarNonlinearFunction(:sin, Any[x])
+    @test (@inferred MOF._write_objective(affine, name_map, MOI.MIN_SENSE)) == (
+        (; :sense => "min", :function => MOF.moi_to_object(affine, name_map)),
+        false,
+    )
+    functions = (
+        x,
+        affine,
+        quadratic,
+        nonlinear,
+        MOI.VectorOfVariables([x, y]),
+        MOI.Utilities.vectorize([affine, 3.0 * y]),
+        MOI.Utilities.vectorize([quadratic, 3.0 * y * y]),
+        MOI.VectorNonlinearFunction([nonlinear]),
+    )
+    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
+    for f in functions
+        MOI.set(model, MOI.ObjectiveFunction{typeof(f)}(), f)
+        expected =
+            (; :sense => "min", :function => MOF.moi_to_object(f, name_map))
+        has_scalar_nonlinear = f isa MOI.ScalarNonlinearFunction
+        @test model.objective === nothing
+        @test MOF._write_objective(model, name_map, MOI.MIN_SENSE) ==
+              (expected, has_scalar_nonlinear)
+        @test MOF._write_objective(model, name_map) ==
+              (expected, has_scalar_nonlinear)
+    end
+    return
+end
+
+function test_write_objective_default_and_feasibility()
+    for T in (Float64, Int)
+        model = MOF.Model(; coefficient_type = T)
+        name_map = Dict{MOI.VariableIndex,String}()
+        expected =
+            MOF.moi_to_object(zero(MOI.ScalarAffineFunction{T}), name_map)
+        for (sense, text) in ((MOI.MIN_SENSE, "min"), (MOI.MAX_SENSE, "max"))
+            MOI.set(model, MOI.ObjectiveSense(), sense)
+            @test MOF._write_objective(model, name_map) ==
+                  ((; :sense => text, :function => expected), false)
+        end
+        MOI.set(model, MOI.ObjectiveSense(), MOI.FEASIBILITY_SENSE)
+        x = MOI.add_variable(model)
+        # A stored objective must not be serialized for feasibility sense. The
+        # empty name map would make serialization of either objective fail.
+        for f in (
+            MOI.ScalarNonlinearFunction(:sin, Any[x]),
+            MOF.Nonlinear(:(sin($x))),
+        )
+            MOI.set(model, MOI.ObjectiveFunction{typeof(f)}(), f)
+            @test MOF._write_objective(model, name_map) ==
+                  ((; sense = "feasibility"), false)
+        end
+    end
+    return
+end
+
+function test_write_objective_universal_fallback()
+    model = MOF.Model()
+    x = MOI.add_variable(model)
+    name_map = Dict(x => "x")
+    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
+    affine = 2.0 * x + 3.0
+    # Integer coefficients and legacy nonlinear expressions are stored by the
+    # fallback, whereas Float64 affine objectives use the inner model.
+    for f in (affine, 4 * x + 5, affine, MOF.Nonlinear(:(sin($x))), affine)
+        MOI.set(model, MOI.ObjectiveFunction{typeof(f)}(), f)
+        @test (model.objective === nothing) ==
+              (f isa MOI.ScalarAffineFunction{Float64})
+        expected = MOF.moi_to_object(f, name_map)
+        @test MOF._write_objective(model, name_map) ==
+              ((; :sense => "max", :function => expected), false)
+    end
+    return
+end
+
+function test_write_objective_generic_storage()
+    inner = MOI.Utilities.Model{Float64}()
+    x = MOI.add_variable(inner)
+    name_map = Dict(x => "x")
+    # GenericModel permits custom objective storage, even without ModelLike.
+    objective = ObjectiveStorage(inner.objective)
+    model = MOI.Utilities.GenericModel{Float64}(
+        objective,
+        inner.variables,
+        inner.constraints,
+    )
+    for f in (2.0 * x + 3.0, MOI.ScalarNonlinearFunction(:sin, Any[x]))
+        MOI.set(inner, MOI.ObjectiveFunction{typeof(f)}(), f)
+        expected = (
+            (; :sense => "min", :function => MOF.moi_to_object(f, name_map)),
+            f isa MOI.ScalarNonlinearFunction,
+        )
+        @test MOF._write_objective(model, name_map, MOI.MIN_SENSE) == expected
+        @test MOF._write_objective(
+            MOI.Utilities.UniversalFallback(model),
+            name_map,
+            MOI.MIN_SENSE,
+        ) == expected
+    end
+    return
+end
+
 function test_min_scalaraffine()
     return _test_model_equality(
         """