Fix for #511 (#512)

* fix for #511

* revert to previous behavior but with recursive application of
replace_returns in return statements

* use Val in the tests for return-values rather than something like
missing and nothing

* Update src/compiler.jl

* Update src/compiler.jl

* bumped patch version

Author: torfjelde

Project.toml

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

src/compiler.jl

@@ -541,9 +541,7 @@ definitions. This is checked using [`isfuncdef`](@ref).
 """
 replace_returns(e) = e
 function replace_returns(e::Expr)
-    if isfuncdef(e)
-        return e
-    end
+    isfuncdef(e) && return e
 
     if Meta.isexpr(e, :return)
         # We capture the original return-value in `retval` and return
@@ -554,7 +552,7 @@ function replace_returns(e::Expr)
         # and is not our intent).
         @gensym retval
         return quote
-            $retval = $(e.args...)
+            $retval = $(map(replace_returns, e.args)...)
             return $retval, __varinfo__
         end
     end
@@ -563,8 +561,8 @@ 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`.
-make_returns_explicit!(body) = Expr(:return, body)
-function make_returns_explicit!(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)
@@ -626,7 +624,7 @@ function build_output(modeldef, linenumbernode)
     # See the docstrings of `replace_returns` for more info.
     evaluatordef[:body] = MacroTools.@q begin
         $(linenumbernode)
-        $(replace_returns(make_returns_explicit!(modeldef[:body])))
+        $(replace_returns(add_return_to_last_statment!(modeldef[:body])))
     end
 
     ## Build the model function.

test/compiler.jl

@@ -599,6 +599,33 @@ end
         # With assignment.
         @model outer() = @submodel x = inner()
         @test outer()() isa Real
+
+        # Edge-cases.
+        # `return` in the last statement.
+        # Ref: issue #511.
+        @model function demo_ret_in_last_stmt(x::Bool)
+            # Two different values not supporting `iterate`.
+            if x
+                return Val(1)
+            else
+                return Val(2)
+            end
+        end
+
+        model_true = demo_ret_in_last_stmt(true)
+        @test model_true() === Val(1)
+
+        model_false = demo_ret_in_last_stmt(false)
+        @test model_false() === Val(2)
+
+        # `return` with `return`
+        @model function demo_ret_with_ret()
+            return begin
+                return Val(1)
+                Val(2)
+            end
+        end
+        @test demo_ret_with_ret()() === Val(1)
     end
 
     @testset "issue #368: hasmissing dispatch" begin