From d20de103854a7b0ce2e83fd481287ebf220e7c93 Mon Sep 17 00:00:00 2001
From: Oscar Dowson <odow@users.noreply.github.com>
Date: Wed, 12 Jun 2024 12:48:48 +1200
Subject: [PATCH] [FileFormats.NL] read ScalarAffineFunction where possible
 (#2512)

---
 src/FileFormats/NL/read.jl  | 38 ++++++++++++++++++++++++++++++-----
 test/FileFormats/NL/read.jl | 40 ++++++++++++++++++++++++++++++++-----
 2 files changed, 68 insertions(+), 10 deletions(-)

diff --git a/src/FileFormats/NL/read.jl b/src/FileFormats/NL/read.jl
index 359ca12ab8..dc399d06ad 100644
--- a/src/FileFormats/NL/read.jl
+++ b/src/FileFormats/NL/read.jl
@@ -208,12 +208,12 @@ function _to_model(data::_CacheModel; use_nlp_block::Bool)
         MOI.set(model, MOI.NLPBlock(), block)
     else
         if data.objective != :()
-            obj = _to_scalar_nonlinear_function(data.objective)
+            obj = _expr_to_function(data.objective)
             MOI.set(model, MOI.ObjectiveFunction{typeof(obj)}(), obj)
         end
         for (i, expr) in enumerate(data.constraints)
             lb, ub = data.constraint_lower[i], data.constraint_upper[i]
-            f = _to_scalar_nonlinear_function(expr)::MOI.ScalarNonlinearFunction
+            f = _expr_to_function(expr)
             if lb == ub
                 MOI.add_constraint(model, f, MOI.EqualTo(lb))
             elseif -Inf == lb && ub < Inf
@@ -228,16 +228,44 @@ function _to_model(data::_CacheModel; use_nlp_block::Bool)
     return model
 end
 
-_to_scalar_nonlinear_function(expr) = expr
+_expr_to_function(expr) = expr
 
-function _to_scalar_nonlinear_function(expr::Expr)
+function _expr_to_function(expr::Expr)
     @assert Meta.isexpr(expr, :call)
+    f = _try_scalar_affine_function(expr)
+    if f !== nothing
+        return convert(MOI.ScalarAffineFunction{Float64}, f)
+    end
     return MOI.ScalarNonlinearFunction(
         expr.args[1],
-        Any[_to_scalar_nonlinear_function(arg) for arg in expr.args[2:end]],
+        Any[_expr_to_function(arg) for arg in expr.args[2:end]],
     )
 end
 
+_try_scalar_affine_function(x::Float64) = x
+
+_try_scalar_affine_function(x::MOI.VariableIndex) = x
+
+function _try_scalar_affine_function(expr::Expr)
+    if expr.args[1] == :+
+        args = _try_scalar_affine_function.(expr.args[2:end])
+        if !any(isnothing, args)
+            return MOI.Utilities.operate(+, Float64, args...)
+        end
+    elseif expr.args[1] == :*
+        args = _try_scalar_affine_function.(expr.args[2:end])
+        n_affine_terms = 0
+        for arg in args
+            n_affine_terms += arg isa MOI.VariableIndex
+            n_affine_terms += arg isa MOI.ScalarAffineFunction{Float64}
+        end
+        if n_affine_terms <= 1
+            return MOI.Utilities.operate(*, Float64, args...)
+        end
+    end
+    return nothing
+end
+
 function _parse_header(io::IO, model::_CacheModel)
     # Line 1
     # We don't support the binary format.
diff --git a/test/FileFormats/NL/read.jl b/test/FileFormats/NL/read.jl
index 1243e6fe9d..eed46db26a 100644
--- a/test/FileFormats/NL/read.jl
+++ b/test/FileFormats/NL/read.jl
@@ -8,7 +8,7 @@ module TestNonlinearRead
 
 using Test
 import MathOptInterface as MOI
-const NL = MOI.FileFormats.NL
+import MathOptInterface.FileFormats: NL
 
 function runtests()
     for name in names(@__MODULE__; all = true)
@@ -720,10 +720,10 @@ function test_hs071_free_constraint_nlexpr()
     open(joinpath(@__DIR__, "data", "hs071_free_constraint.nl"), "r") do io
         return read!(io, model)
     end
-    @test MOI.get(model, MOI.ListOfConstraintTypesPresent()) == [
-        (MOI.ScalarNonlinearFunction, MOI.GreaterThan{Float64}),
-        (MOI.ScalarNonlinearFunction, MOI.Interval{Float64}),
-    ]
+    types = MOI.get(model, MOI.ListOfConstraintTypesPresent())
+    @test length(types) == 2
+    @test (MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) in types
+    @test (MOI.ScalarNonlinearFunction, MOI.Interval{Float64}) in types
     for (F, S) in MOI.get(model, MOI.ListOfConstraintTypesPresent())
         @test MOI.get(model, MOI.NumberOfConstraints{F,S}()) == 1
     end
@@ -783,6 +783,36 @@ function test_mac_minlp()
     return
 end
 
+function test_nl_read_scalar_affine_function()
+    for obj_fn in (x -> x, x -> 1.0 * x, x -> 2.0 * x + 3.0)
+        src = MOI.Utilities.Model{Float64}()
+        x = MOI.add_variable(src)
+        MOI.set(src, MOI.ObjectiveSense(), MOI.MIN_SENSE)
+        f = obj_fn(x)
+        MOI.set(src, MOI.ObjectiveFunction{typeof(f)}(), f)
+        MOI.add_constraint(src, f, MOI.LessThan(1.0))
+        dest = MOI.FileFormats.NL.Model()
+        MOI.copy_to(dest, src)
+        io = IOBuffer()
+        write(io, dest)
+        input = MOI.FileFormats.NL.Model(; use_nlp_block = false)
+        seekstart(io)
+        read!(io, input)
+        model = MOI.Utilities.Model{Float64}()
+        MOI.copy_to(model, input)
+        y = only(MOI.get(model, MOI.ListOfVariableIndices()))
+        g = MOI.Utilities.substitute_variables(
+            _ -> y,
+            convert(MOI.ScalarAffineFunction{Float64}, f),
+        )
+        obj = MOI.get(model, MOI.ObjectiveFunction{typeof(g)}())
+        @test ≈(obj, g)
+        @test MOI.get(model, MOI.ListOfConstraintTypesPresent()) ==
+              [(typeof(f), MOI.LessThan{Float64})]
+    end
+    return
+end
+
 end
 
 TestNonlinearRead.runtests()