Make return statement part of the body and not a statement

1 files changed, 49 insertions, 30 deletions

diff --git a/boot/semantic.cc b/boot/semantic.cc
index 3486830..97c316a 100644
--- a/boot/semantic.cc
+++ b/boot/semantic.cc
@@ -139,7 +139,6 @@ namespace elna::boot
     void type_analysis_visitor::visit(procedure_declaration *declaration)
     {
         this->current_procedure = this->bag.lookup(declaration->identifier.name)->is_procedure();
-        this->returns = false;
 
         if (declaration->body.has_value())
         {
@@ -149,41 +148,53 @@ namespace elna::boot
 
         if (declaration->body.has_value())
         {
-            this->bag.leave();
-            if (!this->returns && declaration->heading().return_type.proper_type != nullptr)
+            if (declaration->body.value().return_expression != nullptr)
+            {
+                expression *return_expr = declaration->body.value().return_expression;
+                type return_type = this->current_procedure->symbol.return_type.proper_type;
+
+                if (!return_type.empty())
+                {
+                    if (!is_assignable_from(return_type, return_expr->type_decoration))
+                    {
+                        add_error<type_expectation_error>(type_expectation_error::kind::result, return_expr->position());
+                    }
+                }
+                else
+                {
+                    add_error<type_expectation_error>(type_expectation_error::kind::result, return_expr->position());
+                }
+            }
+            else if (declaration->heading().return_type.proper_type != nullptr)
             {
                 add_error<return_error>(declaration->identifier.name, declaration->position());
             }
+            this->bag.leave();
         }
         this->current_procedure.reset();
     }
 
-    void type_analysis_visitor::visit(return_statement *statement)
+    void type_analysis_visitor::visit(unit *unit)
     {
-        walking_visitor::visit(statement);
-        type return_type;
+        walking_visitor::visit(unit);
 
-        if (this->current_procedure == nullptr)
+        if (unit->has_body())
         {
-            // In the main function.
-            return_type = this->bag.lookup("Int")->is_type()->symbol;
-        }
-        else
-        {
-            return_type = this->current_procedure->symbol.return_type.proper_type;
-        }
-        if (!return_type.empty())
-        {
-            if (!is_assignable_from(return_type, statement->return_expression().type_decoration))
+            if (unit->return_expression != nullptr)
             {
-                add_error<type_expectation_error>(type_expectation_error::kind::result, statement->position());
+                type return_type = this->bag.lookup("Int")->is_type()->symbol;
+
+                if (!is_assignable_from(return_type, unit->return_expression->type_decoration))
+                {
+                    add_error<type_expectation_error>(type_expectation_error::kind::result,
+                            unit->return_expression->position());
+                }
+            }
+            else
+            {
+                add_error<return_error>("module", unit->position());
             }
         }
-        else
-        {
-            add_error<type_expectation_error>(type_expectation_error::kind::result, statement->position());
-        }
-        this->returns = true;
     }
 
     void type_analysis_visitor::visit(assign_statement *statement)
@@ -641,18 +652,22 @@ namespace elna::boot
                 ++name_iterator;
                 ++type_iterator;
             }
-            for (constant_declaration *const constant : declaration->body.value().constants())
+            for (constant_declaration *const constant : declaration->body.value().constants)
             {
                 constant->accept(this);
             }
-            for (variable_declaration *const variable : declaration->body.value().variables())
+            for (variable_declaration *const variable : declaration->body.value().variables)
             {
                 variable->accept(this);
             }
-            for (statement *const statement : declaration->body.value().statements())
+            for (statement *const statement : declaration->body.value().entry_point)
             {
                 statement->accept(this);
             }
+            if (declaration->body.value().return_expression != nullptr)
+            {
+                declaration->body.value().return_expression->accept(this);
+            }
             this->bag.leave();
         }
         else
@@ -698,7 +713,7 @@ namespace elna::boot
         {
             procedure->accept(this);
         }
-        if (unit->entry_point.has_value())
+        if (unit->has_body())
         {
             this->bag.enter();
             auto variable_type = lookup_primitive_type("Int");
@@ -709,10 +724,14 @@ namespace elna::boot
             variable_type = type(std::make_shared<pointer_type>(variable_type));
             this->bag.enter("parameters", std::make_shared<variable_info>(variable_type, false));
 
-            for (statement *const statement : unit->entry_point.value())
+            for (statement *const statement : unit->entry_point)
             {
                 statement->accept(this);
             }
+            if (unit->return_expression != nullptr)
+            {
+                unit->return_expression->accept(this);
+            }
             this->bag.leave();
         }
     }
@@ -942,11 +961,11 @@ namespace elna::boot
         {
             return;
         }
-        for (constant_declaration *const constant : declaration->body.value().constants())
+        for (constant_declaration *const constant : declaration->body.value().constants)
         {
             constant->accept(this);
         }
-        for (variable_declaration *const variable : declaration->body.value().variables())
+        for (variable_declaration *const variable : declaration->body.value().variables)
         {
             variable->accept(this);
         }