From 92c2835f5b5bd6c182172ed2756d8ab93df568aa Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 4 Aug 2023 17:38:35 +0100
Subject: [PATCH] fix for #511

---
 src/compiler.jl  | 15 ++++++---------
 test/compiler.jl | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/src/compiler.jl b/src/compiler.jl
index 96e98938b..cf3aaf7a9 100644
--- a/src/compiler.jl
+++ b/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
@@ -562,13 +560,12 @@ function replace_returns(e::Expr)
     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])
+        body.args[end] = replace_returns(Expr(:return, body.args[end]))
     end
     return body
 end
@@ -626,7 +623,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])))
+        $(add_return_to_last_statment!(replace_returns(modeldef[:body])))
     end
 
     ## Build the model function.
diff --git a/test/compiler.jl b/test/compiler.jl
index 53c071b2e..39bd5e414 100644
--- a/test/compiler.jl
+++ b/test/compiler.jl
@@ -599,6 +599,24 @@ 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 nothing
+            else
+                return missing
+            end
+        end
+
+        model_true = demo_ret_in_last_stmt(true)
+        @test model_true() === nothing
+
+        model_false = demo_ret_in_last_stmt(false)
+        @test model_false() === missing
     end
 
     @testset "issue #368: hasmissing dispatch" begin