checker/cgen: enable `IfGuardExpr` for `a[k]` and `<-ch` (#9065)

cmd/tools/vtest-self.v

@@ -108,6 +108,7 @@ const (
 		'vlib/v/tests/go_call_generic_fn_test.v',
 		'vlib/v/tests/generics_test.v',
 		'vlib/v/tests/go_wait_2_test.v',
+		'vlib/v/tests/if_guard_test.v',
 		'vlib/v/tests/in_expression_test.v',
 		'vlib/v/tests/interface_edge_cases/assign_to_interface_field_test.v',
 		'vlib/v/tests/interface_fields_test.v',

vlib/v/ast/ast.v

@@ -599,12 +599,13 @@ pub mut:
 // See: token.Kind.is_prefix
 pub struct PrefixExpr {
 pub:
-	op    token.Kind
+	op  token.Kind
-	right Expr
+	pos token.Position
-	pos   token.Position
 pub mut:
 	right_type table.Type
+	right      Expr
 	or_block   OrExpr
+	is_option  bool // IfGuard
 }
 
 pub struct IndexExpr {
@@ -619,6 +620,7 @@ pub mut:
 	is_map    bool
 	is_array  bool
 	is_farray bool
+	is_option bool // IfGuard
 }
 
 pub struct IfExpr {
@@ -1035,9 +1037,9 @@ pub mut:
 pub struct IfGuardExpr {
 pub:
 	var_name string
-	expr     Expr
 	pos      token.Position
 pub mut:
+	expr      Expr
 	expr_type table.Type
 }
 

vlib/v/checker/checker.v

@@ -3752,7 +3752,25 @@ pub fn (mut c Checker) expr(node ast.Expr) table.Type {
 		ast.IfGuardExpr {
 			node.expr_type = c.expr(node.expr)
 			if !node.expr_type.has_flag(.optional) {
-				c.error('expression should return an option', node.expr.position())
+				mut no_opt := true
+				match mut node.expr {
+					ast.IndexExpr {
+						no_opt = false
+						node.expr_type = node.expr_type.set_flag(.optional)
+						node.expr.is_option = true
+					}
+					ast.PrefixExpr {
+						if node.expr.op == .arrow {
+							no_opt = false
+							node.expr_type = node.expr_type.set_flag(.optional)
+							node.expr.is_option = true
+						}
+					}
+					else {}
+				}
+				if no_opt {
+					c.error('expression should return an option', node.expr.position())
+				}
 			}
 			return table.bool_type
 		}

vlib/v/gen/c/cgen.v

@@ -2768,7 +2768,7 @@ fn (mut g Gen) expr(node ast.Expr) {
 			}
 		}
 		ast.PrefixExpr {
-			gen_or := node.op == .arrow && node.or_block.kind != .absent
+			gen_or := node.op == .arrow && (node.or_block.kind != .absent || node.is_option)
 			if node.op == .amp {
 				g.is_amp = true
 			}
@@ -2796,7 +2796,9 @@ fn (mut g Gen) expr(node ast.Expr) {
 				g.expr(node.right)
 				g.write(')')
 				if gen_or {
-					g.or_block(tmp_opt, node.or_block, elem_type)
+					if !node.is_option {
+						g.or_block(tmp_opt, node.or_block, elem_type)
+					}
 					if is_gen_or_and_assign_rhs {
 						elem_styp := g.typ(elem_type)
 						g.write('\n$cur_line*($elem_styp*)${tmp_opt}.data')
@@ -4150,11 +4152,16 @@ fn (mut g Gen) if_expr(node ast.IfExpr) {
 				is_guard = true
 				guard_idx = i
 				guard_vars = []string{len: node.branches.len}
+				g.writeln(';')
 				g.writeln('{ /* if guard */ ')
 			}
-			var_name := g.new_tmp_var()
+			if cond.expr !is ast.IndexExpr && cond.expr !is ast.PrefixExpr {
-			guard_vars[i] = var_name
+				var_name := g.new_tmp_var()
-			g.writeln('${g.typ(cond.expr_type)} $var_name;')
+				guard_vars[i] = var_name
+				g.writeln('${g.typ(cond.expr_type)} $var_name;')
+			} else {
+				guard_vars[i] = ''
+			}
 		}
 	}
 	for i, branch in node.branches {
@@ -4172,13 +4179,28 @@ fn (mut g Gen) if_expr(node ast.IfExpr) {
 		} else {
 			match branch.cond {
 				ast.IfGuardExpr {
-					var_name := guard_vars[i]
+					mut var_name := guard_vars[i]
-					g.write('if ($var_name = ')
+					mut short_opt := false
-					g.expr(branch.cond.expr)
+					if var_name == '' {
-					g.writeln(', ${var_name}.state == 0) {')
+						short_opt = true // we don't need a further tmp, so use the one we'll get later
+						var_name = g.new_tmp_var()
+						guard_vars[i] = var_name // for `else`
+						g.tmp_count--
+						g.writeln('if (${var_name}.state == 0) {')
+					} else {
+						g.write('if ($var_name = ')
+						g.expr(branch.cond.expr)
+						g.writeln(', ${var_name}.state == 0) {')
+					}
 					if branch.cond.var_name != '_' {
 						base_type := g.base_type(branch.cond.expr_type)
-						g.writeln('\t$base_type $branch.cond.var_name = *($base_type*)${var_name}.data;')
+						if short_opt {
+							g.write('\t$base_type $branch.cond.var_name = ')
+							g.expr(branch.cond.expr)
+							g.writeln(';')
+						} else {
+							g.writeln('\t$base_type $branch.cond.var_name = *($base_type*)${var_name}.data;')
+						}
 					}
 				}
 				else {

vlib/v/gen/c/index.v

@@ -92,7 +92,7 @@ fn (mut g Gen) range_expr(node ast.IndexExpr, range ast.RangeExpr) {
 }
 
 fn (mut g Gen) index_of_array(node ast.IndexExpr, sym table.TypeSymbol) {
-	gen_or := node.or_expr.kind != .absent
+	gen_or := node.or_expr.kind != .absent || node.is_option
 	left_is_ptr := node.left_type.is_ptr()
 	info := sym.info as table.Array
 	elem_type_str := g.typ(info.elem_type)
@@ -248,7 +248,9 @@ fn (mut g Gen) index_of_array(node ast.IndexExpr, sym table.TypeSymbol) {
 			g.writeln('} else {')
 			g.writeln('\t${tmp_opt}.state = 2; ${tmp_opt}.err = (Error){.msg=_SLIT("array index out of range"), .code=0};')
 			g.writeln('}')
-			g.or_block(tmp_opt, node.or_expr, elem_type)
+			if !node.is_option {
+				g.or_block(tmp_opt, node.or_expr, elem_type)
+			}
 			g.write('\n$cur_line*($elem_type_str*)${tmp_opt}.data')
 		}
 	}
@@ -287,7 +289,7 @@ fn (mut g Gen) index_of_fixed_array(node ast.IndexExpr, sym table.TypeSymbol) {
 }
 
 fn (mut g Gen) index_of_map(node ast.IndexExpr, sym table.TypeSymbol) {
-	gen_or := node.or_expr.kind != .absent
+	gen_or := node.or_expr.kind != .absent || node.is_option
 	left_is_ptr := node.left_type.is_ptr()
 	info := sym.info as table.Map
 	key_type_str := g.typ(info.key_type)
@@ -414,7 +416,9 @@ fn (mut g Gen) index_of_map(node ast.IndexExpr, sym table.TypeSymbol) {
 			g.writeln('\t${tmp_opt}.state = 2; ${tmp_opt}.err = (Error){.msg=_SLIT("array index out of range"), .code=0};')
 
 			g.writeln('}')
-			g.or_block(tmp_opt, node.or_expr, elem_type)
+			if !node.is_option {
+				g.or_block(tmp_opt, node.or_expr, elem_type)
+			}
 			g.write('\n$cur_line*($elem_type_str*)${tmp_opt}.data')
 		}
 	}

vlib/v/tests/if_guard_test.v

@@ -0,0 +1,61 @@
+fn f(n int) ?f64 {
+	if n < 0 {
+		return error('negative')
+	}
+	return 1.5 * f64(n)
+}
+
+fn test_fn_return() {
+	mut res := []f64{cap:2}
+	for m in [-3, 5] {
+		if x := f(m) {
+			res << x
+		} else {
+			res << 31.0
+		}
+	}
+	assert res == [31.0, 7.5]
+}
+
+fn test_map_get() {
+	mut m := {'xy': 5, 'zu': 7}
+	mut res := []int{cap:2}
+	for k in ['jk', 'zu'] {
+		if x := m[k] {
+			res << x
+		} else {
+			res << -17
+		}
+	}
+	assert res == [-17, 7]
+}
+
+fn test_array_get() {
+	mut a := [12.5, 6.5, -17.25]
+	mut res := []f64{cap:2}
+	for i in [1, 4] {
+		if x := a[i] {
+			res << x
+		} else {
+			res << -23.0
+		}
+	}
+	assert res == [6.5, -23.0]
+}
+
+fn test_chan_pop() {
+	mut res := []f64{cap:3}
+	ch := chan f64{cap: 10}
+	ch <- 6.75
+	ch <- -3.25
+	ch.close()
+	for _ in 0 .. 3 {
+		if x:= <-ch {
+			res << x
+		} else {
+			res << -37.5
+		}
+	}
+	assert res == [6.75, -3.25, -37.5]
+}
+