1 files changed, 27 insertions, 9 deletions

diff --git a/codegen.c b/codegen.c
index f126bda..79deb46 100644
--- a/codegen.c
+++ b/codegen.c
@@ -36,6 +36,21 @@ static struct scope* scope;
 static const struct fn_decl_node* active_fn;
 static unsigned long long branch_counter = 0;
 
+static void enter_scope(struct scope* child_scope) {
+    if (child_scope == NULL ||
+            child_scope->next_out != scope)
+        CGEN_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)
+        CGEN_PANIC("exit_scope: scopes are misaligned");
+
+    scope = child_scope->next_out;
+}
+
 static struct lval_def allocate_register(const struct type_def* type) {
     return (struct lval_def) {
         .loc = {
@@ -547,20 +562,19 @@ static void emit_group_contents(FILE* outfile, const struct group_node* node) {
 }
 
 static void emit_group(FILE* outfile, const struct group_node* node) {
-    if (node->scope->next_out != scope) CGEN_PANIC("scopes are misaligned");
-
-    scope = node->scope;
+    enter_scope(node->scope);
     scope->bp_offset = scope->next_out->bp_offset; /* don't reset bp */
 
     emit_group_contents(outfile, node);
 
     /* don't reset sp because alloca needs to work */
     scope->next_out->bp_offset = scope->bp_offset;
-    scope = scope->next_out;
+    exit_scope(node->scope);
 }
 
 static void emit_if(FILE* outfile, const struct if_node* node) {
-    /* TODO: make type checking real so we aren't passing a null pointer here */
+    enter_scope(node->scope);
+
     struct lval_def cond_result =
         allocate_temporary(outfile, node->cond->resolved_type);
     emit_expr(outfile, node->cond, &cond_result);
@@ -579,6 +593,8 @@ static void emit_if(FILE* outfile, const struct if_node* node) {
         emit_stmt(outfile, node->false_branch);
         fprintf(outfile, "branch%lld:\n", end_branch_num);
     }
+
+    exit_scope(node->scope);
 }
 
 static void emit_stmt(FILE* outfile, const struct stmt_node* node) {
@@ -604,6 +620,8 @@ static void emit_stmt(FILE* outfile, const struct stmt_node* node) {
 }
 
 static void emit_fn_decl(FILE* outfile, const struct fn_decl_node* node) {
+    enter_scope(node->scope);
+
     if (active_fn != NULL)
         CGEN_PANIC(
             "can't define function %s inside function %s",
@@ -615,8 +633,6 @@ static void emit_fn_decl(FILE* outfile, const struct fn_decl_node* node) {
     fprintf(outfile, "\tpush rbp\n");
     fprintf(outfile, "\tmov rbp, rsp\n");
 
-    if (node->scope->next_out != scope) CGEN_PANIC("scopes are misaligned");
-    scope = node->scope;
     scope->bp_offset = 0;
 
     long long spilled_bp_ofs = -16; // return address + old bp
@@ -644,15 +660,17 @@ static void emit_fn_decl(FILE* outfile, const struct fn_decl_node* node) {
         emit_mov(outfile, &arg_dst, &arg_src);
     }
 
+    enter_scope(node->body.scope);
     emit_group_contents(outfile, &node->body);
-
-    scope = scope->next_out;
+    exit_scope(node->body.scope);
 
     fprintf(outfile, RETURN_LABEL_FMT ":\n", node->name);
     fprintf(outfile, "\tmov rsp, rbp\n");
     fprintf(outfile, "\tpop rbp\n");
     fprintf(outfile, "\tret\n");
+
     active_fn = NULL;
+    exit_scope(node->scope);
 }
 
 static void emit_root_node(FILE* outfile, const struct root_node* node) {