diff --git a/cmd/tools/vtest-self.v b/cmd/tools/vtest-self.v index 9e6257573d..ba952637f4 100644 --- a/cmd/tools/vtest-self.v +++ b/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', diff --git a/vlib/v/ast/ast.v b/vlib/v/ast/ast.v index 19cb2fec35..5165a5e214 100644 --- a/vlib/v/ast/ast.v +++ b/vlib/v/ast/ast.v @@ -599,12 +599,13 @@ pub mut: // See: token.Kind.is_prefix pub struct PrefixExpr { pub: - op token.Kind - right Expr - pos token.Position + op token.Kind + 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 } diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v index 060e2cf0e0..4cddda3091 100644 --- a/vlib/v/checker/checker.v +++ b/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 } diff --git a/vlib/v/gen/c/cgen.v b/vlib/v/gen/c/cgen.v index d9a6df9e24..a8fabe1437 100644 --- a/vlib/v/gen/c/cgen.v +++ b/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() - guard_vars[i] = var_name - g.writeln('${g.typ(cond.expr_type)} $var_name;') + if cond.expr !is ast.IndexExpr && cond.expr !is ast.PrefixExpr { + var_name := g.new_tmp_var() + 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] - g.write('if ($var_name = ') - g.expr(branch.cond.expr) - g.writeln(', ${var_name}.state == 0) {') + mut var_name := guard_vars[i] + mut short_opt := false + if var_name == '' { + 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 { diff --git a/vlib/v/gen/c/index.v b/vlib/v/gen/c/index.v index 1b82bad8d4..7f94420575 100644 --- a/vlib/v/gen/c/index.v +++ b/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') } } diff --git a/vlib/v/tests/if_guard_test.v b/vlib/v/tests/if_guard_test.v new file mode 100644 index 0000000000..bbd5e2624b --- /dev/null +++ b/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] +} +