1 files changed, 51 insertions, 29 deletions

diff --git a/boot/type_check.cc b/boot/type_check.cc
index da60a21..cdc78df 100644
--- a/boot/type_check.cc
+++ b/boot/type_check.cc
@@ -114,6 +114,12 @@ namespace elna::boot
         }, this->payload);
     }
 
+    /// The program body is the only callable without a name.
+    static std::string describe_callable(const std::string& identifier)
+    {
+        return identifier.empty() ? "Program body" : "Procedure '" + identifier + "'";
+    }
+
     type_mismatch_error::type_mismatch_error(const source_position position,
             type actual, payload_type payload)
         : diagnostic(position), actual(std::move(actual)), payload(std::move(payload))
@@ -140,12 +146,12 @@ namespace elna::boot
             {
                 if (!this->actual.empty())
                 {
-                    return "Procedure '" + payload.identifier
-                        + "' does not return a value, but return expression has type '"
+                    return describe_callable(payload.identifier)
+                        + " does not return a value, but return expression has type '"
                         + this->actual.to_string() + "'";
                 }
-                return "Procedure '" + payload.identifier
-                    + "' is expected to return, but does not have a return statement";
+                return describe_callable(payload.identifier)
+                    + " is expected to return, but does not have a return statement";
 
             }
             else if constexpr (std::is_same_v<T, unary>)
@@ -660,35 +666,41 @@ namespace elna::boot
         }
         if (declaration->body.has_value())
         {
-            if (declaration->body.value().return_expression != nullptr)
-            {
-                const expression *return_expr = declaration->body.value().return_expression;
-                type const return_type = this->current_procedure->symbol.return_type.proper_type;
+            check_return(declaration->body.value(), declaration->position(),
+                    declaration->identifier.name());
+            this->bag.leave();
+        }
+        this->current_procedure.reset();
+    }
 
-                if (!return_type.empty())
-                {
-                    if (!is_assignable_from(return_type, return_expr->type_decoration))
-                    {
-                        add_error<type_mismatch_error>(return_expr->position(),
-                                return_expr->type_decoration,
-                                type_mismatch_error::expected_type{ return_type });
-                    }
-                }
-                else
-                {
-                    add_error<type_mismatch_error>(return_expr->position(),
-                            return_expr->type_decoration,
-                            type_mismatch_error::return_type{ .identifier = declaration->identifier.name() });
-                }
-            }
-            else if (declaration->heading().return_type.proper_type != nullptr)
+    void type_analysis_visitor::check_return(const procedure_body& body,
+            const source_position position, const std::string& name)
+    {
+        const type return_type = this->current_procedure->symbol.return_type.proper_type;
+
+        if (body.return_expression == nullptr)
+        {
+            if (!return_type.empty())
             {
-                add_error<type_mismatch_error>(declaration->position(), type(),
-                        type_mismatch_error::return_type{ .identifier = declaration->identifier.name() });
+                add_error<type_mismatch_error>(position, type(),
+                        type_mismatch_error::return_type{ .identifier = name });
             }
-            this->bag.leave();
+            return;
+        }
+        const expression *return_expr = body.return_expression;
+
+        if (return_type.empty())
+        {
+            add_error<type_mismatch_error>(return_expr->position(),
+                    return_expr->type_decoration,
+                    type_mismatch_error::return_type{ .identifier = name });
+        }
+        else if (!is_assignable_from(return_type, return_expr->type_decoration))
+        {
+            add_error<type_mismatch_error>(return_expr->position(),
+                    return_expr->type_decoration,
+                    type_mismatch_error::expected_type{ return_type });
         }
-        this->current_procedure.reset();
     }
 
     void type_analysis_visitor::visit(assign_statement *statement)
@@ -714,6 +726,16 @@ namespace elna::boot
         }
     }
 
+    void type_analysis_visitor::visit_entry_point(unit *unit)
+    {
+        this->current_procedure = this->bag.lookup("")->is_procedure();
+        this->bag.enter(this->current_procedure->scope);
+        walking_visitor::visit_entry_point(unit);
+        check_return(unit->entry_point.value(), unit->entry_position.value(), "");
+        this->bag.leave();
+        this->current_procedure.reset();
+    }
+
     void type_analysis_visitor::visit(variable_declaration *declaration)
     {
         walking_visitor::visit(declaration);