Make return statement part of the body and not a statement

3 files changed, 108 insertions, 161 deletions

diff --git a/boot/ast.cc b/boot/ast.cc
index fbecf11..118d92c 100644
--- a/boot/ast.cc
+++ b/boot/ast.cc
@@ -84,11 +84,6 @@ namespace elna::boot
         __builtin_unreachable();
     }
 
-    void empty_visitor::visit(return_statement *)
-    {
-        __builtin_unreachable();
-    }
-
     void empty_visitor::visit(defer_statement *)
     {
         __builtin_unreachable();
@@ -205,18 +200,22 @@ namespace elna::boot
         }
         if (declaration->body.has_value())
         {
-            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);
+            }
         }
     }
 
@@ -269,11 +268,6 @@ namespace elna::boot
         }
     }
 
-    void walking_visitor::visit(return_statement *statement)
-    {
-        statement->return_expression().accept(this);
-    }
-
     void walking_visitor::visit(defer_statement *statement)
     {
         for (struct statement *const block_statement : statement->statements)
@@ -340,12 +334,13 @@ namespace elna::boot
         {
             procedure->accept(this);
         }
-        if (unit->entry_point.has_value())
+        for (statement *const entry_statement : unit->entry_point)
         {
-            for (statement *const entry_statement : unit->entry_point.value())
-            {
-                entry_statement->accept(this);
-            }
+            entry_statement->accept(this);
+        }
+        if (unit->return_expression != nullptr)
+        {
+            unit->return_expression->accept(this);
         }
     }
 
@@ -773,70 +768,54 @@ namespace elna::boot
     }
 
     procedure_body::procedure_body(std::vector<constant_declaration *>&& constants,
-            std::vector<variable_declaration *>&& variables, std::vector<statement *>&& statements)
-        : m_variables(std::move(variables)), m_constants(std::move(constants)), m_statements(std::move(statements))
-    {
-    }
-
-    procedure_body::procedure_body(std::vector<constant_declaration*>&& constants,
-            std::vector<variable_declaration *>&& variables)
-        : procedure_body(std::move(constants), std::move(variables), std::vector<statement *>{})
+            std::vector<variable_declaration *>&& variables, std::vector<statement *>&& entry_point,
+            expression *return_expr)
+        : constants(std::move(constants)), variables(std::move(variables)),
+        entry_point(std::move(entry_point)), return_expression(return_expr)
     {
     }
 
     procedure_body::procedure_body(procedure_body&& that)
-        : m_variables(std::move(that.m_variables)), m_constants(std::move(that.m_constants)),
-        m_statements(std::move(that.m_statements))
-    {
-    }
-
-    procedure_body& procedure_body::operator=(procedure_body&& that)
-    {
-        std::swap(m_variables, that.m_variables);
-        std::swap(m_constants, that.m_constants);
-        std::swap(m_statements, that.m_statements);
-
-        return *this;
-    }
-
-    const std::vector<variable_declaration *>& procedure_body::variables()
-    {
-        return m_variables;
-    }
-
-    const std::vector<constant_declaration *>& procedure_body::constants()
-    {
-        return m_constants;
-    }
-
-    const std::vector<statement *>& procedure_body::statements()
+        : constants(std::move(const_cast<std::vector<constant_declaration *>&>(that.constants))),
+        variables(std::move(const_cast<std::vector<variable_declaration *>&>(that.variables))),
+        entry_point(std::move(const_cast<std::vector<statement *>&>(that.entry_point))),
+        return_expression(that.return_expression)
     {
-        return m_statements;
     }
 
     procedure_body::~procedure_body()
     {
-        for (statement *body_statement : this->statements())
+        for (statement *body_statement : this->entry_point)
         {
             delete body_statement;
         }
-        for (variable_declaration *variable : this->variables())
+        for (variable_declaration *variable : this->variables)
         {
             delete variable;
         }
-        for (constant_declaration *constant : this->constants())
+        for (constant_declaration *constant : this->constants)
         {
             delete constant;
         }
     }
 
     unit::unit(const struct position position)
-        : node(position)
+        : node(position), procedure_body(std::vector<constant_declaration *>{},
+        std::vector<variable_declaration *>{}, std::vector<statement *>{}, nullptr)
     {
     }
 
-    unit::unit(const struct position position, std::vector<statement *>&& entry_point)
-        : node(position), entry_point(std::make_optional<std::vector<statement *>>(std::move(entry_point)))
+    unit::unit(const struct position position,
+            std::vector<import_declaration *>&& imports,
+            std::vector<constant_declaration *>&& constants,
+            std::vector<type_declaration *>&& types,
+            std::vector<variable_declaration *>&& variables,
+            std::vector<procedure_declaration *>&& procedures,
+            std::vector<statement *>&& entry_point,
+            expression *const return_expression)
+        : node(position),
+        procedure_body(std::move(constants), std::move(variables), std::move(entry_point), return_expression),
+        imports(std::move(imports)), types(std::move(types)), procedures(std::move(procedures))
     {
     }
 
@@ -845,24 +824,21 @@ namespace elna::boot
         visitor->visit(this);
     }
 
+    bool unit::has_body() const
+    {
+        return !this->entry_point.empty() || this->return_expression != nullptr;
+    }
+
     unit::~unit()
     {
         for (procedure_declaration *procedure : this->procedures)
         {
             delete procedure;
         }
-        for (variable_declaration *variable : this->variables)
-        {
-            delete variable;
-        }
         for (type_declaration *type : this->types)
         {
             delete type;
         }
-        for (constant_declaration *constant : this->constants)
-        {
-            delete constant;
-        }
         for (import_declaration *declaration : this->imports)
         {
             delete declaration;
@@ -1206,26 +1182,6 @@ namespace elna::boot
         }
     }
 
