From f96a40b2e9798cc8ad7f01c631f59818b903c45a Mon Sep 17 00:00:00 2001
From: Oscar Dowson <odow@users.noreply.github.com>
Date: Thu, 7 Mar 2024 01:19:56 +1300
Subject: [PATCH] Fix _to_polynomial for ScalarAffineFunction and
 ScalarQuadraticFunction (#112)

* Fix _to_polynomial for ScalarAffineFunction and ScalarQuadraticFunction

* Update test/qcqp.jl
---
 src/nl_to_polynomial.jl |  8 ++++++++
 test/qcqp.jl            | 16 ++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/src/nl_to_polynomial.jl b/src/nl_to_polynomial.jl
index bb799e6..71f0a21 100644
--- a/src/nl_to_polynomial.jl
+++ b/src/nl_to_polynomial.jl
@@ -57,6 +57,14 @@ function _to_polynomial!(d, ::Type{T}, vi::MOI.VariableIndex) where {T}
     return d[vi]
 end
 
+function _to_polynomial!(
+    d,
+    ::Type{T},
+    f::Union{MOI.ScalarAffineFunction,MOI.ScalarQuadraticFunction},
+) where {T}
+    return _to_polynomial!(d, T, convert(MOI.ScalarNonlinearFunction, f))
+end
+
 function _to_polynomial!(
     d,
     ::Type{T},
diff --git a/test/qcqp.jl b/test/qcqp.jl
index 46b0a2f..6f48007 100644
--- a/test/qcqp.jl
+++ b/test/qcqp.jl
@@ -199,6 +199,22 @@ function test_unbound_polynomial(x, y, T)
     return
 end
 
+function test_scalar_nonlinear_function(x, y, T)
+    inner = Model{T}()
+    model = PolyJuMP.JuMP.GenericModel{T}() do
+        return PolyJuMP.QCQP.Optimizer{T}(MOI.Utilities.MockOptimizer(inner))
+    end
+    PolyJuMP.@variable(model, 0 <= x <= 1)
+    PolyJuMP.@expression(model, f, 0 + x)
+    PolyJuMP.@expression(model, g, x^2)
+    PolyJuMP.@constraint(model, f * g == 0)
+    PolyJuMP.optimize!(model)
+    F, S = ScalarQuadraticFunction{T}, EqualTo{T}
+    @test MOI.get(inner, MOI.NumberOfConstraints{F,S}()) == 2
+    @test MOI.get(inner, MOI.NumberOfVariables()) == 2
+    return
+end
+
 function runtests(x, y)
     for name in names(@__MODULE__; all = true)
         if startswith("$name", "test_")