From 8db23006d21a55f97d86bc9536ac6442671c6e43 Mon Sep 17 00:00:00 2001
From: Felipe Pena <felipensp@gmail.com>
Date: Sat, 14 Sep 2024 11:04:55 -0300
Subject: [PATCH] cgen: fix code generated for indexexpr with complex assigning
 (#22203)

---
 vlib/v/gen/c/assign.v                     | 18 +++++++++---
 vlib/v/gen/c/cgen.v                       |  1 +
 vlib/v/gen/c/index.v                      |  2 ++
 vlib/v/tests/indexexpr_with_assign_test.v | 34 +++++++++++++++++++++++
 4 files changed, 51 insertions(+), 4 deletions(-)
 create mode 100644 vlib/v/tests/indexexpr_with_assign_test.v

diff --git a/vlib/v/gen/c/assign.v b/vlib/v/gen/c/assign.v
index 5e58f97eba9bf2..60c872a8c8d864 100644
--- a/vlib/v/gen/c/assign.v
+++ b/vlib/v/gen/c/assign.v
@@ -228,6 +228,7 @@ fn (mut g Gen) assign_stmt(node_ ast.AssignStmt) {
 		mut ident := ast.Ident{
 			scope: unsafe { nil }
 		}
+		mut cur_indexexpr := -1
 		left_sym := g.table.sym(g.unwrap_generic(var_type))
 		is_va_list = left_sym.language == .c && left_sym.name == 'C.va_list'
 		if mut left is ast.Ident {
@@ -696,6 +697,9 @@ fn (mut g Gen) assign_stmt(node_ ast.AssignStmt) {
 				g.expr(left)
 			}
 			g.is_assign_lhs = false
+			if left is ast.IndexExpr && g.cur_indexexpr.len > 0 {
+				cur_indexexpr = g.cur_indexexpr.index(left.pos().pos)
+			}
 			if is_fixed_array_var || is_va_list {
 				if is_decl {
 					g.writeln(';')
@@ -703,7 +707,7 @@ fn (mut g Gen) assign_stmt(node_ ast.AssignStmt) {
 						continue
 					}
 				}
-			} else if !g.is_arraymap_set && !str_add && !op_overloaded {
+			} else if cur_indexexpr == -1 && !str_add && !op_overloaded {
 				g.write(' ${op} ')
 			} else if str_add || op_overloaded {
 				g.write(', ')
@@ -813,9 +817,10 @@ fn (mut g Gen) assign_stmt(node_ ast.AssignStmt) {
 			if str_add || op_overloaded {
 				g.write(')')
 			}
-			if g.is_arraymap_set {
+			if cur_indexexpr != -1 {
+				g.cur_indexexpr.delete(cur_indexexpr)
 				g.write(' })')
-				g.is_arraymap_set = false
+				g.is_arraymap_set = g.cur_indexexpr.len > 0
 			}
 			g.is_shared = false
 		}
@@ -841,6 +846,7 @@ fn (mut g Gen) gen_multi_return_assign(node &ast.AssignStmt, return_type ast.Typ
 	g.writeln(';')
 	mr_types := (return_sym.info as ast.MultiReturn).types
 	for i, lx in node.left {
+		mut cur_indexexpr := -1
 		mut is_auto_heap := false
 		mut ident := ast.Ident{
 			scope: unsafe { nil }
@@ -854,6 +860,9 @@ fn (mut g Gen) gen_multi_return_assign(node &ast.AssignStmt, return_type ast.Typ
 				is_auto_heap = lx.obj.is_auto_heap
 			}
 		}
+		if lx is ast.IndexExpr && g.cur_indexexpr.len > 0 {
+			cur_indexexpr = g.cur_indexexpr.index(lx.pos.pos)
+		}
 		styp := if ident.name in g.defer_vars { '' } else { g.typ(node.left_types[i]) }
 		if node.op == .decl_assign {
 			g.write('${styp} ')
@@ -889,7 +898,7 @@ fn (mut g Gen) gen_multi_return_assign(node &ast.AssignStmt, return_type ast.Typ
 				g.writeln(';')
 				g.writeln('memcpy(&${g.expr_string(lx)}, &${mr_var_name}.arg${i}, sizeof(${styp}));')
 			} else {
-				if g.is_arraymap_set {
+				if cur_indexexpr != -1 {
 					if is_auto_heap {
 						g.writeln('HEAP${noscan}(${styp}, ${mr_var_name}.arg${i}) });')
 					} else if is_option {
@@ -897,6 +906,7 @@ fn (mut g Gen) gen_multi_return_assign(node &ast.AssignStmt, return_type ast.Typ
 					} else {
 						g.writeln('${mr_var_name}.arg${i} });')
 					}
+					g.cur_indexexpr.delete(cur_indexexpr)
 				} else {
 					if is_auto_heap {
 						g.writeln(' = HEAP${noscan}(${styp}, ${mr_var_name}.arg${i});')
diff --git a/vlib/v/gen/c/cgen.v b/vlib/v/gen/c/cgen.v
index c066140bff0314..29e62ff7270f07 100644
--- a/vlib/v/gen/c/cgen.v
+++ b/vlib/v/gen/c/cgen.v
@@ -120,6 +120,7 @@ mut:
 	mtxs                      string            // array of mutexes if the `lock` has multiple variables
 	labeled_loops             map[string]&ast.Stmt
 	inner_loop                &ast.Stmt = unsafe { nil }
+	cur_indexexpr             []int          // list of nested indexexpr which generates array_set/map_set
 	shareds                   map[int]string // types with hidden mutex for which decl has been emitted
 	coverage_files            map[u64]&CoverageInfo
 	inside_ternary            int  // ?: comma separated statements on a single line
diff --git a/vlib/v/gen/c/index.v b/vlib/v/gen/c/index.v
index 7f86fb36a15c54..eca18dd164c5a9 100644
--- a/vlib/v/gen/c/index.v
+++ b/vlib/v/gen/c/index.v
@@ -175,6 +175,7 @@ fn (mut g Gen) index_of_array(node ast.IndexExpr, sym ast.TypeSymbol) {
 				g.write('*')
 			}
 		} else {
+			g.cur_indexexpr << node.pos.pos
 			g.is_arraymap_set = true // special handling of assign_op and closing with '})'
 			g.write('array_set(')
 			if !left_is_ptr || node.left_type.has_flag(.shared_f) {
@@ -398,6 +399,7 @@ fn (mut g Gen) index_of_map(node ast.IndexExpr, sym ast.TypeSymbol) {
 	get_and_set_types := val_sym.kind in [.struct_, .map, .array, .array_fixed]
 	if g.is_assign_lhs && !g.is_arraymap_set && !get_and_set_types {
 		if g.assign_op == .assign || info.value_type == ast.string_type {
+			g.cur_indexexpr << node.pos.pos
 			g.is_arraymap_set = true
 			g.write('map_set(')
 		} else {
diff --git a/vlib/v/tests/indexexpr_with_assign_test.v b/vlib/v/tests/indexexpr_with_assign_test.v
new file mode 100644
index 00000000000000..3fdc9c56c31597
--- /dev/null
+++ b/vlib/v/tests/indexexpr_with_assign_test.v
@@ -0,0 +1,34 @@
+fn test_1() {
+	mut test := []f64{}
+	test << 0
+	test[0] = if true {
+		sum := 1.2
+		sum
+	} else {
+		0
+	}
+	assert test == [1.2]
+}
+
+fn test_2() {
+	check := true
+	mut test := []f64{len: 100}
+	for i in 0 .. 100 {
+		test[i] = if check {
+			mut sum := 0.0
+			for x in 0 .. 100 {
+				if check {
+					sum += 1
+				}
+				if i <= x {
+					break
+				}
+			}
+			sum
+		} else {
+			0
+		}
+	}
+	assert test[0] == 1.0
+	assert test[99] == 100.0
+}