1 files changed, 25 insertions, 8 deletions

diff --git a/bootstrap/parse.c b/bootstrap/parse.c
index 31469df..4b56c2b 100644
--- a/bootstrap/parse.c
+++ b/bootstrap/parse.c
@@ -2443,6 +2443,7 @@ parsestmt(stmt_yielder_t yield, void *yarg, struct parser *P) {
    struct stmt st = {0};
    struct tok tok;
    const char *label = NULL;
+   static int deferage;
 
    if (lexmatch(P, &tok, '{')) {
       st.span = tok.span;
@@ -2520,15 +2521,16 @@ parsestmt(stmt_yielder_t yield, void *yarg, struct parser *P) {
       if (!P->curfn)
          fatal(P, tok.span, "return disallowed here");
       st.t = Sreturn;
+      st.ret.deferage = deferage;
       st.span = tok.span;
       if (!lexmatch(P, &tok, ';')) {
          P->targty = P->curfn->retty;
-         st.retex = exprdup(parseexpr(P));
+         st.ret.ex = exprdup(parseexpr(P));
          lexexpect(P, ';');
-         if (!typematchestarg(P->curfn->retty, st.retex->ty))
-            fatal(P, st.retex->span,
+         if (!typematchestarg(P->curfn->retty, st.ret.ex->ty))
+            fatal(P, st.ret.ex->span,
                   "incompatible type in return statement (%t, expected %t)",
-                  st.retex->ty, P->curfn->retty);
+                  st.ret.ex->ty, P->curfn->retty);
       } else if (P->curfn->retty != ty_void) {
          fatal(P, tok.span,
                "return statement in non-void function must return a value");
@@ -2559,6 +2561,18 @@ parsestmt(stmt_yielder_t yield, void *yarg, struct parser *P) {
          st.brkcon.id = decl->loopid;
       }
       lexexpect(P, ';');
+   } else if (lexmatch(P, &tok, TKkw_defer)) {
+      struct expr ex = parseexpr(P);
+      struct blockstmt *block = P->curblock;
+      struct defer *defer = malloc(sizeof *defer);
+      lexexpect(P, ';');
+      assert(block != NULL && "defer outside block");
+      defer->age = deferage++;
+      defer->ex = ex;
+      defer->next = block->defers;
+      block->defers = defer;
+      st.t = Sblock;
+      st.block = (struct blockstmt){0};
    } else if (isdecltokt((tok = lexpeek(P)).t)) {
       st.t = Sdecl;
       st.span = tok.span;
@@ -2606,10 +2620,13 @@ parseblock0(struct parser *P) {
 
    block.env.parent = P->curenv;
    pushenv(P, &block.env);
-   while ((tokt = lexpeek(P).t) != '}' && tokt != TKkw_case && tokt != ')') {
-      if (lexmatch(P, NULL, ';'))
-         continue;
-      parsestmt(parseblockstyield, &stmts, P);
+   block.defers = P->curblock ? P->curblock->defers : NULL;
+   WITH_TMPCHANGE(struct blockstmt *, P->curblock, &block) {
+      while ((tokt = lexpeek(P).t) != '}' && tokt != TKkw_case && tokt != ')') {
+         if (lexmatch(P, NULL, ';'))
+            continue;
+         parsestmt(parseblockstyield, &stmts, P);
+      }
    }
 
    vec_slice_cpy(&block.stmts, &stmts);