Fix for #537 (#538)

* add_return_to_statement is no longer mutating, thus fixing #537

* bumped patch version

* formatting

Author: torfjelde

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.23.17"
+version = "0.23.18"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/compiler.jl

@@ -561,14 +561,15 @@ function replace_returns(e::Expr)
 end
 
 # If it's just a symbol, e.g. `f(x) = 1`, then we make it `f(x) = return 1`.
-add_return_to_last_statment!(body) = Expr(:return, body)
-function add_return_to_last_statment!(body::Expr)
+add_return_to_last_statment(body) = Expr(:return, body)
+function add_return_to_last_statment(body::Expr)
     # If the last statement is a return-statement, we don't do anything.
     # Otherwise we replace the last statement with a `return` statement.
-    if !Meta.isexpr(body.args[end], :return)
-        body.args[end] = Expr(:return, body.args[end])
-    end
-    return body
+    Meta.isexpr(body.args[end], :return) && return body
+    # We need to copy the arguments since we are modifying them.
+    new_args = copy(body.args)
+    new_args[end] = Expr(:return, body.args[end])
+    return Expr(body.head, new_args...)
 end
 
 const FloatOrArrayType = Type{<:Union{AbstractFloat,AbstractArray}}
@@ -602,7 +603,7 @@ function build_output(modeldef, linenumbernode)
     kwargs = modeldef[:kwargs]
 
     ## Build the anonymous evaluator from the user-provided model definition.
-    evaluatordef = deepcopy(modeldef)
+    evaluatordef = copy(modeldef)
 
     # Add the internal arguments to the user-specified arguments (positional + keywords).
     evaluatordef[:args] = vcat(
@@ -624,7 +625,7 @@ function build_output(modeldef, linenumbernode)
     # See the docstrings of `replace_returns` for more info.
     evaluatordef[:body] = MacroTools.@q begin
         $(linenumbernode)
-        $(replace_returns(add_return_to_last_statment!(modeldef[:body])))
+        $(replace_returns(add_return_to_last_statment(modeldef[:body])))
     end
 
     ## Build the model function.

test/compiler.jl

@@ -33,6 +33,8 @@ struct MyCoolStruct{T}
     a::T
 end
 
+module Issue537 end
+
 @testset "compiler.jl" begin
     @testset "model macro" begin
         @model function testmodel_comp(x, y)
@@ -675,4 +677,16 @@ end
         res = f_splat_test_2(1)()
         @test res == (1, (), 1, Int, NamedTuple())
     end
+
+    @testset "issue #537: model with logging" begin
+        # Make sure `Module` is valid to put in a model.
+        @model demo_with_module() = Issue537
+        model = demo_with_module()
+        @test model() === Issue537
+
+        # And one explicit test for logging so know that is working.
+        @model demo_with_logging() = @info "hi"
+        model = demo_with_logging()
+        @test model() == nothing
+    end
 end