From 7fdce50718b205cddec812eb00fe23e35ef769a7 Mon Sep 17 00:00:00 2001
From: Alexander Medvednikov <alexander@medvednikov.com>
Date: Fri, 27 Mar 2020 11:21:00 +0100
Subject: [PATCH] cgen: write defer statements at the end of the function

---
 vlib/builtin/array_test.v        |  2 +-
 vlib/compiler/tests/defer_test.v | 32 ++++++++++++++++++++++----------
 vlib/v/gen/cgen.v                | 29 ++++++++++++++++++-----------
 3 files changed, 41 insertions(+), 22 deletions(-)

diff --git a/vlib/builtin/array_test.v b/vlib/builtin/array_test.v
index 98fef99ee8..3932e65f54 100644
--- a/vlib/builtin/array_test.v
+++ b/vlib/builtin/array_test.v
@@ -372,7 +372,7 @@ fn test_struct_print() {
 	a.b << b
 	a.b << b
 	println('QTODO: v2 struct .str() methods')
-	// assert a.str() == '{Test [{1 2}, {1 2}] }'
+	//assert a.str() == '{Test [{1 2}, {1 2}] }'
 	// assert b.str() == '{1 2}'
 	assert a.b.str() == '[{1 2}, {1 2}]'
 }
diff --git a/vlib/compiler/tests/defer_test.v b/vlib/compiler/tests/defer_test.v
index cc6063d808..1b2b7b4419 100644
--- a/vlib/compiler/tests/defer_test.v
+++ b/vlib/compiler/tests/defer_test.v
@@ -5,8 +5,12 @@ fn foo() string {
 
 fn foo2() string {
 	println('start')
-	defer { println('defer') }
-	defer { println('defer2') }
+	defer {
+		println('defer')
+	}
+	defer {
+		println('defer2')
+	}
 	println('end')
 	return foo()
 }
@@ -16,17 +20,23 @@ fn test_defer() {
 }
 
 fn set_num(i int, n mut Num) {
-	defer { n.val++ }
-	println("Hi")
+	defer {
+		println('exiting')
+		n.val++
+	}
+	println('Hi')
 	if i < 5 {
 		return
-	} else {
+	}
+	else {
 		n.val++
 	}
 }
 
 fn set_num_opt(n mut Num) ?int {
-	defer { n.val = 1 }
+	defer {
+		n.val = 1
+	}
 	return 99
 }
 
@@ -36,16 +46,18 @@ mut:
 }
 
 fn test_defer_early_exit() {
-	mut sum := Num{0}
-	for i in 0..10 {
+	mut sum := Num{
+		0}
+	for i in 0 .. 10 {
 		set_num(i, mut sum)
 	}
-	println("sum: $sum.val")
+	println('sum: $sum.val')
 	assert sum.val == 15
 }
 
 fn test_defer_option() {
-	mut ok := Num{0}
+	mut ok := Num{
+		0}
 	set_num_opt(mut ok) or {
 		assert false
 	}
diff --git a/vlib/v/gen/cgen.v b/vlib/v/gen/cgen.v
index 26dfff23cc..79a4fb87c0 100644
--- a/vlib/v/gen/cgen.v
+++ b/vlib/v/gen/cgen.v
@@ -404,17 +404,7 @@ fn (g mut Gen) stmt(node ast.Stmt) {
 		ast.Import {}
 		ast.Return {
 			if g.defer_stmts.len > 0 {
-				for defer_stmt in g.defer_stmts {
-					g.writeln('// defer')
-					if defer_stmt.ifdef.len > 0 {
-						g.writeln(defer_stmt.ifdef)
-						g.stmts(defer_stmt.stmts)
-						g.writeln('#endif')
-					}
-					else {
-						g.stmts(defer_stmt.stmts)
-					}
-				}
+				g.write_defer_stmts()
 			}
 			g.return_statement(it)
 		}
@@ -442,6 +432,20 @@ fn (g mut Gen) stmt(node ast.Stmt) {
 	}
 }
 
+fn (g mut Gen) write_defer_stmts() {
+	for defer_stmt in g.defer_stmts {
+		g.writeln('// defer')
+		if defer_stmt.ifdef.len > 0 {
+			g.writeln(defer_stmt.ifdef)
+			g.stmts(defer_stmt.stmts)
+			g.writeln('#endif')
+		}
+		else {
+			g.stmts(defer_stmt.stmts)
+		}
+	}
+}
+
 fn (g mut Gen) for_in(it ast.ForInStmt) {
 	if it.is_range {
 		// `for x in 1..10 {`
@@ -755,6 +759,9 @@ fn (g mut Gen) gen_fn_decl(it ast.FnDecl) {
 		}
 		g.writeln('return 0;')
 	}
+	if g.defer_stmts.len > 0 {
+		g.write_defer_stmts()
+	}
 	g.writeln('}')
 	g.defer_stmts = []
 	g.fn_decl = 0