checker,cgen: fix if expressions in lock expression (#14384)

Author: crthpl

vlib/v/checker/checker.v

@@ -3537,6 +3537,7 @@ pub fn (mut c Checker) select_expr(mut node ast.SelectExpr) ast.Type {
 }
 
 pub fn (mut c Checker) lock_expr(mut node ast.LockExpr) ast.Type {
+	expected_type := c.expected_type
 	if c.rlocked_names.len > 0 || c.locked_names.len > 0 {
 		c.error('nested `lock`/`rlock` not allowed', node.pos)
 	}
@@ -3560,16 +3561,17 @@ pub fn (mut c Checker) lock_expr(mut node ast.LockExpr) ast.Type {
 		}
 	}
 	c.stmts(node.stmts)
-	c.rlocked_names = []
-	c.locked_names = []
 	// handle `x := rlock a { a.getval() }`
 	mut ret_type := ast.void_type
 	if node.stmts.len > 0 {
 		last_stmt := node.stmts.last()
 		if last_stmt is ast.ExprStmt {
-			ret_type = last_stmt.typ
+			c.expected_type = expected_type
+			ret_type = c.expr(last_stmt.expr)
 		}
 	}
+	c.rlocked_names = []
+	c.locked_names = []
 	if ret_type != ast.void_type {
 		node.is_expr = true
 	}

vlib/v/checker/tests/lock_already_locked.out

@@ -5,6 +5,13 @@ vlib/v/checker/tests/lock_already_locked.vv:11:3: error: nested `lock`/`rlock` n
       |         ~~~~~
    12 |             a.x++
    13 |         }
+vlib/v/checker/tests/lock_already_locked.vv:12:4: error: a has an `rlock` but needs a `lock`
+   10 |     lock a {
+   11 |         rlock a {
+   12 |             a.x++
+      |             ^
+   13 |         }
+   14 |     }
 vlib/v/checker/tests/lock_already_locked.vv:15:10: error: `a` is `shared` and must be `rlock`ed or `lock`ed to be used as non-mut argument to print
    13 |         }
    14 |     }

vlib/v/gen/c/assign.v

@@ -456,6 +456,8 @@ fn (mut g Gen) gen_assign_stmt(node_ ast.AssignStmt) {
 						}
 						if val is ast.ArrayInit {
 							g.array_init(val, ident.name)
+						} else if val_type.has_flag(.shared_f) {
+							g.expr_with_cast(val, val_type, var_type)
 						} else {
 							g.expr(val)
 						}

vlib/v/gen/c/if.v

@@ -78,7 +78,8 @@ fn (mut g Gen) if_expr(node ast.IfExpr) {
 				g.write(' ? ')
 			}
 			prev_expected_cast_type := g.expected_cast_type
-			if node.is_expr && g.table.sym(node.typ).kind == .sum_type {
+			if node.is_expr
+				&& (g.table.sym(node.typ).kind == .sum_type || node.typ.has_flag(.shared_f)) {
 				g.expected_cast_type = node.typ
 			}
 			g.stmts(branch.stmts)
@@ -204,7 +205,8 @@ fn (mut g Gen) if_expr(node ast.IfExpr) {
 		}
 		if needs_tmp_var {
 			prev_expected_cast_type := g.expected_cast_type
-			if node.is_expr && g.table.sym(node.typ).kind == .sum_type {
+			if node.is_expr
+				&& (g.table.sym(node.typ).kind == .sum_type || node.typ.has_flag(.shared_f)) {
 				g.expected_cast_type = node.typ
 			}
 			g.stmts_with_tmp_var(branch.stmts, tmp)

vlib/v/tests/shared_if_expr_test.v

@@ -0,0 +1,19 @@
+type AA = bool | int
+
+fn test_shared_if_expr() {
+	shared a := [1, 2, 3]
+	b := [4, 5, 6]
+	c := lock a {
+		if a == b { a } else { b }
+	}
+	assert c == [4, 5, 6]
+	d := lock a {
+		if a != b {
+			a << 5
+			a
+		} else {
+			b
+		}
+	}
+	assert d == [1, 2, 3, 5]
+}