save & load scratch in parse_stmt instead of reset in parse_proc

2 files changed, 8 insertions, 6 deletions

diff --git a/main.c b/main.c
index 9b22751..25469bc 100644
--- a/main.c
+++ b/main.c
@@ -347,8 +347,8 @@ void parse_if(Lexer *l, Proc *p) {
 }
 
 void parse_stmt(Lexer *l, Proc *p) {
-	/* TODO */
-	(void)l;
+	ArenaMark m;
+	arena_save(&p->scratch, &m);
 	switch (l->tok) {
 	case TOK_RETURN:
 		parse_return(l, p);
@@ -369,6 +369,7 @@ void parse_stmt(Lexer *l, Proc *p) {
 		lex_expected(l, TM_RBRACE);
 		break;
 	}
+	arena_load(&p->scratch, &m);
 }
 
 void parse_args_list(Lexer *l, Proc *proc) {
@@ -504,7 +505,6 @@ Proc *parse_proc(Lexer *l, Unit *u) {
 	while (l->tok != TOK_RBRACE) {
 		lex_expected_not(l, TM_EOF);
 		parse_stmt(l, proc);
-		arena_reset(&proc->scratch);
 	}
 	scope_pop(&proc->scope, &proc->graph);
 	lex_expected(l, TM_RBRACE);
diff --git a/test.lang b/test.lang
index 52e100b..40c3ff7 100644
--- a/test.lang
+++ b/test.lang
@@ -1,9 +1,11 @@
 func main(a, b i64) i64 {
 	let x = 2
 	if a < b {
-		x := 3 + (a / x)
+		x := 3
 	} else {
-		x := 4 + (x * b)
+		let t = a
+		a := b
+		b := t
 	}
-	return x
+	return a + x
 }