From 7815a5495c95a45ac59386339952d78b8f617c44 Mon Sep 17 00:00:00 2001
From: Tanel Liiv <204019+dkull@users.noreply.github.com>
Date: Fri, 8 May 2020 23:49:45 +0300
Subject: [PATCH] checker: handle multireturn fn calls as if/match last
 expressions

---
 vlib/v/checker/checker.v           | 40 +++++++++++++--------
 vlib/v/gen/cgen.v                  | 14 ++++++--
 vlib/v/tests/complex_assign_test.v | 58 ++++++++++++++++++++++++++++++
 3 files changed, 95 insertions(+), 17 deletions(-)
 create mode 100644 vlib/v/tests/complex_assign_test.v

diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v
index a404e017c2..6f3a2f4e8f 100644
--- a/vlib/v/checker/checker.v
+++ b/vlib/v/checker/checker.v
@@ -1099,32 +1099,38 @@ pub fn (mut c Checker) enum_decl(decl ast.EnumDecl) {
 
 pub fn (mut c Checker) assign_stmt(assign_stmt mut ast.AssignStmt) {
 	c.expected_type = table.none_type // TODO a hack to make `x := if ... work`
-	if assign_stmt.right[0] is ast.CallExpr {
-		call_expr := assign_stmt.right[0] as ast.CallExpr
+	right_first := assign_stmt.right[0]
+	mut right_len := assign_stmt.right.len
+	if right_first is ast.CallExpr || right_first is ast.IfExpr || right_first is ast.MatchExpr {
 		right_type0 := c.expr(assign_stmt.right[0])
 		assign_stmt.right_types = [right_type0]
 		right_type_sym0 := c.table.get_type_symbol(right_type0)
-		mut right_len := if right_type0 == table.void_type { 0 } else { assign_stmt.right.len }
+		right_len = if right_type0 == table.void_type { 0 } else { right_len }
 		if right_type_sym0.kind == .multi_return {
 			assign_stmt.right_types = right_type_sym0.mr_info().types
 			right_len = assign_stmt.right_types.len
 		}
 		if assign_stmt.left.len != right_len {
-			c.error('assignment mismatch: $assign_stmt.left.len variable(s) but `${call_expr.name}()` returns $right_len value(s)',
-				assign_stmt.pos)
-			return
-		}
-	} else {
-		if assign_stmt.left.len != assign_stmt.right.len {
-			c.error('assignment mismatch: $assign_stmt.left.len variable(s) $assign_stmt.right.len value(s)',
-				assign_stmt.pos)
-			return
+			if right_first is ast.CallExpr {
+				call_expr := assign_stmt.right[0] as ast.CallExpr
+				c.error('assignment mismatch: $assign_stmt.left.len variable(s) but `${call_expr.name}()` returns $right_len value(s)',
+					assign_stmt.pos)
+				return
+			} else {
+				c.error('assignment mismatch: $assign_stmt.left.len variable(s) $right_len value(s)',
+					assign_stmt.pos)
+				return
+			}
 		}
+	} else if assign_stmt.left.len != right_len {
+		c.error('assignment mismatch: $assign_stmt.left.len variable(s) $assign_stmt.right.len value(s)',
+			assign_stmt.pos)
+		return
 	}
 	mut scope := c.file.scope.innermost(assign_stmt.pos.pos)
 	for i, _ in assign_stmt.left {
 		mut ident := assign_stmt.left[i]
-		if assign_stmt.right_types.len < assign_stmt.right.len {
+		if assign_stmt.right_types.len < right_len {
 			assign_stmt.right_types << c.expr(assign_stmt.right[i])
 		}
 		val_type := assign_stmt.right_types[i]
@@ -1993,7 +1999,7 @@ pub fn (mut c Checker) if_expr(node mut ast.IfExpr) table.Type {
 	}
 	if node.has_else && node.is_expr {
 		last_branch := node.branches[node.branches.len - 1]
-		if last_branch.stmts.len > 0 {
+		if last_branch.stmts.len > 0 && node.branches[0].stmts.len > 0 {
 			match last_branch.stmts[last_branch.stmts.len - 1] {
 				ast.ExprStmt {
 					// type_sym := p.table.get_type_symbol(it.typ)
@@ -2008,8 +2014,14 @@ pub fn (mut c Checker) if_expr(node mut ast.IfExpr) table.Type {
 				}
 				else {}
 			}
+		} else {
+			c.error('`if` expression needs returns in both branches', node.pos)
 		}
 	}
+	// won't yet work due to eg: if true { println('foo') }
+	/*if node.is_expr && !node.has_else {
+		c.error('`if` expression needs `else` clause. remove return values or add `else`', node.pos)
+	}*/
 	return table.bool_type
 }
 
diff --git a/vlib/v/gen/cgen.v b/vlib/v/gen/cgen.v
index c3d38a4455..cc73b79c93 100644
--- a/vlib/v/gen/cgen.v
+++ b/vlib/v/gen/cgen.v
@@ -837,9 +837,17 @@ fn (mut g Gen) gen_assign_stmt(assign_stmt ast.AssignStmt) {
 		g.write('static ')
 	}
 	mut return_type := table.void_type
-	if assign_stmt.right[0] is ast.CallExpr {
-		it := assign_stmt.right[0] as ast.CallExpr
-		return_type = it.return_type
+	match assign_stmt.right[0] {
+		ast.CallExpr {
+			return_type = it.return_type
+		}
+		ast.IfExpr {
+			return_type = it.typ
+		}
+		ast.MatchExpr {
+			return_type = it.return_type
+		}
+		else {}
 	}
 	mut is_multi := false
 	// json_test failed w/o this check
diff --git a/vlib/v/tests/complex_assign_test.v b/vlib/v/tests/complex_assign_test.v
new file mode 100644
index 0000000000..c675960d02
--- /dev/null
+++ b/vlib/v/tests/complex_assign_test.v
@@ -0,0 +1,58 @@
+
+fn multireturner(n int, s string) (int, string) {
+	return n + 1, s
+}
+
+fn test_assign_multireturn_expression() {
+	a, b := if true {
+		multireturner(13, 'awesome')
+	} else {
+		multireturner(-1, 'notawesome')
+	}
+	assert a == 14
+	assert b == 'awesome'
+
+	c, d := if false {
+		multireturner(-1, 'notawesome')
+	} else if true {
+		multireturner(17, 'awesomer')
+	} else {
+		multireturner(-1, 'notawesome')
+	}
+	assert c == 18
+	assert d == 'awesomer'
+
+	e, f := if false {
+		multireturner(-1, 'notawesome')
+	} else if false {
+		multireturner(-1, 'notawesome')
+	} else {
+		multireturner(17, 'awesomer')
+	}
+	assert e == 18
+	assert f == 'awesomer'
+
+	g, h := match true {
+		true { multireturner(0, 'good') }
+		false { multireturner(100, 'bad') }
+		else { multireturner(200, 'bad') }
+	}
+	assert g == 1
+	assert h == 'good'
+
+	i, j := match true {
+		false { multireturner(100, 'bad') }
+		else { multireturner(0, 'good') }
+	}
+	assert i == 1
+	assert j == 'good'
+
+	// TODO: returning non-function calls does not work yet due to parsing issues
+	/*
+	e, f := if true {
+		1, 'awesome'
+	} else {
+		0, 'bad'
+	}
+	*/
+}