From f496535f560ea1a6bbf5df19031997bdcc1e4022 Mon Sep 17 00:00:00 2001
From: Oscar Dowson <odow@users.noreply.github.com>
Date: Tue, 5 Sep 2023 18:57:20 +1200
Subject: [PATCH] Fix promotion rules for GenericNonlinearExpr (#3483)

---
 src/aff_expr.jl       | 21 +++++++++++++++++++++
 src/nlp_expr.jl       | 27 +++++++++++++++++++--------
 test/test_nlp_expr.jl | 39 +++++++++++++++++++++++++++++++++------
 3 files changed, 73 insertions(+), 14 deletions(-)

diff --git a/src/aff_expr.jl b/src/aff_expr.jl
index 75634bee377..f6b3db65361 100644
--- a/src/aff_expr.jl
+++ b/src/aff_expr.jl
@@ -649,6 +649,27 @@ See also: [`jump_function`](@ref).
 """
 function moi_function end
 
+function moi_function(x::AbstractArray{AbstractJuMPScalar})
+    return error(
+        "Unable to convert array of type `::$(typeof(x))` to an equivalent " *
+        "function in MathOptInterface because the array has the abstract " *
+        "element type `AbstractJuMPScalar`. To fix this error, convert every " *
+        "element in the array to the same concrete element type.\n\n" *
+        """For example, instead of:
+        ```julia
+        model = Model();
+        @variable(model, x);
+        y = AbstractJuMPScalar[x, sin(x)]
+        @objective(model, Min, y)
+        ```
+        do
+        ```julia
+        @objective(model, Min, convert.(NonlinearExpr, y))
+        ```
+        """,
+    )
+end
+
 """
     moi_function_type(::Type{T}) where {T}
 
diff --git a/src/nlp_expr.jl b/src/nlp_expr.jl
index 13dfca3e847..60763939f6a 100644
--- a/src/nlp_expr.jl
+++ b/src/nlp_expr.jl
@@ -675,11 +675,25 @@ function _evaluate_expr(
     end
 end
 
-# MutableArithmetics.jl
+# MutableArithmetics.jl and promotion
 
-# These converts are used in the {add,sub}mul definition for AbstractJuMPScalar.
+function Base.promote_rule(
+    ::Type{GenericNonlinearExpr{V}},
+    ::Type{V},
+) where {V<:AbstractVariableRef}
+    return GenericNonlinearExpr{V}
+end
+
+function Base.promote_rule(
+    ::Type{GenericNonlinearExpr{V}},
+    ::Type{<:Union{GenericAffExpr{C,V},GenericQuadExpr{C,V}}},
+) where {C,V<:AbstractVariableRef}
+    return GenericNonlinearExpr{V}
+end
 
-Base.convert(::Type{<:GenericNonlinearExpr}, x::AbstractVariableRef) = x
+function Base.convert(::Type{GenericNonlinearExpr{V}}, x::V) where {V}
+    return GenericNonlinearExpr{V}(:+, Any[x])
+end
 
 function Base.convert(
     ::Type{<:GenericNonlinearExpr},
@@ -1099,14 +1113,11 @@ function jump_function(
     ]
 end
 
-# We use `AbstractJuMPScalar` as a catch-all fallback for any mix of JuMP
-# scalars that have not been dispatched by some other method.
-
-function moi_function_type(::Type{<:Vector{<:AbstractJuMPScalar}})
+function moi_function_type(::Type{<:AbstractVector{<:GenericNonlinearExpr}})
     return MOI.VectorNonlinearFunction
 end
 
-function moi_function(f::Vector{<:AbstractJuMPScalar})
+function moi_function(f::AbstractVector{<:GenericNonlinearExpr})
     return MOI.VectorNonlinearFunction(f)
 end
 
diff --git a/test/test_nlp_expr.jl b/test/test_nlp_expr.jl
index f7579db2fc7..2afde2d2e2e 100644
--- a/test/test_nlp_expr.jl
+++ b/test/test_nlp_expr.jl
@@ -744,26 +744,45 @@ function test_VectorNonlinearFunction_moi_function()
     return
 end
 
-function test_VectorNonlinearFunction_moi_function_AbstractJuMPScalar()
+function test_VectorNonlinearFunction_moi_function_conversion()
     model = Model()
     @variable(model, x)
-    F = [sin(x), x]
-    @test F isa Vector{AbstractJuMPScalar}
+    F = [sin(x), x, x + 1, x^2]
+    @test F isa Vector{NonlinearExpr}
     @test moi_function_type(typeof(F)) == MOI.VectorNonlinearFunction
     @test isapprox(
         moi_function(F),
         MOI.VectorNonlinearFunction([
             MOI.ScalarNonlinearFunction(:sin, Any[index(x)]),
             MOI.ScalarNonlinearFunction(:+, Any[index(x)]),
+            MOI.ScalarNonlinearFunction(:+, Any[index(x), 1.0]),
+            MOI.ScalarNonlinearFunction(:*, Any[index(x), index(x)]),
         ]),
     )
     @test MOI.VectorNonlinearFunction(F) ≈ moi_function(F)
     @test jump_function_type(model, MOI.VectorNonlinearFunction) ==
           Vector{NonlinearExpr}
-    @test isequal_canonical(
-        jump_function(model, moi_function(F)),
-        [sin(x), NonlinearExpr(:+, x)],
+    @test isequal_canonical(jump_function(model, moi_function(F)), F)
+    return
+end
+
+function test_VectorNonlinearFunction_moi_function_conversion_variable()
+    model = Model()
+    @variable(model, x)
+    F = [sin(x), x]
+    @test F isa Vector{NonlinearExpr}
+    @test moi_function_type(typeof(F)) == MOI.VectorNonlinearFunction
+    @test isapprox(
+        moi_function(F),
+        MOI.VectorNonlinearFunction([
+            MOI.ScalarNonlinearFunction(:sin, Any[index(x)]),
+            MOI.ScalarNonlinearFunction(:+, Any[index(x)]),
+        ]),
     )
+    @test MOI.VectorNonlinearFunction(F) ≈ moi_function(F)
+    @test jump_function_type(model, MOI.VectorNonlinearFunction) ==
+          Vector{NonlinearExpr}
+    @test isequal_canonical(jump_function(model, moi_function(F)), F)
     return
 end
 
@@ -831,6 +850,14 @@ function test_redefinition_of_function()
     return
 end
 
+function test_moi_function_abstract_jump_scalar()
+    model = Model()
+    @variable(model, x)
+    y = AbstractJuMPScalar[x, sin(x)]
+    @test_throws ErrorException moi_function(y)
+    return
+end
+
 function test_linear_algebra_errors()
     model = Model()
     @variable(model, x[1:2, 1:2])