Merge pull request #110 from hdavid16/jump_update

Update Macro and Printing with JuMP 1.18
hdavid16 · Feb 7, 2024 · eb1ce6c · eb1ce6c
2 parents 9ab25ae + ca2ae98
commit eb1ce6c
Show file tree

Hide file tree

Showing 6 changed files with 93 additions and 178 deletions.
diff --git a/Project.toml b/Project.toml
@@ -9,7 +9,7 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 
 [compat]
 Aqua = "0.8"
-JuMP = "1.17"
+JuMP = "1.18"
 Reexport = "1"
 julia = "1.6"
 

diff --git a/src/DisjunctiveProgramming.jl b/src/DisjunctiveProgramming.jl
@@ -10,6 +10,7 @@ using Base.Meta
 # Create aliases
 import JuMP.MOI as _MOI
 import JuMP.MOIU.CleverDicts as _MOIUC
+import JuMP.Containers as _JuMPC
 
 # Load in the source files
 include("datatypes.jl")

diff --git a/src/macros.jl b/src/macros.jl
@@ -1,109 +1,29 @@
 ################################################################################
 #                                BASIC HELPERS
 ################################################################################
-# Macro error function
-# inspired from https://github.com/jump-dev/JuMP.jl/blob/709d41b78e56efb4f2c54414266b30932010bd5a/src/macros.jl#L923-L928
-function _macro_error(macroname, args, source, str...)
-    error("At $(source.file):$(source.line): `@$macroname($(join(args, ", ")))`: ", 
-          str...)
-end
-
 # Escape when needed
 # taken from https://github.com/jump-dev/JuMP.jl/blob/709d41b78e56efb4f2c54414266b30932010bd5a/src/macros.jl#L895-L897
 _esc_non_constant(x::Number) = x
 _esc_non_constant(x::Expr) = isexpr(x,:quote) ? x : esc(x)
 _esc_non_constant(x) = esc(x)
 
-# Extract the name from a macro expression 
-# Inspired from https://github.com/jump-dev/JuMP.jl/blob/45ce630b51fb1d72f1ff8fed35a887d84ef3edf7/src/Containers/macro.jl#L8-L17
-_get_name(c::Symbol) = c
-_get_name(c::Nothing) = ()
-_get_name(c::AbstractString) = c
-function _get_name(c::Expr)
-    if isexpr(c, :string)
-        return c
-    else
-        return c.args[1]
-    end
-end
-
-# Given a base_name and idxvars, returns an expression that constructs the name
-# of the object.
-# Inspired from https://github.com/jump-dev/JuMP.jl/blob/709d41b78e56efb4f2c54414266b30932010bd5a/src/macros.jl#L930-L946
-function _name_call(base_name, idxvars)
-    if isempty(idxvars) || base_name == ""
-        return base_name
-    end
-    ex = Expr(:call, :string, base_name, "[")
-    for i in eachindex(idxvars)
-        # Converting the arguments to strings before concatenating is faster:
-        # https://github.com/JuliaLang/julia/issues/29550.
-        esc_idxvar = esc(idxvars[i])
-        push!(ex.args, :(string($esc_idxvar)))
-        i < length(idxvars) && push!(ex.args, ",")
-    end
-    push!(ex.args, "]")
-    return ex
-end
-
-# Process macro arguments 
-function _extract_kwargs(args)
-    arg_list = collect(args)
-    if !isempty(args) && isexpr(args[1], :parameters)
-        p = popfirst!(arg_list)
-        append!(arg_list, (Expr(:(=), a.args...) for a in p.args))
-    end
-    extra_kwargs = filter(x -> isexpr(x, :(=)) && x.args[1] != :container &&
-                          x.args[1] != :base_name, arg_list)
-    container_type = :Auto
-    base_name = nothing
-    for kw in arg_list
-        if isexpr(kw, :(=)) && kw.args[1] == :container
-            container_type = kw.args[2]
-        elseif isexpr(kw, :(=)) && kw.args[1] == :base_name
-            base_name = esc(kw.args[2])
-        end
-    end
-    pos_args = filter!(x -> !isexpr(x, :(=)), arg_list)
-    return pos_args, extra_kwargs, container_type, base_name
-end
-
-# Add on keyword arguments to a function call expression and escape the expressions
-# Adapted from https://github.com/jump-dev/JuMP.jl/blob/d9cd5fb16c2d0a7e1c06aa9941923492fc9a28b5/src/macros.jl#L11-L36
-function _add_kwargs(call, kwargs)
-    for kw in kwargs
-        push!(call.args, esc(Expr(:kw, kw.args...)))
-    end
-    return
-end
-
-# Add on positional args to a function call and escape
-# Adapted from https://github.com/jump-dev/JuMP.jl/blob/a325eb638d9470204edb2ef548e93e59af56cc19/src/macros.jl#L57C1-L65C4
-function _add_positional_args(call, args)
-    kw_args = filter(arg -> isexpr(arg, :kw), call.args)
-    filter!(arg -> !isexpr(arg, :kw), call.args)
-    for arg in args
-        push!(call.args, esc(arg))
-    end
-    append!(call.args, kw_args)
-    return
-end
-
 # Ensure a model argument is valid
 # Inspired from https://github.com/jump-dev/JuMP.jl/blob/d9cd5fb16c2d0a7e1c06aa9941923492fc9a28b5/src/macros.jl#L38-L44
