parser: enable deferred stms for void and optional functions

compiler/parser.v

@@ -3891,84 +3891,71 @@ 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()
-			ph := p.cgen.add_placeholder()
+		p.inside_return_expr = true
-			p.inside_return_expr = true
+		is_none := p.tok == .key_none
-			is_none := p.tok == .key_none
+		p.expected_type = p.cur_fn.typ
-			p.expected_type = p.cur_fn.typ
+		mut expr_type := p.bool_expression()
-			mut expr_type := p.bool_expression()
+		mut types := []string
-			mut types := []string
+		mut mr_values := [p.cgen.cur_line.right(ph).trim_space()]
-			mut mr_values := [p.cgen.cur_line.right(ph).trim_space()]
+		types << expr_type
-			types << expr_type
+		for p.tok == .comma {
-			for p.tok == .comma {
+			p.check(.comma)
-				p.check(.comma)
+			p.cgen.start_tmp()
-				p.cgen.start_tmp()
+			types << p.bool_expression()
-				types << p.bool_expression()
+			mr_values << p.cgen.end_tmp().trim_space()
-				mr_values << p.cgen.end_tmp().trim_space()
-			}
-			mut cur_fn_typ_chk := p.cur_fn.typ
-			// multiple returns
-			if types.len > 1 {
-				expr_type = types.join(',')
-				cur_fn_typ_chk = cur_fn_typ_chk.replace('_V_MulRet_', '').replace('_PTR_', '*').replace('_V_', ',')
-				mut ret_fields := ''
-				for ret_val_idx, ret_val in mr_values {
-					if ret_val_idx > 0 {
-						ret_fields += ','
-					}
-					ret_fields += '.var_$ret_val_idx=${ret_val}'
-				}
-				p.cgen.resetln('($p.cur_fn.typ){$ret_fields}')
-			}
-			p.inside_return_expr = false
-			// Automatically wrap an object inside an option if the function
-			// returns an option:
-			// `return val` => `return opt_ok(val)`
-			if p.cur_fn.typ.ends_with(expr_type) && !is_none &&
-				p.cur_fn.typ.starts_with('Option_') {
-				tmp := p.get_tmp()
-				ret := p.cgen.cur_line.right(ph)
-				typ := expr_type.replace('Option_', '')
-				p.cgen.resetln('$expr_type $tmp = OPTION_CAST($expr_type)($ret);')
-				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 expr_type == '${p.cur_fn.typ}*' {
-						p.cgen.resetln('return *$ret')
-					} else {
-						p.cgen.resetln('return $ret')
-					}
-				}  else {
-					tmp := p.get_tmp()
-					p.cgen.resetln('$expr_type $tmp = $ret;\n')
-					p.genln(total_text)
-					p.genln('return $tmp;')
-				}
-			}
-			p.check_types(expr_type, cur_fn_typ_chk)
 		}
+		mut cur_fn_typ_chk := p.cur_fn.typ
+		// multiple returns
+		if types.len > 1 {
+			expr_type = types.join(',')
+			cur_fn_typ_chk = cur_fn_typ_chk.replace('_V_MulRet_', '').replace('_PTR_', '*').replace('_V_', ',')
+			mut ret_fields := ''
+			for ret_val_idx, ret_val in mr_values {
+				if ret_val_idx > 0 {
+					ret_fields += ','
+				}
+				ret_fields += '.var_$ret_val_idx=${ret_val}'
+			}
+			p.cgen.resetln('($p.cur_fn.typ){$ret_fields}')
+		}
+		p.inside_return_expr = false
+		// Automatically wrap an object inside an option if the function
+		// returns an option:
+		// `return val` => `return opt_ok(val)`
+		if p.cur_fn.typ.ends_with(expr_type) && !is_none &&
+			p.cur_fn.typ.starts_with('Option_') {
+			tmp := p.get_tmp()
+			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)
+
+			if deferred_text == '' || expr_type == 'void*' {
+				// no defer{} necessary?
+				if expr_type == '${p.cur_fn.typ}*' {
+					p.cgen.resetln('return *$ret')
+				} else {
+					p.cgen.resetln('return $ret')
+				}
+			}  else {
+				tmp := p.get_tmp()
+				p.cgen.resetln('$expr_type $tmp = $ret;\n')
+				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
@@ -3976,6 +3963,7 @@ fn (p mut Parser) return_st() {
 			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
+}