Merge 8d4fae6 into aff74fb

src/macros.jl

@@ -156,25 +156,65 @@ function parse_one_operator_constraint(_error::Function, vectorized::Bool,
     return parse_code, _build_call(_error, vectorized, variable, set)
 end
 
+# To be defined if such a function must be rewritten; otherwise,
+# rewrite_call_expression will never be called.
+expression_to_rewrite(head::Val, args...) = true # false
+
+# When `expression_to_rewrite` called with the same arguments save the first one,
+# returns three things:
+# - code for parsing (which must be called first), if needed; otherwise, :().
+# - code for building the required constraints, if needed; otherwise, :().
+# - the symbol of the variable that replaces the expression (<: VariableRef).
+function rewrite_call_expression(_error::Function, head::Val{F}, args...) where F
+    return :(), :(), Expr(:call, F, args...)
+end
+
 _functionize(v::VariableRef) = convert(AffExpr, v)
 _functionize(v::AbstractArray{VariableRef}) = _functionize.(v)
 _functionize(x) = x
 _functionize(::MutableArithmetics.Zero) = 0.0
 function parse_one_operator_constraint(_error::Function, vectorized::Bool, sense::Val, lhs, rhs)
+    parse_code_rhs, build_code_rhs, new_rhs = :(), :(), rhs
+    parse_code_lhs, build_code_lhs, new_lhs = :(), :(), lhs
+    if isexpr(rhs, :call) && expression_to_rewrite(Val(rhs.args[1]), rhs.args[2:end]...)
+        parse_code_rhs, build_code_rhs, new_rhs = rewrite_call_expression(_error, Val(rhs.args[1]), rhs.args[2:end]...)
+    end
+    if isexpr(lhs, :call) && expression_to_rewrite(Val(lhs.args[1]), lhs.args[2:end]...)
+        parse_code_lhs, build_code_lhs, new_lhs = rewrite_call_expression(_error, Val(lhs.args[1]), lhs.args[2:end]...)
+    end
+
     # Simple comparison - move everything to the LHS.
     #
     # `_functionize` deals with the pathological case where the `lhs` is a `VariableRef`
     # and the `rhs` is a summation with no terms. `_build_call` should be passed a
     # `GenericAffExpr` or a `GenericQuadExpr`, and not a `VariableRef` as would be the case
     # without `_functionize`.
-    if vectorized
-        func = :($lhs .- $rhs)
+
+    # TODO: bug in MutableArithmetics? The returned code does not find $new_rhs (ERROR: UndefVarError: #642###976 not defined) when only using the first code path.
+    parse_code, variable = nothing, nothing
+    if rhs == new_rhs && lhs == new_lhs
+        if vectorized
+            func = :($lhs .- $rhs)
+        else
+            func = :($lhs - $rhs)
+        end
+        variable, parse_code = _MA.rewrite(func)
+    elseif rhs != new_rhs && lhs == new_lhs
+        variable, parse_code = _MA.rewrite(lhs)
+        parse_code = :($parse_code; $variable = _MA.operate!(-, $variable, $new_rhs))
+    elseif lhs != new_lhs && rhs == new_rhs
+        variable, parse_code = _MA.rewrite(rhs)
+        parse_code = :($parse_code; $variable = _MA.operate!(-, $new_lhs, $variable))
     else
-        func = :($lhs - $rhs)
+        @assert rhs != new_rhs
+        @assert lhs != new_lhs
+        variable = gensym()
+        parse_code = :($parse_code; $variable = _MA.operate!(-, $new_lhs, $new_rhs))
     end
+
     set = sense_to_set(_error, sense)
-    variable, parse_code = _MA.rewrite(func)
-    return parse_code, _build_call(_error, vectorized, :(_functionize($variable)), set)
+    build_code = _build_call(_error, vectorized, :(_functionize($variable)), set)
+    return :($parse_code_lhs; $parse_code_rhs; $parse_code), :($build_code_lhs; $build_code_rhs; $build_code)
 end
 
 function parse_constraint_expr(_error::Function, expr::Expr)

test/macros.jl

@@ -176,6 +176,44 @@ function custom_expression_test(ModelType::Type{<:JuMP.AbstractModel})
     end
 end
 
+JuMP.expression_to_rewrite(head::Val{:donothing}, var) = true
+function JuMP.rewrite_call_expression(errorf::Function, head::Val{:donothing}, var)
+    return :(), :(), esc(var)
+end
+function build_constraint_test(ModelType::Type{<:JuMP.AbstractModel})
+    @testset "Extension of @constraint with rewrite_call_expression #2229" begin
+        # RHS
+        model = ModelType()
+        @variable(model, x)
+        @variable(model, y)
+        cref = @constraint(model, x == donothing(y))
+
+        c = JuMP.constraint_object(cref)
+        @test JuMP.isequal_canonical(c.func, x - y)
+        @test c.set == MOI.EqualTo(0.0)
+
+        # LHS
+        model = ModelType()
+        @variable(model, x)
+        @variable(model, y)
+        cref = @constraint(model, donothing(x) == y)
+
+        c = JuMP.constraint_object(cref)
+        @test JuMP.isequal_canonical(c.func, x - y)
+        @test c.set == MOI.EqualTo(0.0)
+
+        # Both sides.
+        model = ModelType()
+        @variable(model, x)
+        @variable(model, y)
+        cref = @constraint(model, donothing(x) == donothing(y))
+
+        c = JuMP.constraint_object(cref)
+        @test JuMP.isequal_canonical(c.func, x - y)
+        @test c.set == MOI.EqualTo(0.0)
+    end
+end
+
 function macros_test(ModelType::Type{<:JuMP.AbstractModel}, VariableRefType::Type{<:JuMP.AbstractVariableRef})
     @testset "build_constraint on variable" begin
         m = ModelType()
@@ -357,8 +395,8 @@ function macros_test(ModelType::Type{<:JuMP.AbstractModel}, VariableRefType::Typ
     end
 
     build_constraint_keyword_test(ModelType)
-
     custom_expression_test(ModelType)
+    build_constraint_test(ModelType)
 end
 
 @testset "Macros for JuMP.Model" begin