Merge check_belongs_to_model with moi_function

Author: blegat

src/constraints.jl

@@ -760,6 +760,10 @@ function moi_function(constraint::AbstractConstraint)
     return moi_function(jump_function(constraint))
 end
 
+function moi_function(constraint::AbstractConstraint, model)
+    return moi_function(jump_function(constraint), model)
+end
+
 """
     moi_set(constraint::AbstractConstraint)
 
@@ -1016,6 +1020,11 @@ function _moi_add_constraint(
     return MOI.add_constraint(model, f, s)
 end
 
+function moi_function(f, model)
+    check_belongs_to_model(f, model)
+    return moi_function(f)
+end
+
 """
     add_constraint(
         model::GenericModel,
@@ -1032,10 +1041,9 @@ function add_constraint(
     name::String = "",
 )
     con = model_convert(model, con)
+    func, set = moi_function(con, model), moi_set(con)
     # The type of backend(model) is unknown so we directly redirect to another
     # function.
-    check_belongs_to_model(con, model)
-    func, set = moi_function(con), moi_set(con)
     cindex = _moi_add_constraint(
         backend(model),
         func,

src/nlp_expr.jl

@@ -552,24 +552,25 @@ end
 
 moi_function(x::Number) = x
 
-function moi_function(f::GenericNonlinearExpr{V}) where {V}
-    model = owner_model(f)
-    cache = isnothing(model) ? nothing : model.subexpressions
-    if !isnothing(cache) && haskey(cache, f)
+function moi_function(f::GenericNonlinearExpr{V}, model::JuMP.GenericModel) where {V}
+    cache = model.subexpressions
+    if haskey(cache, f)
         return cache[f]
     end
     ret = MOI.ScalarNonlinearFunction(f.head, similar(f.args))
     stack = Tuple{MOI.ScalarNonlinearFunction,Int,GenericNonlinearExpr{V}}[]
     for i in length(f.args):-1:1
         if f.args[i] isa GenericNonlinearExpr{V}
             push!(stack, (ret, i, f.args[i]))
+        elseif f.args[i] isa AbstractJuMPScalar
+            ret.args[i] = moi_function(model, f.args[i])
         else
             ret.args[i] = moi_function(f.args[i])
         end
     end
     while !isempty(stack)
         parent, i, arg = pop!(stack)
-        if !isnothing(cache) && haskey(cache, arg)
+        if haskey(cache, arg)
             parent.args[i] = cache[arg]
             continue
         end
@@ -582,13 +583,9 @@ function moi_function(f::GenericNonlinearExpr{V}) where {V}
                 child.args[j] = moi_function(arg.args[j])
             end
         end
-        if !isnothing(cache)
-            cache[arg] = child
-        end
-    end
-    if !isnothing(cache)
-        cache[f] = ret
+        cache[arg] = child
     end
+    cache[f] = ret
     return ret
 end