Fix for #511

src/compiler.jl

@@ -541,9 +541,7 @@
 """
 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,19 +552,18 @@
         # and is not our intent).
         @gensym retval
         return quote
-            $retval = $(e.args...)
+            $retval = $(map(replace_returns, e.args)...)
             return $retval, __varinfo__
         end
     end
 
     return Expr(e.head, map(replace_returns, e.args)...)
 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.
+    # Otherwise we replace the last statement with an altered `return`.
     if !Meta.isexpr(body.args[end], :return)
         body.args[end] = Expr(:return, body.args[end])
     end
@@ -626,7 +623,7 @@
     # 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