Fix MA.add_mul for GenericNonlinearExpr arguments (#3488)

src/mutable_arithmetics.jl

@@ -348,3 +348,21 @@ function _MA.sub_mul(
     end
     return _MA.operate!(_MA.sub_mul, expr, x, y, args...)
 end
+
+# The default implementation of `add_mul` falls back to Base.muladd, which ends
+# up rewriting `x + y * z` as `y * z + x`. Intercept some methods, ensuring that
+# we don't cause ambiguities.
+
+for F in (:_Scalar, :AbstractJuMPScalar)
+    @eval begin
+        _MA.add_mul(x::$F, y::GenericNonlinearExpr, z::_Scalar) = x + y * z
+        _MA.add_mul(x::$F, y::_Scalar, z::GenericNonlinearExpr) = x + y * z
+        function _MA.add_mul(
+            x::$F,
+            y::GenericNonlinearExpr,
+            z::GenericNonlinearExpr,
+        )
+            return x + y * z
+        end
+    end
+end

test/test_nlp_expr.jl

@@ -174,6 +174,22 @@ function test_extension_expression_addmul(
     return
 end
 
+function test_extension_expression_explicit_add_mul(
+    ModelType = Model,
+    VariableRefType = VariableRef,
+)
+    model = ModelType()
+    @variable(model, x)
+    f = sin(x)
+    @test string(MA.operate!!(MA.add_mul, 1, 2, f)) == "1.0 + (2.0 * $f)"
+    @test string(MA.operate!!(MA.add_mul, 1, f, 2)) == "1.0 + ($f * 2.0)"
+    @test string(MA.operate!!(MA.add_mul, 1, f, f)) == "1.0 + ($f * $f)"
+    @test string(MA.operate!!(MA.add_mul, f, 2, f)) == "$f + (2.0 * $f)"
+    @test string(MA.operate!!(MA.add_mul, f, f, 2)) == "$f + ($f * 2.0)"
+    @test string(MA.operate!!(MA.add_mul, f, f, f)) == "$f + ($f * $f)"
+    return
+end
+
 function test_extension_expression_submul(
     ModelType = Model,
     VariableRefType = VariableRef,