-    return_statement::return_statement(const struct position position, expression *return_expression)
-        : node(position), m_return_expression(return_expression)
-    {
-    }
-
-    void return_statement::accept(parser_visitor *visitor)
-    {
-        visitor->visit(this);
-    }
-
-    expression& return_statement::return_expression()
-    {
-        return *m_return_expression;
-    }
-
-    return_statement::~return_statement()
-    {
-        delete m_return_expression;
-    }
-
     case_statement::case_statement(const struct position position,
             expression *condition, std::vector<switch_case>&& cases, std::vector<statement *> *alternative)
         : node(position), m_condition(condition), cases(std::move(cases)), alternative(alternative)
diff --git a/boot/parser.yy b/boot/parser.yy
index aace9a1..f1bf915 100644
--- a/boot/parser.yy
+++ b/boot/parser.yy
@@ -131,16 +131,14 @@ along with GCC; see the file COPYING3.  If not see
 %type <elna::boot::type_expression *> type_expression;
 %type <std::vector<elna::boot::type_expression *>> type_expressions;
 %type <elna::boot::traits_expression *> traits_expression;
-%type <elna::boot::expression *> expression operand simple_expression;
+%type <elna::boot::expression *> expression operand simple_expression return_statement;
 %type <elna::boot::unary_expression *> unary_expression;
 %type <elna::boot::binary_expression *> binary_expression;
 %type <std::vector<elna::boot::expression *>> expressions actual_parameter_list;
 %type <elna::boot::designator_expression *> designator_expression;
 %type <elna::boot::procedure_call*> call_expression;
-%type <elna::boot::return_statement *> return_statement;
 %type <elna::boot::statement *> statement;
-%type <std::vector<elna::boot::statement *>> statements;
-%type <std::unique_ptr<std::vector<elna::boot::statement *>>> statement_part;
+%type <std::vector<elna::boot::statement *>> statements statement_part;
 %type <elna::boot::procedure_declaration *> procedure_declaration;
 %type <elna::boot::procedure_type_expression *> procedure_heading;
 %type <elna::boot::procedure_type_expression::return_t> return_declaration;
@@ -159,49 +157,22 @@ along with GCC; see the file COPYING3.  If not see
 %type <std::vector<elna::boot::import_declaration *>> import_declarations import_part;
 %%
 program:
-    import_part constant_part type_part variable_part procedure_part statement_part "end" "."
-        {
-            boot::unit *tree;
-            if ($6)
-            {
-                tree = new boot::unit(boot::make_position(@1),
-				        std::move(*$6.release()));
-            }
-            else
-            {
-                tree = new boot::unit(boot::make_position(@1));
-            }
-			driver.tree.reset(tree);
-
-            std::swap(driver.tree->imports, $1);
-            std::swap(driver.tree->constants, $2);
-            std::swap(driver.tree->types , $3);
-            std::swap(driver.tree->variables, $4);
-            std::swap(driver.tree->procedures, $5);
-        }
-procedure_body: constant_part variable_part statement_part "end"
-        {
-            if ($3)
-            {
-                $$ = std::make_unique<boot::procedure_body>(std::move($1), std::move($2), std::move(*$3.release()));
-            }
-            else
-            {
-                $$ = std::make_unique<boot::procedure_body>(std::move($1), std::move($2));
-            }
-        }
-statement_part:
-    /* no statements */ {}
-    | "begin" statements { $$ = std::make_unique<std::vector<boot::statement *>>(std::move($2));; }
-    | return_statement
+    import_part constant_part type_part variable_part procedure_part statement_part return_statement "end" "."
         {
-            $$ = std::make_unique<std::vector<boot::statement *>>(std::vector<boot::statement *>{ $1 });
+            boot::unit *tree = new boot::unit(boot::make_position(@1),
+                    std::move($1), std::move($2), std::move($3), std::move($4), std::move($5),
+                    std::move($6), $7);
+			driver.tree.reset(tree);
         }
-    | "begin" statements ";" return_statement
+procedure_body:
+    constant_part variable_part statement_part return_statement "end"
         {
-            $$ = std::make_unique<std::vector<boot::statement *>>(std::move($2));
-            $$->push_back($4);
+            $$ = std::make_unique<boot::procedure_body>(std::move($1), std::move($2), std::move($3), $4);
         }
+
+statement_part:
+    /* no statements */ {}
+    | "begin" statements { std::swap($$, $2); }
 identifier_definition:
     IDENTIFIER "*" { $$ = boot::identifier_definition{ $1, true }; }
     | IDENTIFIER { $$ = boot::identifier_definition{ $1, false }; }
@@ -263,8 +234,9 @@ elsif_then_statements:
             $$.emplace($$.begin(), branch);
         }
     | /* no branches */ {}
-return_statement: "return" expression
-        { $$ = new boot::return_statement(boot::make_position(@1), $2); }
+return_statement:
+    "return" expression { $$ = $2; }
+    | /* no return statement */ { $$ = nullptr; }
 literal:
     INTEGER { $$ = new boot::literal<std::int32_t>(boot::make_position(@1), $1); }
     | WORD { $$ = new boot::literal<std::uint32_t>(boot::make_position(@1), $1); }
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);
         }