1 files changed, 32 insertions, 27 deletions

diff --git a/boot/ast.cc b/boot/ast.cc
index d4f015b..b8647f4 100644
--- a/boot/ast.cc
+++ b/boot/ast.cc
@@ -232,18 +232,7 @@ namespace elna::boot
         }
         if (declaration->body.has_value())
         {
-            for (variable_declaration *variable : declaration->body.value().variables)
-            {
-                variable->accept(this);
-            }
-            for (auto *statement : declaration->body.value().entry_point)
-            {
-                statement->accept(this);
-            }
-            if (declaration->body.value().return_expression != nullptr)
-            {
-                declaration->body.value().return_expression->accept(this);
-            }
+            traverse_body(this, declaration->body.value());
         }
     }
 
@@ -376,9 +365,9 @@ namespace elna::boot
         {
             procedure->accept(this);
         }
-        for (auto *entry_statement : unit->entry_point)
+        if (unit->entry_point.has_value())
         {
-            entry_statement->accept(this);
+            traverse_body(this, unit->entry_point.value());
         }
     }
 
@@ -989,20 +978,20 @@ namespace elna::boot
     procedure_body::procedure_body(std::vector<variable_declaration *>&& variables,
             std::vector<statement *>&& entry_point, expression *return_expression)
         : variables(std::move(variables)),
-        entry_point(std::move(entry_point)), return_expression(return_expression)
+        statements(std::move(entry_point)), return_expression(return_expression)
     {
     }
 
     procedure_body::procedure_body(procedure_body&& that) noexcept
         : variables(std::move(const_cast<std::vector<variable_declaration *>&>(that.variables))),
-        entry_point(std::move(const_cast<std::vector<statement *>&>(that.entry_point))),
+        statements(std::move(const_cast<std::vector<statement *>&>(that.statements))),
         return_expression(that.return_expression)
     {
     }
 
     procedure_body::~procedure_body()
     {
-        for (const statement *body_statement : this->entry_point)
+        for (const statement *body_statement : this->statements)
         {
             delete body_statement;
         }
@@ -1012,9 +1001,30 @@ namespace elna::boot
         }
     }
 
-    unit::unit(const source_position position)
+    void traverse_body(parser_visitor *visitor, const procedure_body& body)
+    {
+        for (variable_declaration *variable : body.variables)
+        {
+            variable->accept(visitor);
+        }
+        for (statement *body_statement : body.statements)
+        {
+            body_statement->accept(visitor);
+        }
+        if (body.return_expression != nullptr)
+        {
+            body.return_expression->accept(visitor);
+        }
+    }
+
+    unit::unit(const source_position position,
+            std::vector<import_declaration *>&& imports,
+            std::vector<type_declaration *>&& types,
+            std::vector<variable_declaration *>&& variables,
+            std::vector<procedure_declaration *>&& procedures)
         : node(position),
-        procedure_body(std::vector<variable_declaration *>{}, std::vector<statement *>{}, nullptr)
+        imports(std::move(imports)), types(std::move(types)), procedures(std::move(procedures)),
+        variables(std::move(variables))
     {
     }
 
@@ -1023,10 +1033,10 @@ namespace elna::boot
             std::vector<type_declaration *>&& types,
             std::vector<variable_declaration *>&& variables,
             std::vector<procedure_declaration *>&& procedures,
-            std::vector<statement *>&& entry_point)
+            std::optional<procedure_body>&& body)
         : node(position),
-        procedure_body(std::move(variables), std::move(entry_point)),
-        imports(std::move(imports)), types(std::move(types)), procedures(std::move(procedures))
+        imports(std::move(imports)), types(std::move(types)), procedures(std::move(procedures)),
+        variables(std::move(variables)), entry_point(std::move(body))
     {
     }
 
@@ -1035,11 +1045,6 @@ namespace elna::boot
         visitor->visit(this);
     }
 
-    bool unit::has_body() const
-    {
-        return !this->entry_point.empty();
-    }
-
     unit::~unit()
     {
         for (const procedure_declaration *procedure : this->procedures)