defer seems to be fully working

7 files changed, 126 insertions, 16 deletions

diff --git a/cgen.c b/cgen.c
index d9dc189..ad0c673 100644
--- a/cgen.c
+++ b/cgen.c
@@ -672,19 +672,25 @@ static void cgen_fn_header(CGenerator *g, FnExpr *f, U64 which_are_const) {
 	}
 }
 
-static inline void cgen_deferred_stmts_from_block(CGenerator *g, Block *from) {
-	arr_foreach(from->deferred, StatementPtr, s) {
+static inline void cgen_deferred_from_block(CGenerator *g, Block *from) {
+	arr_foreach_reversed(from->deferred, StatementPtr, s) {
 		cgen_stmt(g, *s);
 	}
 }
 
 /* generates deferred statements in g->block, g->block->parent, ..., to) */
-static void cgen_deferred_stmts_up_to(CGenerator *g, Block *to) {
+static inline void cgen_deferred_up_to(CGenerator *g, Block *to) {
 	for (Block *b = g->block; b; b = b == to ? NULL : b->parent) {
-		cgen_deferred_stmts_from_block(g, b);
+		cgen_deferred_from_block(g, b);
 	}
 }
 
+/* same as cgen_deferred_up_to but doesn't generate to->deferred */
+static inline void cgen_deferred_up_to_not_including(CGenerator *g, Block *to) {
+	for (Block *b = g->block; b != to; b = b->parent)
+		cgen_deferred_from_block(g, b);
+}
+
 /* 
    Either set_expr or set_str should be NULL and either to_expr or to_str should be NULL 
    Also, set_str and/or to_str should be NULL
@@ -1551,7 +1557,7 @@ static void cgen_expr(CGenerator *g, Expression *e) {
 			}
 		}
 		cgen_block(g, &fo->body, NULL, CGEN_BLOCK_NOBRACES);
-		
+		cgen_deferred_from_block(g, &fo->body);
 		cgen_write(g, "}}");
 		if (fo->body.c.break_lbl) {
 			cgen_lbl(g, fo->body.c.break_lbl);
@@ -1710,7 +1716,7 @@ static void cgen_block(CGenerator *g, Block *b, const char *ret_name, U16 flags)
 	}
 	--g->indent_lvl;
 	if (!(flags & CGEN_BLOCK_NOBRACES)) {
-		cgen_deferred_stmts_from_block(g, b);
+		cgen_deferred_from_block(g, b);
 		arr_clear(&b->deferred);
 		cgen_write(g, "}");
 		if (b->c.break_lbl) {
@@ -1931,7 +1937,7 @@ static void cgen_ret(CGenerator *g, Block *returning_from, Expression *ret_expr)
 		cgen_set(g, NULL, "ret_", ret_expr, NULL);
 		cgen_nl(g);
 	}
-	cgen_deferred_stmts_up_to(g, returning_from);
+	cgen_deferred_up_to(g, returning_from);
 	if (f->ret_decls) {
 		if (f->ret_type.kind == TYPE_TUPLE) {
 			Expression tuple_expr = {0};
@@ -2019,16 +2025,20 @@ static void cgen_stmt(CGenerator *g, Statement *s) {
 				cgen_stmt(g, sub);
 		}
 		break;
-	case STMT_BREAK:
+	case STMT_BREAK: {
+		Block *b = s->referring_to;
+		cgen_deferred_up_to(g, b);
 		cgen_write(g, "goto ");
-		cgen_lbl(g, s->referring_to->c.break_lbl);
+		cgen_lbl(g, b->c.break_lbl);
 		cgen_writeln(g, ";");
-		break;
-	case STMT_CONT:
+	} break;
+	case STMT_CONT: {
+		Block *b = s->referring_to;
+		cgen_deferred_up_to_not_including(g, b);
 		cgen_write(g, "goto ");
-		cgen_lbl(g, s->referring_to->c.cont_lbl);
+		cgen_lbl(g, b->c.cont_lbl);
 		cgen_writeln(g, ";");
-		break;
+	} break;
 	case STMT_MESSAGE:
 		break;
 	case STMT_DEFER:
diff --git a/data_structures.c b/data_structures.c
index 494aaf1..06bbc90 100644
--- a/data_structures.c
+++ b/data_structures.c
@@ -214,7 +214,8 @@ You shouldn't rely on this, though, e.g. by doing
 #define arr_last(arr) arr_last_((void *)(arr), sizeof *(arr))
 /* one past last, or NULL if empty */
 #define arr_end(arr) arr_end_((void *)(arr), sizeof *(arr))
-#define arr_foreach(arr, type, var) for (type *var = arr, *var##_foreach_end = arr_end(arr); var < var##_foreach_end; ++var) /* NOTE: < is useful here because currently it's possible for var_foreach_end to be NULL but var could start out not null */
+#define arr_foreach(arr, type, var) for (type *var = (arr), *var##_foreach_end = arr_end(arr); var != var##_foreach_end; ++var)
+#define arr_foreach_reversed(arr, type, var) for (type *var = arr_last(arr), *var##_foreach_last = arr; var; var = var == var##_foreach_last ? NULL : (var-1))
 #define arr_remove_last(arr) arr_remove_last_((void **)(arr)), (void)sizeof **(arr)
 #define arr_remove_lasta(arr, a) arr_remove_lasta_((void **)(arr), sizeof **(arr), (a))
 #define arr_copya(out, in, a) do { assert(sizeof *(in) == sizeof **(out)); arr_copya_((void **)(out), (in), sizeof **(out), (a)); } while(0)
diff --git a/main.c b/main.c
index d1caf33..b24c794 100644
--- a/main.c
+++ b/main.c
@@ -8,7 +8,6 @@
 
 /* 
 TODO:
-make sure defer works with for, break, continue
 make sure you can't return a #C() (because of the current defer system)
 switch to:
 	static void 
diff --git a/test.toc b/test.toc
index abeebd1..28e639e 100644
--- a/test.toc
+++ b/test.toc
@@ -6,9 +6,31 @@ main ::= fn() {
 		defer io.puts("deferred from for");
 		io.puti(i);
 		if i == 2 {
+			defer io.puts("deferred from if1");
+			defer io.puts("deferred from if2");
+			defer io.puts("deferred from if3");
+			defer io.puts("deferred from if4");
+			defer io.puts("deferred from if5");
+			defer io.puts("deferred from if6");
+			defer io.puts("deferred from if7");
+			defer io.puts("deferred from if8");
+			continue;
+		}
+		if i == 8 {
+			break;
+		}
+	}
+	i := 0;
+	while {
+		defer io.puts("deferred from while");
+		i += 1;
+		io.puti(i);
+		if i % 2 == 0 { continue; }
+		if i == 7 {
 			defer io.puts("deferred from if");
-			return;
+			break;
 		}
+
 	}
 	io.puts("end of main()");
 }
diff --git a/tests/defer.toc b/tests/defer.toc
new file mode 100644
index 0000000..77a4c4e
--- /dev/null
+++ b/tests/defer.toc
@@ -0,0 +1,36 @@
+#include "io.toc", io;
+
+main ::= fn() {
+	defer io.puts("deferred from main()");
+	for i := 1..10 {
+		defer io.puts("deferred from for");
+		io.puti(i);
+		if i == 2 {
+			defer io.puts("deferred from if1");
+			defer io.puts("deferred from if2");
+			defer io.puts("deferred from if3");
+			defer io.puts("deferred from if4");
+			defer io.puts("deferred from if5");
+			defer io.puts("deferred from if6");
+			defer io.puts("deferred from if7");
+			defer io.puts("deferred from if8");
+			continue;
+		}
+		if i == 8 {
+			break;
+		}
+	}
+	i := 0;
+	while {
+		defer io.puts("deferred from while");
+		i += 1;
+		io.puti(i);
+		if i % 2 == 0 { continue; }
+		if i == 7 {
+			defer io.puts("deferred from if");
+			break;
+		}
+
+	}
+	io.puts("end of main()");
+}
diff --git a/tests/defer_expected b/tests/defer_expected
new file mode 100644
index 0000000..0a6daa7
--- /dev/null
+++ b/tests/defer_expected
@@ -0,0 +1,41 @@
+1
+deferred from for
+2
+deferred from if8
+deferred from if7
+deferred from if6
+deferred from if5
+deferred from if4
+deferred from if3
+deferred from if2
+deferred from if1
+deferred from for
+3
+deferred from for
+4
+deferred from for
+5
+deferred from for
+6
+deferred from for
+7
+deferred from for
+8
+deferred from for
+1
+deferred from while
+2
+deferred from while
+3
+deferred from while
+4
+deferred from while
+5
+deferred from while
+6
+deferred from while
+7
+deferred from if
+deferred from while
+end of main()
+deferred from main()
diff --git a/tests/test.sh b/tests/test.sh
index 4cb04dc..3d728d3 100755
--- a/tests/test.sh
+++ b/tests/test.sh
@@ -2,6 +2,7 @@
 
 tests='bf
 control_flow
+defer
 sizeof
 new
 arr