ast, checker, cgen: implement if guard with multi return optional (#13273)

cmd/tools/vast/vast.v

@@ -1613,9 +1613,17 @@ fn (t Tree) par_expr(node ast.ParExpr) &Node {
 fn (t Tree) if_guard_expr(node ast.IfGuardExpr) &Node {
 	mut obj := new_object()
 	obj.add_terse('ast_type', t.string_node('IfGuardExpr'))
-	obj.add_terse('var_name', t.string_node(node.var_name))
+	obj.add_terse('vars', t.array_node_if_guard_var(node.vars))
 	obj.add_terse('expr', t.expr(node.expr))
 	obj.add_terse('expr_type', t.type_node(node.expr_type))
+	return obj
+}
+
+fn (t Tree) if_guard_var(node ast.IfGuardVar) &Node {
+	mut obj := new_object()
+	obj.add_terse('ast_type', t.string_node('IfGuardVar'))
+	obj.add_terse('name', t.string_node(node.name))
+	obj.add_terse('is_mut', t.bool_node(node.is_mut))
 	obj.add('pos', t.position(node.pos))
 	return obj
 }
@@ -2224,6 +2232,14 @@ fn (t Tree) array_node_struct_init_field(nodes []ast.StructInitField) &Node {
 	return arr
 }
 
+fn (t Tree) array_node_if_guard_var(nodes []ast.IfGuardVar) &Node {
+	mut arr := new_array()
+	for node in nodes {
+		arr.add_item(t.if_guard_var(node))
+	}
+	return arr
+}
+
 fn (t Tree) array_node_struct_init_embed(nodes []ast.StructInitEmbed) &Node {
 	mut arr := new_array()
 	for node in nodes {

vlib/v/ast/ast.v

@@ -1417,12 +1417,17 @@ pub mut:
 	is_used bool // asserts are used in _test.v files, as well as in non -prod builds of all files
 }
 
-// `if [x := opt()] {`
+pub struct IfGuardVar {
-pub struct IfGuardExpr {
+pub mut:
-pub:
+	name   string
-	var_name string
 	is_mut bool
 	pos    token.Position
+}
+
+// `if x := opt() {`
+pub struct IfGuardExpr {
+pub:
+	vars []IfGuardVar
 pub mut:
 	expr      Expr
 	expr_type Type

vlib/v/ast/str.v

@@ -456,7 +456,14 @@ pub fn (x Expr) str() string {
 			} + ')'
 		}
 		IfGuardExpr {
-			return x.var_name + ' := ' + x.expr.str()
+			mut s := ''
+			for i, var in x.vars {
+				s += var.name
+				if i != x.vars.len - 1 {
+					s += ', '
+				}
+			}
+			return s + ' := ' + x.expr.str()
 		}
 		StructInit {
 			sname := global_table.sym(x.typ).name

vlib/v/checker/checker.v

@@ -2962,7 +2962,19 @@ pub fn (mut c Checker) ident(mut node ast.Ident) ast.Type {
 						if mut obj.expr is ast.IfGuardExpr {
 							// new variable from if guard shouldn't have the optional flag for further use
 							// a temp variable will be generated which unwraps it
+							sym := c.table.sym(obj.expr.expr_type)
+							if sym.kind == .multi_return {
+								mr_info := sym.info as ast.MultiReturn
+								if mr_info.types.len == obj.expr.vars.len {
+									for vi, var in obj.expr.vars {
+										if var.name == node.name {
+											typ = mr_info.types[vi]
+										}
+									}
+								}
+							} else {
 								typ = obj.expr.expr_type.clear_flag(.optional)
+							}
 						} else {
 							typ = c.expr(obj.expr)
 						}

vlib/v/checker/if.v

@@ -206,6 +206,20 @@ pub fn (mut c Checker) if_expr(mut node ast.IfExpr) ast.Type {
 				st.check_c_expr() or { c.error('`if` expression branch has $err.msg', st.pos) }
 			}
 		}
+		if mut branch.cond is ast.IfGuardExpr {
+			sym := c.table.sym(branch.cond.expr_type)
+			if sym.kind == .multi_return {
+				mr_info := sym.info as ast.MultiReturn
+				if branch.cond.vars.len != mr_info.types.len {
+					c.error('if guard expects $mr_info.types.len variables, but got $branch.cond.vars.len',
+						branch.pos)
+				} else {
+					for vi, var in branch.cond.vars {
+						branch.scope.update_var_type(var.name, mr_info.types[vi])
+					}
+				}
+			}
+		}
 		// Also check for returns inside a comp.if's statements, even if its contents aren't parsed
 		if has_return := c.has_return(branch.stmts) {
 			if has_return {

vlib/v/checker/tests/if_guard_variables_err.out

@@ -0,0 +1,14 @@
+vlib/v/checker/tests/if_guard_variables_err.vv:6:2: error: if guard expects 3 variables, but got 1
+    4 |
+    5 | fn main() {
+    6 |     if r1 := create() {
+      |     ~~~~~~~~~~~~~~~~~
+    7 |         println(r1)
+    8 |     }
+vlib/v/checker/tests/if_guard_variables_err.vv:10:2: error: if guard expects 3 variables, but got 4
+    8 |     }
+    9 |
+   10 |     if r1, r2, r3, r4 := create() {
+      |     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+   11 |         println(r1)
+   12 |         println(r2)

vlib/v/checker/tests/if_guard_variables_err.vv

@@ -0,0 +1,16 @@
+fn create() ?(int, string, bool) {
+	return 5, 'aa', true
+}
+
+fn main() {
+	if r1 := create() {
+		println(r1)
+	}
+
+	if r1, r2, r3, r4 := create() {
+		println(r1)
+		println(r2)
+		println(r3)
+		println(r4)
+	}
+}

vlib/v/fmt/fmt.v

@@ -1889,10 +1889,16 @@ fn branch_is_single_line(b ast.IfBranch) bool {
 }
 
 pub fn (mut f Fmt) if_guard_expr(node ast.IfGuardExpr) {
-	if node.is_mut {
+	for i, var in node.vars {
+		if var.is_mut {
 			f.write('mut ')
 		}
-	f.write(node.var_name + ' := ')
+		f.write(var.name)
+		if i != node.vars.len - 1 {
+			f.write(', ')
+		}
+	}
+	f.write(' := ')
 	f.expr(node.expr)
 }
 

vlib/v/gen/c/cgen.v

@@ -4677,13 +4677,13 @@ fn (mut g Gen) if_expr(node ast.IfExpr) {
 						g.expr(branch.cond.expr)
 						g.writeln(', ${var_name}.state == 0) {')
 					}
-					if short_opt || branch.cond.var_name != '_' {
+					if short_opt || branch.cond.vars[0].name != '_' {
 						base_type := g.base_type(branch.cond.expr_type)
 						if short_opt {
-							cond_var_name := if branch.cond.var_name == '_' {
+							cond_var_name := if branch.cond.vars[0].name == '_' {
 								'_dummy_${g.tmp_count + 1}'
 							} else {
-								branch.cond.var_name
+								branch.cond.vars[0].name
 							}
 							g.write('\t$base_type $cond_var_name = ')
 							g.expr(branch.cond.expr)
@@ -4692,16 +4692,34 @@ fn (mut g Gen) if_expr(node ast.IfExpr) {
 							mut is_auto_heap := false
 							if branch.stmts.len > 0 {
 								scope := g.file.scope.innermost(ast.Node(branch.stmts[branch.stmts.len - 1]).position().pos)
-								if v := scope.find_var(branch.cond.var_name) {
+								if v := scope.find_var(branch.cond.vars[0].name) {
 									is_auto_heap = v.is_auto_heap
 								}
 							}
-							left_var_name := c_name(branch.cond.var_name)
+							if branch.cond.vars.len == 1 {
+								left_var_name := c_name(branch.cond.vars[0].name)
 								if is_auto_heap {
 									g.writeln('\t$base_type* $left_var_name = HEAP($base_type, *($base_type*)${var_name}.data);')
 								} else {
 									g.writeln('\t$base_type $left_var_name = *($base_type*)${var_name}.data;')
 								}
+							} else if branch.cond.vars.len > 1 {
+								for vi, var in branch.cond.vars {
+									left_var_name := c_name(var.name)
+									sym := g.table.sym(branch.cond.expr_type)
+									if sym.kind == .multi_return {
+										mr_info := sym.info as ast.MultiReturn
+										if mr_info.types.len == branch.cond.vars.len {
+											var_typ := g.typ(mr_info.types[vi])
+											if is_auto_heap {
+												g.writeln('\t$var_typ* $left_var_name = (HEAP($base_type, *($base_type*)${var_name}.data).arg$vi);')
+											} else {
+												g.writeln('\t$var_typ $left_var_name = (*($base_type*)${var_name}.data).arg$vi;')
+											}
+										}
+									}
+								}
+							}
 						}
 					}
 				}

vlib/v/gen/js/js.v

@@ -2833,18 +2833,18 @@ fn (mut g JsGen) gen_if_expr(node ast.IfExpr) {
 						g.expr(branch.cond.expr)
 						g.writeln(', ${var_name}.state == 0) {')
 					}
-					if short_opt || branch.cond.var_name != '_' {
+					if short_opt || branch.cond.vars[0].name != '_' {
 						if short_opt {
-							cond_var_name := if branch.cond.var_name == '_' {
+							cond_var_name := if branch.cond.vars[0].name == '_' {
 								'_dummy_${g.tmp_count + 1}'
 							} else {
-								branch.cond.var_name
+								branch.cond.vars[0].name
 							}
 							g.write('\tlet $cond_var_name = ')
 							g.expr(branch.cond.expr)
 							g.writeln(';')
 						} else {
-							g.writeln('\tlet $branch.cond.var_name = ${var_name}.data;')
+							g.writeln('\tlet $branch.cond.vars[0].name = ${var_name}.data;')
 						}
 					}
 				}

vlib/v/parser/if_match.v

@@ -82,36 +82,48 @@ fn (mut p Parser) if_expr(is_comptime bool) ast.IfExpr {
 		mut cond := ast.empty_expr()
 		mut is_guard := false
 
-		// `if x := opt() {`
+		// if guard `if x,y := opt() {`
-		if !is_comptime && (p.peek_tok.kind == .decl_assign
+		if !is_comptime && p.peek_token_after_var_list().kind == .decl_assign {
-			|| (p.tok.kind == .key_mut && p.peek_token(2).kind == .decl_assign)) {
 			p.open_scope()
 			is_guard = true
+			mut vars := []ast.IfGuardVar{}
+			for {
+				mut var := ast.IfGuardVar{}
 				mut is_mut := false
 				if p.tok.kind == .key_mut {
 					is_mut = true
-				p.check(.key_mut)
+					p.next()
 				}
-			var_pos := p.tok.position()
+				var.is_mut = is_mut
-			var_name := p.check_name()
+				var.pos = p.tok.position()
-			if p.scope.known_var(var_name) {
+				var.name = p.check_name()
-				p.error_with_pos('redefinition of `$var_name`', var_pos)
+
+				if p.scope.known_var(var.name) {
+					p.error_with_pos('redefinition of `$var.name`', var.pos)
+				}
+				vars << var
+				if p.tok.kind != .comma {
+					break
+				}
+				p.next()
 			}
 			comments << p.eat_comments()
 			p.check(.decl_assign)
 			comments << p.eat_comments()
 			expr := p.expr(0)
+
 			cond = ast.IfGuardExpr{
-				var_name: var_name
+				vars: vars
-				is_mut: is_mut
 				expr: expr
 			}
+			for var in vars {
 				p.scope.register(ast.Var{
-				name: var_name
+					name: var.name
-				is_mut: is_mut
+					is_mut: var.is_mut
 					expr: cond
-				pos: var_pos
+					pos: var.pos
 				})
+			}
 			prev_guard = true
 		} else {
 			prev_guard = false

vlib/v/parser/parser.v

@@ -438,6 +438,27 @@ pub fn (p &Parser) peek_token(n int) token.Token {
 	return p.scanner.peek_token(n - 2)
 }
 
+// peek token in if guard `if x,y := opt()` after var_list `x,y`
+pub fn (p &Parser) peek_token_after_var_list() token.Token {
+	mut n := 0
+	mut tok := p.tok
+	for {
+		if tok.kind == .key_mut {
+			n += 2
+		} else {
+			n++
+		}
+		tok = p.scanner.peek_token(n - 2)
+		if tok.kind != .comma {
+			break
+		} else {
+			n++
+			tok = p.scanner.peek_token(n - 2)
+		}
+	}
+	return tok
+}
+
 pub fn (mut p Parser) open_scope() {
 	p.scope = &ast.Scope{
 		parent: p.scope

vlib/v/tests/if_guard_with_multi_return_test.v

@@ -0,0 +1,18 @@
+fn create() ?(int, string, bool) {
+	return 5, 'aa', true
+}
+
+fn test_if_guard_with_multi_return() {
+	if r1, r2, r3 := create() {
+		println(r1)
+		assert r1 == 5
+
+		println(r2)
+		assert r2 == 'aa'
+
+		println(r3)
+		assert r3 == true
+	} else {
+		assert false
+	}
+}