From 64e9bdc213c2c75ceb0d476834e83a4ca46f0d63 Mon Sep 17 00:00:00 2001
From: yuyi <yuyi98@163.com>
Date: Wed, 8 Jul 2020 22:01:17 +0800
Subject: [PATCH] parser: fix non-void optional fn missing return value (fix
 #5736) (#5741)

---
 vlib/v/ast/ast.v           |  2 +-
 vlib/v/checker/checker.v   | 25 +++++++++++++++++++++----
 vlib/v/parser/fn.v         |  8 --------
 vlib/v/tests/option_test.v | 12 ++++++++++++
 4 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/vlib/v/ast/ast.v b/vlib/v/ast/ast.v
index 7d70520f5a..51fb3bcd1a 100644
--- a/vlib/v/ast/ast.v
+++ b/vlib/v/ast/ast.v
@@ -222,7 +222,6 @@ pub struct FnDecl {
 pub:
 	name          string
 	mod           string
-	stmts         []Stmt
 	args          []table.Arg
 	is_deprecated bool
 	is_pub        bool
@@ -242,6 +241,7 @@ pub:
 	file          string
 	is_generic    bool
 pub mut:
+	stmts         []Stmt
 	return_type   table.Type
 }
 
diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v
index feb54a8fde..14fe732621 100644
--- a/vlib/v/checker/checker.v
+++ b/vlib/v/checker/checker.v
@@ -1899,7 +1899,7 @@ fn (mut c Checker) stmt(node ast.Stmt) {
 			// node.typ = c.check_expr_opt_call(node.expr, table.void_type)
 		}
 		ast.FnDecl {
-			c.fn_decl(node)
+			c.fn_decl(mut node)
 		}
 		ast.ForCStmt {
 			c.in_for_count++
@@ -2988,14 +2988,14 @@ fn (c &Checker) fetch_and_verify_orm_fields(info table.Struct, pos token.Positio
 	return fields
 }
 
-fn (mut c Checker) fn_decl(node ast.FnDecl) {
+fn (mut c Checker) fn_decl(mut node ast.FnDecl) {
 	if node.is_generic && c.cur_generic_type == 0 { // need the cur_generic_type check to avoid inf. recursion
 		// loop thru each generic type and generate a function
 		for gen_type in c.table.fn_gen_types[node.name] {
 			c.cur_generic_type = gen_type
 			// sym:=c.table.get_type_symbol(gen_type)
 			// println('\ncalling check for $node.name for type $sym.name')
-			c.fn_decl(node)
+			c.fn_decl(mut node)
 		}
 		c.cur_generic_type = 0
 		return
@@ -3046,7 +3046,24 @@ fn (mut c Checker) fn_decl(node ast.FnDecl) {
 		}
 	}
 	c.expected_type = table.void_type
-	c.cur_fn = &node
+	c.cur_fn = node
+
+	// Add return if `fn(...) ? {...}` have no return at end
+	if node.return_type != table.void_type && node.return_type.has_flag(.optional) &&
+	(node.stmts.len == 0 || node.stmts[node.stmts.len - 1] !is ast.Return) {
+		sym := c.table.get_type_symbol(node.return_type)
+		if sym.kind == .void {
+			node.stmts << ast.Return{
+				pos: node.pos
+			}
+		} else {
+			node.stmts << ast.Return{
+				pos: node.pos
+				exprs: [ast.Expr(ast.None{pos: node.pos})]
+			}
+		}
+	}
+
 	c.stmts(node.stmts)
 	if node.language == .v && !node.no_body &&
 		node.return_type != table.void_type && !c.returns &&
diff --git a/vlib/v/parser/fn.v b/vlib/v/parser/fn.v
index 9e7bf4f70f..db5a11b0f9 100644
--- a/vlib/v/parser/fn.v
+++ b/vlib/v/parser/fn.v
@@ -284,14 +284,6 @@ fn (mut p Parser) fn_decl() ast.FnDecl {
 	body_start_pos := p.peek_tok.position()
 	if p.tok.kind == .lcbr {
 		stmts = p.parse_block_no_scope(true)
-		// Add return if `fn(...) ? {...}` have no return at end
-		if return_type != table.void_type &&
-			p.table.get_type_symbol(return_type).kind == .void && return_type.has_flag(.optional) &&
-			(stmts.len == 0 || stmts[stmts.len - 1] !is ast.Return) {
-			stmts << ast.Return{
-				pos: p.tok.position()
-			}
-		}
 	}
 	p.close_scope()
 	if !no_body && are_args_type_only {
diff --git a/vlib/v/tests/option_test.v b/vlib/v/tests/option_test.v
index ba1e432bf6..fe639a83d2 100644
--- a/vlib/v/tests/option_test.v
+++ b/vlib/v/tests/option_test.v
@@ -316,3 +316,15 @@ fn test_option_void_return_types_of_anon_fn_in_struct() {
 	}
 }
 
+fn get_string(param bool) ?string {
+	if param {
+		return 'Hello World'
+	}
+}
+
+fn test_option_auto_add_return_none() {
+	r := get_string(false) or {
+		'test'
+	}
+	assert r == 'test'
+}