-function _valid_model(_error::Function, model, name)
-    is_gdp_model(model) || _error("$name is not a `GDPModel`.")
+_valid_model(error_fn::Function, model::JuMP.AbstractModel, name) = nothing
+function _valid_model(error_fn::Function, model, name)
+    error_fn("Expected $name to be an `JuMP.AbstractModel`, but it has type ", 
+           typeof(model))
 end
 
 # Check if a macro julia variable can be registered 
 # Adapted from https://github.com/jump-dev/JuMP.jl/blob/d9cd5fb16c2d0a7e1c06aa9941923492fc9a28b5/src/macros.jl#L66-L86
 function _error_if_cannot_register(
-    _error::Function, 
-    model, 
+    error_fn::Function, 
+    model::JuMP.AbstractModel, 
     name::Symbol
     )
     if haskey(JuMP.object_dictionary(model), name)
-        _error("An object of name $name is already attached to this model. If ",
+       error_fn("An object of name $name is already attached to this model. If ",
                "this is intended, consider using the anonymous construction ",
                "syntax, e.g., `x = @macro_name(model, ...)` where the name ",
                "of the object does not appear inside the macro. Alternatively, ",
@@ -114,22 +34,35 @@ function _error_if_cannot_register(
     end
     return
 end
-
-# Update the creation code to register and assign the object to the name
-# Inspired from https://github.com/jump-dev/JuMP.jl/blob/d9cd5fb16c2d0a7e1c06aa9941923492fc9a28b5/src/macros.jl#L88-L120
-function _macro_assign_and_return(_error, code, name, model)
-    return quote
-        _error_if_cannot_register($_error, $model, $(quot(name)))
-        $(esc(name)) = $code
-        $model[$(quot(name))] = $(esc(name))
-    end
+function _error_if_cannot_register(error_fn::Function, ::JuMP.AbstractModel, name)
+    error_fn("Invalid name `$name`.")
 end
 
-# Wrap the macro generated code for better stacttraces (assumes model is escaped)
-# Inspired from https://github.com/jump-dev/JuMP.jl/blob/d9cd5fb16c2d0a7e1c06aa9941923492fc9a28b5/src/macros.jl#L46-L64
-function _finalize_macro(_error, model, code, source::LineNumberNode)
-    return Expr(:block, source, 
-                :(_valid_model($_error, $model, $(quot(model.args[1])))), code)
+# Inspired from https://github.com/jump-dev/JuMP.jl/blob/246cccb0d3167d5ed3df72fba97b1569476d46cf/src/macros.jl#L332-L377
+function _finalize_macro(
+    error_fn::Function,
+    model::Expr,
+    code::Any,
+    source::LineNumberNode,
+    register_name::Union{Nothing,Symbol}
+    )
+    @assert Meta.isexpr(model, :escape)
+    if model.args[1] isa Symbol
+        code = quote
+            let $model = $model
+                $code
+            end
+        end
+    end
+    if register_name !== nothing
+        sym_name = Meta.quot(register_name)
+        code = quote
+            _error_if_cannot_register($error_fn, $model, $sym_name)
+            $(esc(register_name)) = $model[$sym_name] = $code
+        end
+    end
+    is_valid_code = :(_valid_model($error_fn, $model, $(Meta.quot(model.args[1]))))
+    return Expr(:block, source, is_valid_code, code)
 end
 
 ################################################################################
@@ -162,21 +95,23 @@ To create disjunctions without macros, see [`disjunction`](@ref).
 """
 macro disjunction(args...)
     # define error message function
-
-    _error(str...) = _macro_error(:disjunction, (args...,),
-                                  __source__, str...)
+    error_fn = _JuMPC.build_error_fn(:disjunction, args, __source__)
 
-    # parse the arguments
-    pos_args, extra_kwargs, container_type, base_name = _extract_kwargs(args)
+    # process the inputs
+    pos_args, kwargs = _JuMPC.parse_macro_arguments(
+        error_fn, 
+        args,
+        num_positional_args = 2:4,
+        valid_kwargs = [:container, :base_name, :exactly1]
+    )
 
     # initial processing of positional arguments
-    length(pos_args) >= 2 || _error("Not enough arguments, please see docs for accepted `@disjunction` syntax..")
-    model = popfirst!(pos_args)
-    esc_model = esc(model)
+    model_sym = popfirst!(pos_args)
+    model = esc(model_sym)
     y = first(pos_args)
     extra = pos_args[2:end]
     if isexpr(args[2], :block)
-        _error("Invalid syntax. Did you mean to use `@disjunctions`?")
+        error_fn("Invalid syntax. Did you mean to use `@disjunctions`?")
     end
 
     # TODO: three cases lead to problems when julia variables are used for Disjunct tags
@@ -201,72 +136,50 @@ macro disjunction(args...)
 
     # Case 8
     if isexpr(y, :ref) && (isempty(extra) || isexpr(extra[1], :call)) 
-        c = gensym()
+        c = nothing
         x = _esc_non_constant(y)
-        is_anon = true
     # Cases 2, 3, 5
     elseif isexpr(y, (:vcat, :ref, :typed_vcat))
-        length(extra) >= 1 || _error("No disjunction expression was given, please see docs for accepted `@disjunction` syntax..")
+        length(extra) >= 1 || error_fn("No disjunction expression was given, please see docs for accepted `@disjunction` syntax..")
         c = y
         x = _esc_non_constant(popfirst!(extra))
-        is_anon = isexpr(y, :vcat) # check for Case 5
     # Cases 1, 4
     elseif (isa(y, Symbol) || isexpr(y, :vect)) && 
         !isempty(extra) && 
         (isa(extra[1], Symbol) || isexpr(extra[1], (:vect, :comprehension, :ref)))
         c = y
         x = _esc_non_constant(popfirst!(extra))
-        is_anon = isexpr(y, :vect) # check for Case 4
     # Cases 6, 7
     elseif isa(y, Symbol) || isexpr(y, (:vect, :comprehension))
-        c = gensym()
+        c = nothing
         x = _esc_non_constant(y)
-        is_anon = true
     # Case 9
     else
-        _error("Unrecognized syntax, please see docs for accepted `@disjunction` syntax.")
+        error_fn("Unrecognized syntax, please see docs for accepted `@disjunction` syntax.")
     end
 
+    # make sure param is something reasonable (I don't think this is needed)
+    # if !(c isa Union{Nothing, Symbol, Expr})
+    #     error_fn("Expected `$c` to be a disjunction name.")
+    # end
+
+    # process the container input
+    name, idxvars, inds = Containers.parse_ref_sets(
+        error_fn,
+        c;
+        invalid_index_variables = [model_sym],
+    )
+
     # process the name
-    name = _get_name(c)
-    if isnothing(base_name)
-        base_name = is_anon ? "" : string(name)
-    end
-    if !isa(name, Symbol) && !is_anon
-        _error("Expression $name should not be used as a disjunction name. Use " *
-               "the \"anonymous\" syntax $name = @disjunction(model, " *
-               "...) instead.")
-    end
+    name_expr = _JuMPC.build_name_expr(name, idxvars, kwargs)
 
     # make the creation code
-    if isa(c, Symbol)
-        # easy case with single parameter
-        creation_code = :( _disjunction($_error, $esc_model, $x, $base_name) )
-        _add_positional_args(creation_code, extra)
-        _add_kwargs(creation_code, extra_kwargs)
-    else
-        # we have a container of parameters
-        idxvars, inds = JuMP.Containers.build_ref_sets(_error, c)
-        if model in idxvars
-            _error("Index $(model) is the same symbol as the model. Use a ",
-                   "different name for the index.")
-        end
-        name_code = _name_call(base_name, idxvars)
-        disjunction_call = :( _disjunction($_error, $esc_model, $x, $name_code) )
-        _add_positional_args(disjunction_call, extra)
-        _add_kwargs(disjunction_call, extra_kwargs)
-        creation_code = JuMP.Containers.container_code(idxvars, inds, disjunction_call,
-                                                       container_type)
-    end
+    creation_code = :( _disjunction($error_fn, $model, $x, $name_expr) )
+    _JuMPC.add_additional_args(creation_code, extra, kwargs; kwarg_exclude = [:container, :base_name])
+    code = _JuMPC.container_code(idxvars, inds, creation_code, kwargs)
 
     # finalize the macro
-    if is_anon
-        macro_code = creation_code
-    else
-        macro_code = _macro_assign_and_return(_error, creation_code, name,
-                                              esc_model)
-    end
-    return _finalize_macro(_error, esc_model, macro_code, __source__)
+    return _finalize_macro(error_fn, model, code, __source__, name)
 end
 
 # Pluralize the @disjunction macro

diff --git a/src/print.jl b/src/print.jl
@@ -20,6 +20,19 @@ _card_func_str(::MIME"text/latex", ::_MOIAtMost) = "\\text{atmost}"
 _card_func_str(::MIME"text/plain", ::_MOIExactly) = "exactly"
 _card_func_str(::MIME"text/latex", ::_MOIExactly) = "\\text{exactly}"
 
+################################################################################
+#                               LOGICAL PRINTING
+################################################################################
+# Extend op_string to correct print _LogicalExpr
+JuMP.op_string(::MIME"text/latex", ::_LogicalExpr, ::Val{:!}) =  "\\neg"
+JuMP.op_string(::MIME"text/latex", ::_LogicalExpr, ::Val{:(==)}) =  "\\iff"
+JuMP.op_string(::MIME"text/latex", ::_LogicalExpr, ::Val{:(=>)}) =  "\\implies"
+JuMP.op_string(::MIME"text/plain", ::_LogicalExpr, ::Val{:&&}) =  Sys.iswindows() ? "and" : "∧"
+JuMP.op_string(::MIME"text/plain", ::_LogicalExpr, ::Val{:||}) =  Sys.iswindows() ? "or" : "∨"
+JuMP.op_string(::MIME"text/plain", ::_LogicalExpr, ::Val{:!}) =  Sys.iswindows() ? "!" : "¬"
+JuMP.op_string(::MIME"text/plain", ::_LogicalExpr, ::Val{:(==)}) =  Sys.iswindows() ? "<-->" : "⟺"
+JuMP.op_string(::MIME"text/plain", ::_LogicalExpr, ::Val{:(=>)}) =  Sys.iswindows() ? "-->" : "⟹"
+
 ################################################################################
 #                               CONSTRAINT PRINTING
 ################################################################################
@@ -76,27 +89,12 @@ function JuMP.constraint_string(
     cref::LogicalConstraintRef; 
     in_math_mode = false
     )
-    constr_str = JuMP.constraint_string(
+    return JuMP.constraint_string(
         mode, 
         JuMP.name(cref), 
         JuMP.constraint_object(cref); 
         in_math_mode = in_math_mode
     )
-    # temporary hack until JuMP provides a better solution for operator printing
-    # TODO improve the printing of implications (not recognized by JuMP as two-sided operators)
-    if mode == MIME("text/latex")
-        for p in ("&&" => "\\wedge", "||" => "\\vee", "\\textsf{!}" => "\\neg", 
-                  "==" => "\\iff", "\\textsf{=>}" => "\\implies")
-            constr_str = replace(constr_str, p)
-        end
-        return constr_str
-    else
-        constr_str = replace(constr_str, "&&" => Sys.iswindows() ? "and" : "∧")
-        constr_str = replace(constr_str, "||" => Sys.iswindows() ? "or" : "∨")
-        constr_str = replace(constr_str, "!" => Sys.iswindows() ? "!" : "¬")
-        constr_str = replace(constr_str, "==" => Sys.iswindows() ? "<-->" : "⟺")
-        return replace(constr_str, "=>" => Sys.iswindows() ? "-->" : "⟹")
-    end
 end
 
 # Return the string of a Disjunction for plain printing

diff --git a/test/constraints/disjunction.jl b/test/constraints/disjunction.jl
@@ -1,11 +1,6 @@
 function test_macro_helpers()
     @test DP._esc_non_constant(1) == 1
-    @test DP._get_name(:x) == :x
-    @test DP._get_name("x") == "x"
-    @test DP._get_name(nothing) == ()
-    @test DP._get_name(Expr(:string,"x")) == Expr(:string,"x")
-    @test DP._name_call("",[]) == ""
-    @test DP._name_call("name",[]) == "name"
+    @test_throws ErrorException DP._error_if_cannot_register(error, Model(), 42)
 end
 
 function test_disjunction_add_fail()
@@ -15,9 +10,11 @@ function test_disjunction_add_fail()
     @constraint(model, x == 5, Disjunct(y[1]))
 
     @test_macro_throws ErrorException @disjunction(model) #not enough arguments
+    @test_macro_throws ErrorException @disjunction(42, y) #a model was not given
+    @test_macro_throws ErrorException @disjunction(model, "bob"[i = 1:1], y[1:1]) #bad name
     @test_macro_throws UndefVarError @disjunction(model, y) #unassociated indicator
     @test_macro_throws UndefVarError @disjunction(GDPModel(), y) #wrong model
-    @test_macro_throws ErrorException @disjunction(Model(), y) #not a GDPModel
+    @test_throws ErrorException disjunction(Model(), y) #not a GDPModel
     @test_macro_throws UndefVarError @disjunction(model, [y[1], y[1]]) #duplicate indicator
     @test_macro_throws UndefVarError @disjunction(model, y[1]) #unrecognized disjunction expression
     @test_throws ErrorException disjunction(model, y[1]) #unrecognized disjunction expression
@@ -30,6 +27,7 @@ function test_disjunction_add_fail()
 
     @constraint(model, x == 10, Disjunct(y[2]))
     @disjunction(model, disj, y)
+    @test model[:disj] isa DisjunctionRef
     @test_macro_throws UndefVarError @disjunction(model, disj, y) #duplicate name
     @test_throws ErrorException DP._error_if_cannot_register(error, model, :disj) #duplicate name