Reuse container visiting logic

1 files changed, 14 insertions, 11 deletions

diff --git a/boot/type_check.cc b/boot/type_check.cc
index 71fd8ab..15a01dc 100644
--- a/boot/type_check.cc
+++ b/boot/type_check.cc
@@ -800,13 +800,19 @@ namespace elna::boot
         }
     }
 
-    void type_analysis_visitor::visit_entry_point(unit *unit)
+    void type_analysis_visitor::visit(unit *unit)
     {
-        this->current_procedure = this->bag.lookup("")->is_procedure();
-        this->bag.enter(this->current_procedure->scope);
-        walking_visitor::visit_entry_point(unit);
-        this->bag.leave();
-        this->current_procedure.reset();
+        visit_all(this, unit->imports);
+        visit_all(this, unit->declarations);
+
+        if (unit->entry_point.has_value())
+        {
+            this->current_procedure = this->bag.lookup("")->is_procedure();
+            this->bag.enter(this->current_procedure->scope);
+            visit_all(this, unit->entry_point.value());
+            this->bag.leave();
+            this->current_procedure.reset();
+        }
     }
 
     void type_analysis_visitor::visit(variable_declaration *declaration)
@@ -885,7 +891,7 @@ namespace elna::boot
     void type_analysis_visitor::traverse_block(block& body)
     {
         this->bag.enter(body.symbols);
-        walking_visitor::traverse_block(body);
+        boot::traverse_block(this, body);
         this->bag.leave();
     }
 
@@ -1409,10 +1415,7 @@ namespace elna::boot
             {
                 // Inference reads the argument types, so they are decorated
                 // before the parameter types they will be checked against exist.
-                for (expression *const argument : call->arguments)
-                {
-                    argument->accept(this);
-                }
+                visit_all(this, call->arguments);
                 inferred = procedure;
                 callable_type = infer_call(*call, *reference, *procedure);