parser: enable deferred stms for void and optional functions

compiler/parser.v

@@ -3891,12 +3891,12 @@ if (!$tmp) {
 fn (p mut Parser) return_st() {
 	p.check(.key_return)
 	p.fgen(' ')
+	deferred_text := p.get_deferred_text()
 	fn_returns := p.cur_fn.typ != 'void'
 	if fn_returns {
 		if p.tok == .rcbr {
 			p.error('`$p.cur_fn.name` needs to return `$p.cur_fn.typ`')
 		}
-		else {
 		ph := p.cgen.add_placeholder()
 		p.inside_return_expr = true
 		is_none := p.tok == .key_none
@@ -3935,26 +3935,14 @@ fn (p mut Parser) return_st() {
 			ret := p.cgen.cur_line.right(ph)
 			typ := expr_type.replace('Option_', '')
 			p.cgen.resetln('$expr_type $tmp = OPTION_CAST($expr_type)($ret);')
+			p.genln(deferred_text)
 			p.gen('return opt_ok(&$tmp, sizeof($typ))')
 		}
 		else {
 			ret := p.cgen.cur_line.right(ph)
 
-				// @emily33901: Scoped defer
-				// Check all of our defer texts to see if there is one at a higher scope level
-				// The one for our current scope would be the last so any before that need to be
-				// added.
-
-				mut total_text := ''
-
-				for text in p.cur_fn.defer_text {
-					if text != '' {
-						// In reverse order
-						total_text = text + total_text
-					}
-				}
-
-				if total_text == '' || expr_type == 'void*' {
+			if deferred_text == '' || expr_type == 'void*' {
+				// no defer{} necessary?
 				if expr_type == '${p.cur_fn.typ}*' {
 					p.cgen.resetln('return *$ret')
 				} else {
@@ -3963,19 +3951,19 @@ fn (p mut Parser) return_st() {
 			}  else {
 				tmp := p.get_tmp()
 				p.cgen.resetln('$expr_type $tmp = $ret;\n')
-					p.genln(total_text)
+				p.genln(deferred_text)
 				p.genln('return $tmp;')
 			}
 		}
 		p.check_types(expr_type, cur_fn_typ_chk)
 	}
-	}
 	else {
 		// Don't allow `return val` in functions that don't return anything
 		if !p.is_vweb && (p.tok == .name || p.tok == .number || p.tok == .str) {
 			p.error_with_token_index('function `$p.cur_fn.name` should not return a value', p.cur_fn.fn_name_token_idx)
 		}
 
+		p.genln(deferred_text)
 		if p.cur_fn.name == 'main' {
 			p.gen('return 0')
 		}
@@ -3986,6 +3974,21 @@ fn (p mut Parser) return_st() {
 	p.returns = true
 }
 
+fn (p Parser) get_deferred_text() string {
+	// @emily33901: Scoped defer
+	// Check all of our defer texts to see if there is one at a higher scope level
+	// The one for our current scope would be the last so any before that need to be
+	// added.
+	mut deferred_text := ''
+	for text in p.cur_fn.defer_text {
+		if text != '' {
+			// In reverse order
+			deferred_text = text + deferred_text
+		}
+	}
+	return deferred_text
+}
+
 fn prepend_mod(mod, name string) string {
 	return '${mod}__${name}'
 }

compiler/tests/defer_test.v

@@ -14,3 +14,38 @@ fn foo2() string {
 fn test_defer() {
 	assert foo2() == 'foo'
 }
+
+fn set_num(i int, n mut Num) {
+	defer { n.val+=1 }
+	println("Hi")
+	if i < 5 {
+		return
+	} else {
+		n.val+=1
+	}
+}
+
+fn set_num_opt(n mut Num) ?int {
+	defer { n.val = 1 }
+	return 99
+}
+
+struct Num {
+mut:
+	val int
+}
+
+fn test_defer_early_exit() {
+	mut sum := Num{0}
+	for i in 0..10 {
+		set_num(i, mut sum)
+	}
+	println("sum: $sum.val")
+	assert sum.val == 15
+}
+
+fn test_defer_option() {
+	mut ok := Num{0}
+	set_num_opt(mut ok)
+	assert ok.val == 1
+}