Add basic `defer` implementation.

We don't have any closures yet, so it's essentially the same as just
moving the statement after the `defer` keyword to the end of the block/
right before returning from the function.

4 files changed, 27 insertions, 0 deletions

diff --git a/src/ast.h b/src/ast.h
index 39d1a4e..09d15d1 100644
--- a/src/ast.h
+++ b/src/ast.h
@@ -28,6 +28,7 @@
   F(AST_CONDITIONAL, "conditional expression")                                 \
   F(AST_IF, "if statement")                                                    \
   F(AST_WHILE, "while statement")                                              \
+  F(AST_DEFER, "defer statement")                                              \
   F(AST_FOR, "for statement")                                                  \
   F(AST_VARDECL, "variable decl")                                              \
   F(AST_LOCAL_VAR, "local variable")                                           \
diff --git a/src/generator.c b/src/generator.c
index c7a50bc..403bf4a 100644
--- a/src/generator.c
+++ b/src/generator.c
@@ -12,6 +12,10 @@
 static int label_counter = 0;
 static Node *current_function = NULL;
 
+#define DEFER_STACK_SIZE 1024
+static Node *defer_stack[DEFER_STACK_SIZE];
+static i64 defer_stack_count = 0;
+
 void make_syscall(i64 syscall_no, FILE *out) {
 #if __APPLE__
     syscall_no += 0x2000000;
@@ -244,7 +248,14 @@ void generate_statement(Node *stmt, FILE *out)
 {
     if (stmt->type == AST_RETURN) {
         generate_expr_into_rax(stmt->unary_expr, out);
+        fprintf(out, "    push rax\n"); // Save the return value
+
+        // Run all the defer statements
+        for (int i = defer_stack_count - 1; i >= 0; i--)
+            generate_statement(defer_stack[i], out);
+
         // TODO: Only do this if we have local variables
+        fprintf(out, "    pop rax\n");
         fprintf(out, "    mov rsp, rbp\n");
         fprintf(out, "    pop rbp\n");
         fprintf(out, "    ret\n");
@@ -312,6 +323,9 @@ void generate_statement(Node *stmt, FILE *out)
 
     } else if (stmt->type == AST_BLOCK) {
         generate_block(stmt, out);
+    } else if (stmt->type == AST_DEFER) {
+        assert(defer_stack_count < DEFER_STACK_SIZE);
+        defer_stack[defer_stack_count++] = stmt->unary_expr;
     } else {
         // Once again, default to an expression here...
         generate_expr_into_rax(stmt, out);
@@ -320,9 +334,16 @@ void generate_statement(Node *stmt, FILE *out)
 
 void generate_block(Node *block, FILE *out)
 {
+    int cur_defer_pos = defer_stack_count;
     assert(block->type == AST_BLOCK);
     for (int i = 0; i < block->block.num_children; i++)
         generate_statement(block->block.children[i], out);
+    
+    assert(defer_stack_count - cur_defer_pos >= 0);
+    while (defer_stack_count > cur_defer_pos) {
+        Node *deferred = defer_stack[--defer_stack_count];
+        generate_statement(deferred, out);
+    }
 }
 
 void generate_function_header(Node *func, FILE *out)
diff --git a/src/parser.c b/src/parser.c
index a0ec7a1..0cd0e5f 100644
--- a/src/parser.c
+++ b/src/parser.c
@@ -507,6 +507,10 @@ Node *parse_statement(Lexer *lexer)
         node->loop.body = parse_statement(lexer);
     } else if (token.type == TOKEN_OPEN_BRACE) {
         node = parse_block(lexer);
+    } else if (token.type == TOKEN_DEFER) {
+        Lexer_next(lexer);
+        node = Node_new(AST_DEFER);
+        node->unary_expr = parse_statement(lexer);
     } else {
         // Default to trying to handle it as an expression
         node = parse_expression(lexer);
diff --git a/src/tokens.h b/src/tokens.h
index 7af4780..42f4647 100644
--- a/src/tokens.h
+++ b/src/tokens.h
@@ -44,6 +44,7 @@
 
 #define ENUM_KEYWORDS(F)                                                       \
   F(TOKEN_ELSE, "else")                                                        \
+  F(TOKEN_DEFER, "defer")                                                      \
   F(TOKEN_FN, "fn")                                                            \
   F(TOKEN_FOR, "for")                                                          \
   F(TOKEN_IF, "if")                                                            \