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.

7 files changed, 104 insertions, 1 deletions

diff --git a/examples/defer.cup b/examples/defer.cup
new file mode 100644
index 0000000..e8e4778
--- /dev/null
+++ b/examples/defer.cup
@@ -0,0 +1,13 @@
+fn main() {
+    defer print(1);
+    {
+        defer print(2);
+        {
+            defer print(3);
+            print(4);
+        }
+        print(5);
+        return 0;
+    }
+    print(10);
+}
\ No newline at end of file
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")                                                            \
diff --git a/tests/common.sh b/tests/common.sh
index 1047708..965975c 100644
--- a/tests/common.sh
+++ b/tests/common.sh
@@ -57,7 +57,7 @@ function assert_stdout_text() {
         echo "----------------------------------------------" 
         echo "Test failed: executable returned non-0 code"
         echo "----------------------------------------------"
-        echo "$code"
+        echo "$1"
         exit 1
     fi
     if [[ "$output" != $2 ]]
diff --git a/tests/basics.sh b/tests/core.sh
index fc43847..71c539d 100755
--- a/tests/basics.sh
+++ b/tests/core.sh
@@ -123,4 +123,67 @@ fn main() {
     return x + y - 1;  
 }
 EOF
+echo " OK"
+
+echo -n "- Defer: "
+assert_stdout_text \
+"fn main() {
+    defer print(5);
+    print(4);
+}" \
+"4
+5"
+
+assert_stdout_text \
+"fn main() {
+    defer print(1);
+    {
+        defer print(2);
+        {
+            defer print(3);
+            print(4);
+        }
+        print(5);
+        return 0;
+    }
+    print(10);
+}" \
+"4
+3
+5
+2
+1"
+
+assert_stdout_text \
+"fn test() {
+    defer {
+        defer print(1);
+        print(2);
+    }
+    print(3);
+}
+fn main() {
+    defer print(4);
+    defer test();
+    print(10);
+}" \
+"10
+3
+2
+1
+4"
+
+assert_stdout_text \
+"let g: int;
+fn test(): int {
+    g = 5;
+    defer g = 10;
+    return g;
+}
+fn main() {
+    print(test());
+    print(g);
+}" \
+"5
+10"
 echo " OK"
\ No newline at end of file