defer fix

5 files changed, 22 insertions, 15 deletions

diff --git a/bootstrap/all.h b/bootstrap/all.h
index 64ab9b1..cd51499 100644
--- a/bootstrap/all.h
+++ b/bootstrap/all.h
@@ -346,6 +346,7 @@ struct blockstmt {
    struct env env;
    slice_t(struct stmt) stmts;
    struct defer *defers;
+   int id;
 };
 
 struct expr {
@@ -496,7 +497,7 @@ struct stmt {
 };
 struct defer {
    struct defer *next;
-   int age;
+   int blockid, age;
    struct expr ex;
 } ;
 
diff --git a/bootstrap/cgen.c b/bootstrap/cgen.c
index 94ae420..5fae699 100644
--- a/bootstrap/cgen.c
+++ b/bootstrap/cgen.c
@@ -488,6 +488,14 @@ genblock(struct blockstmt block) {
    pri("{\n");
    for (int i = 0; i < block.stmts.n; ++i)
       genstmt(&block, &block.stmts.d[i]);
+   for (struct defer* defer = block.defers; defer; defer = defer->next) {
+      if (defer->blockid != block.id)
+         continue;
+      const char *p = fileid2path(defer->ex.span.fileid);
+      pri("\n#line %d %S\n", defer->ex.span.line, p, (u64)strlen(p));
+      genexpr(&defer->ex);
+      pri(";\n");
+   }
    pri("}\n");
 }
 
@@ -699,15 +707,7 @@ genfn(bool externp, const char *cname, struct fn *fn) {
    if (!fn->body) {
       pri(";\n");
    } else {
-      pri("{");
       genblock(fn->body->block);
-      for (struct defer* defer = fn->body->block.defers; defer; defer = defer->next) {
-         const char *p = fileid2path(defer->ex.span.fileid);
-         pri("\n#line %d %S\n", defer->ex.span.line, p, (u64)strlen(p));
-         genexpr(&defer->ex);
-         pri(";\n");
-      }
-      pri("}");
    }
 }
 
diff --git a/bootstrap/parse.c b/bootstrap/parse.c
index 4b56c2b..371289f 100644
--- a/bootstrap/parse.c
+++ b/bootstrap/parse.c
@@ -2569,6 +2569,7 @@ parsestmt(stmt_yielder_t yield, void *yarg, struct parser *P) {
       assert(block != NULL && "defer outside block");
       defer->age = deferage++;
       defer->ex = ex;
+      defer->blockid = P->curblock->id;
       defer->next = block->defers;
       block->defers = defer;
       st.t = Sblock;
@@ -2617,10 +2618,12 @@ parseblock0(struct parser *P) {
    vec_t(struct stmt) stmts = {0};
    struct blockstmt block = {0};
    int tokt;
+   static int id;
 
    block.env.parent = P->curenv;
    pushenv(P, &block.env);
    block.defers = P->curblock ? P->curblock->defers : NULL;
+   block.id = id++;
    WITH_TMPCHANGE(struct blockstmt *, P->curblock, &block) {
       while ((tokt = lexpeek(P).t) != '}' && tokt != TKkw_case && tokt != ')') {
          if (lexmatch(P, NULL, ';'))
diff --git a/examples/defer.cff b/examples/defer.cff
index f91b8dc..7833717 100644
--- a/examples/defer.cff
+++ b/examples/defer.cff
@@ -9,6 +9,9 @@ extern fn main() void {
       }
       return;
    }
+   {
+      defer 0;
+   }
    defer printf("bye\n");
    if 0 > 0 {
       return;
diff --git a/src/parse.cff b/src/parse.cff
index c3831ef..f2bbb41 100644
--- a/src/parse.cff
+++ b/src/parse.cff
@@ -922,6 +922,7 @@ fn parseexpandtepl(P *Parser, tepl *Tepl) *const Type {
       return cache.ty;
    }
    let env = mkenv(tepl.env, P.tlalloc);
+   defer envfree(env);
    with_tmpchange(P.curenv, env) {
       let decl *Decl #?;
       let expan Expan = {
@@ -932,7 +933,6 @@ fn parseexpandtepl(P *Parser, tepl *Tepl) *const Type {
       ty = parseagg(P, loc, tepl.kind, tname, &decl);
    }
    P.peektok = :None;
-   envfree(env);
    (as(*Type)ty).u.Agg.tpargs = tpargs;
    cache = P.tlalloc->alloc(sizeof(*cache), alignof(*cache));
    cache.next = tepl.cache;
@@ -2039,13 +2039,13 @@ fn pstwhile(P *Parser, loc Loc, label *const u8) Stmt {
    let body [#]Stmt #?;
    with_tmpchange(P.curloop, ++P.loopid) {
       let env = mkenv(P.curenv, P.alloc);
+      defer envfree(env);
       pushenv(P, env);
       if label {
          putdecl(P, loc, { label, loc, .u: :Label(P.loopid) });
       }
       body = parseblock(P);
       popenv(P);
-      envfree(env);
    }
    return { loc, :While { test, body, P.loopid }};
 }
@@ -2088,13 +2088,13 @@ fn pstfor(P *Parser, loc Loc, label *const u8) Stmt {
    let body [#]Stmt #?;
    with_tmpchange(P.curloop, ++P.loopid) {
       let env = mkenv(P.curenv, P.alloc);
+      defer envfree(env);
       pushenv(P, env);
       if label {
          putdecl(P, loc, { label, loc, .u: :Label(P.loopid) });
       }
       body = parseblock(P);
       popenv(P);
-      envfree(env);
    }
    return { loc, :For { ini, test, next, body, P.loopid }};
 }
@@ -2169,6 +2169,7 @@ fn psteuswitch(P *Parser, loc Loc, test Expr) Stmt {
          }
       }
       let env = mkenv(P.curenv, P.alloc);
+      defer envfree(env);
       if c.fld.ty != #null and lexmatch(P, &tok, :ident) {
          let ty = c.fld.ty;
          c.capt = tok.u.ident;
@@ -2185,7 +2186,6 @@ fn psteuswitch(P *Parser, loc Loc, test Expr) Stmt {
       if c.capt {
          popenv(P);
       }
-      envfree(env);
    }
 
 
@@ -2320,6 +2320,7 @@ fn parseblock0(P *Parser) [#]Stmt {
    let sts Vec<Stmt> = {};
    let tok Tok = {};
    let env = mkenv(P.curenv, P.alloc);
+   defer envfree(env);
    pushenv(P, env);
    while (tok = lexpeek(P)).t != '}' and tok.t != :kw_case and tok.t != ')' {
       if lexmatch(P, #null, ';') { continue; }
@@ -2330,7 +2331,6 @@ fn parseblock0(P *Parser) [#]Stmt {
       parsestmts(P, &pushstmt, &sts);
    }
    popenv(P);
-   envfree(env);
    return sts->move(P.alloc);
 }
 
@@ -2433,6 +2433,7 @@ fn parsefn(P *Parser, loc Loc, externp bool, name *const u8) *Decl {
    with_tmpchange(P.alloc, &Allocator { &Arena {}, &Arena:allocf }) {
       with_tmpchange(P.curfn, Fn) {
          let env = mkenv(P.curenv, P.alloc);
+         defer envfree(env);
          pushenv(P, env);
          foreach(name, i, Fn.paramnames) {
             if name {
@@ -2443,7 +2444,6 @@ fn parsefn(P *Parser, loc Loc, externp bool, name *const u8) *Decl {
          }
          Fn.body = :Some(parseblock(P));
          popenv(P);
-         envfree(env);
          (as(*Arena)P.alloc.a)->destroy();
       }
    }