From 82018034ef00a8178d258c6a30289c3071f458bb Mon Sep 17 00:00:00 2001
From: crthpl <56052645+crthpl@users.noreply.github.com>
Date: Sun, 15 May 2022 02:31:07 -0700
Subject: [PATCH] checker,cgen: fix if expressions in lock expression (#14384)

---
 vlib/v/checker/checker.v                     |  8 +++++---
 vlib/v/checker/tests/lock_already_locked.out |  7 +++++++
 vlib/v/gen/c/assign.v                        |  2 ++
 vlib/v/gen/c/if.v                            |  6 ++++--
 vlib/v/tests/shared_if_expr_test.v           | 19 +++++++++++++++++++
 5 files changed, 37 insertions(+), 5 deletions(-)
 create mode 100644 vlib/v/tests/shared_if_expr_test.v

diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v
index 174e2fe931..f2c475b720 100644
--- a/vlib/v/checker/checker.v
+++ b/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
 	}
diff --git a/vlib/v/checker/tests/lock_already_locked.out b/vlib/v/checker/tests/lock_already_locked.out
index 29696aa076..ae6152c3dd 100644
--- a/vlib/v/checker/tests/lock_already_locked.out
+++ b/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 |     }
diff --git a/vlib/v/gen/c/assign.v b/vlib/v/gen/c/assign.v
index 55056936f6..9296065456 100644
--- a/vlib/v/gen/c/assign.v
+++ b/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)
 						}
diff --git a/vlib/v/gen/c/if.v b/vlib/v/gen/c/if.v
index 3868d55c77..61503c8b85 100644
--- a/vlib/v/gen/c/if.v
+++ b/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)
diff --git a/vlib/v/tests/shared_if_expr_test.v b/vlib/v/tests/shared_if_expr_test.v
new file mode 100644
index 0000000000..02197558c3
--- /dev/null
+++ b/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]
+}