1 files changed, 25 insertions, 6 deletions

diff --git a/type_checker.c b/type_checker.c
index 2842ea2..aabd7c9 100644
--- a/type_checker.c
+++ b/type_checker.c
@@ -33,6 +33,21 @@ static void assert_cast_compatible(
     }
 }
 
+static void enter_scope(struct scope* child_scope) {
+    if (child_scope == NULL ||
+            child_scope->next_out != scope)
+        TYPE_PANIC("enter_scope: scopes are misaligned");
+
+    scope = child_scope;
+}
+
+static void exit_scope(struct scope* child_scope) {
+    if (child_scope != scope || child_scope->next_out == NULL)
+        TYPE_PANIC("exit_scope: scopes are misaligned");
+
+    scope = child_scope->next_out;
+}
+
 static const struct type_def* resolve_type_ref(struct type_ref_node* node) {
     /* TODO: support pointers and such */
     if (node->def_ref == NULL) TYPE_PANIC("use of unresolved type");
@@ -160,9 +175,13 @@ static void type_check_return(struct return_node* node) {
 }
 
 static void type_check_if(struct if_node* node) {
+    enter_scope(node->scope);
+
     type_check_expr(node->cond);
     type_check_stmt(node->true_branch);
     if (node->false_branch != NULL) type_check_stmt(node->false_branch);
+
+    exit_scope(node->scope);
 }
 
 static void type_check_stmt(struct stmt_node* node) {
@@ -188,19 +207,19 @@ static void type_check_stmt(struct stmt_node* node) {
 }
 
 static void type_check_group(struct group_node* node) {
-    if (node->scope->next_out != scope) TYPE_PANIC("scopes are borked");
-    scope = node->scope;
+    enter_scope(node->scope);
+
     struct stmt_node* stmt = node->head;
     while (stmt != NULL) {
         type_check_stmt(stmt);
         stmt = stmt->next;
     }
-    scope = scope->next_out;
+
+    exit_scope(node->scope);
 }
 
 static void type_check_fn_decl(struct fn_decl_node* node) {
-    if (node->scope->next_out != scope) TYPE_PANIC("scopes are borked");
-    scope = node->scope;
+    enter_scope(node->scope);
 
     node->return_type.def_ref = resolve_type_ref(&node->return_type);
 
@@ -211,7 +230,7 @@ static void type_check_fn_decl(struct fn_decl_node* node) {
 
     type_check_group(&node->body);
 
-    scope = scope->next_out;
+    exit_scope(node->scope);
 }
 
 static void type_check_root(struct root_node* node) {