From 34952ab5e556407f507688cfad481521ad4eda3e Mon Sep 17 00:00:00 2001 From: odow Date: Tue, 12 Sep 2023 08:59:21 +1200 Subject: [PATCH 1/5] Use let model=model in variable macro to improve type stability --- src/macros.jl | 8 +++++++- test/test_hygiene.jl | 6 ++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/src/macros.jl b/src/macros.jl index 7660ab6576d..19bb34e427a 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -3016,7 +3016,13 @@ macro variable(args...) requested_container, ) end - + # Wrap the entire code block in a let statement to make the model act as + # a type stable local variable. + creation_code = quote + let $model = $model + $creation_code + end + end if anonvar # Anonymous variable, no need to register it in the model-level # dictionary nor to assign it to a variable in the user scope. diff --git a/test/test_hygiene.jl b/test/test_hygiene.jl index a7bc87a8ed5..09b82cb11ff 100644 --- a/test/test_hygiene.jl +++ b/test/test_hygiene.jl @@ -57,4 +57,10 @@ Test.@test ex[3] == 6 Test.@test i == 10 Test.@test j == 10 +# Test that `model` is inferred correctly inside macros, despite being a +# non-const global. +model = JuMP.Model() +JuMP.@variable(model, x[1:0]) +Test.@test x == JuMP.VariableRef[] + end From be6beb38326e2fc7c22e13aca148c5be4bfdae52 Mon Sep 17 00:00:00 2001 From: odow Date: Thu, 14 Sep 2023 10:11:49 +1200 Subject: [PATCH 2/5] Add to other macros --- src/macros.jl | 15 ++++++++++++++- test/test_hygiene.jl | 4 ++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/src/macros.jl b/src/macros.jl index 19bb34e427a..a74c7eb70e1 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -1334,7 +1334,13 @@ function _constraint_macro( creation_code = Containers.container_code(idxvars, indices, code, requested_container) - + # Wrap the entire code block in a let statement to make the model act as + # a type stable local variable. + creation_code = quote + let $model = $model + $creation_code + end + end if anonvar # Anonymous constraint, no need to register it in the model-level # dictionary nor to assign it to a variable in the user scope. @@ -1939,6 +1945,13 @@ macro expression(args...) end code = Containers.container_code(idxvars, indices, code, requested_container) + # Wrap the entire code block in a let statement to make the model act as + # a type stable local variable. + creation_code = quote + let $m = $m + $creation_code + end + end # don't do anything with the model, but check that it's valid anyway if anonvar macro_code = code diff --git a/test/test_hygiene.jl b/test/test_hygiene.jl index 09b82cb11ff..9ff08bb6edc 100644 --- a/test/test_hygiene.jl +++ b/test/test_hygiene.jl @@ -59,8 +59,8 @@ Test.@test j == 10 # Test that `model` is inferred correctly inside macros, despite being a # non-const global. -model = JuMP.Model() -JuMP.@variable(model, x[1:0]) +m = JuMP.Model() +JuMP.@variable(m, x[1:0]) Test.@test x == JuMP.VariableRef[] end From 17429fdfc4387dbf435108cf1900b2d30e7fd535 Mon Sep 17 00:00:00 2001 From: Oscar Dowson Date: Thu, 14 Sep 2023 10:37:26 +1200 Subject: [PATCH 3/5] Apply suggestions from code review --- src/macros.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/macros.jl b/src/macros.jl index a74c7eb70e1..4c94d09d3b6 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -1947,9 +1947,9 @@ macro expression(args...) Containers.container_code(idxvars, indices, code, requested_container) # Wrap the entire code block in a let statement to make the model act as # a type stable local variable. - creation_code = quote + code = quote let $m = $m - $creation_code + $code end end # don't do anything with the model, but check that it's valid anyway From 45a39c5aa1254c76cff71070cad73434eb56c157 Mon Sep 17 00:00:00 2001 From: odow Date: Thu, 14 Sep 2023 13:34:08 +1200 Subject: [PATCH 4/5] Update --- src/macros.jl | 30 +++++++++++++--------------- test/test_macros.jl | 48 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 16 deletions(-) diff --git a/src/macros.jl b/src/macros.jl index 4c94d09d3b6..afbd777b460 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -1204,6 +1204,17 @@ end # TODO: update 3-argument @constraint macro to pass through names like @variable +function _wrap_let(model, code) + if Meta.isexpr(model, :escape) && model.args[1] isa Symbol + return quote + let $model = $model + $code + end + end + end + return code +end + """ _constraint_macro( args, macro_name::Symbol, parsefun::Function, source::LineNumberNode @@ -1331,16 +1342,11 @@ function _constraint_macro( $parsecode $constraintcall end - creation_code = Containers.container_code(idxvars, indices, code, requested_container) # Wrap the entire code block in a let statement to make the model act as # a type stable local variable. - creation_code = quote - let $model = $model - $creation_code - end - end + creation_code = _wrap_let(model, creation_code) if anonvar # Anonymous constraint, no need to register it in the model-level # dictionary nor to assign it to a variable in the user scope. @@ -1947,11 +1953,7 @@ macro expression(args...) Containers.container_code(idxvars, indices, code, requested_container) # Wrap the entire code block in a let statement to make the model act as # a type stable local variable. - code = quote - let $m = $m - $code - end - end + code = _wrap_let(m, code) # don't do anything with the model, but check that it's valid anyway if anonvar macro_code = code @@ -3031,11 +3033,7 @@ macro variable(args...) end # Wrap the entire code block in a let statement to make the model act as # a type stable local variable. - creation_code = quote - let $model = $model - $creation_code - end - end + creation_code = _wrap_let(model, creation_code) if anonvar # Anonymous variable, no need to register it in the model-level # dictionary nor to assign it to a variable in the user scope. diff --git a/test/test_macros.jl b/test/test_macros.jl index 3c58375340b..62841c7e644 100644 --- a/test/test_macros.jl +++ b/test/test_macros.jl @@ -2034,4 +2034,52 @@ function test_unsupported_ternary_operator() return end +# This code needs to be evaluated in a top-level scope to prevent inference from +# knowing the type of `model`. +function test_wrap_let_non_symbol_models() + module_name = @eval module $(gensym()) + using JuMP, Test + data = (; model = Model()) + end + @eval module_name begin + @variable(data.model, x) + @test x isa VariableRef + @objective(data.model, Min, x^2) + @test isequal_canonical(objective_function(data.model), x^2) + @expression(data.model, expr[i = 1:2], x + i) + @test expr == [x + 1, x + 2] + @constraint(data.model, c[i = 1:2], i * expr[i] <= i) + @test c isa Vector{<:ConstraintRef} + @variable(data.model, bad_var[1:0]) + @test bad_var isa Vector{Any} + @expression(data.model, bad_expr[i = 1:0], x + i) + @test bad_expr isa Vector{Any} + end + return +end + +# This code needs to be evaluated in a top-level scope to prevent inference from +# knowing the type of `model`. +function test_wrap_let_symbol_models() + module_name = @eval module $(gensym()) + using JuMP, Test + model = Model() + end + @eval module_name begin + @variable(model, x) + @test x isa VariableRef + @objective(model, Min, x^2) + @test isequal_canonical(objective_function(model), x^2) + @expression(model, expr[i = 1:2], x + i) + @test expr == [x + 1, x + 2] + @constraint(model, c[i = 1:2], i * expr[i] <= i) + @test c isa Vector{<:ConstraintRef} + @variable(model, bad_var[1:0]) + @test bad_var isa Vector{VariableRef} + @expression(model, bad_expr[i = 1:0], x + i) + @test bad_expr isa Vector{Any} + end + return +end + end # module From 2cd6bca416d788807c7fdb971a4455f377c8f508 Mon Sep 17 00:00:00 2001 From: odow Date: Thu, 14 Sep 2023 14:08:58 +1200 Subject: [PATCH 5/5] Update --- test/test_macros.jl | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/test/test_macros.jl b/test/test_macros.jl index 62841c7e644..340a69f73df 100644 --- a/test/test_macros.jl +++ b/test/test_macros.jl @@ -2038,8 +2038,8 @@ end # knowing the type of `model`. function test_wrap_let_non_symbol_models() module_name = @eval module $(gensym()) - using JuMP, Test - data = (; model = Model()) + using JuMP, Test + data = (; model = Model()) end @eval module_name begin @variable(data.model, x) @@ -2051,7 +2051,8 @@ function test_wrap_let_non_symbol_models() @constraint(data.model, c[i = 1:2], i * expr[i] <= i) @test c isa Vector{<:ConstraintRef} @variable(data.model, bad_var[1:0]) - @test bad_var isa Vector{Any} + @test bad_var isa Vector{<:Any} + @test !(bad_var isa Vector{VariableRef}) # Cannot prove type @expression(data.model, bad_expr[i = 1:0], x + i) @test bad_expr isa Vector{Any} end @@ -2062,8 +2063,8 @@ end # knowing the type of `model`. function test_wrap_let_symbol_models() module_name = @eval module $(gensym()) - using JuMP, Test - model = Model() + using JuMP, Test + model = Model() end @eval module_name begin @variable(model, x)